sql_splitter/redactor/strategy/
shuffle.rs1use super::{RedactValue, Strategy, StrategyKind};
10use rand::seq::SliceRandom;
11
12#[derive(Debug, Clone, Default)]
14pub struct ShuffleStrategy {
15 values: Vec<RedactValue>,
17 shuffled: Vec<RedactValue>,
19 index: usize,
21}
22
23impl ShuffleStrategy {
24 pub fn new() -> Self {
25 Self::default()
26 }
27
28 pub fn collect(&mut self, value: RedactValue) {
30 self.values.push(value);
31 }
32
33 pub fn shuffle(&mut self, rng: &mut impl rand::Rng) {
35 self.shuffled = self.values.clone();
36 self.shuffled.shuffle(rng);
37 self.index = 0;
38 }
39
40 pub fn next_value(&mut self) -> Option<RedactValue> {
42 if self.index < self.shuffled.len() {
43 let value = self.shuffled[self.index].clone();
44 self.index += 1;
45 Some(value)
46 } else {
47 None
48 }
49 }
50
51 pub fn has_values(&self) -> bool {
53 !self.values.is_empty()
54 }
55
56 pub fn len(&self) -> usize {
58 self.values.len()
59 }
60
61 pub fn is_empty(&self) -> bool {
63 self.values.is_empty()
64 }
65}
66
67impl Strategy for ShuffleStrategy {
68 fn apply(&self, value: &RedactValue, _rng: &mut dyn rand::RngCore) -> RedactValue {
69 value.clone()
80 }
81
82 fn kind(&self) -> StrategyKind {
83 StrategyKind::Shuffle
84 }
85}
86
87#[cfg(test)]
88mod tests {
89 use super::*;
90 use rand::SeedableRng;
91
92 #[test]
93 fn test_shuffle_strategy() {
94 let mut strategy = ShuffleStrategy::new();
95 let mut rng = rand::rngs::StdRng::seed_from_u64(42);
96
97 strategy.collect(RedactValue::Integer(1));
99 strategy.collect(RedactValue::Integer(2));
100 strategy.collect(RedactValue::Integer(3));
101 strategy.collect(RedactValue::Integer(4));
102 strategy.collect(RedactValue::Integer(5));
103
104 assert_eq!(strategy.len(), 5);
105
106 strategy.shuffle(&mut rng);
108
109 let mut results = Vec::new();
111 while let Some(v) = strategy.next_value() {
112 results.push(v);
113 }
114
115 assert_eq!(results.len(), 5);
116
117 let mut ints: Vec<i64> = results
119 .iter()
120 .map(|v| match v {
121 RedactValue::Integer(i) => *i,
122 _ => panic!("Expected Integer"),
123 })
124 .collect();
125 ints.sort();
126 assert_eq!(ints, vec![1, 2, 3, 4, 5]);
127 }
128
129 #[test]
130 fn test_shuffle_with_nulls() {
131 let mut strategy = ShuffleStrategy::new();
132 let mut rng = rand::rngs::StdRng::seed_from_u64(42);
133
134 strategy.collect(RedactValue::String("a".to_string()));
135 strategy.collect(RedactValue::Null);
136 strategy.collect(RedactValue::String("b".to_string()));
137
138 strategy.shuffle(&mut rng);
139
140 let mut null_count = 0;
141 let mut string_count = 0;
142
143 while let Some(v) = strategy.next_value() {
144 match v {
145 RedactValue::Null => null_count += 1,
146 RedactValue::String(_) => string_count += 1,
147 _ => {}
148 }
149 }
150
151 assert_eq!(null_count, 1);
152 assert_eq!(string_count, 2);
153 }
154}