verit-core 0.2.0

//! Encoder: `Schema` + `Value` tree → message bytes.
//!
//! Envelope: 24-byte header, optional inline canonical schema, padding to 8,
//! then the payload arena. All offsets in the message are absolute u32 byte
//! positions within the buffer. See the architecture docs.

use crate::error::{Error, Result};
use crate::layout::{align_up, slot_size_align};
use crate::schema::{Schema, Type};
use crate::value::Value;

pub const MESSAGE_MAGIC: &[u8; 4] = b"VRT2";
pub const FLAG_INLINE_SCHEMA: u16 = 1;
/// Envelope: magic(4) flags(2) reserved(2) schema_id(16) root_off(4) schema_len(4).
pub const HEADER_LEN: usize = 32;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SchemaMode {
    /// Embed the full canonical schema: the message alone is self-describing.
    Inline,
    /// Carry only the 128-bit schema id (streams / registries amortize the
    /// schema out of band).
    HashOnly,
}

pub fn encode(schema: &Schema, root: &Value, mode: SchemaMode) -> Result<Vec<u8>> {
    let root_fields = match root {
        Value::Struct(fields) => fields,
        other => {
            return Err(Error::TypeMismatch {
                expected: "struct".into(),
                got: other.kind().into(),
            })
        }
    };
    let mut enc = Encoder {
        schema,
        buf: Vec::with_capacity(256),
    };
    enc.buf.extend_from_slice(MESSAGE_MAGIC);
    let inline = mode == SchemaMode::Inline;
    let flags: u16 = if inline { FLAG_INLINE_SCHEMA } else { 0 };
    enc.buf.extend_from_slice(&flags.to_le_bytes());
    enc.buf.extend_from_slice(&0u16.to_le_bytes()); // reserved
    enc.buf.extend_from_slice(&schema.id().to_le_bytes()); // 16-byte schema id
    enc.buf.extend_from_slice(&0u32.to_le_bytes()); // root offset, patched below
    let schema_len = if inline {
        schema.canonical_bytes().len()
    } else {
        0
    };
    let schema_len_u32 = u32::try_from(schema_len).map_err(|_| Error::MessageTooLarge)?;
    enc.buf.extend_from_slice(&schema_len_u32.to_le_bytes());
    if inline {
        enc.buf.extend_from_slice(schema.canonical_bytes());
    }
    enc.pad_to(8)?;
    let root_off = enc.write_struct(schema.root_index(), root_fields)?;
    enc.buf[24..28].copy_from_slice(&root_off.to_le_bytes());
    Ok(enc.buf)
}

enum ElemKind {
    Scalar,
    Ref,
    InlineStruct(u16),
}

fn elem_info(schema: &Schema, elem: &Type) -> (u32, u32, ElemKind) {
    match elem {
        Type::Struct(i) => match schema.layout_unchecked(*i) {
            // Fixed-size struct elements are stored inline at a constant
            // stride. Packed struct elements are variable-size, so they ride
            // in the list by u32 offset (the Ref discipline) instead.
            crate::layout::StructLayout::Fixed(f) => (f.size, f.align, ElemKind::InlineStruct(*i)),
            crate::layout::StructLayout::Packed(_) => (4, 4, ElemKind::Ref),
        },
        Type::String | Type::Bytes | Type::List(_) | Type::Map(_, _) | Type::Union(_) => {
            (4, 4, ElemKind::Ref)
        }
        other => {
            let (size, align) = slot_size_align(other);
            (size, align, ElemKind::Scalar)
        }
    }
}

/// Total canonical order over map keys of a single key type: integers by value
/// (sign-aware), `bool` false < true, enums by their u32, strings by UTF-8
/// bytes. The encoder sorts entries by this so one logical map has exactly one
/// byte encoding, regardless of insertion order (crucial for languages whose
/// maps iterate in an unspecified order).
fn key_cmp(a: &Value, b: &Value) -> std::cmp::Ordering {
    use std::cmp::Ordering::Equal;
    use Value::*;
    match (a, b) {
        (Bool(x), Bool(y)) => x.cmp(y),
        (U8(x), U8(y)) => x.cmp(y),
        (U16(x), U16(y)) => x.cmp(y),
        (U32(x), U32(y)) => x.cmp(y),
        (U64(x), U64(y)) => x.cmp(y),
        (I8(x), I8(y)) => x.cmp(y),
        (I16(x), I16(y)) => x.cmp(y),
        (I32(x), I32(y)) => x.cmp(y),
        (I64(x), I64(y)) => x.cmp(y),
        (Enum(x), Enum(y)) => x.cmp(y),
        (Str(x), Str(y)) => x.as_bytes().cmp(y.as_bytes()),
        // Mixed or non-key variants: leave order unspecified here; the per-entry
        // type check in `write_map` rejects a key of the wrong type.
        _ => Equal,
    }
}

fn type_mismatch(schema: &Schema, expected: &Type, got: &Value) -> Error {
    Error::TypeMismatch {
        expected: expected.describe(schema),
        got: got.kind().into(),
    }
}

struct Encoder<'s> {
    schema: &'s Schema,
    buf: Vec<u8>,
}

impl<'s> Encoder<'s> {
    fn pos(&self) -> Result<u32> {
        u32::try_from(self.buf.len()).map_err(|_| Error::MessageTooLarge)
    }

    fn pad_to(&mut self, align: u32) -> Result<u32> {
        let pos = self.pos()?;
        let target = align_up(pos, align);
        self.buf.resize(self.buf.len() + (target - pos) as usize, 0);
        Ok(target)
    }

    fn patch_u32(&mut self, at: u32, v: u32) {
        let at = at as usize;
        self.buf[at..at + 4].copy_from_slice(&v.to_le_bytes());
    }

    /// Reserve a zeroed, aligned fixed struct block and return its offset.
    fn alloc_struct_block(&mut self, type_index: u16) -> Result<u32> {
        let lay = self.schema.layout_unchecked(type_index).as_fixed();
        let (align, size) = (lay.align, lay.size);
        let base = self.pad_to(align)?;
        self.buf.resize(self.buf.len() + size as usize, 0);
        self.pos()?; // enforce the 4 GiB cap after growing
        Ok(base)
    }

    fn write_struct(&mut self, type_index: u16, values: &[(u16, Value)]) -> Result<u32> {
        if self.schema.struct_def_unchecked(type_index).is_packed() {
            return self.write_packed_struct(type_index, values);
        }
        let base = self.alloc_struct_block(type_index)?;
        self.fill_struct_at(base, type_index, values)?;
        Ok(base)
    }

    /// Fill an already-reserved struct block: presence bits and scalar slots
    /// in place, heap children appended afterwards with their offsets patched
    /// into the slots.
    fn fill_struct_at(
        &mut self,
        base: u32,
        type_index: u16,
        values: &[(u16, Value)],
    ) -> Result<()> {
        let schema = self.schema;
        let sd = schema.struct_def_unchecked(type_index);
        let lay = schema.layout_unchecked(type_index).as_fixed();
        let mut heap_jobs: Vec<(u32, &Value, &'s Type)> = Vec::new();
        let mut seen = vec![false; sd.fields.len()];
        for (id, value) in values {
            let pos = sd
                .fields
                .binary_search_by_key(id, |f| f.id)
                .map_err(|_| Error::UnknownFieldId(*id))?;
            if seen[pos] {
                return Err(Error::DuplicateField(*id));
            }
            seen[pos] = true;
            if !sd.is_dense() {
                let bit_at = (base + pos as u32 / 8) as usize;
                self.buf[bit_at] |= 1 << (pos % 8);
            }
            let slot = base + lay.slots[pos];
            let ty = &sd.fields[pos].ty;
            match ty {
                Type::String
                | Type::Bytes
                | Type::List(_)
                | Type::Struct(_)
                | Type::Map(_, _)
                | Type::Union(_) => {
                    // Cheap shape check now for a good error; full check when written.
                    let ok = matches!(
                        (ty, value),
                        (Type::String, Value::Str(_))
                            | (Type::Bytes, Value::Bytes(_))
                            | (Type::List(_), Value::List(_))
                            | (Type::List(_), Value::Scalars(_))
                            | (Type::Struct(_), Value::Struct(_))
                            | (Type::Map(_, _), Value::Map(_))
                            | (Type::Union(_), Value::Union(..))
                    );
                    if !ok {
                        return Err(type_mismatch(schema, ty, value));
                    }
                    heap_jobs.push((slot, value, ty));
                }
                _ => self.store_scalar_at(slot, value, ty)?,
            }
        }
        if sd.is_dense() {
            if let Some(pos) = seen.iter().position(|s| !s) {
                return Err(Error::MissingField(sd.fields[pos].id));
            }
        }
        for (slot, value, ty) in heap_jobs {
            let off = self.write_heap(value, ty)?;
            self.patch_u32(slot, off);
        }
        Ok(())
    }

    /// Encode a packed struct: presence bitmap, then present fields packed by
    /// size class. Block size and slot offsets come from the popcount layout.
    fn write_packed_struct(&mut self, type_index: u16, values: &[(u16, Value)]) -> Result<u32> {
        let schema = self.schema;
        let sd = schema.struct_def_unchecked(type_index);
        let lay = schema.layout_unchecked(type_index).as_packed();

        // Resolve ids to positions, detect duplicates, build the bitmap.
        let mut bitmap = 0u64;
        let mut present: Vec<(usize, &Value, &'s Type)> = Vec::with_capacity(values.len());
        for (id, value) in values {
            let pos = sd
                .fields
                .binary_search_by_key(id, |f| f.id)
                .map_err(|_| Error::UnknownFieldId(*id))?;
            let bit = 1u64 << pos;
            if bitmap & bit != 0 {
                return Err(Error::DuplicateField(*id));
            }
            bitmap |= bit;
            present.push((pos, value, &sd.fields[pos].ty));
        }

        let size = lay.block_size(bitmap);
        let base = self.pad_to(lay.align)?;
        self.buf.resize(self.buf.len() + size as usize, 0);
        self.pos()?;
        // Write the presence bitmap (only the used low bytes).
        let bmap_bytes = bitmap.to_le_bytes();
        let bstart = base as usize;
        self.buf[bstart..bstart + lay.bitmap_bytes as usize]
            .copy_from_slice(&bmap_bytes[..lay.bitmap_bytes as usize]);

        let mut heap_jobs: Vec<(u32, &Value, &'s Type)> = Vec::new();
        for (pos, value, ty) in present {
            let slot = base + lay.field_offset(bitmap, pos);
            match ty {
                Type::String
                | Type::Bytes
                | Type::List(_)
                | Type::Struct(_)
                | Type::Map(_, _)
                | Type::Union(_) => {
                    let ok = matches!(
                        (ty, value),
                        (Type::String, Value::Str(_))
                            | (Type::Bytes, Value::Bytes(_))
                            | (Type::List(_), Value::List(_))
                            | (Type::List(_), Value::Scalars(_))
                            | (Type::Struct(_), Value::Struct(_))
                            | (Type::Map(_, _), Value::Map(_))
                            | (Type::Union(_), Value::Union(..))
                    );
                    if !ok {
                        return Err(type_mismatch(schema, ty, value));
                    }
                    heap_jobs.push((slot, value, ty));
                }
                _ => self.store_scalar_at(slot, value, ty)?,
            }
        }
        for (slot, value, ty) in heap_jobs {
            let off = self.write_heap(value, ty)?;
            self.patch_u32(slot, off);
        }
        Ok(base)
    }

    /// Write a scalar into an existing (zeroed) slot.
    fn store_scalar_at(&mut self, at: u32, value: &Value, ty: &Type) -> Result<()> {
        let at = at as usize;
        let buf = &mut self.buf;
        macro_rules! put {
            ($bytes:expr) => {{
                let b = $bytes;
                buf[at..at + b.len()].copy_from_slice(&b);
            }};
        }
        match (ty, value) {
            (Type::Bool, Value::Bool(x)) => buf[at] = *x as u8,
            (Type::U8, Value::U8(x)) => buf[at] = *x,
            (Type::U16, Value::U16(x)) => put!(x.to_le_bytes()),
            (Type::U32, Value::U32(x)) => put!(x.to_le_bytes()),
            (Type::U64, Value::U64(x)) => put!(x.to_le_bytes()),
            (Type::I8, Value::I8(x)) => buf[at] = *x as u8,
            (Type::I16, Value::I16(x)) => put!(x.to_le_bytes()),
            (Type::I32, Value::I32(x)) => put!(x.to_le_bytes()),
            (Type::I64, Value::I64(x)) => put!(x.to_le_bytes()),
            (Type::F32, Value::F32(x)) => put!(x.to_le_bytes()),
            (Type::F64, Value::F64(x)) => put!(x.to_le_bytes()),
            (Type::Enum(_), Value::Enum(x)) => put!(x.to_le_bytes()),
            _ => return Err(type_mismatch(self.schema, ty, value)),
        }
        Ok(())
    }

    /// Append a scalar list element at the current (aligned) position.
    fn push_scalar(&mut self, value: &Value, ty: &Type) -> Result<()> {
        let buf = &mut self.buf;
        match (ty, value) {
            (Type::Bool, Value::Bool(x)) => buf.push(*x as u8),
            (Type::U8, Value::U8(x)) => buf.push(*x),
            (Type::U16, Value::U16(x)) => buf.extend_from_slice(&x.to_le_bytes()),
            (Type::U32, Value::U32(x)) => buf.extend_from_slice(&x.to_le_bytes()),
            (Type::U64, Value::U64(x)) => buf.extend_from_slice(&x.to_le_bytes()),
            (Type::I8, Value::I8(x)) => buf.push(*x as u8),
            (Type::I16, Value::I16(x)) => buf.extend_from_slice(&x.to_le_bytes()),
            (Type::I32, Value::I32(x)) => buf.extend_from_slice(&x.to_le_bytes()),
            (Type::I64, Value::I64(x)) => buf.extend_from_slice(&x.to_le_bytes()),
            (Type::F32, Value::F32(x)) => buf.extend_from_slice(&x.to_le_bytes()),
            (Type::F64, Value::F64(x)) => buf.extend_from_slice(&x.to_le_bytes()),
            (Type::Enum(_), Value::Enum(x)) => buf.extend_from_slice(&x.to_le_bytes()),
            _ => return Err(type_mismatch(self.schema, ty, value)),
        }
        Ok(())
    }

    /// Write a heap object (string, bytes, list, struct block) and return its
    /// absolute offset.
    fn write_heap(&mut self, value: &Value, ty: &Type) -> Result<u32> {
        match (ty, value) {
            (Type::String, Value::Str(s)) => self.write_blob(s.as_bytes()),
            (Type::Bytes, Value::Bytes(b)) => self.write_blob(b),
            (Type::Struct(i), Value::Struct(fields)) => self.write_struct(*i, fields),
            (Type::List(elem), Value::List(items)) => self.write_list(elem, items),
            (Type::List(elem), Value::Scalars(runs)) => self.write_scalar_list(elem, runs),
            (Type::Map(k, v), Value::Map(entries)) => self.write_map(k, v, entries),
            (Type::Union(variants), Value::Union(tag, inner)) => {
                self.write_union(variants, *tag, inner)
            }
            _ => Err(type_mismatch(self.schema, ty, value)),
        }
    }

    /// Write a `union` value: `u32 tag`, then the selected variant's value at
    /// its natural-aligned payload slot (scalar in place, or a u32 heap offset).
    fn write_union(&mut self, variants: &[Type], tag: u32, inner: &Value) -> Result<u32> {
        let vty = variants.get(tag as usize).ok_or(Error::BadUnionTag(tag))?;
        let (pslot, palign) = slot_size_align(vty);
        let payload_off = crate::layout::align_up(4, palign);
        let block_align = palign.max(4);
        let block_size = crate::layout::align_up(payload_off + pslot, block_align);
        let base = self.pad_to(block_align)?;
        self.buf.resize(self.buf.len() + block_size as usize, 0);
        self.pos()?;
        self.patch_u32(base, tag);
        let slot = base + payload_off;
        match vty {
            Type::String
            | Type::Bytes
            | Type::List(_)
            | Type::Struct(_)
            | Type::Map(_, _)
            | Type::Union(_) => {
                let off = self.write_heap(inner, vty)?;
                self.patch_u32(slot, off);
            }
            _ => self.store_scalar_at(slot, inner, vty)?,
        }
        Ok(base)
    }

    /// Write a `map<K, V>`: `count`, then `count` entry blocks (each a dense
    /// 2-field {key, value} struct), entries sorted by key in canonical order.
    /// Returns the map's absolute offset.
    fn write_map(
        &mut self,
        key_ty: &Type,
        val_ty: &Type,
        entries: &[(Value, Value)],
    ) -> Result<u32> {
        let lay = crate::layout::map_entry_layout(key_ty, val_ty);
        let (key_slot, val_slot) = (lay.slots[0], lay.slots[1]);

        // Sort a view of the entries by key into canonical order (leaving the
        // caller's Vec untouched), then reject duplicate keys.
        let mut order: Vec<usize> = (0..entries.len()).collect();
        order.sort_by(|&i, &j| key_cmp(&entries[i].0, &entries[j].0));
        for w in order.windows(2) {
            if key_cmp(&entries[w[0]].0, &entries[w[1]].0) == std::cmp::Ordering::Equal {
                return Err(Error::DuplicateMapKey);
            }
        }

        let count = u32::try_from(entries.len()).map_err(|_| Error::MessageTooLarge)?;
        let off = self.pad_to(4)?;
        self.buf.extend_from_slice(&count.to_le_bytes());
        let base = self.pad_to(lay.align)?;
        let total = (lay.size as usize)
            .checked_mul(entries.len())
            .ok_or(Error::MessageTooLarge)?;
        self.buf.resize(self.buf.len() + total, 0);
        self.pos()?;

        // Two passes so all in-block scalars land before any heap object is
        // appended (mirrors struct filling), keeping the entry region contiguous.
        let mut heap_jobs: Vec<(u32, &Value, &Type)> = Vec::new();
        for (rank, &i) in order.iter().enumerate() {
            let entry_base = base + rank as u32 * lay.size;
            for (slot_off, ty, val) in [
                (key_slot, key_ty, &entries[i].0),
                (val_slot, val_ty, &entries[i].1),
            ] {
                let slot = entry_base + slot_off;
                match ty {
                    Type::String
                    | Type::Bytes
                    | Type::List(_)
                    | Type::Struct(_)
                    | Type::Map(_, _) => {
                        heap_jobs.push((slot, val, ty));
                    }
                    _ => self.store_scalar_at(slot, val, ty)?,
                }
            }
        }
        for (slot, val, ty) in heap_jobs {
            let child = self.write_heap(val, ty)?;
            self.patch_u32(slot, child);
        }
        Ok(off)
    }

    fn write_blob(&mut self, bytes: &[u8]) -> Result<u32> {
        let len = u32::try_from(bytes.len()).map_err(|_| Error::MessageTooLarge)?;
        let off = self.pad_to(4)?;
        self.buf.extend_from_slice(&len.to_le_bytes());
        self.buf.extend_from_slice(bytes);
        self.pos()?;
        Ok(off)
    }

    /// Write a `list<scalar>` from native values — the bulk peer of
    /// [`write_list`](Self::write_list).
    ///
    /// Produces exactly the bytes `write_list` would for the equivalent
    /// `Value::List`: `u32 count`, alignment padding, then the elements packed
    /// at their natural stride. The difference is only that the caller did not
    /// have to build a `Value` per element, and the run is appended in one pass
    /// the optimiser can vectorize.
    fn write_scalar_list(&mut self, elem: &Type, runs: &crate::value::Scalars) -> Result<u32> {
        use crate::value::Scalars;
        let count = u32::try_from(runs.len()).map_err(|_| Error::MessageTooLarge)?;
        let (_, elem_align, kind) = elem_info(self.schema, elem);
        if !matches!(kind, ElemKind::Scalar) {
            return Err(type_mismatch(
                self.schema,
                elem,
                &Value::Scalars(runs.clone()),
            ));
        }
        let off = self.pad_to(4)?;
        self.buf.extend_from_slice(&count.to_le_bytes());
        self.pad_to(elem_align)?;

        // Each arm must match `push_scalar`'s bytes for that (Type, Value) pair
        // exactly — the golden vectors are the proof that it does.
        macro_rules! run {
            ($values:expr) => {{
                self.buf
                    .reserve($values.len() * std::mem::size_of_val(&$values[0]));
                for x in $values {
                    self.buf.extend_from_slice(&x.to_le_bytes());
                }
            }};
        }
        match (elem, runs) {
            (Type::Bool, Scalars::Bool(v)) => {
                self.buf.reserve(v.len());
                for x in v {
                    self.buf.push(*x as u8);
                }
            }
            (Type::U8, Scalars::U8(v)) => self.buf.extend_from_slice(v),
            (Type::I8, Scalars::I8(v)) => {
                self.buf.reserve(v.len());
                for x in v {
                    self.buf.push(*x as u8);
                }
            }
            (Type::U16, Scalars::U16(v)) => run!(v),
            (Type::U32, Scalars::U32(v)) => run!(v),
            (Type::U64, Scalars::U64(v)) => run!(v),
            (Type::I16, Scalars::I16(v)) => run!(v),
            (Type::I32, Scalars::I32(v)) => run!(v),
            (Type::I64, Scalars::I64(v)) => run!(v),
            (Type::F32, Scalars::F32(v)) => run!(v),
            (Type::F64, Scalars::F64(v)) => run!(v),
            _ => {
                return Err(type_mismatch(
                    self.schema,
                    elem,
                    &Value::Scalars(runs.clone()),
                ))
            }
        }
        self.pos()?;
        Ok(off)
    }

    fn write_list(&mut self, elem: &Type, items: &[Value]) -> Result<u32> {
        let schema = self.schema;
        let (stride, elem_align, kind) = elem_info(schema, elem);
        let count = u32::try_from(items.len()).map_err(|_| Error::MessageTooLarge)?;
        let off = self.pad_to(4)?;
        self.buf.extend_from_slice(&count.to_le_bytes());
        self.pad_to(elem_align)?;
        match kind {
            ElemKind::Scalar => {
                for item in items {
                    self.push_scalar(item, elem)?;
                }
            }
            ElemKind::Ref => {
                let slots_base = self.pos()?;
                self.buf.resize(
                    self.buf.len() + items.len().checked_mul(4).ok_or(Error::MessageTooLarge)?,
                    0,
                );
                self.pos()?;
                for (i, item) in items.iter().enumerate() {
                    let child = self.write_heap(item, elem)?;
                    self.patch_u32(slots_base + (i as u32) * 4, child);
                }
            }
            ElemKind::InlineStruct(type_index) => {
                let base = self.pos()?;
                let total = (stride as usize)
                    .checked_mul(items.len())
                    .ok_or(Error::MessageTooLarge)?;
                self.buf.resize(self.buf.len() + total, 0);
                self.pos()?;
                for (i, item) in items.iter().enumerate() {
                    let fields = match item {
                        Value::Struct(fields) => fields,
                        other => return Err(type_mismatch(schema, elem, other)),
                    };
                    self.fill_struct_at(base + (i as u32) * stride, type_index, fields)?;
                }
            }
        }
        self.pos()?;
        Ok(off)
    }
}