sbrd_gen/writer/
writer_base.rs

1use crate::error::SchemaResult;
2use crate::generator::Randomizer;
3use crate::Schema;
4use serde::ser::{Error, SerializeMap, SerializeSeq};
5use serde::{Serialize, Serializer};
6use std::sync::Mutex;
7
8/// Trait of Writer for keys and generated values
9pub trait GeneratedValueWriterBase<W: std::io::Write> {
10    /// Create from the writer
11    fn from_writer(writer: W) -> Self;
12
13    /// Take writer
14    fn into_inner(self) -> W;
15
16    /// Flush buffers that have not been written yet
17    fn flush(&mut self) -> SchemaResult<()>;
18
19    /// Generate all values and then write keys and the values
20    fn write_after_all_generated<R: Randomizer + ?Sized>(
21        &mut self,
22        use_key_header: bool,
23        schema: &Schema<R>,
24        rng: &mut R,
25        count: u64,
26    ) -> SchemaResult<()>;
27
28    /// Write keys and generated values while generating values
29    fn write_with_generate<R: Randomizer + ?Sized>(
30        &mut self,
31        use_key_header: bool,
32        schema: &Schema<R>,
33        rng: &mut R,
34        count: u64,
35    ) -> SchemaResult<()>;
36}
37
38/// dummy key of `keys` for a KVS
39pub const DUMMY_KEYS_NAME: &str = "keys";
40/// dummy key of `values` for a KVS
41pub const DUMMY_VALUES_NAME: &str = "values";
42
43/// A Sequence of generated values
44pub struct GeneratedDisplayValues<K: Serialize, V: Serialize> {
45    key_values: Vec<(K, V)>,
46}
47
48impl<K: Serialize, V: Serialize> GeneratedDisplayValues<K, V> {
49    /// Constructor
50    pub fn new(key_values: Vec<(K, V)>) -> Self {
51        Self { key_values }
52    }
53}
54
55impl<K: Serialize, V: Serialize> Serialize for GeneratedDisplayValues<K, V> {
56    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
57    where
58        S: Serializer,
59    {
60        let mut map_state = serializer.serialize_map(Some(self.key_values.len()))?;
61        for (k, v) in self.key_values.iter() {
62            map_state.serialize_entry(k, v)?;
63        }
64        map_state.end()
65    }
66}
67
68/// A Sequence of generated values while generating it
69pub struct SerializeWithGenerate<'a, R: Randomizer + ?Sized> {
70    schema: &'a Schema<R>,
71    rng: Mutex<&'a mut R>,
72    count: &'a u64,
73}
74
75impl<'a, R: Randomizer + ?Sized> SerializeWithGenerate<'a, R> {
76    /// Constructor
77    pub fn new(schema: &'a Schema<R>, rng: &'a mut R, count: &'a u64) -> Self {
78        Self {
79            schema,
80            rng: Mutex::new(rng),
81            count,
82        }
83    }
84}
85
86impl<'a, R: Randomizer + ?Sized> Serialize for SerializeWithGenerate<'a, R> {
87    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
88    where
89        S: Serializer,
90    {
91        let mut seq_state = serializer.serialize_seq(Some(*self.count as usize))?;
92
93        for _ in 0..*self.count {
94            let generated = {
95                let mut rng = self.rng.try_lock().map_err(S::Error::custom)?;
96                self.schema.generate(*rng).map_err(S::Error::custom)?
97            };
98
99            let values = generated.into_values_with_key().map_err(S::Error::custom)?;
100
101            let json_map = GeneratedDisplayValues::new(values);
102            seq_state.serialize_element(&json_map)?;
103        }
104
105        seq_state.end()
106    }
107}
sbrd_gen/writer/writer_base.rs

sbrd_gen/writer/
writer_base.rs
