wasmgc 0.0.0

//! Arranges schema declarations for Wasm recursive type encoding.
//!
//! Array equivalence is refined through element references, while distinct
//! struct declarations retain their identity. Strongly connected components
//! become recursion groups, ordered after their external dependencies. Every
//! schema reference in operations, exports, and function signatures is remapped
//! whenever the schema sequence changes.

use std::{cmp::Ordering, collections::BTreeMap, ops::Range};

use super::{Input, Schema};
use crate::schema::{FieldTag, SchemaKind};

fn remap(input: &mut Input, indices: &[u32]) {
    for schema in &mut input.schemas {
        for field in &mut schema.fields {
            field.target = field.target.map(|target| indices[target as usize]);
        }
    }
    for operation in input.operations.values_mut() {
        operation.schema = indices[operation.schema as usize];
    }
    for export in &mut input.exports {
        export.schema = indices[export.schema as usize];
    }
    for bridge in &mut input.bridges {
        for field in bridge.params.iter_mut().chain(&mut bridge.result) {
            field.target = field.target.map(|target| indices[target as usize]);
        }
    }
}

/// Merges equivalent array types and remaps all references to their schemas.
///
/// Rust may promote the same array descriptor at several addresses. Array
/// equivalence follows element references; struct declarations keep their identity.
pub(super) fn merge_array_types(input: &mut Input) {
    let schemas = core::mem::take(&mut input.schemas);
    let arrays: Vec<_> = schemas
        .iter()
        .enumerate()
        .filter_map(|(index, schema)| (schema.kind == SchemaKind::Array).then_some(index))
        .collect();
    if arrays.is_empty() {
        input.schemas = schemas;
        return;
    }

    let mut classes = vec![0; schemas.len()];
    let mut previous_count = 1;
    loop {
        let mut groups = BTreeMap::new();
        let mut next = classes.clone();
        for &index in &arrays {
            let field = &schemas[index].fields[0];
            let target = field.target.map(|target| {
                let array = schemas[target as usize].kind == SchemaKind::Array;
                (
                    array,
                    if array {
                        classes[target as usize]
                    } else {
                        target
                    },
                )
            });
            let count = groups.len() as u32;
            next[index] = *groups.entry((field.tag, target)).or_insert(count);
        }
        let stable = groups.len() == previous_count;
        previous_count = groups.len();
        classes = next;
        if stable {
            break;
        }
    }

    let mut canonical = BTreeMap::new();
    let mut indices = vec![0; schemas.len()];
    let mut merged = Vec::new();
    for (index, schema) in schemas.into_iter().enumerate() {
        if schema.kind == SchemaKind::Array {
            if let Some(&target) = canonical.get(&classes[index]) {
                indices[index] = target;
                continue;
            }
            canonical.insert(classes[index], merged.len() as u32);
        }
        indices[index] = merged.len() as u32;
        merged.push(schema);
    }
    input.schemas = merged;
    remap(input, &indices);
}

/// Orders schema dependencies into canonical Wasm recursion groups.
pub(super) fn group_recursive_types(input: &mut Input) {
    let components = components(&input.schemas);
    let mut group_of = vec![0; input.schemas.len()];
    for (group, members) in components.iter().enumerate() {
        for &index in members {
            group_of[index] = group;
        }
    }

    let mut keys = vec![None; input.schemas.len()];
    let mut canonical_keys = BTreeMap::new();
    let mut arena = Vec::new();
    let mut positions = vec![usize::MAX; input.schemas.len()];
    let mut order = Vec::with_capacity(input.schemas.len());
    let mut groups: Vec<Range<usize>> = Vec::with_capacity(components.len());
    for members in components {
        let (members, key) = canonical_order(
            &members,
            &input.schemas,
            &group_of,
            &keys,
            &arena,
            &mut positions,
        );
        // Keys refer to previously interned groups by arena index instead of owning
        // nested keys. Long type chains therefore need no recursive comparison/drop.
        let group = *canonical_keys.entry(key.clone()).or_insert_with(|| {
            let group = arena.len();
            arena.push(key);
            group
        });
        for (member, &index) in members.iter().enumerate() {
            keys[index] = Some(TypeKey { group, member });
        }

        let start = order.len();
        order.extend(members);
        groups.push(start..order.len());
    }

    let mut indices = vec![0; input.schemas.len()];
    for (new, &old) in order.iter().enumerate() {
        indices[old] = new as u32;
    }
    input.schemas = order
        .into_iter()
        .map(|index| {
            let schema = &mut input.schemas[index];
            Schema {
                kind: schema.kind,
                fields: core::mem::take(&mut schema.fields),
            }
        })
        .collect();
    remap(input, &indices);
    input.schema_groups = groups;
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct TypeKey {
    group: usize,
    member: usize,
}

#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
struct SchemaKey {
    kind: SchemaKind,
    fields: Vec<(FieldTag, Option<TargetKey>)>,
}

#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
enum TargetKey {
    Internal(usize),
    External(TypeKey),
}

// Arena indices establish equality, but their allocation order is incidental.
// Compare referenced structures with an explicit stack so unrelated discovery
// order cannot choose a different recursion-group traversal.
fn compare(left: &[SchemaKey], right: &[SchemaKey], arena: &[Vec<SchemaKey>]) -> Ordering {
    enum Task<'a> {
        Schemas(&'a [SchemaKey], &'a [SchemaKey]),
        Fields(
            &'a [(FieldTag, Option<TargetKey>)],
            &'a [(FieldTag, Option<TargetKey>)],
        ),
        Order(Ordering),
    }

    let mut stack = vec![Task::Schemas(left, right)];
    while let Some(task) = stack.pop() {
        let order = match task {
            Task::Schemas(left, right) => {
                if let (Some((left, rest_left)), Some((right, rest_right))) =
                    (left.split_first(), right.split_first())
                {
                    stack.push(Task::Schemas(rest_left, rest_right));
                    stack.push(Task::Fields(&left.fields, &right.fields));
                    left.kind.cmp(&right.kind)
                } else {
                    left.len().cmp(&right.len())
                }
            }
            Task::Fields(left, right) => {
                if let (Some((left, rest_left)), Some((right, rest_right))) =
                    (left.split_first(), right.split_first())
                {
                    stack.push(Task::Fields(rest_left, rest_right));
                    let order = left.0.cmp(&right.0);
                    if order != Ordering::Equal {
                        return order;
                    }
                    match (&left.1, &right.1) {
                        (Some(TargetKey::External(left)), Some(TargetKey::External(right)))
                            if left.group != right.group =>
                        {
                            stack.push(Task::Order(left.member.cmp(&right.member)));
                            stack.push(Task::Schemas(&arena[left.group], &arena[right.group]));
                            Ordering::Equal
                        }
                        _ => left.1.cmp(&right.1),
                    }
                } else {
                    left.len().cmp(&right.len())
                }
            }
            Task::Order(order) => order,
        };
        if order != Ordering::Equal {
            return order;
        }
    }
    Ordering::Equal
}

fn components(schemas: &[Schema]) -> Vec<Vec<usize>> {
    let mut seen = vec![false; schemas.len()];
    let mut finished = Vec::new();

    for root in 0..schemas.len() {
        if seen[root] {
            continue;
        }

        let mut stack = vec![(root, false)];
        while let Some((index, finish)) = stack.pop() {
            if finish {
                finished.push(index);
                continue;
            }
            if core::mem::replace(&mut seen[index], true) {
                continue;
            }

            stack.push((index, true));
            stack.extend(
                schemas[index]
                    .fields
                    .iter()
                    .rev()
                    .filter_map(|field| field.target)
                    .map(|target| (target as usize, false)),
            );
        }
    }

    let mut incoming = vec![Vec::new(); schemas.len()];
    for (index, schema) in schemas.iter().enumerate() {
        for target in schema.fields.iter().filter_map(|field| field.target) {
            incoming[target as usize].push(index);
        }
    }

    seen.fill(false);
    let mut groups = Vec::new();
    for root in finished.into_iter().rev() {
        if seen[root] {
            continue;
        }

        let mut members = Vec::new();
        let mut stack = vec![root];
        while let Some(index) = stack.pop() {
            if core::mem::replace(&mut seen[index], true) {
                continue;
            }

            members.push(index);
            stack.extend(&incoming[index]);
        }
        groups.push(members);
    }

    // Edges run from a declaration to its field dependencies. Reverse the SCC
    // traversal order so references outside a group always point to earlier types.
    groups.reverse();
    groups
}

fn canonical_order(
    members: &[usize],
    schemas: &[Schema],
    group_of: &[usize],
    keys: &[Option<TypeKey>],
    arena: &[Vec<SchemaKey>],
    positions: &mut [usize],
) -> (Vec<usize>, Vec<SchemaKey>) {
    let group = group_of[members[0]];
    let mut best: Option<(Vec<usize>, Vec<SchemaKey>)> = None;

    // Try each member as the traversal root and select the least structural key.
    // Field order breaks traversal ties; completely symmetric members have no
    // structural distinction, so their tie keeps the first traversal encountered.
    for &root in members {
        let mut order = Vec::with_capacity(members.len());
        let mut stack = vec![root];
        while let Some(index) = stack.pop() {
            if positions[index] != usize::MAX {
                continue;
            }

            positions[index] = order.len();
            order.push(index);
            stack.extend(
                schemas[index]
                    .fields
                    .iter()
                    .rev()
                    .filter_map(|field| field.target)
                    .map(|target| target as usize)
                    .filter(|&target| group_of[target] == group),
            );
        }

        let key: Vec<_> = order
            .iter()
            .map(|&index| SchemaKey {
                kind: schemas[index].kind,
                fields: schemas[index]
                    .fields
                    .iter()
                    .map(|field| {
                        let target = field.target.map(|target| {
                            let target = target as usize;
                            if group_of[target] == group {
                                TargetKey::Internal(positions[target])
                            } else {
                                TargetKey::External(keys[target].unwrap())
                            }
                        });
                        (field.tag, target)
                    })
                    .collect(),
            })
            .collect();
        if best
            .as_ref()
            .is_none_or(|(_, best)| compare(&key, best, arena).is_lt())
        {
            best = Some((order, key));
        }
        for &index in members {
            positions[index] = usize::MAX;
        }
    }

    best.unwrap()
}