Skip to main content

scripting/
serializer.rs

1use std::collections::HashMap;
2
3use finance::split::Split;
4use finance::transaction::Transaction;
5
6use crate::format::{
7    AccountData, BASE_OFFSET, CommodityData, ContextType, ENTITY_HEADER_SIZE, EntityFlags,
8    EntityHeader, EntityType, GLOBAL_HEADER_SIZE, GlobalHeader, OUTPUT_HEADER_SIZE, Operation,
9    OutputHeader, SplitData, TAG_DATA_SIZE, TagData, TransactionData,
10};
11
12fn transaction_to_data(tx: &Transaction) -> TransactionData {
13    TransactionData {
14        post_date: tx.post_date.timestamp_millis(),
15        enter_date: tx.enter_date.timestamp_millis(),
16        split_count: 0,
17        tag_count: 0,
18        is_multi_currency: 0,
19        reserved: [0; 23],
20    }
21}
22
23fn split_to_data(split: &Split) -> SplitData {
24    SplitData {
25        account_id: *split.account_id.as_bytes(),
26        commodity_id: *split.commodity_id.as_bytes(),
27        value_num: split.value_num,
28        value_denom: split.value_denom,
29        reconcile_state: split.reconcile_state.map_or(0, u8::from),
30        reserved: [0; 7],
31        reconcile_date: split
32            .reconcile_date
33            .map_or(0, |d: chrono::DateTime<chrono::Utc>| d.timestamp_millis()),
34    }
35}
36
37pub struct MemorySerializer {
38    context_type: ContextType,
39    primary_entity_type: EntityType,
40    primary_entity_idx: u32,
41    entities: Vec<SerializedEntity>,
42    strings_pool: Vec<u8>,
43    string_cache: HashMap<String, (u32, u16)>,
44}
45
46struct SerializedEntity {
47    header: EntityHeader,
48    data: Vec<u8>,
49}
50
51impl Default for MemorySerializer {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57impl MemorySerializer {
58    #[must_use]
59    pub fn new() -> Self {
60        Self {
61            context_type: ContextType::EntityCreate,
62            primary_entity_type: EntityType::Transaction,
63            primary_entity_idx: 0,
64            entities: Vec::new(),
65            strings_pool: Vec::new(),
66            string_cache: HashMap::new(),
67        }
68    }
69
70    pub fn set_context(&mut self, context_type: ContextType, primary_entity_type: EntityType) {
71        self.context_type = context_type;
72        self.primary_entity_type = primary_entity_type;
73    }
74
75    pub fn set_primary(&mut self, entity_idx: u32) {
76        self.primary_entity_idx = entity_idx;
77    }
78
79    pub fn add_string(&mut self, s: &str) -> (u32, u16) {
80        if let Some(&cached) = self.string_cache.get(s) {
81            return cached;
82        }
83        let offset = self.strings_pool.len() as u32;
84        let len = s.len() as u16;
85        self.strings_pool.extend_from_slice(s.as_bytes());
86        self.string_cache.insert(s.to_string(), (offset, len));
87        (offset, len)
88    }
89
90    pub fn add_transaction(
91        &mut self,
92        id: [u8; 16],
93        parent_idx: i32,
94        is_primary: bool,
95        is_context: bool,
96        post_date: i64,
97        enter_date: i64,
98        split_count: u32,
99        tag_count: u32,
100        is_multi_currency: bool,
101    ) -> u32 {
102        let flags = EntityFlags::make(is_primary, is_context);
103        let data = TransactionData {
104            post_date,
105            enter_date,
106            split_count,
107            tag_count,
108            is_multi_currency: u8::from(is_multi_currency),
109            reserved: [0; 23],
110        };
111        let header = EntityHeader::new(
112            EntityType::Transaction,
113            Operation::Nop,
114            flags,
115            id,
116            parent_idx,
117            0,
118            data.to_bytes().len() as u32,
119        );
120        let idx = self.entities.len() as u32;
121        self.entities.push(SerializedEntity {
122            header,
123            data: data.to_bytes().to_vec(),
124        });
125        idx
126    }
127
128    pub fn add_split(
129        &mut self,
130        id: [u8; 16],
131        parent_idx: i32,
132        is_primary: bool,
133        is_context: bool,
134        account_id: [u8; 16],
135        commodity_id: [u8; 16],
136        value_num: i64,
137        value_denom: i64,
138        reconcile_state: u8,
139        reconcile_date: i64,
140    ) -> u32 {
141        let flags = EntityFlags::make(is_primary, is_context);
142        let data = SplitData {
143            account_id,
144            commodity_id,
145            value_num,
146            value_denom,
147            reconcile_state,
148            reserved: [0; 7],
149            reconcile_date,
150        };
151        let header = EntityHeader::new(
152            EntityType::Split,
153            Operation::Nop,
154            flags,
155            id,
156            parent_idx,
157            0,
158            data.to_bytes().len() as u32,
159        );
160        let idx = self.entities.len() as u32;
161        self.entities.push(SerializedEntity {
162            header,
163            data: data.to_bytes().to_vec(),
164        });
165        idx
166    }
167
168    pub fn add_tag(
169        &mut self,
170        id: [u8; 16],
171        parent_idx: i32,
172        is_primary: bool,
173        is_context: bool,
174        name: &str,
175        value: &str,
176    ) -> u32 {
177        let flags = EntityFlags::make(is_primary, is_context);
178        let (name_offset, name_len) = self.add_string(name);
179        let (value_offset, value_len) = self.add_string(value);
180        let data = TagData {
181            name_offset,
182            value_offset,
183            name_len,
184            value_len,
185            reserved: [0; 4],
186        };
187        let header = EntityHeader::new(
188            EntityType::Tag,
189            Operation::Nop,
190            flags,
191            id,
192            parent_idx,
193            0,
194            TAG_DATA_SIZE as u32,
195        );
196        let idx = self.entities.len() as u32;
197        self.entities.push(SerializedEntity {
198            header,
199            data: data.to_bytes().to_vec(),
200        });
201        idx
202    }
203
204    pub fn add_account(
205        &mut self,
206        id: [u8; 16],
207        parent_idx: i32,
208        is_primary: bool,
209        is_context: bool,
210        parent_account_id: [u8; 16],
211        name: &str,
212        path: &str,
213        tag_count: u32,
214    ) -> u32 {
215        let flags = EntityFlags::make(is_primary, is_context);
216        let (name_offset, name_len) = self.add_string(name);
217        let (path_offset, path_len) = self.add_string(path);
218        let data = AccountData {
219            parent_account_id,
220            name_offset,
221            path_offset,
222            tag_count,
223            name_len,
224            path_len,
225            reserved: [0; 16],
226        };
227        let header = EntityHeader::new(
228            EntityType::Account,
229            Operation::Nop,
230            flags,
231            id,
232            parent_idx,
233            0,
234            data.to_bytes().len() as u32,
235        );
236        let idx = self.entities.len() as u32;
237        self.entities.push(SerializedEntity {
238            header,
239            data: data.to_bytes().to_vec(),
240        });
241        idx
242    }
243
244    pub fn add_commodity(
245        &mut self,
246        id: [u8; 16],
247        parent_idx: i32,
248        is_primary: bool,
249        is_context: bool,
250        symbol: &str,
251        name: &str,
252        tag_count: u32,
253    ) -> u32 {
254        let flags = EntityFlags::make(is_primary, is_context);
255        let (symbol_offset, symbol_len) = self.add_string(symbol);
256        let (name_offset, name_len) = self.add_string(name);
257        let data = CommodityData {
258            symbol_offset,
259            name_offset,
260            tag_count,
261            symbol_len,
262            name_len,
263            reserved: [0; 16],
264        };
265        let header = EntityHeader::new(
266            EntityType::Commodity,
267            Operation::Nop,
268            flags,
269            id,
270            parent_idx,
271            0,
272            data.to_bytes().len() as u32,
273        );
274        let idx = self.entities.len() as u32;
275        self.entities.push(SerializedEntity {
276            header,
277            data: data.to_bytes().to_vec(),
278        });
279        idx
280    }
281
282    #[must_use]
283    pub fn entity_count(&self) -> u32 {
284        self.entities.len() as u32
285    }
286
287    pub fn add_transaction_from(
288        &mut self,
289        tx: &Transaction,
290        is_primary: bool,
291        split_count: u32,
292        tag_count: u32,
293        is_multi_currency: bool,
294    ) -> u32 {
295        let flags = EntityFlags::make(is_primary, false);
296        let mut data = transaction_to_data(tx);
297        data.split_count = split_count;
298        data.tag_count = tag_count;
299        data.is_multi_currency = u8::from(is_multi_currency);
300        let header = EntityHeader::new(
301            EntityType::Transaction,
302            Operation::Nop,
303            flags,
304            *tx.id.as_bytes(),
305            -1,
306            0,
307            data.to_bytes().len() as u32,
308        );
309        let idx = self.entities.len() as u32;
310        self.entities.push(SerializedEntity {
311            header,
312            data: data.to_bytes().to_vec(),
313        });
314        idx
315    }
316
317    pub fn add_split_from(&mut self, split: &Split, parent_idx: i32) -> u32 {
318        let data = split_to_data(split);
319        let header = EntityHeader::new(
320            EntityType::Split,
321            Operation::Nop,
322            0,
323            *split.id.as_bytes(),
324            parent_idx,
325            0,
326            data.to_bytes().len() as u32,
327        );
328        let idx = self.entities.len() as u32;
329        self.entities.push(SerializedEntity {
330            header,
331            data: data.to_bytes().to_vec(),
332        });
333        idx
334    }
335
336    #[must_use]
337    pub fn finalize(mut self, output_size: u32) -> Vec<u8> {
338        let entity_count = self.entities.len() as u32;
339        let entities_offset = BASE_OFFSET + GLOBAL_HEADER_SIZE as u32;
340
341        let mut entities_total_size = 0u32;
342        for entity in &self.entities {
343            entities_total_size += ENTITY_HEADER_SIZE as u32 + entity.data.len() as u32;
344        }
345
346        let strings_pool_offset = entities_offset + entities_total_size;
347        let strings_pool_size = self.strings_pool.len() as u32;
348        let output_offset = strings_pool_offset + strings_pool_size;
349
350        let output_header = OutputHeader::new(entity_count);
351
352        let mut global_header = GlobalHeader::new(
353            self.context_type,
354            self.primary_entity_type,
355            entity_count,
356            self.primary_entity_idx,
357        );
358        global_header.entities_offset = entities_offset;
359        global_header.strings_pool_offset = strings_pool_offset;
360        global_header.strings_pool_size = strings_pool_size;
361        global_header.output_offset = output_offset;
362        global_header.output_size = output_size;
363
364        let total_size = GLOBAL_HEADER_SIZE
365            + entities_total_size as usize
366            + strings_pool_size as usize
367            + output_size as usize;
368        let mut buffer = vec![0u8; total_size];
369
370        buffer[..GLOBAL_HEADER_SIZE].copy_from_slice(global_header.as_bytes());
371
372        let headers_total = entity_count as usize * ENTITY_HEADER_SIZE;
373        let mut data_offset = entities_offset + headers_total as u32;
374        let mut write_pos = GLOBAL_HEADER_SIZE;
375
376        for entity in &mut self.entities {
377            entity.header.data_offset = data_offset;
378            data_offset += entity.data.len() as u32;
379        }
380
381        for entity in &self.entities {
382            let header_bytes = entity.header.to_bytes();
383            buffer[write_pos..write_pos + ENTITY_HEADER_SIZE].copy_from_slice(&header_bytes);
384            write_pos += ENTITY_HEADER_SIZE;
385        }
386
387        for entity in &self.entities {
388            buffer[write_pos..write_pos + entity.data.len()].copy_from_slice(&entity.data);
389            write_pos += entity.data.len();
390        }
391
392        buffer[write_pos..write_pos + self.strings_pool.len()].copy_from_slice(&self.strings_pool);
393        write_pos += self.strings_pool.len();
394
395        buffer[write_pos..write_pos + OUTPUT_HEADER_SIZE]
396            .copy_from_slice(&output_header.to_bytes());
397
398        buffer
399    }
400}
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405    use crate::format::MAGIC_NOMI;
406
407    #[test]
408    fn test_serializer_basic() {
409        let mut serializer = MemorySerializer::new();
410        serializer.set_context(ContextType::EntityCreate, EntityType::Transaction);
411
412        let tx_id = [1u8; 16];
413        let tx_idx = serializer.add_transaction(tx_id, -1, true, false, 1000, 2000, 2, 1, false);
414        serializer.set_primary(tx_idx);
415
416        let split_id = [2u8; 16];
417        let account_id = [3u8; 16];
418        let commodity_id = [4u8; 16];
419        serializer.add_split(
420            split_id,
421            tx_idx as i32,
422            false,
423            false,
424            account_id,
425            commodity_id,
426            -5000,
427            100,
428            0,
429            0,
430        );
431
432        let tag_id = [5u8; 16];
433        serializer.add_tag(
434            tag_id,
435            tx_idx as i32,
436            false,
437            false,
438            "note",
439            "test transaction",
440        );
441
442        assert_eq!(serializer.entity_count(), 3);
443
444        let buffer = serializer.finalize(1024);
445
446        let header = GlobalHeader::from_bytes(&buffer).unwrap();
447        assert_eq!(header.magic, MAGIC_NOMI);
448        assert_eq!(header.input_entity_count, 3);
449        assert_eq!(header.context_type, ContextType::EntityCreate as u8);
450        assert_eq!(header.primary_entity_type, EntityType::Transaction as u8);
451    }
452
453    #[test]
454    fn test_string_deduplication() {
455        let mut serializer = MemorySerializer::new();
456        let (offset1, len1) = serializer.add_string("test");
457        let (offset2, len2) = serializer.add_string("test");
458        let (offset3, _) = serializer.add_string("other");
459
460        assert_eq!(offset1, offset2);
461        assert_eq!(len1, len2);
462        assert_ne!(offset1, offset3);
463    }
464}