rdgen_lib/
lib.rs

1use std::{io::Cursor, num::NonZeroUsize};
2
3use blake2::{Blake2b, Digest};
4
5#[derive(thiserror::Error, Debug)]
6pub enum Error {
7    #[error("Error while reading data stream: `{0}`")]
8    DataStreamError(String),
9}
10
11#[must_use]
12pub struct InfiniteDataWriter {
13    seed: [u8; 64],
14}
15
16impl InfiniteDataWriter {
17    /// Create a new instance with the given seed.
18    pub fn new(seed: impl AsRef<[u8]>) -> Self {
19        Self::new_from_stream(Cursor::new(seed.as_ref())).expect("Cannot fail")
20    }
21
22    /// Create a new instance with the given stream of data.
23    pub fn new_from_stream(mut source: impl std::io::Read) -> Result<Self, Error> {
24        let mut seed_hasher = Blake2b::new();
25
26        let mut buffer = [0; 4096];
27
28        loop {
29            let bytes_read = source
30                .read(&mut buffer)
31                .map_err(|e| Error::DataStreamError(e.to_string()))?;
32
33            if bytes_read == 0 {
34                break;
35            }
36
37            seed_hasher.update(&buffer[..bytes_read]);
38        }
39
40        let seed = seed_hasher.finalize().into();
41        Ok(Self { seed })
42    }
43
44    /// Pull a batch of data, and generate new data in seed
45    pub fn pull(&mut self) -> [u8; 64] {
46        let mut hasher = Blake2b::new();
47        hasher.update(self.seed.as_ref());
48        let mut seed = hasher.finalize().into();
49        std::mem::swap(&mut seed, &mut self.seed);
50        seed
51    }
52
53    pub const fn batch_size(&self) -> NonZeroUsize {
54        match NonZeroUsize::new(self.seed.len()) {
55            Some(r) => r,
56            None => panic!("Size must be larger than zero"),
57        }
58    }
59}
60
61impl Iterator for InfiniteDataWriter {
62    type Item = [u8; 64];
63
64    fn next(&mut self) -> Option<Self::Item> {
65        Some(self.pull())
66    }
67}
68
69#[must_use]
70pub struct FiniteDataWriter {
71    writer: InfiniteDataWriter,
72    desired_length: Option<usize>,
73    pulled_length: usize,
74}
75
76impl FiniteDataWriter {
77    /// Create a new instance with the given seed.
78    /// If `desired length` is Some(), the output will be limited to that length. If None, the output will never have an end.
79    pub fn new(seed: impl AsRef<[u8]>, desired_length: Option<usize>) -> Self {
80        Self {
81            writer: InfiniteDataWriter::new(seed),
82            desired_length,
83            pulled_length: 0,
84        }
85    }
86
87    /// Create a new instance with the given stream of data.
88    /// If `desired length` is Some(), the output will be limited to that length. If None, the output will never have an end.
89    pub fn new_from_stream(
90        source: impl std::io::Read,
91        desired_length: Option<usize>,
92    ) -> Result<Self, Error> {
93        Ok(Self {
94            writer: InfiniteDataWriter::new_from_stream(source)?,
95            desired_length,
96            pulled_length: 0,
97        })
98    }
99
100    /// Pull a batch of data, and generate new data in seed
101    pub fn pull(&mut self) -> Vec<u8> {
102        let data = self.writer.pull();
103
104        let desired_length = match self.desired_length {
105            Some(l) => l,
106            None => return data.to_vec(),
107        };
108
109        let max_length_to_push = desired_length - self.pulled_length;
110
111        if max_length_to_push > self.writer.batch_size().get() {
112            self.pulled_length += self.writer.batch_size().get();
113            data.to_vec()
114        } else {
115            self.pulled_length += max_length_to_push;
116            data.split_at(max_length_to_push).0.to_vec()
117        }
118    }
119}
120
121impl Iterator for FiniteDataWriter {
122    type Item = Vec<u8>;
123
124    fn next(&mut self) -> Option<Self::Item> {
125        let data = self.pull();
126        if data.is_empty() {
127            None
128        } else {
129            Some(data)
130        }
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137
138    #[test]
139    fn basic() {
140        let mut writer = FiniteDataWriter::new("abc", Some(256));
141        assert_eq!(hex::encode(writer.pull()), "ba80a53f981c4d0d6a2797b69f12f6e94c212f14685ac4b74b12bb6fdbffa2d17d87c5392aab792dc252d5de4533cc9518d38aa8dbf1925ab92386edd4009923");
142        assert_eq!(hex::encode(writer.pull()), "66cb547665e462bbdd51d9b6ce1221116e9cfc6711c78d8798158349d12fa8ca513efb14bd84edf4e7cd3551355f14c1cf54dd203669b95675e52d72d3ec00d9");
143        assert_eq!(hex::encode(writer.pull()), "2ddda015a6b31d39fa9e6d54bb55bab1999a224d23b094fb1f77c41a1ea597c485e10bc721dd5531f1cddc52fdafa09c03ac4fbaaac9271241bd1da64dbd390c");
144        assert_eq!(hex::encode(writer.pull()), "50f4b533357084ec5a41ff26dfd36e069a1bf23ed6fd17ee341cf082d409854480332831399565d3f6fa0bed4cab0fad7c81c62b66c2b328ab880f139a094e1c");
145        // All subsequent pulls, should yield nothing
146        for _ in 0..1000 {
147            assert_eq!(hex::encode(writer.pull()), "");
148        }
149    }
150
151    #[test]
152    fn non_multiple_len() {
153        let mut writer = FiniteDataWriter::new("abc", Some(100));
154        assert_eq!(hex::encode(writer.pull()), "ba80a53f981c4d0d6a2797b69f12f6e94c212f14685ac4b74b12bb6fdbffa2d17d87c5392aab792dc252d5de4533cc9518d38aa8dbf1925ab92386edd4009923");
155        assert_eq!(
156            hex::encode(writer.pull()),
157            "66cb547665e462bbdd51d9b6ce1221116e9cfc6711c78d8798158349d12fa8ca513efb14"
158        );
159        // All subsequent pulls, should yield nothing
160        for _ in 0..1000 {
161            assert_eq!(hex::encode(writer.pull()), "");
162        }
163    }
164
165    #[test]
166    fn empty() {
167        let mut writer = FiniteDataWriter::new("abc", Some(0));
168        // All subsequent pulls, should yield nothing
169        for _ in 0..1000 {
170            assert_eq!(hex::encode(writer.pull()), "");
171        }
172    }
173
174    #[test]
175    fn no_data() {
176        let mut writer = FiniteDataWriter::new("abc", None);
177        assert_eq!(hex::encode(writer.pull()), "ba80a53f981c4d0d6a2797b69f12f6e94c212f14685ac4b74b12bb6fdbffa2d17d87c5392aab792dc252d5de4533cc9518d38aa8dbf1925ab92386edd4009923");
178        assert_eq!(hex::encode(writer.pull()), "66cb547665e462bbdd51d9b6ce1221116e9cfc6711c78d8798158349d12fa8ca513efb14bd84edf4e7cd3551355f14c1cf54dd203669b95675e52d72d3ec00d9");
179        assert_eq!(hex::encode(writer.pull()), "2ddda015a6b31d39fa9e6d54bb55bab1999a224d23b094fb1f77c41a1ea597c485e10bc721dd5531f1cddc52fdafa09c03ac4fbaaac9271241bd1da64dbd390c");
180        assert_eq!(hex::encode(writer.pull()), "50f4b533357084ec5a41ff26dfd36e069a1bf23ed6fd17ee341cf082d409854480332831399565d3f6fa0bed4cab0fad7c81c62b66c2b328ab880f139a094e1c");
181        assert_eq!(hex::encode(writer.pull()), "500cb0c9c086a7d65309a6e1d792501f811812411dc22f557c687af44428b68ce19f15ffe1f469cad0fe1180182151ac86f7f406f97e35f943bb084f1f51462b");
182
183        // The data should never end, since size is None
184        for _ in 0..100 {
185            assert!(writer.pull().len() > 0);
186        }
187    }
188
189    #[test]
190    fn all_sizes_homomorphism() {
191        const MAX_SIZE: usize = 2000;
192        const SEED: &str = "abc";
193
194        // Generate data with MAX_SIZE limit.
195        let writer = FiniteDataWriter::new(SEED, Some(MAX_SIZE));
196        let expected = writer.into_iter().fold(Vec::new(), |mut so_far, curr| {
197            so_far.extend(curr);
198            so_far
199        });
200
201        assert_eq!(expected.len(), MAX_SIZE);
202
203        // Make sure that any data generated with size < MAX_SIZE is a subset of the previous result.
204        for curr_size in 0..MAX_SIZE {
205            let writer = FiniteDataWriter::new(SEED, Some(curr_size));
206            let actual = writer.into_iter().fold(Vec::new(), |mut so_far, curr| {
207                so_far.extend(curr);
208                so_far
209            });
210            assert_eq!(actual.len(), curr_size);
211            assert_eq!(actual, expected[0..curr_size]);
212        }
213    }
214}
rdgen_lib/lib.rs

rdgen_lib/
lib.rs
