scripting/

serializer.rs

use std::collections::HashMap;

use finance::split::Split;
use finance::transaction::Transaction;

use crate::format::{
    AccountData, BASE_OFFSET, CommodityData, ContextType, ENTITY_HEADER_SIZE, EntityFlags,
    EntityHeader, EntityType, GLOBAL_HEADER_SIZE, GlobalHeader, OUTPUT_HEADER_SIZE, Operation,
    OutputHeader, SplitData, TAG_DATA_SIZE, TagData, TransactionData,
};

fn transaction_to_data(tx: &Transaction) -> TransactionData {
    TransactionData {
        post_date: tx.post_date.timestamp_millis(),
        enter_date: tx.enter_date.timestamp_millis(),
        split_count: 0,
        tag_count: 0,
        is_multi_currency: 0,
        reserved: [0; 23],
    }
}

fn split_to_data(split: &Split, account_name_offset: u32, account_name_len: u32) -> SplitData {
    SplitData {
        account_id: *split.account_id.as_bytes(),
        commodity_id: *split.commodity_id.as_bytes(),
        value_num: split.value_num,
        value_denom: split.value_denom,
        reconcile_state: split.reconcile_state.map_or(0, u8::from),
        reserved: [0; 7],
        reconcile_date: split
            .reconcile_date
            .map_or(0, |d: chrono::DateTime<chrono::Utc>| d.timestamp_millis()),
        account_name_offset,
        account_name_len,
    }
}

pub struct MemorySerializer {
    context_type: ContextType,
    primary_entity_type: EntityType,
    primary_entity_idx: u32,
    entities: Vec<SerializedEntity>,
    strings_pool: Vec<u8>,
    string_cache: HashMap<String, (u32, u16)>,
}

struct SerializedEntity {
    header: EntityHeader,
    data: Vec<u8>,
}

/// Header fields shared by every `add_*` entry: identity, parenting, and
/// the two flag bits (`is_primary`, `is_context`). Pulled out so each
/// add_* fn signature stays under clippy's `too_many_arguments` cap.
#[derive(Debug, Clone, Copy)]
pub struct EntityHeaderArgs {
    pub id: [u8; 16],
    pub parent_idx: i32,
    pub is_primary: bool,
    pub is_context: bool,
}

/// Body fields specific to a transaction entity. Combined with
/// [`EntityHeaderArgs`] by [`MemorySerializer::add_transaction`].
#[derive(Debug, Clone, Copy)]
pub struct TransactionArgs {
    pub post_date: i64,
    pub enter_date: i64,
    pub split_count: u32,
    pub tag_count: u32,
    pub is_multi_currency: bool,
}

/// Body fields specific to a split entity. `account_name` borrows from the
/// caller and is interned into the strings pool by
/// [`MemorySerializer::add_split`] so a trigger script can read the posting
/// account's display name via `SPLIT-ACCOUNT-NAME`.
#[derive(Debug, Clone, Copy)]
pub struct SplitArgs<'a> {
    pub account_id: [u8; 16],
    pub commodity_id: [u8; 16],
    pub value_num: i64,
    pub value_denom: i64,
    pub reconcile_state: u8,
    pub reconcile_date: i64,
    pub account_name: &'a str,
}

/// Body fields specific to an account entity. The `name` and `path`
/// borrow from the caller's scope and are interned into the strings
/// pool by [`MemorySerializer::add_account`].
#[derive(Debug, Clone, Copy)]
pub struct AccountArgs<'a> {
    pub parent_account_id: [u8; 16],
    pub name: &'a str,
    pub path: &'a str,
    pub tag_count: u32,
}

/// Body fields specific to a commodity entity.
#[derive(Debug, Clone, Copy)]
pub struct CommodityArgs<'a> {
    pub symbol: &'a str,
    pub name: &'a str,
    pub tag_count: u32,
}

/// Aggregate args for `add_transaction_from`: the source `Transaction`
/// plus the runtime-supplied counters and primary flag the test harness
/// can't infer from the row.
#[derive(Debug, Clone, Copy)]
pub struct TransactionFromArgs<'a> {
    pub transaction: &'a Transaction,
    pub is_primary: bool,
    pub split_count: u32,
    pub tag_count: u32,
    pub is_multi_currency: bool,
}

impl Default for MemorySerializer {
    fn default() -> Self {
        Self::new()
    }
}

impl MemorySerializer {
    #[must_use]
    pub fn new() -> Self {
        Self {
            context_type: ContextType::EntityCreate,
            primary_entity_type: EntityType::Transaction,
            primary_entity_idx: 0,
            entities: Vec::new(),
            strings_pool: Vec::new(),
            string_cache: HashMap::new(),
        }
    }

    pub fn set_context(&mut self, context_type: ContextType, primary_entity_type: EntityType) {
        self.context_type = context_type;
        self.primary_entity_type = primary_entity_type;
    }

    pub fn set_primary(&mut self, entity_idx: u32) {
        self.primary_entity_idx = entity_idx;
    }

    pub fn add_string(&mut self, s: &str) -> (u32, u16) {
        if let Some(&cached) = self.string_cache.get(s) {
            return cached;
        }
        let offset = self.strings_pool.len() as u32;
        let len = s.len() as u16;
        self.strings_pool.extend_from_slice(s.as_bytes());
        self.string_cache.insert(s.to_string(), (offset, len));
        (offset, len)
    }

    pub fn add_transaction(&mut self, header: EntityHeaderArgs, args: TransactionArgs) -> u32 {
        let flags = EntityFlags::make(header.is_primary, header.is_context);
        let data = TransactionData {
            post_date: args.post_date,
            enter_date: args.enter_date,
            split_count: args.split_count,
            tag_count: args.tag_count,
            is_multi_currency: u8::from(args.is_multi_currency),
            reserved: [0; 23],
        };
        let entity_header = EntityHeader::new(
            EntityType::Transaction,
            Operation::Nop,
            flags,
            header.id,
            header.parent_idx,
            0,
            data.to_bytes().len() as u32,
        );
        let idx = self.entities.len() as u32;
        self.entities.push(SerializedEntity {
            header: entity_header,
            data: data.to_bytes().to_vec(),
        });
        idx
    }

    pub fn add_split(&mut self, header: EntityHeaderArgs, args: SplitArgs) -> u32 {
        let flags = EntityFlags::make(header.is_primary, header.is_context);
        let (name_offset, _) = self.add_string(args.account_name);
        let data = SplitData {
            account_id: args.account_id,
            commodity_id: args.commodity_id,
            value_num: args.value_num,
            value_denom: args.value_denom,
            reconcile_state: args.reconcile_state,
            reserved: [0; 7],
            reconcile_date: args.reconcile_date,
            account_name_offset: name_offset,
            // The format gives the split account name a u32 length; take the full
            // byte length rather than `add_string`'s u16 (which would wrap a name
            // longer than 65535 bytes while the field claims u32 capacity).
            account_name_len: args.account_name.len() as u32,
        };
        let entity_header = EntityHeader::new(
            EntityType::Split,
            Operation::Nop,
            flags,
            header.id,
            header.parent_idx,
            0,
            data.to_bytes().len() as u32,
        );
        let idx = self.entities.len() as u32;
        self.entities.push(SerializedEntity {
            header: entity_header,
            data: data.to_bytes().to_vec(),
        });
        idx
    }

    pub fn add_tag(
        &mut self,
        id: [u8; 16],
        parent_idx: i32,
        is_primary: bool,
        is_context: bool,
        name: &str,
        value: &str,
    ) -> u32 {
        let flags = EntityFlags::make(is_primary, is_context);
        let (name_offset, name_len) = self.add_string(name);
        let (value_offset, value_len) = self.add_string(value);
        let data = TagData {
            name_offset,
            value_offset,
            name_len,
            value_len,
            reserved: [0; 4],
        };
        let header = EntityHeader::new(
            EntityType::Tag,
            Operation::Nop,
            flags,
            id,
            parent_idx,
            0,
            TAG_DATA_SIZE as u32,
        );
        let idx = self.entities.len() as u32;
        self.entities.push(SerializedEntity {
            header,
            data: data.to_bytes().to_vec(),
        });
        idx
    }

    pub fn add_account(&mut self, header: EntityHeaderArgs, args: AccountArgs<'_>) -> u32 {
        let flags = EntityFlags::make(header.is_primary, header.is_context);
        let (name_offset, name_len) = self.add_string(args.name);
        let (path_offset, path_len) = self.add_string(args.path);
        let data = AccountData {
            parent_account_id: args.parent_account_id,
            name_offset,
            path_offset,
            tag_count: args.tag_count,
            name_len,
            path_len,
            reserved: [0; 16],
        };
        let entity_header = EntityHeader::new(
            EntityType::Account,
            Operation::Nop,
            flags,
            header.id,
            header.parent_idx,
            0,
            data.to_bytes().len() as u32,
        );
        let idx = self.entities.len() as u32;
        self.entities.push(SerializedEntity {
            header: entity_header,
            data: data.to_bytes().to_vec(),
        });
        idx
    }

    pub fn add_commodity(&mut self, header: EntityHeaderArgs, args: CommodityArgs<'_>) -> u32 {
        let flags = EntityFlags::make(header.is_primary, header.is_context);
        let (symbol_offset, symbol_len) = self.add_string(args.symbol);
        let (name_offset, name_len) = self.add_string(args.name);
        let data = CommodityData {
            symbol_offset,
            name_offset,
            tag_count: args.tag_count,
            symbol_len,
            name_len,
            reserved: [0; 16],
        };
        let entity_header = EntityHeader::new(
            EntityType::Commodity,
            Operation::Nop,
            flags,
            header.id,
            header.parent_idx,
            0,
            data.to_bytes().len() as u32,
        );
        let idx = self.entities.len() as u32;
        self.entities.push(SerializedEntity {
            header: entity_header,
            data: data.to_bytes().to_vec(),
        });
        idx
    }

    #[must_use]
    pub fn entity_count(&self) -> u32 {
        self.entities.len() as u32
    }

    pub fn add_transaction_from(&mut self, args: TransactionFromArgs<'_>) -> u32 {
        let flags = EntityFlags::make(args.is_primary, false);
        let mut data = transaction_to_data(args.transaction);
        data.split_count = args.split_count;
        data.tag_count = args.tag_count;
        data.is_multi_currency = u8::from(args.is_multi_currency);
        let header = EntityHeader::new(
            EntityType::Transaction,
            Operation::Nop,
            flags,
            *args.transaction.id.as_bytes(),
            -1,
            0,
            data.to_bytes().len() as u32,
        );
        let idx = self.entities.len() as u32;
        self.entities.push(SerializedEntity {
            header,
            data: data.to_bytes().to_vec(),
        });
        idx
    }

    pub fn add_split_from(&mut self, split: &Split, parent_idx: i32, account_name: &str) -> u32 {
        let (name_offset, _) = self.add_string(account_name);
        let data = split_to_data(split, name_offset, account_name.len() as u32);
        let header = EntityHeader::new(
            EntityType::Split,
            Operation::Nop,
            0,
            *split.id.as_bytes(),
            parent_idx,
            0,
            data.to_bytes().len() as u32,
        );
        let idx = self.entities.len() as u32;
        self.entities.push(SerializedEntity {
            header,
            data: data.to_bytes().to_vec(),
        });
        idx
    }

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

        let mut entities_total_size = 0u32;
        for entity in &self.entities {
            entities_total_size += ENTITY_HEADER_SIZE as u32 + entity.data.len() as u32;
        }

        let strings_pool_offset = entities_offset + entities_total_size;
        let strings_pool_size = self.strings_pool.len() as u32;
        let output_offset = strings_pool_offset + strings_pool_size;

        let output_header = OutputHeader::new(entity_count);

        let mut global_header = GlobalHeader::new(
            self.context_type,
            self.primary_entity_type,
            entity_count,
            self.primary_entity_idx,
        );
        global_header.entities_offset = entities_offset;
        global_header.strings_pool_offset = strings_pool_offset;
        global_header.strings_pool_size = strings_pool_size;
        global_header.output_offset = output_offset;
        global_header.output_size = output_size;

        let total_size = GLOBAL_HEADER_SIZE
            + entities_total_size as usize
            + strings_pool_size as usize
            + output_size as usize;
        let mut buffer = vec![0u8; total_size];

        buffer[..GLOBAL_HEADER_SIZE].copy_from_slice(global_header.as_bytes());

        let headers_total = entity_count as usize * ENTITY_HEADER_SIZE;
        let mut data_offset = entities_offset + headers_total as u32;
        let mut write_pos = GLOBAL_HEADER_SIZE;

        for entity in &mut self.entities {
            entity.header.data_offset = data_offset;
            data_offset += entity.data.len() as u32;
        }

        for entity in &self.entities {
            let header_bytes = entity.header.to_bytes();
            buffer[write_pos..write_pos + ENTITY_HEADER_SIZE].copy_from_slice(&header_bytes);
            write_pos += ENTITY_HEADER_SIZE;
        }

        for entity in &self.entities {
            buffer[write_pos..write_pos + entity.data.len()].copy_from_slice(&entity.data);
            write_pos += entity.data.len();
        }

        buffer[write_pos..write_pos + self.strings_pool.len()].copy_from_slice(&self.strings_pool);
        write_pos += self.strings_pool.len();

        buffer[write_pos..write_pos + OUTPUT_HEADER_SIZE]
            .copy_from_slice(&output_header.to_bytes());

        buffer
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::format::MAGIC_NOMI;

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

        let tx_id = [1u8; 16];
        let tx_idx = serializer.add_transaction(
            EntityHeaderArgs {
                id: tx_id,
                parent_idx: -1,
                is_primary: true,
                is_context: false,
            },
            TransactionArgs {
                post_date: 1000,
                enter_date: 2000,
                split_count: 2,
                tag_count: 1,
                is_multi_currency: false,
            },
        );
        serializer.set_primary(tx_idx);

        let split_id = [2u8; 16];
        let account_id = [3u8; 16];
        let commodity_id = [4u8; 16];
        serializer.add_split(
            EntityHeaderArgs {
                id: split_id,
                parent_idx: tx_idx as i32,
                is_primary: false,
                is_context: false,
            },
            SplitArgs {
                account_id,
                commodity_id,
                value_num: -5000,
                value_denom: 100,
                reconcile_state: 0,
                reconcile_date: 0,
                account_name: "Assets:Test",
            },
        );

        let tag_id = [5u8; 16];
        serializer.add_tag(
            tag_id,
            tx_idx as i32,
            false,
            false,
            "note",
            "test transaction",
        );

        assert_eq!(serializer.entity_count(), 3);

        let buffer = serializer.finalize(1024);

        let header = GlobalHeader::from_bytes(&buffer).unwrap();
        assert_eq!(header.magic, MAGIC_NOMI);
        assert_eq!(header.input_entity_count, 3);
        assert_eq!(header.context_type, ContextType::EntityCreate as u8);
        assert_eq!(header.primary_entity_type, EntityType::Transaction as u8);
    }

    #[test]
    fn split_account_name_round_trips_through_strings_pool() {
        let mut ser = MemorySerializer::new();
        ser.set_context(ContextType::EntityCreate, EntityType::Transaction);
        let tx_idx = ser.add_transaction(
            EntityHeaderArgs {
                id: [1u8; 16],
                parent_idx: -1,
                is_primary: true,
                is_context: false,
            },
            TransactionArgs {
                post_date: 0,
                enter_date: 0,
                split_count: 1,
                tag_count: 0,
                is_multi_currency: false,
            },
        );
        ser.set_primary(tx_idx);
        ser.add_split(
            EntityHeaderArgs {
                id: [2u8; 16],
                parent_idx: tx_idx as i32,
                is_primary: false,
                is_context: false,
            },
            SplitArgs {
                account_id: [3u8; 16],
                commodity_id: [4u8; 16],
                value_num: -5000,
                value_denom: 100,
                reconcile_state: 0,
                reconcile_date: 0,
                account_name: "Metro",
            },
        );
        let buf = ser.finalize(1024);

        let header = GlobalHeader::from_bytes(&buf).unwrap();
        // Header offsets are absolute wasm addresses (BASE_OFFSET-relative); the
        // raw buffer starts at BASE_OFFSET, so index it after subtracting that.
        let base = BASE_OFFSET as usize;
        // Entity 1 is the split (entity 0 is the transaction).
        let split_header_off = header.entities_offset as usize - base + ENTITY_HEADER_SIZE;
        let split_header = EntityHeader::from_bytes(&buf[split_header_off..]).unwrap();
        let split =
            SplitData::from_bytes(&buf[split_header.data_offset as usize - base..]).unwrap();

        let name_start =
            header.strings_pool_offset as usize - base + split.account_name_offset as usize;
        let name = &buf[name_start..name_start + split.account_name_len as usize];
        assert_eq!(name, b"Metro");
    }

    #[test]
    fn test_string_deduplication() {
        let mut serializer = MemorySerializer::new();
        let (offset1, len1) = serializer.add_string("test");
        let (offset2, len2) = serializer.add_string("test");
        let (offset3, _) = serializer.add_string("other");

        assert_eq!(offset1, offset2);
        assert_eq!(len1, len2);
        assert_ne!(offset1, offset3);
    }
}