rustywallet_batch/
stream.rs

1//! Memory-efficient key streaming.
2//!
3//! This module provides [`KeyStream`] for generating keys on-demand
4//! without storing all keys in memory.
5
6use crate::error::BatchError;
7use rustywallet_keys::private_key::PrivateKey;
8
9/// A memory-efficient stream of generated private keys.
10///
11/// `KeyStream` implements the `Iterator` trait, allowing keys to be
12/// generated on-demand without storing all keys in memory.
13///
14/// # Example
15///
16/// ```rust
17/// use rustywallet_batch::prelude::*;
18///
19/// let stream = BatchGenerator::new()
20///     .count(1_000_000)
21///     .generate()
22///     .unwrap();
23///
24/// // Process keys one at a time
25/// for key in stream.take(100) {
26///     println!("{}", key.unwrap().to_hex());
27/// }
28/// ```
29pub struct KeyStream {
30    /// The inner iterator that generates keys.
31    inner: Box<dyn Iterator<Item = Result<PrivateKey, BatchError>> + Send>,
32
33    /// Total number of keys to generate (if known).
34    total_count: Option<usize>,
35
36    /// Number of keys generated so far.
37    generated_count: usize,
38}
39
40impl KeyStream {
41    /// Create a new key stream from an iterator.
42    pub fn new<I>(iter: I, total_count: Option<usize>) -> Self
43    where
44        I: Iterator<Item = Result<PrivateKey, BatchError>> + Send + 'static,
45    {
46        Self {
47            inner: Box::new(iter),
48            total_count,
49            generated_count: 0,
50        }
51    }
52
53    /// Get the progress as a percentage (0.0 to 1.0).
54    ///
55    /// Returns `None` if the total count is unknown.
56    pub fn progress(&self) -> Option<f64> {
57        self.total_count.map(|total| {
58            if total == 0 {
59                1.0
60            } else {
61                self.generated_count as f64 / total as f64
62            }
63        })
64    }
65
66    /// Get the number of keys generated so far.
67    pub fn generated(&self) -> usize {
68        self.generated_count
69    }
70
71    /// Get the total number of keys to generate (if known).
72    pub fn total(&self) -> Option<usize> {
73        self.total_count
74    }
75
76    /// Get the number of remaining keys (if known).
77    pub fn remaining(&self) -> Option<usize> {
78        self.total_count.map(|total| total.saturating_sub(self.generated_count))
79    }
80
81    /// Collect a chunk of keys from the stream.
82    ///
83    /// This is useful for processing keys in batches while still
84    /// maintaining memory efficiency.
85    pub fn collect_chunk(&mut self, size: usize) -> Vec<Result<PrivateKey, BatchError>> {
86        let mut chunk = Vec::with_capacity(size);
87        for _ in 0..size {
88            match self.next() {
89                Some(key) => chunk.push(key),
90                None => break,
91            }
92        }
93        chunk
94    }
95}
96
97impl Iterator for KeyStream {
98    type Item = Result<PrivateKey, BatchError>;
99
100    fn next(&mut self) -> Option<Self::Item> {
101        let result = self.inner.next();
102        if result.is_some() {
103            self.generated_count += 1;
104        }
105        result
106    }
107
108    fn size_hint(&self) -> (usize, Option<usize>) {
109        match self.total_count {
110            Some(total) => {
111                let remaining = total.saturating_sub(self.generated_count);
112                (remaining, Some(remaining))
113            }
114            None => (0, None),
115        }
116    }
117}
118
119impl ExactSizeIterator for KeyStream {
120    fn len(&self) -> usize {
121        self.total_count
122            .map(|total| total.saturating_sub(self.generated_count))
123            .unwrap_or(0)
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn test_stream_progress() {
133        let keys: Vec<Result<PrivateKey, BatchError>> = (0..10)
134            .map(|_| Ok(PrivateKey::random()))
135            .collect();
136        
137        let mut stream = KeyStream::new(keys.into_iter(), Some(10));
138        
139        assert_eq!(stream.progress(), Some(0.0));
140        assert_eq!(stream.generated(), 0);
141        assert_eq!(stream.total(), Some(10));
142        assert_eq!(stream.remaining(), Some(10));
143
144        // Consume 5 keys
145        for _ in 0..5 {
146            stream.next();
147        }
148
149        assert_eq!(stream.progress(), Some(0.5));
150        assert_eq!(stream.generated(), 5);
151        assert_eq!(stream.remaining(), Some(5));
152    }
153
154    #[test]
155    fn test_stream_collect_chunk() {
156        let keys: Vec<Result<PrivateKey, BatchError>> = (0..100)
157            .map(|_| Ok(PrivateKey::random()))
158            .collect();
159        
160        let mut stream = KeyStream::new(keys.into_iter(), Some(100));
161        
162        let chunk = stream.collect_chunk(25);
163        assert_eq!(chunk.len(), 25);
164        assert_eq!(stream.generated(), 25);
165        assert_eq!(stream.remaining(), Some(75));
166    }
167
168    #[test]
169    fn test_stream_size_hint() {
170        let keys: Vec<Result<PrivateKey, BatchError>> = (0..50)
171            .map(|_| Ok(PrivateKey::random()))
172            .collect();
173        
174        let mut stream = KeyStream::new(keys.into_iter(), Some(50));
175        
176        assert_eq!(stream.size_hint(), (50, Some(50)));
177        
178        for _ in 0..20 {
179            stream.next();
180        }
181        
182        assert_eq!(stream.size_hint(), (30, Some(30)));
183    }
184}
rustywallet_batch/stream.rs

rustywallet_batch/
stream.rs
