Skip to main content

hyperopt_storage/
memory.rs

1use hyperopt_core::{Storage, StorageError, StudyMetadata, Trial};
2use std::collections::{BTreeMap, HashMap};
3use std::sync::Mutex;
4
5/// In-memory storage: trials kept in a `BTreeMap` per study (ordered by trial
6/// number) behind a `Mutex`. Fast and dependency-free, but everything is lost
7/// when the process exits. This is the Phase 1–2 default.
8#[derive(Default)]
9pub struct InMemoryStorage {
10    inner: Mutex<Inner>,
11}
12
13#[derive(Default)]
14struct Inner {
15    trials: HashMap<String, BTreeMap<usize, Trial>>,
16    meta: HashMap<String, StudyMetadata>,
17}
18
19impl InMemoryStorage {
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    fn lock(&self) -> std::sync::MutexGuard<'_, Inner> {
25        self.inner.lock().unwrap_or_else(|p| p.into_inner())
26    }
27}
28
29impl Storage for InMemoryStorage {
30    fn save_trial(&self, study_name: &str, trial: &Trial) -> Result<(), StorageError> {
31        let mut inner = self.lock();
32        inner
33            .trials
34            .entry(study_name.to_string())
35            .or_default()
36            .insert(trial.number, trial.clone());
37        Ok(())
38    }
39
40    fn load_trials(&self, study_name: &str) -> Result<Vec<Trial>, StorageError> {
41        let inner = self.lock();
42        Ok(inner
43            .trials
44            .get(study_name)
45            .map(|m| m.values().cloned().collect())
46            .unwrap_or_default())
47    }
48
49    fn save_study_metadata(&self, meta: &StudyMetadata) -> Result<(), StorageError> {
50        let mut inner = self.lock();
51        inner.meta.insert(meta.study_name.clone(), meta.clone());
52        Ok(())
53    }
54
55    fn load_study_metadata(
56        &self,
57        study_name: &str,
58    ) -> Result<Option<StudyMetadata>, StorageError> {
59        let inner = self.lock();
60        Ok(inner.meta.get(study_name).cloned())
61    }
62}