waves_rust/util/
byte_writer.rs

1use crate::model::{Arg, ByteString, Function};
2
3pub struct ByteWriter {
4    buf: Vec<u8>,
5}
6
7impl ByteWriter {
8    pub fn new() -> Self {
9        Self { buf: vec![] }
10    }
11
12    pub fn push_byte(&mut self, byte: u8) -> &mut Self {
13        self.buf.push(byte);
14        self
15    }
16
17    pub fn push_bytes(&mut self, bytes: &mut Vec<u8>) -> &mut Self {
18        self.buf.append(bytes);
19        self
20    }
21
22    pub fn push_int(&mut self, int: i32) -> &mut Self {
23        let mut int_bytes: Vec<u8> = int.to_be_bytes().into();
24        self.buf.append(&mut int_bytes);
25        self
26    }
27
28    pub fn push_long(&mut self, long: i64) -> &mut Self {
29        let mut long_bytes: Vec<u8> = long.to_be_bytes().into();
30        self.buf.append(&mut long_bytes);
31        self
32    }
33
34    pub fn bytes(&self) -> Vec<u8> {
35        self.buf.clone()
36    }
37
38    pub fn bytes_from_function(function: &Function) -> Vec<u8> {
39        let mut byte_writer = ByteWriter::new();
40        if function.is_default() {
41            return byte_writer.push_byte(0).bytes();
42        } else {
43            byte_writer
44                .push_bytes(&mut vec![1, 9, 1])
45                .push_int(function.name().len() as i32)
46                .push_bytes(&mut function.name().into_bytes())
47                .push_func_args(function.args())
48                .bytes()
49        }
50    }
51
52    pub fn push_func_args(&mut self, args: Vec<Arg>) -> &mut Self {
53        self.push_int(args.len() as i32);
54        for arg in args {
55            match arg {
56                Arg::Integer(integer) => {
57                    self.push_byte(0).push_long(integer);
58                }
59                Arg::Binary(binary) => {
60                    let bytes_len = binary.bytes().len() as i32;
61                    self.push_byte(1)
62                        .push_int(bytes_len)
63                        .push_bytes(&mut binary.bytes());
64                }
65                Arg::String(string) => {
66                    let mut string_bytes = string.into_bytes();
67                    self.push_byte(2)
68                        .push_int(string_bytes.len() as i32)
69                        .push_bytes(&mut string_bytes);
70                }
71                Arg::Boolean(boolean) => {
72                    if boolean {
73                        self.push_byte(6);
74                    } else {
75                        self.push_byte(7);
76                    }
77                }
78                Arg::List(list) => {
79                    self.push_byte(11).push_func_args(list);
80                }
81            }
82        }
83        self
84    }
85}
86
87impl Default for ByteWriter {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use crate::model::{Arg, Base64String, Function};
96    use crate::util::ByteWriter;
97
98    #[test]
99    fn test_push_int() {
100        assert_eq!(
101            vec![0, 18, 213, 59],
102            ByteWriter::new().push_int(1234235).bytes()
103        );
104    }
105
106    #[test]
107    fn test_push_long() {
108        assert_eq!(
109            vec![0, 0, 0, 0, 0, 18, 213, 59],
110            ByteWriter::new().push_long(1234235).bytes()
111        );
112    }
113
114    #[test]
115    fn test_bytes_from_function() {
116        let function = Function::new(
117            "storeData".to_owned(),
118            vec![
119                Arg::Boolean(true),
120                Arg::String("some string".to_owned()),
121                Arg::Integer(123),
122                Arg::Binary(Base64String::from_bytes(vec![3, 5, 2, 11, 15])),
123                Arg::List(vec![Arg::Integer(123), Arg::Integer(543)]),
124            ],
125        );
126
127        let function_bytes = ByteWriter::bytes_from_function(&function);
128        assert_eq!(
129            vec![
130                1, 9, 1, 0, 0, 0, 9, 115, 116, 111, 114, 101, 68, 97, 116, 97, 0, 0, 0, 5, 6, 2, 0,
131                0, 0, 11, 115, 111, 109, 101, 32, 115, 116, 114, 105, 110, 103, 0, 0, 0, 0, 0, 0,
132                0, 0, 123, 1, 0, 0, 0, 5, 3, 5, 2, 11, 15, 11, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0,
133                123, 0, 0, 0, 0, 0, 0, 0, 2, 31
134            ],
135            function_bytes
136        )
137    }
138}