refactor: move next_trial_id counter from Study into Storage trait

src/storage/journal.rs

@@ -115,12 +115,19 @@ impl<V: Serialize + DeserializeOwned + Send + Sync> Storage<V> for JournalStorag
         self.memory.trials_arc()
     }
 
+    fn next_trial_id(&self) -> u64 {
+        self.memory.next_trial_id()
+    }
+
     fn refresh(&self) -> bool {
         let Ok(loaded) = load_trials_from_file::<V>(&self.path) else {
             return false;
         };
         let mut guard = self.memory.trials_arc().write();
         if loaded.len() > guard.len() {
+            if let Some(max_id) = loaded.iter().map(|t| t.id).max() {
+                self.memory.bump_next_id(max_id + 1);
+            }
             *guard = loaded;
             true
         } else {

src/storage/memory.rs

@@ -1,3 +1,4 @@
+use core::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
 use parking_lot::RwLock;
@@ -10,6 +11,7 @@ use crate::sampler::CompletedTrial;
 /// This is a thin wrapper around `Arc<RwLock<Vec<CompletedTrial<V>>>>`.
 pub struct MemoryStorage<V> {
     trials: Arc<RwLock<Vec<CompletedTrial<V>>>>,
+    next_id: AtomicU64,
 }
 
 impl<V> MemoryStorage<V> {
@@ -18,16 +20,24 @@ impl<V> MemoryStorage<V> {
     pub fn new() -> Self {
         Self {
             trials: Arc::new(RwLock::new(Vec::new())),
+            next_id: AtomicU64::new(0),
         }
     }
 
     /// Creates an in-memory store pre-populated with `trials`.
     #[must_use]
     pub fn with_trials(trials: Vec<CompletedTrial<V>>) -> Self {
+        let next_id = trials.iter().map(|t| t.id).max().map_or(0, |id| id + 1);
         Self {
             trials: Arc::new(RwLock::new(trials)),
+            next_id: AtomicU64::new(next_id),
         }
     }
+
+    /// Ensures the ID counter is at least `min_value`.
+    pub(crate) fn bump_next_id(&self, min_value: u64) {
+        self.next_id.fetch_max(min_value, Ordering::SeqCst);
+    }
 }
 
 impl<V> Default for MemoryStorage<V> {
@@ -44,4 +54,8 @@ impl<V: Send + Sync> Storage<V> for MemoryStorage<V> {
     fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>> {
         &self.trials
     }
+
+    fn next_trial_id(&self) -> u64 {
+        self.next_id.fetch_add(1, Ordering::SeqCst)
+    }
 }

src/storage/mod.rs

@@ -42,6 +42,12 @@ pub trait Storage<V>: Send + Sync {
     /// lock for efficient, allocation-free access.
     fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>>;
 
+    /// Atomically returns the next unique trial ID.
+    ///
+    /// Each call increments an internal counter so that consecutive
+    /// calls always produce distinct IDs.
+    fn next_trial_id(&self) -> u64;
+
     /// Reload from an external source (e.g. a file written by another
     /// process).  Returns `true` if the in-memory buffer was updated.
     ///

src/storage/sqlite.rs

@@ -98,6 +98,10 @@ impl<V: Serialize + DeserializeOwned + Send + Sync> Storage<V> for SqliteStorage
         self.memory.trials_arc()
     }
 
+    fn next_trial_id(&self) -> u64 {
+        self.memory.next_trial_id()
+    }
+
     fn refresh(&self) -> bool {
         let conn = self.conn.lock();
         let Ok(loaded) = load_all::<V>(&conn) else {
@@ -105,6 +109,9 @@ impl<V: Serialize + DeserializeOwned + Send + Sync> Storage<V> for SqliteStorage
         };
         let mut guard = self.memory.trials_arc().write();
         if loaded.len() > guard.len() {
+            if let Some(max_id) = loaded.iter().map(|t| t.id).max() {
+                self.memory.bump_next_id(max_id + 1);
+            }
             *guard = loaded;
             true
         } else {

src/study.rs

@@ -6,7 +6,6 @@ use core::fmt;
 use core::future::Future;
 use core::marker::PhantomData;
 use core::ops::ControlFlow;
-use core::sync::atomic::{AtomicU64, Ordering};
 use core::time::Duration;
 use std::collections::{HashMap, VecDeque};
 use std::sync::Arc;
@@ -52,8 +51,6 @@ where
     pruner: Arc<dyn Pruner>,
     /// Trial storage backend (default: [`MemoryStorage`](crate::storage::MemoryStorage)).
     storage: Arc<dyn crate::storage::Storage<V>>,
-    /// Counter for generating unique trial IDs.
-    next_trial_id: AtomicU64,
     /// Optional factory for creating sampler-aware trials.
     /// Set automatically for `Study<f64>` so that `create_trial()` and all
     /// optimization methods use the sampler without requiring `_with_sampler` suffixes.
@@ -240,26 +237,18 @@ where
         let storage: Arc<dyn crate::storage::Storage<V>> = Arc::new(storage);
         let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner);
 
-        let next_id = storage
-            .trials_arc()
-            .read()
-            .iter()
-            .map(|t| t.id)
-            .max()
-            .map_or(0, |id| id + 1);
-
         Self {
             direction,
             sampler,
             pruner,
             storage,
-            next_trial_id: AtomicU64::new(next_id),
             trial_factory,
             enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
         }
     }
 
     /// Returns the optimization direction.
+    #[must_use]
     pub fn direction(&self) -> Direction {
         self.direction
     }
@@ -316,7 +305,6 @@ where
             sampler,
             pruner,
             storage,
-            next_trial_id: AtomicU64::new(0),
             trial_factory,
             enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
         }
@@ -344,6 +332,7 @@ where
     }
 
     /// Returns a reference to the study's pruner.
+    #[must_use]
     pub fn pruner(&self) -> &dyn Pruner {
         &*self.pruner
     }
@@ -423,7 +412,7 @@ where
 
     /// Generates the next unique trial ID.
     pub(crate) fn next_trial_id(&self) -> u64 {
-        self.next_trial_id.fetch_add(1, Ordering::SeqCst)
+        self.storage.next_trial_id()
     }
 
     /// Creates a new trial with a unique ID.
@@ -448,13 +437,9 @@ where
     /// let trial2 = study.create_trial();
     /// assert_eq!(trial2.id(), 1);
     /// ```
+    #[must_use]
     pub fn create_trial(&self) -> Trial {
-        if self.storage.refresh() {
+        self.storage.refresh();
-            let trials = self.storage.trials_arc().read();
-            if let Some(max_id) = trials.iter().map(|t| t.id).max() {
-                self.next_trial_id.fetch_max(max_id + 1, Ordering::SeqCst);
-            }
-        }
 
         let id = self.next_trial_id();
         let mut trial = if let Some(factory) = &self.trial_factory {
@@ -565,6 +550,7 @@ where
     /// let value = x_val * x_val;
     /// study.tell(trial, Ok::<_, &str>(value));
     /// ```
+    #[must_use]
     pub fn ask(&self) -> Trial {
         self.create_trial()
     }
@@ -649,6 +635,7 @@ where
     ///     println!("Trial {} has value {:?}", completed.id, completed.value);
     /// }
     /// ```
+    #[must_use]
     pub fn trials(&self) -> Vec<CompletedTrial<V>>
     where
         V: Clone,
@@ -675,11 +662,13 @@ where
     /// study.complete_trial(trial, 0.5);
     /// assert_eq!(study.n_trials(), 1);
     /// ```
+    #[must_use]
     pub fn n_trials(&self) -> usize {
         self.storage.trials_arc().read().len()
     }
 
     /// Returns the number of pruned trials.
+    #[must_use]
     pub fn n_pruned_trials(&self) -> usize {
         self.storage
             .trials_arc()
@@ -821,6 +810,7 @@ where
     /// Only includes completed trials (not failed or pruned).
     ///
     /// If fewer than `n` completed trials exist, returns all of them.
+    #[must_use]
     pub fn top_trials(&self, n: usize) -> Vec<CompletedTrial<V>>
     where
         V: Clone,
@@ -1987,6 +1977,7 @@ where
     ///
     /// This clones the internal trial list, so it is suitable for
     /// analysis and iteration but not for hot paths.
+    #[must_use]
     pub fn iter(&self) -> std::vec::IntoIter<CompletedTrial<V>> {
         self.trials().into_iter()
     }
@@ -2234,6 +2225,7 @@ impl Study<f64> {
         since = "0.2.0",
         note = "use `create_trial()` instead — it now uses the sampler automatically for Study<f64>"
     )]
+    #[must_use]
     pub fn create_trial_with_sampler(&self) -> Trial {
         self.create_trial()
     }
@@ -2336,20 +2328,11 @@ impl<V: PartialOrd + Send + Sync + 'static> Study<V> {
         let storage: Arc<dyn crate::storage::Storage<V>> = Arc::new(storage);
         let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner);
 
-        let next_id = storage
-            .trials_arc()
-            .read()
-            .iter()
-            .map(|t| t.id)
-            .max()
-            .map_or(0, |id| id + 1);
-
         Self {
             direction,
             sampler,
             pruner,
             storage,
-            next_trial_id: AtomicU64::new(next_id),
             trial_factory,
             enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
         }
@@ -2451,20 +2434,11 @@ impl<V: PartialOrd> StudyBuilder<V> {
         let storage: Arc<dyn crate::storage::Storage<V>> = Arc::from(storage);
         let trial_factory = Study::make_trial_factory(&sampler, &storage, &pruner);
 
-        let next_id = storage
-            .trials_arc()
-            .read()
-            .iter()
-            .map(|t| t.id)
-            .max()
-            .map_or(0, |id| id + 1);
-
         Study {
             direction: self.direction,
             sampler,
             pruner,
             storage,
-            next_trial_id: AtomicU64::new(next_id),
             trial_factory,
             enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
         }
@@ -2590,11 +2564,13 @@ impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
     /// Returns an I/O error if the file cannot be created or written.
     pub fn save(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
         let path = path.as_ref();
+        let trials = self.trials();
+        let next_trial_id = trials.iter().map(|t| t.id).max().map_or(0, |id| id + 1);
         let snapshot = StudySnapshot {
             version: 1,
             direction: self.direction,
-            trials: self.trials(),
+            trials,
-            next_trial_id: self.next_trial_id.load(Ordering::Relaxed),
+            next_trial_id,
             metadata: HashMap::new(),
         };
 
@@ -2660,12 +2636,12 @@ impl<V: PartialOrd + Send + Sync + Clone + serde::de::DeserializeOwned + 'static
         let file = std::fs::File::open(path)?;
         let snapshot: StudySnapshot<V> = serde_json::from_reader(file)
             .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
-        let study = Study::new(snapshot.direction);
+        let storage = crate::storage::MemoryStorage::with_trials(snapshot.trials);
-        *study.storage.trials_arc().write() = snapshot.trials;
+        Ok(Self::with_sampler_and_storage(
-        study
+            snapshot.direction,
-            .next_trial_id
+            RandomSampler::new(),
-            .store(snapshot.next_trial_id, Ordering::Relaxed);
+            storage,
-        Ok(study)
+        ))
     }
 }