wasmgc 0.0.0

//! Reads linked descriptors without instantiating or interpreting the module.
//!
//! Descriptor addresses come from temporary exported globals or constants in
//! functions that call a marker. Active, constant-offset data segments provide
//! their bytes in memory zero. Schema addresses are interned before their
//! fields are visited so recursive references can be resolved iteratively.
//!
//! All recorded function indices belong to the original module. Reordering
//! and intrinsic removal happen later in the encoder.

use std::collections::{BTreeMap, BTreeSet};

use crate::schema::{
    EXPORT_MAGIC, EXPORT_PREFIX, FUNCTION_MAGIC, FUNCTION_PREFIX, FieldTag, FunctionKind, Layout,
    MAGIC_LEN, OPERATION_MAGIC, OperationKind, SCHEMA_MAGIC, SIGNATURE_PREFIX, SchemaKind,
};
use wasmparser::{ConstExpr, DataKind, ExternalKind, Operator, Parser, Payload, TypeRef};

use super::{Bridge, Export, Field, Input, Operation, Result, Schema};

struct Memory<'a> {
    layout: Layout,
    segments: Vec<(u64, &'a [u8])>,
}

#[derive(Default)]
struct FunctionExport {
    function: Option<u32>,
    signature: Option<u64>,
}

struct Decoder<'a> {
    memory: Memory<'a>,
    addresses: BTreeMap<u64, u32>,
    schemas: Vec<Schema>,
    pending: Vec<u64>,
    byte_len: usize,
}

enum Placeholder {
    Operation(Operation),
    Import(Bridge),
}

/// Section counts, descriptor locations, and marker callers from the linked module.
pub(super) struct ScannedModule<'a> {
    input: Input,
    segments: Vec<(u64, &'a [u8])>,
    bodies: BTreeMap<u32, BTreeSet<u64>>,
    export_descriptors: Vec<u64>,
    function_exports: BTreeMap<String, FunctionExport>,
    byte_len: usize,
}

/// Collects linked addresses and module contents without decoding descriptors.
pub(super) fn scan(bytes: &[u8]) -> Result<ScannedModule<'_>> {
    let mut input = Input::default();
    let mut markers = BTreeSet::new();
    let mut segments = Vec::new();
    let mut bodies = BTreeMap::new();
    let mut global_addresses = BTreeMap::new();
    let mut export_descriptors = Vec::new();
    let mut function_exports = BTreeMap::<String, FunctionExport>::new();
    let mut next_function = 0;

    for payload in Parser::new(0).parse_all(bytes) {
        match payload? {
            Payload::TypeSection(reader) => {
                for group in reader {
                    input.types += group?.types().len() as u32;
                }
            }
            Payload::ImportSection(reader) => {
                for import in reader.into_imports() {
                    let import = import?;
                    let intrinsic = import.module == "wasmgc";

                    match import.ty {
                        TypeRef::Func(ty) | TypeRef::FuncExact(ty) => {
                            if intrinsic && import.name.starts_with("marker_") {
                                markers.insert(next_function);
                            }

                            input
                                .imports
                                .push((intrinsic.then(|| import.name.to_owned()), ty));
                            next_function += 1;
                        }
                        _ if intrinsic => {
                            return Err("wasmgc imports must be functions".into());
                        }
                        TypeRef::Table(_) => input.tables += 1,
                        TypeRef::Global(_) => input.globals += 1,
                        TypeRef::Memory(memory) => {
                            input.memory.get_or_insert(memory);
                        }
                        _ => {}
                    }
                }
            }
            Payload::FunctionSection(reader) => input.functions = reader.count(),
            Payload::TableSection(reader) => input.tables += reader.count(),
            Payload::GlobalSection(reader) => {
                for global in reader {
                    if let Some(address) = offset(global?.init_expr)? {
                        global_addresses.insert(input.globals, address);
                    }
                    input.globals += 1;
                }
            }
            Payload::ExportSection(reader) => {
                for export in reader {
                    let export = export?;
                    if export.kind == ExternalKind::Global && export.name.starts_with(EXPORT_PREFIX)
                    {
                        export_descriptors.push(
                            *global_addresses
                                .get(&export.index)
                                .ok_or("wasmgc export descriptor requires a constant address")?,
                        );
                    } else if export.kind == ExternalKind::Func
                        && let Some(name) = export.name.strip_prefix(FUNCTION_PREFIX)
                    {
                        function_exports
                            .entry(name.to_owned())
                            .or_default()
                            .function = Some(export.index);
                    } else if export.kind == ExternalKind::Global
                        && let Some(name) = export.name.strip_prefix(SIGNATURE_PREFIX)
                    {
                        let address = *global_addresses
                            .get(&export.index)
                            .ok_or("wasmgc function signature requires a constant address")?;
                        function_exports
                            .entry(name.to_owned())
                            .or_default()
                            .signature = Some(address);
                    }
                }
            }
            Payload::MemorySection(reader) => {
                for memory in reader {
                    input.memory.get_or_insert(memory?);
                }
            }
            Payload::CodeSectionEntry(body) => {
                if markers.is_empty() {
                    next_function += 1;
                    continue;
                }

                let mut constants = BTreeSet::new();
                let mut calls_marker = false;
                let mut operators = body.get_operators_reader()?;

                while !operators.eof() {
                    match operators.read()? {
                        Operator::I32Const { value } => {
                            constants.insert(u64::from(value as u32));
                        }
                        Operator::I64Const { value } => {
                            constants.insert(value as u64);
                        }
                        Operator::Call { function_index }
                        | Operator::ReturnCall { function_index }
                            if markers.contains(&function_index) =>
                        {
                            calls_marker = true;
                        }
                        _ => {}
                    }
                }

                if calls_marker {
                    bodies.insert(next_function, constants);
                }

                next_function += 1;
            }
            Payload::DataSection(reader) => {
                for data in reader {
                    let data = data?;

                    if let DataKind::Active {
                        memory_index: 0,
                        offset_expr,
                    } = data.kind
                        && let Some(start) = offset(offset_expr)?
                    {
                        segments.push((start, data.data));
                    }
                }
            }
            _ => {}
        }
    }

    if input.imports.iter().all(|(name, _)| name.is_none())
        && export_descriptors.is_empty()
        && function_exports.is_empty()
    {
        return Err(
            "no wasmgc imports or exports found; expected an unprocessed Rust module using wasmgc"
                .into(),
        );
    }

    Ok(ScannedModule {
        input,
        segments,
        bodies,
        export_descriptors,
        function_exports,
        byte_len: bytes.len(),
    })
}

impl ScannedModule<'_> {
    /// Decodes the located descriptors and every schema reachable from them.
    pub(super) fn decode_descriptors(self) -> Result<Input> {
        let Self {
            mut input,
            segments,
            mut bodies,
            export_descriptors,
            function_exports,
            byte_len,
        } = self;

        let layout = if input.memory.is_some_and(|memory| memory.memory64) {
            Layout::WASM64
        } else {
            Layout::WASM32
        };

        let mut decoder = Decoder {
            memory: Memory { layout, segments },
            addresses: BTreeMap::new(),
            schemas: Vec::new(),
            pending: Vec::new(),
            byte_len,
        };

        // Pair by export name: the linker may share a numeric bridge between
        // exports whose native reference signatures differ.
        for export in function_exports.into_values() {
            let function = export
                .function
                .ok_or("wasmgc function signature has no bridge")?;
            let address = export
                .signature
                .ok_or("wasmgc function bridge has no signature")?;
            let bridge = decoder.bridge(address, function)?;
            if bridge.kind != FunctionKind::Export {
                return Err("wasmgc signature export requires an exported function".into());
            }
            input.bridges.push(bridge);
            bodies.remove(&function);
        }

        // A marker identifies the function to replace. Constants only nominate
        // descriptor addresses; their magic bytes identify valid descriptors.
        // Keeping this association avoids decoding the function-pointer table slot.
        for (function, constants) in bodies {
            match decoder.placeholder(function, constants)? {
                Placeholder::Operation(operation) => {
                    input.operations.insert(function, operation);
                }
                Placeholder::Import(bridge) => input.bridges.push(bridge),
            }
        }

        for address in export_descriptors {
            input.exports.push(decoder.export(address)?);
        }

        input.schemas = decoder.finish()?;
        for operation in input.operations.values() {
            if matches!(
                operation.kind,
                OperationKind::StructGet | OperationKind::StructSet
            ) && operation.field as usize
                >= input.schemas[operation.schema as usize].fields.len()
            {
                return Err("wasmgc operation field index out of bounds".into());
            }
        }

        for export in &input.exports {
            let schema = &input.schemas[export.schema as usize];
            if schema.kind != SchemaKind::Struct
                || (export.kind != OperationKind::StructNew
                    && export.field as usize >= schema.fields.len())
            {
                return Err("wasmgc export requires a valid struct operation".into());
            }
        }

        Ok(input)
    }
}

impl Decoder<'_> {
    fn schema(&mut self, address: u64) -> Result<u32> {
        if let Some(index) = self.addresses.get(&address) {
            return Ok(*index);
        }

        let layout = self.memory.layout;
        if address == 0
            || self
                .memory
                .read_at::<MAGIC_LEN>(address, layout.schema_magic_offset)?
                != SCHEMA_MAGIC
        {
            return Err("invalid wasmgc schema descriptor".into());
        }

        let index = u32::try_from(self.schemas.len()).map_err(|_| "too many wasmgc schemas")?;
        let kind = SchemaKind::from_raw(self.memory.word(address, layout.schema_kind_offset)?)
            .ok_or("invalid wasmgc schema kind")?;

        // Reserve the index before visiting fields so a self-reference or cycle
        // resolves to the same declaration rather than enqueueing it again.
        self.addresses.insert(address, index);
        self.schemas.push(Schema {
            kind,
            fields: Vec::new(),
        });
        self.pending.push(address);

        Ok(index)
    }

    fn field(&mut self, address: u64) -> Result<Field> {
        let layout = self.memory.layout;
        let tag = FieldTag::from_raw(self.memory.word(address, layout.field_kind_tag_offset)?)
            .ok_or("invalid wasmgc field tag")?;
        let target = if matches!(
            tag,
            FieldTag::Ref | FieldTag::NullableRef | FieldTag::Array | FieldTag::NullableArray
        ) {
            Some(
                self.schema(
                    self.memory
                        .pointer(address, layout.field_kind_target_offset)?,
                )?,
            )
        } else {
            None
        };

        Ok(Field { tag, target })
    }

    fn fields(&mut self, address: u64, count: u32) -> Result<Vec<Field>> {
        let size = self.memory.layout.field_kind_size;
        if count as usize > self.byte_len / size as usize {
            return Err("invalid wasmgc field count".into());
        }

        let mut fields = Vec::with_capacity(count as usize);
        for index in 0..count {
            let address = address
                .checked_add(u64::from(index) * u64::from(size))
                .ok_or("wasmgc fields address overflow")?;
            fields.push(self.field(address)?);
        }

        Ok(fields)
    }

    fn bridge(&mut self, address: u64, function: u32) -> Result<Bridge> {
        let layout = self.memory.layout;
        if self
            .memory
            .read_at::<MAGIC_LEN>(address, layout.function_magic_offset)?
            != FUNCTION_MAGIC
        {
            return Err("invalid wasmgc function descriptor".into());
        }

        let kind = FunctionKind::from_raw(self.memory.word(address, layout.function_kind_offset)?)
            .ok_or("invalid wasmgc function kind")?;
        let count = self
            .memory
            .word(address, layout.function_param_count_offset)?;
        let start = self
            .memory
            .pointer(address, layout.function_params_offset)?;
        let params = self.fields(start, count)?;
        let result_address = self
            .memory
            .pointer(address, layout.function_result_offset)?;
        let result = if result_address == 0 {
            None
        } else {
            Some(self.field(result_address)?)
        };

        Ok(Bridge {
            function,
            kind,
            params,
            result,
            module: self.memory.string(
                address,
                layout.function_module_offset,
                layout.function_module_len_offset,
            )?,
            name: self.memory.string(
                address,
                layout.function_name_offset,
                layout.function_name_len_offset,
            )?,
        })
    }

    fn export(&mut self, address: u64) -> Result<Export> {
        let layout = self.memory.layout;
        if self
            .memory
            .read_at::<MAGIC_LEN>(address, layout.export_magic_offset)?
            != EXPORT_MAGIC
        {
            return Err("invalid wasmgc export descriptor".into());
        }

        let kind = OperationKind::from_raw(self.memory.word(address, layout.export_kind_offset)?)
            .filter(|kind| {
                matches!(
                    kind,
                    OperationKind::StructNew | OperationKind::StructGet | OperationKind::StructSet
                )
            })
            .ok_or("invalid wasmgc export kind")?;
        let schema = self.schema(self.memory.pointer(address, layout.export_schema_offset)?)?;

        Ok(Export {
            schema,
            kind,
            field: self.memory.word(address, layout.export_field_offset)?,
            name: self.memory.string(
                address,
                layout.export_name_offset,
                layout.export_name_len_offset,
            )?,
        })
    }

    fn placeholder(&mut self, function: u32, constants: BTreeSet<u64>) -> Result<Placeholder> {
        let layout = self.memory.layout;
        let mut decoded = None;
        for address in constants {
            let Ok(magic) = self.memory.read::<MAGIC_LEN>(address) else {
                continue;
            };
            let placeholder = match magic {
                FUNCTION_MAGIC => {
                    let bridge = self.bridge(address, function)?;
                    if bridge.kind != FunctionKind::Import {
                        return Err("wasmgc marker function requires an imported function".into());
                    }
                    Placeholder::Import(bridge)
                }
                OPERATION_MAGIC => {
                    let kind = OperationKind::from_raw(
                        self.memory.word(address, layout.operation_kind_offset)?,
                    )
                    .ok_or("invalid wasmgc operation kind")?;
                    Placeholder::Operation(Operation {
                        schema: self.schema(
                            self.memory
                                .pointer(address, layout.operation_schema_offset)?,
                        )?,
                        kind,
                        field: self.memory.word(address, layout.operation_field_offset)?,
                    })
                }
                _ => continue,
            };

            if decoded.replace(placeholder).is_some() {
                return Err("multiple wasmgc descriptors identify one function".into());
            }
        }

        decoded.ok_or_else(|| "wasmgc marker call has no valid operation descriptor".into())
    }

    fn finish(mut self) -> Result<Vec<Schema>> {
        let layout = self.memory.layout;

        // Reading a field may discover another schema. Drain this worklist after
        // all roots have been registered, without recursively walking the graph.
        while let Some(address) = self.pending.pop() {
            let count = self
                .memory
                .word(address, layout.schema_field_count_offset)?;
            let fields_address = self.memory.pointer(address, layout.schema_fields_offset)?;
            let fields = self.fields(fields_address, count)?;
            let schema = &mut self.schemas[self.addresses[&address] as usize];
            if schema.kind == SchemaKind::Array && fields.len() != 1 {
                return Err("wasmgc array schema requires one element type".into());
            }
            schema.fields = fields;
        }

        Ok(self.schemas)
    }
}

impl Memory<'_> {
    fn read<const N: usize>(&self, address: u64) -> Result<[u8; N]> {
        let mut result = [0; N];

        for (offset, byte) in result.iter_mut().enumerate() {
            let address = address
                .checked_add(offset as u64)
                .ok_or("wasmgc descriptor address overflow")?;

            // Later active segments overwrite earlier ones, including partial overlaps.
            // Resolve each byte independently rather than assuming one containing segment.
            *byte = self
                .segments
                .iter()
                .rev()
                .find_map(|(start, bytes)| {
                    let offset = usize::try_from(address.checked_sub(*start)?).ok()?;

                    bytes.get(offset).copied()
                })
                .ok_or("wasmgc descriptor points outside initialized data")?;
        }

        Ok(result)
    }

    fn read_at<const N: usize>(&self, address: u64, offset: u32) -> Result<[u8; N]> {
        self.read(
            address
                .checked_add(u64::from(offset))
                .ok_or("wasmgc descriptor address overflow")?,
        )
    }

    fn word(&self, address: u64, offset: u32) -> Result<u32> {
        Ok(u32::from_le_bytes(self.read_at(address, offset)?))
    }

    fn string(&self, address: u64, offset: u32, length_offset: u32) -> Result<String> {
        let start = self.pointer(address, offset)?;
        let length = self.word(address, length_offset)?;
        let mut bytes = Vec::new();
        for offset in 0..length {
            bytes.push(self.read_at::<1>(start, offset)?[0]);
        }
        Ok(String::from_utf8(bytes)?)
    }

    fn pointer(&self, address: u64, offset: u32) -> Result<u64> {
        if self.layout.pointer_size == 8 {
            Ok(u64::from_le_bytes(self.read_at(address, offset)?))
        } else {
            Ok(u64::from(self.word(address, offset)?))
        }
    }
}

// Accept a single linked integer constant, not an expression requiring runtime
// evaluation. Signed Wasm literals preserve the unsigned address bit pattern.
fn offset(expr: ConstExpr<'_>) -> Result<Option<u64>> {
    let mut operators = expr.get_operators_reader();

    let value = match operators.read()? {
        Operator::I32Const { value } => Some(u64::from(value as u32)),
        Operator::I64Const { value } => Some(value as u64),
        _ => None,
    };

    Ok(
        if matches!(operators.read()?, Operator::End) && operators.eof() {
            value
        } else {
            None
        },
    )
}