base_d/encoders/algorithms/schema/
binary_packer.rs

1use crate::encoders::algorithms::schema::types::{
2    FLAG_HAS_NULLS, FLAG_HAS_ROOT_KEY, FieldType, IntermediateRepresentation, SchemaValue,
3};
4
5/// Pack intermediate representation into binary format
6pub fn pack(ir: &IntermediateRepresentation) -> Vec<u8> {
7    let mut buffer = Vec::new();
8
9    // Pack header
10    pack_header(&mut buffer, ir);
11
12    // Pack values
13    pack_values(&mut buffer, ir);
14
15    buffer
16}
17
18/// Pack the schema header
19fn pack_header(buffer: &mut Vec<u8>, ir: &IntermediateRepresentation) {
20    let header = &ir.header;
21
22    // Flags
23    buffer.push(header.flags);
24
25    // Root key (if present)
26    if header.has_flag(FLAG_HAS_ROOT_KEY) && header.root_key.is_some() {
27        let key = header.root_key.as_ref().unwrap();
28        encode_varint(buffer, key.len() as u64);
29        buffer.extend_from_slice(key.as_bytes());
30    }
31
32    // Row count
33    encode_varint(buffer, header.row_count as u64);
34
35    // Field count
36    encode_varint(buffer, header.fields.len() as u64);
37
38    // Field types (4 bits each, packed)
39    pack_field_types(buffer, ir);
40
41    // Field names
42    for field in &header.fields {
43        encode_varint(buffer, field.name.len() as u64);
44        buffer.extend_from_slice(field.name.as_bytes());
45    }
46
47    // Null bitmap (if present)
48    if header.has_flag(FLAG_HAS_NULLS) && header.null_bitmap.is_some() {
49        let bitmap = header.null_bitmap.as_ref().unwrap();
50        buffer.extend_from_slice(bitmap);
51    }
52}
53
54/// Pack field types (4 bits each)
55fn pack_field_types(buffer: &mut Vec<u8>, ir: &IntermediateRepresentation) {
56    let mut type_buffer = Vec::new();
57    let mut nibble_count = 0;
58
59    for field in &ir.header.fields {
60        pack_field_type_recursive(&mut type_buffer, &field.field_type, &mut nibble_count);
61    }
62
63    // Encode length of type buffer
64    encode_varint(buffer, type_buffer.len() as u64);
65    buffer.extend_from_slice(&type_buffer);
66}
67
68/// Pack a field type recursively (handles nested arrays)
69fn pack_field_type_recursive(
70    buffer: &mut Vec<u8>,
71    field_type: &FieldType,
72    nibble_count: &mut usize,
73) {
74    let tag = field_type.type_tag();
75
76    // Pack as 4-bit nibbles (2 per byte)
77    if (*nibble_count).is_multiple_of(2) {
78        // Start new byte with tag in lower nibble
79        buffer.push(tag);
80    } else {
81        // Add tag to upper nibble of last byte
82        let last_idx = buffer.len() - 1;
83        buffer[last_idx] |= tag << 4;
84    }
85    *nibble_count += 1;
86
87    // If array, recursively pack element type
88    if let FieldType::Array(element_type) = field_type {
89        pack_field_type_recursive(buffer, element_type, nibble_count);
90    }
91}
92
93/// Pack values
94fn pack_values(buffer: &mut Vec<u8>, ir: &IntermediateRepresentation) {
95    for value in &ir.values {
96        pack_value(buffer, value);
97    }
98}
99
100/// Pack a single value
101fn pack_value(buffer: &mut Vec<u8>, value: &SchemaValue) {
102    match value {
103        SchemaValue::U64(v) => encode_varint(buffer, *v),
104        SchemaValue::I64(v) => encode_signed_varint(buffer, *v),
105        SchemaValue::F64(v) => buffer.extend_from_slice(&v.to_le_bytes()),
106        SchemaValue::String(s) => {
107            encode_varint(buffer, s.len() as u64);
108            buffer.extend_from_slice(s.as_bytes());
109        }
110        SchemaValue::Bool(b) => buffer.push(if *b { 1 } else { 0 }),
111        SchemaValue::Null => {} // Null encoded in bitmap, no value bytes
112        SchemaValue::Array(arr) => {
113            encode_varint(buffer, arr.len() as u64);
114            for item in arr {
115                pack_value(buffer, item);
116            }
117        }
118    }
119}
120
121/// Encode unsigned varint (LEB128)
122pub(crate) fn encode_varint(buffer: &mut Vec<u8>, mut value: u64) {
123    loop {
124        let mut byte = (value & 0x7F) as u8;
125        value >>= 7;
126        if value != 0 {
127            byte |= 0x80; // More bytes follow
128        }
129        buffer.push(byte);
130        if value == 0 {
131            break;
132        }
133    }
134}
135
136/// Encode signed varint using zigzag encoding
137pub(crate) fn encode_signed_varint(buffer: &mut Vec<u8>, value: i64) {
138    let encoded = ((value << 1) ^ (value >> 63)) as u64;
139    encode_varint(buffer, encoded);
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use crate::encoders::algorithms::schema::types::{FieldDef, SchemaHeader};
146
147    #[test]
148    fn test_encode_varint() {
149        let mut buf = Vec::new();
150        encode_varint(&mut buf, 0);
151        assert_eq!(buf, vec![0]);
152
153        buf.clear();
154        encode_varint(&mut buf, 1);
155        assert_eq!(buf, vec![1]);
156
157        buf.clear();
158        encode_varint(&mut buf, 127);
159        assert_eq!(buf, vec![127]);
160
161        buf.clear();
162        encode_varint(&mut buf, 128);
163        assert_eq!(buf, vec![0x80, 0x01]);
164
165        buf.clear();
166        encode_varint(&mut buf, 16383);
167        assert_eq!(buf, vec![0xFF, 0x7F]);
168
169        buf.clear();
170        encode_varint(&mut buf, 16384);
171        assert_eq!(buf, vec![0x80, 0x80, 0x01]);
172    }
173
174    #[test]
175    fn test_encode_signed_varint() {
176        let mut buf = Vec::new();
177        encode_signed_varint(&mut buf, 0);
178        assert_eq!(buf, vec![0]);
179
180        buf.clear();
181        encode_signed_varint(&mut buf, -1);
182        assert_eq!(buf, vec![1]);
183
184        buf.clear();
185        encode_signed_varint(&mut buf, 1);
186        assert_eq!(buf, vec![2]);
187
188        buf.clear();
189        encode_signed_varint(&mut buf, -64);
190        assert_eq!(buf, vec![127]);
191
192        buf.clear();
193        encode_signed_varint(&mut buf, 64);
194        assert_eq!(buf, vec![128, 1]);
195    }
196
197    #[test]
198    fn test_pack_simple_ir() {
199        let fields = vec![
200            FieldDef::new("id", FieldType::U64),
201            FieldDef::new("name", FieldType::String),
202        ];
203        let header = SchemaHeader::new(1, fields);
204
205        let values = vec![
206            SchemaValue::U64(42),
207            SchemaValue::String("Alice".to_string()),
208        ];
209
210        let ir = IntermediateRepresentation::new(header, values).unwrap();
211        let packed = pack(&ir);
212
213        // Verify it produces some output
214        assert!(!packed.is_empty());
215
216        // First byte should be flags (0 for no flags)
217        assert_eq!(packed[0], 0);
218    }
219
220    #[test]
221    fn test_pack_with_root_key() {
222        let mut header = SchemaHeader::new(1, vec![FieldDef::new("id", FieldType::U64)]);
223        header.root_key = Some("users".to_string());
224        header.set_flag(FLAG_HAS_ROOT_KEY);
225
226        let values = vec![SchemaValue::U64(42)];
227        let ir = IntermediateRepresentation::new(header, values).unwrap();
228        let packed = pack(&ir);
229
230        // First byte should have FLAG_HAS_ROOT_KEY set
231        assert_eq!(packed[0] & FLAG_HAS_ROOT_KEY, FLAG_HAS_ROOT_KEY);
232    }
233
234    #[test]
235    fn test_pack_field_types() {
236        let fields = vec![
237            FieldDef::new("a", FieldType::U64),    // tag 0
238            FieldDef::new("b", FieldType::I64),    // tag 1
239            FieldDef::new("c", FieldType::String), // tag 3
240        ];
241        let header = SchemaHeader::new(1, fields);
242        let values = vec![
243            SchemaValue::U64(1),
244            SchemaValue::I64(-1),
245            SchemaValue::String("x".to_string()),
246        ];
247
248        let ir = IntermediateRepresentation::new(header, values).unwrap();
249        let packed = pack(&ir);
250
251        // Types should be packed as nibbles: 0, 1, 3
252        // In bytes: 0x10 (0 and 1), 0x03 (3)
253        // We need to find the type section in the packed data
254        assert!(!packed.is_empty());
255    }
256
257    #[test]
258    fn test_pack_values() {
259        let mut buffer = Vec::new();
260
261        pack_value(&mut buffer, &SchemaValue::U64(42));
262        assert_eq!(buffer, vec![42]);
263
264        buffer.clear();
265        pack_value(&mut buffer, &SchemaValue::Bool(true));
266        assert_eq!(buffer, vec![1]);
267
268        buffer.clear();
269        pack_value(&mut buffer, &SchemaValue::String("hi".to_string()));
270        assert_eq!(buffer, vec![2, b'h', b'i']);
271    }
272
273    #[test]
274    fn test_pack_array() {
275        let mut buffer = Vec::new();
276        let array = SchemaValue::Array(vec![SchemaValue::U64(1), SchemaValue::U64(2)]);
277        pack_value(&mut buffer, &array);
278
279        // Should be: count (2) + value (1) + value (2)
280        assert_eq!(buffer, vec![2, 1, 2]);
281    }
282}