1 use crate::packets::Builder;
2 
3 pub struct PayloadAccumulator<T: Builder> {
4     curr: usize,
5     lim: usize,
6     elems: Vec<T>,
7 }
8 
9 impl<T: Builder> PayloadAccumulator<T> {
new(size: usize) -> Self10     pub fn new(size: usize) -> Self {
11         Self { curr: 0, lim: size * 8, elems: vec![] }
12     }
13 
14     #[must_use]
push(&mut self, builder: T) -> bool15     pub fn push(&mut self, builder: T) -> bool {
16         // if serialization fails we WANT to continue, to get a clean SerializeError at
17         // the end
18         let elem_size = builder.size_in_bits().unwrap_or(0);
19         if elem_size + self.curr > self.lim {
20             return false;
21         }
22         self.elems.push(builder);
23         self.curr += elem_size;
24         true
25     }
26 
into_boxed_slice(self) -> Box<[T]>27     pub fn into_boxed_slice(self) -> Box<[T]> {
28         self.elems.into_boxed_slice()
29     }
30 
is_empty(&self) -> bool31     pub fn is_empty(&self) -> bool {
32         self.elems.is_empty()
33     }
34 }
35 
36 #[cfg(test)]
37 mod test {
38     use crate::packets::{AttBuilder, AttChild, AttOpcode};
39 
40     use super::PayloadAccumulator;
41 
42     #[test]
test_empty()43     fn test_empty() {
44         let accumulator = PayloadAccumulator::<AttBuilder>::new(0);
45         assert!(accumulator.is_empty())
46     }
47     #[test]
test_nonempty()48     fn test_nonempty() {
49         let mut accumulator = PayloadAccumulator::new(128);
50 
51         let ok = accumulator.push(AttBuilder {
52             opcode: AttOpcode::WRITE_RESPONSE,
53             _child_: AttChild::RawData([1, 2].into()),
54         });
55 
56         assert!(ok);
57         assert!(!accumulator.is_empty())
58     }
59 
60     #[test]
test_push_serialize()61     fn test_push_serialize() {
62         let mut accumulator = PayloadAccumulator::new(128);
63 
64         let ok = accumulator.push(AttBuilder {
65             opcode: AttOpcode::WRITE_RESPONSE,
66             _child_: AttChild::RawData([1, 2].into()),
67         });
68 
69         assert!(ok);
70         assert_eq!(
71             accumulator.into_boxed_slice().as_ref(),
72             [AttBuilder {
73                 opcode: AttOpcode::WRITE_RESPONSE,
74                 _child_: AttChild::RawData([1, 2].into()),
75             }]
76         );
77     }
78 
79     #[test]
test_push_past_capacity()80     fn test_push_past_capacity() {
81         let mut accumulator = PayloadAccumulator::new(5);
82 
83         // each builder is 3 bytes, so the first should succeed, the second should fail
84         let first_ok = accumulator.push(AttBuilder {
85             opcode: AttOpcode::WRITE_RESPONSE,
86             _child_: AttChild::RawData([1, 2].into()),
87         });
88         let second_ok = accumulator.push(AttBuilder {
89             opcode: AttOpcode::WRITE_RESPONSE,
90             _child_: AttChild::RawData([3, 4].into()),
91         });
92 
93         // assert: the first one is pushed and is correctly output, but the second is
94         // dropped
95         assert!(first_ok);
96         assert!(!second_ok);
97         assert_eq!(
98             accumulator.into_boxed_slice().as_ref(),
99             [AttBuilder {
100                 opcode: AttOpcode::WRITE_RESPONSE,
101                 _child_: AttChild::RawData([1, 2].into()),
102             }]
103         );
104     }
105 
106     #[test]
test_push_to_capacity()107     fn test_push_to_capacity() {
108         let mut accumulator = PayloadAccumulator::new(5);
109 
110         // 3 + 2 bytes = the size, so both should push correctly
111         let first_ok = accumulator.push(AttBuilder {
112             opcode: AttOpcode::WRITE_RESPONSE,
113             _child_: AttChild::RawData([1, 2].into()),
114         });
115         let second_ok = accumulator.push(AttBuilder {
116             opcode: AttOpcode::WRITE_RESPONSE,
117             _child_: AttChild::RawData([3].into()),
118         });
119 
120         // assert: both are pushed and output correctly
121         assert!(first_ok);
122         assert!(second_ok);
123         assert_eq!(
124             accumulator.into_boxed_slice().as_ref(),
125             [
126                 AttBuilder {
127                     opcode: AttOpcode::WRITE_RESPONSE,
128                     _child_: AttChild::RawData([1, 2].into()),
129                 },
130                 AttBuilder {
131                     opcode: AttOpcode::WRITE_RESPONSE,
132                     _child_: AttChild::RawData([3].into()),
133                 }
134             ]
135         );
136     }
137 }
138