cellular_raza_core/storage/
concepts.rs

1use std::collections::{BTreeMap, HashMap};
2use std::error::Error;
3use std::fmt::Display;
4
5use serde::{Deserialize, Serialize};
6use uniquevec::UniqueVec;
7
8#[cfg(feature = "tracing")]
9use tracing::instrument;
10
11use super::memory_storage::MemoryStorageInterface;
12use super::ron::RonStorageInterface;
13use super::serde_json::JsonStorageInterface;
14use super::sled_database::SledStorageInterface;
15
16/// Error related to storing and reading elements
17#[derive(Debug)]
18pub enum StorageError {
19    /// Error related to File Io operations.
20    IoError(std::io::Error),
21    /// Occurs during parsing of json structs.
22    SerdeJsonError(serde_json::Error),
23    /// Generic error related to serialization in the [ron] crate.
24    RonError(ron::Error),
25    /// Generic error related to deserialization in the [ron] crate.
26    RonSpannedError(ron::error::SpannedError),
27    /// Generic error related to the [sled] database.
28    SledError(sled::Error),
29    /// Generic serialization error thrown by the [bincode] library.
30    BincodeSeError(bincode::error::EncodeError),
31    /// Generic deserialization error thrown by the [bincode] library.
32    BincodeDeError(bincode::error::DecodeError),
33    /// Initialization error mainly used for initialization of databases such as [sled].
34    InitError(String),
35    /// Error when parsing file/folder names.
36    ParseIntError(std::num::ParseIntError),
37    /// Generic Utf8 error.
38    Utf8Error(std::str::Utf8Error),
39    /// Error during locking of Mutex
40    PoisonError(String),
41}
42
43impl From<serde_json::Error> for StorageError {
44    fn from(err: serde_json::Error) -> Self {
45        StorageError::SerdeJsonError(err)
46    }
47}
48
49impl From<ron::error::SpannedError> for StorageError {
50    fn from(err: ron::error::SpannedError) -> Self {
51        StorageError::RonSpannedError(err)
52    }
53}
54
55impl From<ron::Error> for StorageError {
56    fn from(err: ron::Error) -> Self {
57        StorageError::RonError(err)
58    }
59}
60
61impl From<sled::Error> for StorageError {
62    fn from(err: sled::Error) -> Self {
63        StorageError::SledError(err)
64    }
65}
66
67impl From<bincode::error::EncodeError> for StorageError {
68    fn from(err: bincode::error::EncodeError) -> Self {
69        StorageError::BincodeSeError(err)
70    }
71}
72
73impl From<bincode::error::DecodeError> for StorageError {
74    fn from(err: bincode::error::DecodeError) -> Self {
75        StorageError::BincodeDeError(err)
76    }
77}
78
79impl From<std::io::Error> for StorageError {
80    fn from(err: std::io::Error) -> Self {
81        StorageError::IoError(err)
82    }
83}
84
85impl From<std::str::Utf8Error> for StorageError {
86    fn from(err: std::str::Utf8Error) -> Self {
87        StorageError::Utf8Error(err)
88    }
89}
90
91impl From<std::num::ParseIntError> for StorageError {
92    fn from(err: std::num::ParseIntError) -> Self {
93        StorageError::ParseIntError(err)
94    }
95}
96
97impl<T> From<std::sync::PoisonError<T>> for StorageError {
98    fn from(err: std::sync::PoisonError<T>) -> Self {
99        StorageError::PoisonError(format!("{err}"))
100    }
101}
102
103impl Display for StorageError {
104    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
105        match self {
106            StorageError::SerdeJsonError(message) => write!(f, "{}", message),
107            StorageError::RonError(message) => write!(f, "{}", message),
108            StorageError::RonSpannedError(message) => write!(f, "{}", message),
109            StorageError::SledError(message) => write!(f, "{}", message),
110            StorageError::BincodeSeError(message) => write!(f, "{}", message),
111            StorageError::BincodeDeError(message) => write!(f, "{}", message),
112            StorageError::IoError(message) => write!(f, "{}", message),
113            StorageError::InitError(message) => write!(f, "{}", message),
114            StorageError::Utf8Error(message) => write!(f, "{}", message),
115            StorageError::ParseIntError(message) => write!(f, "{}", message),
116            StorageError::PoisonError(message) => write!(f, "{}", message),
117        }
118    }
119}
120
121impl Error for StorageError {}
122
123/// Define how to store results of the simulation.
124///
125/// We currently support saving results in a [sled] database, or as a json file by using
126/// [serde_json].
127#[cfg_attr(feature = "pyo3", pyo3::pyclass(eq, eq_int))]
128#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)]
129pub enum StorageOption {
130    /// Save results as [sled] database.
131    Sled,
132    /// Save results as [sled] database but remove them when dropping the struct
133    SledTemp,
134    /// Save results as [json](https://www.json.org/json-en.html) file.
135    SerdeJson,
136    /// Store results in the [ron] file format specifically designed for Rust structs.
137    /// This format guarantees round-trips `Rust -> Ron -> Rust` and is thus preferred together
138    /// with the well-established [StorageOption::SerdeJson] format.
139    Ron,
140    /// A [std::collections::HashMap](HashMap) based memory storage.
141    Memory,
142}
143
144impl StorageOption {
145    /// Which storage option should be used by default.
146    pub fn default_priority() -> UniqueVec<Self> {
147        vec![
148            StorageOption::SerdeJson,
149            // TODO fix sled! This is currently not working on multiple threads
150            // StorageOptions::Sled,
151        ]
152        .into()
153    }
154}
155
156/// Define how elements and identifiers are saved when being serialized together.
157#[derive(Clone, Debug, Deserialize, Serialize)]
158pub struct CombinedSaveFormat<Id, Element> {
159    /// Identifier of the element
160    pub identifier: Id,
161    /// Actual element which is being stored
162    pub element: Element,
163}
164
165/// Define how batches of elements and identifiers are saved when being serialized.
166#[derive(Clone, Debug, Deserialize, Serialize)]
167pub struct BatchSaveFormat<Id, Element> {
168    pub(super) data: Vec<CombinedSaveFormat<Id, Element>>,
169}
170
171/// This manager handles if multiple storage options have been specified
172/// It can load resources from one storage aspect and will
173#[derive(Clone, Debug)]
174pub struct StorageManager<Id, Element> {
175    storage_priority: UniqueVec<StorageOption>,
176    builder: StorageBuilder<true>,
177    instance: u64,
178
179    sled_storage: Option<SledStorageInterface<Id, Element>>,
180    sled_temp_storage: Option<SledStorageInterface<Id, Element, true>>,
181    json_storage: Option<StorageWrapper<JsonStorageInterface<Id, Element>>>,
182    ron_storage: Option<StorageWrapper<RonStorageInterface<Id, Element>>>,
183    memory_storage: Option<MemoryStorageInterface<Id, Element>>,
184}
185
186/// Used to construct a [StorageManager]
187///
188/// This builder contains multiple options which can be used to configure the location and type in
189/// which results are stored.
190/// To get an overview over all possible options, we refer to the [module](crate::storage)
191/// documentation.
192///
193/// ```
194/// use cellular_raza_core::storage::{StorageBuilder, StorageOption};
195///
196/// let storage_priority = StorageOption::default_priority();
197/// let storage_builder = StorageBuilder::new()
198///     .priority(storage_priority)
199///     .location("./");
200/// ```
201#[derive(Clone, Debug, Deserialize, Serialize)]
202pub struct StorageBuilder<const INIT: bool = false> {
203    location: std::path::PathBuf,
204    priority: UniqueVec<StorageOption>,
205    suffix: std::path::PathBuf,
206    #[cfg(feature = "timestamp")]
207    add_date: bool,
208    #[cfg(feature = "timestamp")]
209    date: std::path::PathBuf,
210}
211
212impl<const INIT: bool> StorageBuilder<INIT> {
213    /// Define the priority of [StorageOption]. See [StorageOption::default_priority].
214    pub fn priority(self, priority: impl IntoIterator<Item = StorageOption>) -> Self {
215        let (priority, _) = UniqueVec::from_iter(priority);
216        Self { priority, ..self }
217    }
218
219    /// Get the current priority
220    pub fn get_priority(&self) -> UniqueVec<StorageOption> {
221        self.priority.clone()
222    }
223
224    /// Define a suffix which will be appended to the save path
225    pub fn suffix(self, suffix: impl Into<std::path::PathBuf>) -> Self {
226        Self {
227            suffix: suffix.into(),
228            ..self
229        }
230    }
231
232    /// Get the current suffix
233    pub fn get_suffix(&self) -> std::path::PathBuf {
234        self.suffix.clone()
235    }
236
237    /// Store results by their current date inside the specified folder path
238    #[cfg(feature = "timestamp")]
239    pub fn add_date(self, add_date: bool) -> Self {
240        Self { add_date, ..self }
241    }
242
243    /// Get information if the current date should be appended to the storage path
244    #[cfg(feature = "timestamp")]
245    pub fn get_add_date(&self) -> bool {
246        self.add_date
247    }
248}
249
250impl StorageBuilder<false> {
251    /// Constructs a new [StorageBuilder] with default settings.
252    ///
253    /// ```
254    /// use cellular_raza_core::storage::StorageBuilder;
255    /// let storage_builder = StorageBuilder::new();
256    /// ```
257    #[cfg_attr(feature = "tracing", instrument(skip_all))]
258    pub fn new() -> Self {
259        Self {
260            location: "./out".into(),
261            priority: UniqueVec::from_iter([StorageOption::SerdeJson]).0,
262            suffix: "".into(),
263            #[cfg(feature = "timestamp")]
264            add_date: true,
265            #[cfg(feature = "timestamp")]
266            date: "".into(),
267        }
268    }
269
270    /// Initializes the [StorageBuilder] thus filling information about time.
271    #[cfg_attr(feature = "tracing", instrument(skip_all))]
272    pub fn init(self) -> StorageBuilder<true> {
273        #[cfg(feature = "timestamp")]
274        let date: std::path::PathBuf = if self.add_date {
275            format!("{}", chrono::Local::now().format("%Y-%m-%d-T%H-%M-%S")).into()
276        } else {
277            "".into()
278        };
279        self.init_with_date(&date)
280    }
281
282    /// Specify the time at which the results should be saved
283    #[cfg_attr(feature = "tracing", instrument(skip_all))]
284    pub fn init_with_date(self, date: &std::path::Path) -> StorageBuilder<true> {
285        StorageBuilder::<true> {
286            location: self.location,
287            priority: self.priority,
288            suffix: self.suffix,
289            #[cfg(feature = "timestamp")]
290            add_date: self.add_date,
291            #[cfg(feature = "timestamp")]
292            date: date.into(),
293        }
294    }
295
296    /// Define a folder where to store results
297    ///
298    /// Note that this functionality is only available as long as the [StorageBuilder] has not been
299    /// initialized.
300    pub fn location<P>(self, location: P) -> Self
301    where
302        std::path::PathBuf: From<P>,
303    {
304        Self {
305            location: location.into(),
306            ..self
307        }
308    }
309
310    /// Get the current storage_location
311    ///
312    /// Note that this functionality is only available as long as the [StorageBuilder] has not been
313    /// initialized.
314    pub fn get_location(&self) -> std::path::PathBuf {
315        self.location.clone()
316    }
317}
318
319impl StorageBuilder<true> {
320    /// Get the fully constructed path after the Builder has been initialized with the
321    /// [StorageBuilder::init] function.
322    #[cfg_attr(feature = "tracing", instrument(skip_all))]
323    pub fn get_full_path(&self) -> std::path::PathBuf {
324        let mut full_path = self.location.clone();
325        #[cfg(feature = "timestamp")]
326        if self.add_date {
327            full_path.extend(&self.date);
328        }
329        full_path.extend(&self.suffix);
330        full_path
331    }
332
333    #[doc(hidden)]
334    pub fn init(self) -> Self {
335        self
336    }
337
338    /// De-initializes the StorageBuilder, making it possible to edit it again.
339    pub fn de_init(self) -> StorageBuilder<false> {
340        StorageBuilder {
341            location: self.location,
342            priority: self.priority,
343            suffix: self.suffix,
344            #[cfg(feature = "timestamp")]
345            add_date: self.add_date,
346            #[cfg(feature = "timestamp")]
347            date: "".into(),
348        }
349    }
350}
351
352impl<Id, Element> StorageManager<Id, Element> {
353    /// Constructs the [StorageManager] from the instance identifier
354    /// and the settings given by the [StorageBuilder].
355    ///
356    /// ```
357    /// use cellular_raza_core::storage::*;
358    /// let builder = StorageBuilder::new()
359    ///     .location("/tmp")
360    ///     .init();
361    ///
362    /// let manager = StorageManager::<usize, f64>::open_or_create(builder, 0)?;
363    /// # Ok::<(), StorageError>(())
364    /// ```
365    #[cfg_attr(feature = "tracing", instrument(skip_all))]
366    pub fn open_or_create(
367        storage_builder: StorageBuilder<true>,
368        instance: u64,
369    ) -> Result<Self, StorageError> {
370        let location = storage_builder.get_full_path();
371
372        let mut sled_storage = None;
373        let mut sled_temp_storage = None;
374        let mut json_storage = None;
375        let mut ron_storage = None;
376        let mut memory_storage = None;
377        for storage_variant in storage_builder.priority.iter() {
378            match storage_variant {
379                StorageOption::SerdeJson => {
380                    json_storage = Some(StorageWrapper(
381                        JsonStorageInterface::<Id, Element>::open_or_create(
382                            &location
383                                .to_path_buf()
384                                .join(JsonStorageInterface::<Id, Element>::EXTENSION),
385                            instance,
386                        )?,
387                    ));
388                }
389                StorageOption::Sled => {
390                    sled_storage =
391                        Some(SledStorageInterface::<Id, Element, false>::open_or_create(
392                            &location.to_path_buf().join("sled"),
393                            instance,
394                        )?);
395                }
396                StorageOption::SledTemp => {
397                    sled_temp_storage =
398                        Some(SledStorageInterface::<Id, Element, true>::open_or_create(
399                            &location.to_path_buf().join("sled_memory"),
400                            instance,
401                        )?);
402                }
403                StorageOption::Ron => {
404                    ron_storage = Some(StorageWrapper(
405                        RonStorageInterface::<Id, Element>::open_or_create(
406                            &location
407                                .to_path_buf()
408                                .join(RonStorageInterface::<Id, Element>::EXTENSION),
409                            instance,
410                        )?,
411                    ));
412                }
413                StorageOption::Memory => {
414                    memory_storage = Some(MemoryStorageInterface::<Id, Element>::open_or_create(
415                        &location.to_path_buf(),
416                        instance,
417                    )?);
418                }
419            }
420        }
421        let manager = StorageManager {
422            storage_priority: storage_builder.priority.clone(),
423            builder: storage_builder.clone(),
424            instance,
425
426            sled_storage,
427            sled_temp_storage,
428            json_storage,
429            ron_storage,
430            memory_storage,
431        };
432
433        Ok(manager)
434    }
435
436    /// Uses an existing storage manager to construct a new one.
437    /// ```
438    /// # use cellular_raza_core::storage::*;
439    /// let builder = StorageBuilder::new()
440    ///     .location("/tmp")
441    ///     .init();
442    ///
443    /// let manager = StorageManager::<usize, f64>::open_or_create(builder, 0)?;
444    /// let manager2 = manager.clone_to_new_instance(1);
445    /// # Ok::<(), StorageError>(())
446    ///
447    /// ```
448    pub fn clone_to_new_instance(&self, storage_instance: u64) -> Self {
449        Self {
450            storage_priority: self.storage_priority.clone(),
451            builder: self.builder.clone(),
452            instance: storage_instance,
453
454            sled_storage: self
455                .sled_storage
456                .as_ref()
457                .map(|x| x.clone_to_new_instance(storage_instance)),
458            sled_temp_storage: self
459                .sled_temp_storage
460                .as_ref()
461                .map(|x| x.clone_to_new_instance(storage_instance)),
462            json_storage: self
463                .json_storage
464                .as_ref()
465                .map(|x| StorageWrapper(x.0.clone_to_new_instance(storage_instance))),
466            ron_storage: self
467                .ron_storage
468                .as_ref()
469                .map(|x| StorageWrapper(x.0.clone_to_new_instance(storage_instance))),
470            memory_storage: self
471                .memory_storage
472                .as_ref()
473                .map(|x| x.clone_to_new_instance(storage_instance)),
474        }
475    }
476
477    /// Extracts all information given by the [StorageBuilder] when constructing
478    #[cfg_attr(feature = "tracing", instrument(skip_all))]
479    pub fn extract_builder(&self) -> StorageBuilder<true> {
480        self.builder.clone()
481    }
482
483    /// Get the instance of this object.
484    ///
485    /// These instances should not be overlapping, ie. there should not be two objects existing in
486    /// parallel with the same instance number.
487    pub fn get_instance(&self) -> u64 {
488        self.instance
489    }
490}
491
492macro_rules! exec_for_all_storage_options(
493    (@internal $self:ident, $storage_option:ident, $field:ident, $function:ident, $($args:tt)*) => {
494        {
495            if let Some($field) = &$self.$field {
496                $field.$function($($args)*)
497            } else {
498                Err(StorageError::InitError(
499                    stringify!($storage_option, " storage was not initialized but called").into(),
500                ))?
501            }
502        }
503    };
504    (mut $self:ident, $field:ident, $function:ident, $($args:tt)*) => {
505        if let Some($field) = &mut $self.$field {
506            $field.$function($($args)*)?;
507        }
508    };
509    (all mut $self:ident, $function:ident, $($args:tt)*) => {
510        exec_for_all_storage_options!(mut $self, sled_storage, $function, $($args)*);
511        exec_for_all_storage_options!(mut $self, sled_temp_storage, $function, $($args)*);
512        exec_for_all_storage_options!(mut $self, json_storage, $function, $($args)*);
513        exec_for_all_storage_options!(mut $self, ron_storage, $function, $($args)*);
514        exec_for_all_storage_options!(mut $self, memory_storage, $function, $($args)*);
515    };
516    ($self:ident, $priority:ident, $function:ident, $($args:tt)*) => {
517        match $priority {
518            StorageOption::Sled => exec_for_all_storage_options!(
519                @internal $self, Sled, sled_storage, $function, $($args)*
520            ),
521            StorageOption::SledTemp => exec_for_all_storage_options!(
522                @internal $self, SledTemp, sled_temp_storage, $function, $($args)*
523            ),
524            StorageOption::SerdeJson => exec_for_all_storage_options!(
525                @internal $self, SerdeJson, json_storage, $function, $($args)*
526            ),
527            StorageOption::Ron => exec_for_all_storage_options!(
528                @internal $self, Ron, ron_storage, $function, $($args)*
529            ),
530            StorageOption::Memory => exec_for_all_storage_options!(
531                @internal $self, Memory, memory_storage, $function, $($args)*
532            ),
533        }
534    }
535);
536
537impl<Id, Element> StorageInterfaceStore<Id, Element> for StorageManager<Id, Element>
538where
539    Id: core::hash::Hash + core::cmp::Eq + Clone,
540    Element: Clone,
541{
542    #[allow(unused)]
543    fn store_single_element(
544        &mut self,
545        iteration: u64,
546        identifier: &Id,
547        element: &Element,
548    ) -> Result<(), StorageError>
549    where
550        Id: Serialize,
551        Element: Serialize,
552    {
553        exec_for_all_storage_options!(
554            all mut self,
555            store_single_element,
556            iteration, identifier, element
557        );
558        Ok(())
559    }
560
561    #[allow(unused)]
562    fn store_batch_elements<'a, I>(
563        &'a mut self,
564        iteration: u64,
565        identifiers_elements: I,
566    ) -> Result<(), StorageError>
567    where
568        Id: 'a + Serialize,
569        Element: 'a + Serialize,
570        I: Clone + IntoIterator<Item = (&'a Id, &'a Element)>,
571    {
572        exec_for_all_storage_options!(
573            all mut self,
574            store_batch_elements,
575            iteration,
576            identifiers_elements.clone()
577        );
578        Ok(())
579    }
580}
581
582impl<Id, Element> StorageInterfaceLoad<Id, Element> for StorageManager<Id, Element>
583where
584    Id: core::hash::Hash + core::cmp::Eq + Clone,
585    Element: Clone,
586{
587    #[allow(unused)]
588    fn load_single_element(
589        &self,
590        iteration: u64,
591        identifier: &Id,
592    ) -> Result<Option<Element>, StorageError>
593    where
594        Id: Serialize + for<'a> Deserialize<'a>,
595        Element: for<'a> Deserialize<'a>,
596    {
597        for priority in self.storage_priority.iter() {
598            return exec_for_all_storage_options!(
599                self,
600                priority,
601                load_single_element,
602                iteration,
603                identifier
604            );
605        }
606        Ok(None)
607    }
608
609    #[allow(unused)]
610    fn load_all_elements_at_iteration(
611        &self,
612        iteration: u64,
613    ) -> Result<HashMap<Id, Element>, StorageError>
614    where
615        Id: std::hash::Hash + std::cmp::Eq + for<'a> Deserialize<'a>,
616        Element: for<'a> Deserialize<'a>,
617    {
618        for priority in self.storage_priority.iter() {
619            return exec_for_all_storage_options!(
620                self,
621                priority,
622                load_all_elements_at_iteration,
623                iteration
624            );
625        }
626        Ok(HashMap::new())
627    }
628
629    fn get_all_iterations(&self) -> Result<Vec<u64>, StorageError> {
630        for priority in self.storage_priority.iter() {
631            return exec_for_all_storage_options!(self, priority, get_all_iterations,);
632        }
633        Ok(Vec::new())
634    }
635}
636
637/// The mode in which to generate paths and store results.
638pub enum StorageMode {
639    /// Save one element to a single file
640    Single,
641    /// Save many elements in one file.
642    Batch,
643}
644
645impl StorageMode {
646    fn to_str(&self) -> &str {
647        match self {
648            Self::Single => "single",
649            Self::Batch => "batch",
650        }
651    }
652}
653
654/// Abstraction and simplification of many file-based storage solutions
655pub trait FileBasedStorage<Id, Element> {
656    /// The suffix which is used to distinguish this storage solution from others.
657    const EXTENSION: &'static str;
658
659    /// Get path where results are stored.
660    fn get_path(&self) -> &std::path::Path;
661
662    /// Get the number of this storage instance.
663    /// This value may coincide with the thread number.
664    fn get_storage_instance(&self) -> u64;
665
666    /// Writes either [BatchSaveFormat] or [CombinedSaveFormat] to the disk.
667    fn to_writer_pretty<V, W>(&self, writer: W, value: &V) -> Result<(), StorageError>
668    where
669        V: Serialize,
670        W: std::io::Write;
671
672    /// Deserialize the given value from a string
673    fn from_str<V>(&self, input: &str) -> Result<V, StorageError>
674    where
675        V: for<'a> Deserialize<'a>;
676
677    /// Creates a new iteration file with a predefined naming scheme.
678    ///
679    /// The path which to use is by default determined by the
680    /// [FileBasedStorage::get_iteration_save_path_batch_with_prefix] function.
681    fn create_or_get_iteration_file_with_prefix(
682        &self,
683        iteration: u64,
684        mode: StorageMode,
685    ) -> Result<std::io::BufWriter<std::fs::File>, StorageError> {
686        let save_path = self.get_iteration_save_path_batch_with_prefix(iteration, mode)?;
687
688        // Open+Create a file and wrap it inside a buffer writer
689        let file = std::fs::OpenOptions::new()
690            .read(true)
691            .write(true)
692            .create(true)
693            .open(&save_path)?;
694
695        Ok(std::io::BufWriter::new(file))
696    }
697
698    /// Get the path which holds saved entries if the given iteration.
699    ///
700    /// By default this function joins the path generated by [FileBasedStorage::get_path]
701    /// with a 0-delimited number according to the iteration number.
702    fn get_iteration_path(&self, iteration: u64) -> std::path::PathBuf {
703        self.get_path().join(format!("{:020.0}", iteration))
704    }
705
706    /// Creates the path used by the [FileBasedStorage::create_or_get_iteration_file_with_prefix]
707    /// function.
708    fn get_iteration_save_path_batch_with_prefix(
709        &self,
710        iteration: u64,
711        mode: StorageMode,
712    ) -> Result<std::path::PathBuf, StorageError> {
713        // First we get the folder path of the iteration
714        let iteration_path = self.get_iteration_path(iteration);
715        // If this folder does not exist, we create it
716        std::fs::create_dir_all(&iteration_path)?;
717
718        // Check if other batch files are already existing
719        // If this is the case increase the batch number until we find one where no batch is existing
720        let prefix = mode.to_str();
721        let create_save_path = |i: usize| -> std::path::PathBuf {
722            iteration_path
723                .join(format!(
724                    "{}_{:020.0}_{:020.0}",
725                    prefix,
726                    self.get_storage_instance(),
727                    i
728                ))
729                .with_extension(Self::EXTENSION)
730        };
731        let mut counter = 0;
732        let mut save_path;
733        while {
734            save_path = create_save_path(counter);
735            save_path.exists()
736        } {
737            counter += 1
738        }
739        Ok(save_path)
740    }
741
742    /// Converts a given path of a folder to a iteration number.
743    ///
744    /// This function is used for loading results
745    fn folder_name_to_iteration(
746        &self,
747        file: &std::path::Path,
748    ) -> Result<Option<u64>, StorageError> {
749        match file.file_stem() {
750            Some(filename) => match filename.to_str() {
751                Some(filename_string) => Ok(Some(filename_string.parse::<u64>()?)),
752                None => Ok(None),
753            },
754            None => Ok(None),
755        }
756    }
757}
758
759#[derive(Clone, Debug, Deserialize, Serialize)]
760pub(crate) struct StorageWrapper<T>(pub(crate) T);
761
762impl<T, Id, Element> StorageInterfaceStore<Id, Element> for StorageWrapper<T>
763where
764    T: FileBasedStorage<Id, Element>,
765{
766    fn store_batch_elements<'a, I>(
767        &'a mut self,
768        iteration: u64,
769        identifiers_elements: I,
770    ) -> Result<(), StorageError>
771    where
772        Id: 'a + Serialize,
773        Element: 'a + Serialize,
774        I: Clone + IntoIterator<Item = (&'a Id, &'a Element)>,
775    {
776        let iteration_file = self
777            .0
778            .create_or_get_iteration_file_with_prefix(iteration, StorageMode::Batch)?;
779        let batch = BatchSaveFormat {
780            data: identifiers_elements
781                .into_iter()
782                .map(|(id, element)| CombinedSaveFormat {
783                    identifier: id,
784                    element,
785                })
786                .collect(),
787        };
788        self.0.to_writer_pretty(iteration_file, &batch)?;
789        Ok(())
790    }
791
792    fn store_single_element(
793        &mut self,
794        iteration: u64,
795        identifier: &Id,
796        element: &Element,
797    ) -> Result<(), StorageError>
798    where
799        Id: Serialize,
800        Element: Serialize,
801    {
802        let iteration_file = self
803            .0
804            .create_or_get_iteration_file_with_prefix(iteration, StorageMode::Single)?;
805        let save_format = CombinedSaveFormat {
806            identifier,
807            element,
808        };
809        self.0.to_writer_pretty(iteration_file, &save_format)?;
810        Ok(())
811    }
812}
813
814/// Open or create a new instance of the Storage controller.
815pub trait StorageInterfaceOpen {
816    /// Initializes the current storage device.
817    ///
818    /// In the case of databases, this may already result in an IO operation
819    /// while when saving as files such as json folders might be created.
820    fn open_or_create(
821        location: &std::path::Path,
822        storage_instance: u64,
823    ) -> Result<Self, StorageError>
824    where
825        Self: Sized;
826
827    /// Constructs a new instance from an existing one
828    ///
829    /// For the case of `storage_instance == 0`, an instance with the same value may already exist.
830    fn clone_to_new_instance(&self, storage_instance: u64) -> Self;
831}
832
833/// Handles storing of elements
834pub trait StorageInterfaceStore<Id, Element> {
835    /// Saves a single element at given iteration.
836    fn store_single_element(
837        &mut self,
838        iteration: u64,
839        identifier: &Id,
840        element: &Element,
841    ) -> Result<(), StorageError>
842    where
843        Id: Serialize,
844        Element: Serialize;
845
846    /// Stores a batch of multiple elements with identifiers all at the same iteration.
847    fn store_batch_elements<'a, I>(
848        &'a mut self,
849        iteration: u64,
850        identifiers_elements: I,
851    ) -> Result<(), StorageError>
852    where
853        Id: 'a + Serialize,
854        Element: 'a + Serialize,
855        I: Clone + IntoIterator<Item = (&'a Id, &'a Element)>;
856}
857
858/// Handles loading of elements
859pub trait StorageInterfaceLoad<Id, Element> {
860    // TODO decide if these functions should be &mut self instead of &self
861    // This could be useful when implementing buffers, but maybe unnecessary.
862    /// Loads a single element from the storage solution if the element exists.
863    fn load_single_element(
864        &self,
865        iteration: u64,
866        identifier: &Id,
867    ) -> Result<Option<Element>, StorageError>
868    where
869        Id: Eq + Serialize + for<'a> Deserialize<'a>,
870        Element: for<'a> Deserialize<'a>;
871
872    /// Loads the elements history, meaning every occurrence of the element in the storage.
873    /// This function by default provides the results in ordered fashion such that the time
874    /// direction is retained.
875    /// Furthermore this function assumes that a given index occurs over the course of a complete
876    /// time segment with no interceptions.
877    /// ```
878    /// // All elements (given by Strings) occur over a period of time
879    /// // but do not appear afterwards.
880    /// use std::collections::HashMap;
881    /// let valid_state = HashMap::from([
882    ///     (0, vec!["E1", "E2", "E3"]),
883    ///     (1, vec!["E1", "E2", "E3", "E4"]),
884    ///     (2, vec!["E1", "E2", "E3", "E4"]),
885    ///     (3, vec!["E1", "E2", "E4"]),
886    ///     (4, vec!["E2", "E4"]),
887    ///     (5, vec!["E2", "E4", "E5"]),
888    ///     (6, vec!["E4", "E5"]),
889    /// ]);
890    /// // The entry "E2" is missing in iteration 1 but present afterwards.
891    /// // This is an invalid state but will not be caught.
892    /// // The backend is responsible to avoid this state.
893    /// let invalid_state = HashMap::from([
894    ///     (0, vec!["E1", "E2"]),
895    ///     (1, vec!["E1"]),
896    ///     (2, vec!["E1", "E2"]),
897    /// ]);
898    /// ```
899    fn load_element_history(&self, identifier: &Id) -> Result<HashMap<u64, Element>, StorageError>
900    where
901        Id: Eq + Serialize + for<'a> Deserialize<'a>,
902        Element: for<'a> Deserialize<'a>,
903    {
904        let mut iterations = self.get_all_iterations()?;
905        iterations.sort();
906        let mut started_gathering = false;
907        let mut stop_gathering = false;
908        let results = iterations
909            .iter()
910            .filter_map(|&iteration| {
911                if stop_gathering {
912                    None
913                } else {
914                    match self.load_single_element(iteration, identifier) {
915                        Ok(Some(element)) => {
916                            started_gathering = true;
917                            Some(Ok((iteration, element)))
918                        }
919                        Ok(None) => {
920                            if started_gathering {
921                                stop_gathering = true;
922                            }
923                            None
924                        }
925                        Err(e) => Some(Err(e)),
926                    }
927                }
928            })
929            .collect::<Result<HashMap<u64, _>, StorageError>>()?;
930        Ok(results)
931    }
932
933    /// Gets a snapshot of all elements at a given iteration.
934    ///
935    /// This function might be useful when implementing how simulations can be restored from saved
936    /// results.
937    fn load_all_elements_at_iteration(
938        &self,
939        iteration: u64,
940    ) -> Result<HashMap<Id, Element>, StorageError>
941    where
942        Id: std::hash::Hash + std::cmp::Eq + for<'a> Deserialize<'a>,
943        Element: for<'a> Deserialize<'a>;
944
945    /// Get all iteration values which have been saved.
946    fn get_all_iterations(&self) -> Result<Vec<u64>, StorageError>;
947
948    /// Loads all elements for every iteration.
949    /// This will yield the complete storage and may result in extremely large allocations of
950    /// memory.
951    fn load_all_elements(&self) -> Result<BTreeMap<u64, HashMap<Id, Element>>, StorageError>
952    where
953        Id: std::hash::Hash + std::cmp::Eq + for<'a> Deserialize<'a>,
954        Element: for<'a> Deserialize<'a>,
955    {
956        let iterations = self.get_all_iterations()?;
957        let all_elements = iterations
958            .iter()
959            .map(|iteration| {
960                let elements = self.load_all_elements_at_iteration(*iteration)?;
961                Ok((*iteration, elements))
962            })
963            .collect::<Result<BTreeMap<_, _>, StorageError>>()?;
964        Ok(all_elements)
965    }
966
967    /// Similarly to the [load_all_elements](StorageInterfaceLoad::load_all_elements) function,
968    /// but this function returns all elements as their histories.
969    fn load_all_element_histories(
970        &self,
971    ) -> Result<HashMap<Id, BTreeMap<u64, Element>>, StorageError>
972    where
973        Id: std::hash::Hash + std::cmp::Eq + for<'a> Deserialize<'a>,
974        Element: for<'a> Deserialize<'a>,
975    {
976        let all_elements = self.load_all_elements()?;
977        let reordered_elements: HashMap<Id, BTreeMap<u64, Element>> = all_elements
978            .into_iter()
979            .map(|(iteration, identifier_to_elements)| {
980                identifier_to_elements
981                    .into_iter()
982                    .map(move |(identifier, element)| (identifier, iteration, element))
983            })
984            .flatten()
985            .fold(
986                HashMap::new(),
987                |mut acc, (identifier, iteration, element)| {
988                    let existing_elements = acc.entry(identifier).or_default();
989                    existing_elements.insert(iteration, element);
990                    acc
991                },
992            );
993        Ok(reordered_elements)
994    }
995}
996
997impl<T, Id, Element> StorageInterfaceLoad<Id, Element> for StorageWrapper<T>
998where
999    T: FileBasedStorage<Id, Element>,
1000{
1001    fn load_single_element(
1002        &self,
1003        iteration: u64,
1004        identifier: &Id,
1005    ) -> Result<Option<Element>, StorageError>
1006    where
1007        Id: Eq + Serialize + for<'a> Deserialize<'a>,
1008        Element: for<'a> Deserialize<'a>,
1009    {
1010        let iterations = self.get_all_iterations()?;
1011        if iterations.contains(&iteration) {
1012            // Get the path where the iteration folder is
1013            let iteration_path = self.0.get_iteration_path(iteration);
1014
1015            // Load all elements which are inside this folder from batches and singles
1016            for path in std::fs::read_dir(&iteration_path)? {
1017                let p = path?.path();
1018                let content = std::fs::read_to_string(&p)?;
1019
1020                match p.file_stem() {
1021                    Some(stem) => match stem.to_str() {
1022                        Some(tail) => {
1023                            let first_name_segment = tail.split("_").into_iter().next();
1024                            if first_name_segment == Some("batch") {
1025                                let result: BatchSaveFormat<Id, Element> =
1026                                    self.0.from_str(&content)?;
1027                                for json_save_format in result.data.into_iter() {
1028                                    if &json_save_format.identifier == identifier {
1029                                        return Ok(Some(json_save_format.element));
1030                                    }
1031                                }
1032                            } else if first_name_segment == Some("single") {
1033                                let result: CombinedSaveFormat<Id, Element> =
1034                                    self.0.from_str(&content)?;
1035                                if &result.identifier == identifier {
1036                                    return Ok(Some(result.element));
1037                                }
1038                            }
1039                        }
1040                        None => (),
1041                    },
1042                    None => (),
1043                }
1044            }
1045            return Ok(None);
1046        } else {
1047            return Ok(None);
1048        }
1049    }
1050
1051    fn load_all_elements_at_iteration(
1052        &self,
1053        iteration: u64,
1054    ) -> Result<HashMap<Id, Element>, StorageError>
1055    where
1056        Id: std::hash::Hash + std::cmp::Eq + for<'a> Deserialize<'a>,
1057        Element: for<'a> Deserialize<'a>,
1058    {
1059        let iterations = self.get_all_iterations()?;
1060        if iterations.contains(&iteration) {
1061            // Create a new empty hashmap
1062            let mut all_elements_at_iteration = HashMap::new();
1063
1064            // Get the path where the iteration folder is
1065            let iteration_path = self.0.get_iteration_path(iteration);
1066
1067            // Load all elements which are inside this folder from batches and singles
1068            for path in std::fs::read_dir(&iteration_path)? {
1069                let p = path?.path();
1070                let content = std::fs::read_to_string(&p)?;
1071
1072                match p.file_stem() {
1073                    Some(stem) => match stem.to_str() {
1074                        Some(tail) => {
1075                            let first_name_segment = tail.split("_").into_iter().next();
1076                            if first_name_segment == Some("batch") {
1077                                let result: BatchSaveFormat<Id, Element> =
1078                                    self.0.from_str(&content)?;
1079                                all_elements_at_iteration.extend(result.data.into_iter().map(
1080                                    |json_save_format| {
1081                                        (json_save_format.identifier, json_save_format.element)
1082                                    },
1083                                ));
1084                            } else if first_name_segment == Some("single") {
1085                                let result: CombinedSaveFormat<Id, Element> =
1086                                    self.0.from_str(&content)?;
1087                                all_elements_at_iteration
1088                                    .extend([(result.identifier, result.element)]);
1089                            }
1090                        }
1091                        None => (),
1092                    },
1093                    None => (),
1094                }
1095            }
1096            return Ok(all_elements_at_iteration);
1097        } else {
1098            return Ok(HashMap::new());
1099        }
1100    }
1101
1102    fn get_all_iterations(&self) -> Result<Vec<u64>, StorageError> {
1103        let paths = std::fs::read_dir(&self.0.get_path())?;
1104        paths
1105            .into_iter()
1106            .filter_map(|path| match path {
1107                Ok(p) => match self.0.folder_name_to_iteration(&p.path()) {
1108                    Ok(Some(entry)) => Some(Ok(entry)),
1109                    Ok(None) => None,
1110                    Err(e) => Some(Err(e)),
1111                },
1112                Err(_) => None,
1113            })
1114            .collect::<Result<Vec<_>, _>>()
1115    }
1116}
1117
1118/// Provide methods to initialize, store and load single and multiple elements at iterations.
1119pub trait StorageInterface<Id, Element>:
1120    StorageInterfaceOpen + StorageInterfaceLoad<Id, Element> + StorageInterfaceStore<Id, Element>
1121{
1122}
1123
1124impl<Id, Element, T> StorageInterface<Id, Element> for T
1125where
1126    T: StorageInterfaceLoad<Id, Element>,
1127    T: StorageInterfaceStore<Id, Element>,
1128    T: StorageInterfaceOpen,
1129{
1130}
cellular_raza_core/storage/concepts.rs

cellular_raza_core/storage/
concepts.rs
