1
use std::collections::HashMap;
2

            
3
use finance::split::Split;
4
use finance::transaction::Transaction;
5

            
6
use 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

            
12
588
fn transaction_to_data(tx: &Transaction) -> TransactionData {
13
588
    TransactionData {
14
588
        post_date: tx.post_date.timestamp_millis(),
15
588
        enter_date: tx.enter_date.timestamp_millis(),
16
588
        split_count: 0,
17
588
        tag_count: 0,
18
588
        is_multi_currency: 0,
19
588
        reserved: [0; 23],
20
588
    }
21
588
}
22

            
23
980
fn split_to_data(split: &Split) -> SplitData {
24
    SplitData {
25
980
        account_id: *split.account_id.as_bytes(),
26
980
        commodity_id: *split.commodity_id.as_bytes(),
27
980
        value_num: split.value_num,
28
980
        value_denom: split.value_denom,
29
980
        reconcile_state: split.reconcile_state.map_or(0, u8::from),
30
980
        reserved: [0; 7],
31
980
        reconcile_date: split
32
980
            .reconcile_date
33
980
            .map_or(0, |d: chrono::DateTime<chrono::Utc>| d.timestamp_millis()),
34
    }
35
980
}
36

            
37
pub 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

            
46
struct SerializedEntity {
47
    header: EntityHeader,
48
    data: Vec<u8>,
49
}
50

            
51
impl Default for MemorySerializer {
52
    fn default() -> Self {
53
        Self::new()
54
    }
55
}
56

            
57
impl MemorySerializer {
58
    #[must_use]
59
5994
    pub fn new() -> Self {
60
5994
        Self {
61
5994
            context_type: ContextType::EntityCreate,
62
5994
            primary_entity_type: EntityType::Transaction,
63
5994
            primary_entity_idx: 0,
64
5994
            entities: Vec::new(),
65
5994
            strings_pool: Vec::new(),
66
5994
            string_cache: HashMap::new(),
67
5994
        }
68
5994
    }
69

            
70
5993
    pub fn set_context(&mut self, context_type: ContextType, primary_entity_type: EntityType) {
71
5993
        self.context_type = context_type;
72
5993
        self.primary_entity_type = primary_entity_type;
73
5993
    }
74

            
75
1457
    pub fn set_primary(&mut self, entity_idx: u32) {
76
1457
        self.primary_entity_idx = entity_idx;
77
1457
    }
78

            
79
2973
    pub fn add_string(&mut self, s: &str) -> (u32, u16) {
80
2973
        if let Some(&cached) = self.string_cache.get(s) {
81
85
            return cached;
82
2888
        }
83
2888
        let offset = self.strings_pool.len() as u32;
84
2888
        let len = s.len() as u16;
85
2888
        self.strings_pool.extend_from_slice(s.as_bytes());
86
2888
        self.string_cache.insert(s.to_string(), (offset, len));
87
2888
        (offset, len)
88
2973
    }
89

            
90
841
    pub fn add_transaction(
91
841
        &mut self,
92
841
        id: [u8; 16],
93
841
        parent_idx: i32,
94
841
        is_primary: bool,
95
841
        is_context: bool,
96
841
        post_date: i64,
97
841
        enter_date: i64,
98
841
        split_count: u32,
99
841
        tag_count: u32,
100
841
        is_multi_currency: bool,
101
841
    ) -> u32 {
102
841
        let flags = EntityFlags::make(is_primary, is_context);
103
841
        let data = TransactionData {
104
841
            post_date,
105
841
            enter_date,
106
841
            split_count,
107
841
            tag_count,
108
841
            is_multi_currency: u8::from(is_multi_currency),
109
841
            reserved: [0; 23],
110
841
        };
111
841
        let header = EntityHeader::new(
112
841
            EntityType::Transaction,
113
841
            Operation::Nop,
114
841
            flags,
115
841
            id,
116
841
            parent_idx,
117
            0,
118
841
            data.to_bytes().len() as u32,
119
        );
120
841
        let idx = self.entities.len() as u32;
121
841
        self.entities.push(SerializedEntity {
122
841
            header,
123
841
            data: data.to_bytes().to_vec(),
124
841
        });
125
841
        idx
126
841
    }
127

            
128
1597
    pub fn add_split(
129
1597
        &mut self,
130
1597
        id: [u8; 16],
131
1597
        parent_idx: i32,
132
1597
        is_primary: bool,
133
1597
        is_context: bool,
134
1597
        account_id: [u8; 16],
135
1597
        commodity_id: [u8; 16],
136
1597
        value_num: i64,
137
1597
        value_denom: i64,
138
1597
        reconcile_state: u8,
139
1597
        reconcile_date: i64,
140
1597
    ) -> u32 {
141
1597
        let flags = EntityFlags::make(is_primary, is_context);
142
1597
        let data = SplitData {
143
1597
            account_id,
144
1597
            commodity_id,
145
1597
            value_num,
146
1597
            value_denom,
147
1597
            reconcile_state,
148
1597
            reserved: [0; 7],
149
1597
            reconcile_date,
150
1597
        };
151
1597
        let header = EntityHeader::new(
152
1597
            EntityType::Split,
153
1597
            Operation::Nop,
154
1597
            flags,
155
1597
            id,
156
1597
            parent_idx,
157
            0,
158
1597
            data.to_bytes().len() as u32,
159
        );
160
1597
        let idx = self.entities.len() as u32;
161
1597
        self.entities.push(SerializedEntity {
162
1597
            header,
163
1597
            data: data.to_bytes().to_vec(),
164
1597
        });
165
1597
        idx
166
1597
    }
167

            
168
1457
    pub fn add_tag(
169
1457
        &mut self,
170
1457
        id: [u8; 16],
171
1457
        parent_idx: i32,
172
1457
        is_primary: bool,
173
1457
        is_context: bool,
174
1457
        name: &str,
175
1457
        value: &str,
176
1457
    ) -> u32 {
177
1457
        let flags = EntityFlags::make(is_primary, is_context);
178
1457
        let (name_offset, name_len) = self.add_string(name);
179
1457
        let (value_offset, value_len) = self.add_string(value);
180
1457
        let data = TagData {
181
1457
            name_offset,
182
1457
            value_offset,
183
1457
            name_len,
184
1457
            value_len,
185
1457
            reserved: [0; 4],
186
1457
        };
187
1457
        let header = EntityHeader::new(
188
1457
            EntityType::Tag,
189
1457
            Operation::Nop,
190
1457
            flags,
191
1457
            id,
192
1457
            parent_idx,
193
            0,
194
1457
            TAG_DATA_SIZE as u32,
195
        );
196
1457
        let idx = self.entities.len() as u32;
197
1457
        self.entities.push(SerializedEntity {
198
1457
            header,
199
1457
            data: data.to_bytes().to_vec(),
200
1457
        });
201
1457
        idx
202
1457
    }
203

            
204
28
    pub fn add_account(
205
28
        &mut self,
206
28
        id: [u8; 16],
207
28
        parent_idx: i32,
208
28
        is_primary: bool,
209
28
        is_context: bool,
210
28
        parent_account_id: [u8; 16],
211
28
        name: &str,
212
28
        path: &str,
213
28
        tag_count: u32,
214
28
    ) -> u32 {
215
28
        let flags = EntityFlags::make(is_primary, is_context);
216
28
        let (name_offset, name_len) = self.add_string(name);
217
28
        let (path_offset, path_len) = self.add_string(path);
218
28
        let data = AccountData {
219
28
            parent_account_id,
220
28
            name_offset,
221
28
            path_offset,
222
28
            tag_count,
223
28
            name_len,
224
28
            path_len,
225
28
            reserved: [0; 16],
226
28
        };
227
28
        let header = EntityHeader::new(
228
28
            EntityType::Account,
229
28
            Operation::Nop,
230
28
            flags,
231
28
            id,
232
28
            parent_idx,
233
            0,
234
28
            data.to_bytes().len() as u32,
235
        );
236
28
        let idx = self.entities.len() as u32;
237
28
        self.entities.push(SerializedEntity {
238
28
            header,
239
28
            data: data.to_bytes().to_vec(),
240
28
        });
241
28
        idx
242
28
    }
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
1
    pub fn entity_count(&self) -> u32 {
284
1
        self.entities.len() as u32
285
1
    }
286

            
287
588
    pub fn add_transaction_from(
288
588
        &mut self,
289
588
        tx: &Transaction,
290
588
        is_primary: bool,
291
588
        split_count: u32,
292
588
        tag_count: u32,
293
588
        is_multi_currency: bool,
294
588
    ) -> u32 {
295
588
        let flags = EntityFlags::make(is_primary, false);
296
588
        let mut data = transaction_to_data(tx);
297
588
        data.split_count = split_count;
298
588
        data.tag_count = tag_count;
299
588
        data.is_multi_currency = u8::from(is_multi_currency);
300
588
        let header = EntityHeader::new(
301
588
            EntityType::Transaction,
302
588
            Operation::Nop,
303
588
            flags,
304
588
            *tx.id.as_bytes(),
305
            -1,
306
            0,
307
588
            data.to_bytes().len() as u32,
308
        );
309
588
        let idx = self.entities.len() as u32;
310
588
        self.entities.push(SerializedEntity {
311
588
            header,
312
588
            data: data.to_bytes().to_vec(),
313
588
        });
314
588
        idx
315
588
    }
316

            
317
980
    pub fn add_split_from(&mut self, split: &Split, parent_idx: i32) -> u32 {
318
980
        let data = split_to_data(split);
319
980
        let header = EntityHeader::new(
320
980
            EntityType::Split,
321
980
            Operation::Nop,
322
            0,
323
980
            *split.id.as_bytes(),
324
980
            parent_idx,
325
            0,
326
980
            data.to_bytes().len() as u32,
327
        );
328
980
        let idx = self.entities.len() as u32;
329
980
        self.entities.push(SerializedEntity {
330
980
            header,
331
980
            data: data.to_bytes().to_vec(),
332
980
        });
333
980
        idx
334
980
    }
335

            
336
    #[must_use]
337
5993
    pub fn finalize(mut self, output_size: u32) -> Vec<u8> {
338
5993
        let entity_count = self.entities.len() as u32;
339
5993
        let entities_offset = BASE_OFFSET + GLOBAL_HEADER_SIZE as u32;
340

            
341
5993
        let mut entities_total_size = 0u32;
342
5995
        for entity in &self.entities {
343
5491
            entities_total_size += ENTITY_HEADER_SIZE as u32 + entity.data.len() as u32;
344
5491
        }
345

            
346
5993
        let strings_pool_offset = entities_offset + entities_total_size;
347
5993
        let strings_pool_size = self.strings_pool.len() as u32;
348
5993
        let output_offset = strings_pool_offset + strings_pool_size;
349

            
350
5993
        let output_header = OutputHeader::new(entity_count);
351

            
352
5993
        let mut global_header = GlobalHeader::new(
353
5993
            self.context_type,
354
5993
            self.primary_entity_type,
355
5993
            entity_count,
356
5993
            self.primary_entity_idx,
357
        );
358
5993
        global_header.entities_offset = entities_offset;
359
5993
        global_header.strings_pool_offset = strings_pool_offset;
360
5993
        global_header.strings_pool_size = strings_pool_size;
361
5993
        global_header.output_offset = output_offset;
362
5993
        global_header.output_size = output_size;
363

            
364
5993
        let total_size = GLOBAL_HEADER_SIZE
365
5993
            + entities_total_size as usize
366
5993
            + strings_pool_size as usize
367
5993
            + output_size as usize;
368
5993
        let mut buffer = vec![0u8; total_size];
369

            
370
5993
        buffer[..GLOBAL_HEADER_SIZE].copy_from_slice(global_header.as_bytes());
371

            
372
5993
        let headers_total = entity_count as usize * ENTITY_HEADER_SIZE;
373
5993
        let mut data_offset = entities_offset + headers_total as u32;
374
5993
        let mut write_pos = GLOBAL_HEADER_SIZE;
375

            
376
5995
        for entity in &mut self.entities {
377
5491
            entity.header.data_offset = data_offset;
378
5491
            data_offset += entity.data.len() as u32;
379
5491
        }
380

            
381
5995
        for entity in &self.entities {
382
5491
            let header_bytes = entity.header.to_bytes();
383
5491
            buffer[write_pos..write_pos + ENTITY_HEADER_SIZE].copy_from_slice(&header_bytes);
384
5491
            write_pos += ENTITY_HEADER_SIZE;
385
5491
        }
386

            
387
5995
        for entity in &self.entities {
388
5491
            buffer[write_pos..write_pos + entity.data.len()].copy_from_slice(&entity.data);
389
5491
            write_pos += entity.data.len();
390
5491
        }
391

            
392
5993
        buffer[write_pos..write_pos + self.strings_pool.len()].copy_from_slice(&self.strings_pool);
393
5993
        write_pos += self.strings_pool.len();
394

            
395
5993
        buffer[write_pos..write_pos + OUTPUT_HEADER_SIZE]
396
5993
            .copy_from_slice(&output_header.to_bytes());
397

            
398
5993
        buffer
399
5993
    }
400
}
401

            
402
#[cfg(test)]
403
mod tests {
404
    use super::*;
405
    use crate::format::MAGIC_NOMI;
406

            
407
    #[test]
408
1
    fn test_serializer_basic() {
409
1
        let mut serializer = MemorySerializer::new();
410
1
        serializer.set_context(ContextType::EntityCreate, EntityType::Transaction);
411

            
412
1
        let tx_id = [1u8; 16];
413
1
        let tx_idx = serializer.add_transaction(tx_id, -1, true, false, 1000, 2000, 2, 1, false);
414
1
        serializer.set_primary(tx_idx);
415

            
416
1
        let split_id = [2u8; 16];
417
1
        let account_id = [3u8; 16];
418
1
        let commodity_id = [4u8; 16];
419
1
        serializer.add_split(
420
1
            split_id,
421
1
            tx_idx as i32,
422
            false,
423
            false,
424
1
            account_id,
425
1
            commodity_id,
426
            -5000,
427
            100,
428
            0,
429
            0,
430
        );
431

            
432
1
        let tag_id = [5u8; 16];
433
1
        serializer.add_tag(
434
1
            tag_id,
435
1
            tx_idx as i32,
436
            false,
437
            false,
438
1
            "note",
439
1
            "test transaction",
440
        );
441

            
442
1
        assert_eq!(serializer.entity_count(), 3);
443

            
444
1
        let buffer = serializer.finalize(1024);
445

            
446
1
        let header = GlobalHeader::from_bytes(&buffer).unwrap();
447
1
        assert_eq!(header.magic, MAGIC_NOMI);
448
1
        assert_eq!(header.input_entity_count, 3);
449
1
        assert_eq!(header.context_type, ContextType::EntityCreate as u8);
450
1
        assert_eq!(header.primary_entity_type, EntityType::Transaction as u8);
451
1
    }
452

            
453
    #[test]
454
1
    fn test_string_deduplication() {
455
1
        let mut serializer = MemorySerializer::new();
456
1
        let (offset1, len1) = serializer.add_string("test");
457
1
        let (offset2, len2) = serializer.add_string("test");
458
1
        let (offset3, _) = serializer.add_string("other");
459

            
460
1
        assert_eq!(offset1, offset2);
461
1
        assert_eq!(len1, len2);
462
1
        assert_ne!(offset1, offset3);
463
1
    }
464
}