submilli-engine 0.3.0

//! Primitive operators for live values whose runtime type can differ from a
//! retained TypeScript control-flow refinement.

use crate::runtime::host::{abi_arg, abi_result};
use num_traits::{FromPrimitive, ToPrimitive, Zero};
use std::cmp::Ordering;
use wasmtime::{Caller, FuncType, HeapType, Linker, RefType, Val, ValType};

use crate::runtime::StoreData;
use crate::runtime::host::{
    range_error, register_host_fn_async, type_error, write_boxed_number_struct,
    write_submilli_string_struct_units,
};
use crate::runtime::intrinsic_types::{build_intrinsic_types, intrinsic_types};
use crate::runtime::number::{format_number_js, string_to_number_js};
use crate::runtime::prelude::bigint::ops::{
    limbs_to_bigint, make_bigint_struct, read_bigint_struct,
};
use crate::runtime::prelude::{MODULE_NAME, declare_method};
use crate::{PackageDeclaration, Param, Type};

const ARITHMETIC: [&str; 6] = ["add", "sub", "mul", "div", "rem", "pow"];
const RELATIONAL: [&str; 4] = ["lt", "gt", "le", "ge"];

pub(super) fn install(linker: &mut Linker<StoreData>) -> wasmtime::Result<()> {
    let engine = linker.engine().clone();
    let intr = build_intrinsic_types(&engine)?;
    let value = ValType::Ref(RefType::new(true, HeapType::ConcreteStruct(intr.object)));
    for operation in ARITHMETIC.into_iter().chain(super::bitwise::OPERATIONS) {
        register_host_fn_async(
            linker,
            MODULE_NAME,
            crate::mangle::prelude(&format!("__value_{operation}")),
            FuncType::new(&engine, [value.clone(), value.clone()], [value.clone()]),
            true,
            move |caller, params, results| {
                Box::pin(async move {
                    let lhs = primitive(caller, abi_arg(params, 0)?).await?;
                    let rhs = primitive(caller, abi_arg(params, 1)?).await?;
                    *abi_result(results, 0)? = arithmetic(caller, operation, lhs, rhs)?;
                    Ok(())
                })
            },
        )?;
    }
    for operation in RELATIONAL {
        register_host_fn_async(
            linker,
            MODULE_NAME,
            crate::mangle::prelude(&format!("__value_{operation}")),
            FuncType::new(&engine, [value.clone(), value.clone()], [ValType::I32]),
            true,
            move |caller, params, results| {
                Box::pin(async move {
                    let lhs = primitive(caller, abi_arg(params, 0)?).await?;
                    let rhs = primitive(caller, abi_arg(params, 1)?).await?;
                    let ordering = compare(caller, &lhs, &rhs)?;
                    let answer = match operation {
                        "lt" => ordering == Some(Ordering::Less),
                        "gt" => ordering == Some(Ordering::Greater),
                        "le" => matches!(ordering, Some(Ordering::Less | Ordering::Equal)),
                        "ge" => matches!(ordering, Some(Ordering::Greater | Ordering::Equal)),
                        _ => {
                            return Err(crate::runtime::host::invariant_trap(
                                "value comparison: unknown operation",
                            ));
                        }
                    };
                    *abi_result(results, 0)? = Val::I32(i32::from(answer));
                    Ok(())
                })
            },
        )?;
    }
    for operation in ["neg", "pos", "numeric", "inc", "dec", "bitnot"] {
        register_host_fn_async(
            linker,
            MODULE_NAME,
            crate::mangle::prelude(&format!("__value_{operation}")),
            FuncType::new(&engine, [value.clone()], [value.clone()]),
            true,
            move |caller, params, results| {
                Box::pin(async move {
                    let value = primitive(caller, abi_arg(params, 0)?).await?;
                    *abi_result(results, 0)? = unary(caller, operation, value)?;
                    Ok(())
                })
            },
        )?;
    }
    register_host_fn_async(
        linker,
        MODULE_NAME,
        crate::mangle::prelude("__value_to_number"),
        FuncType::new(&engine, [value.clone()], [ValType::F64]),
        true,
        |caller, params, results| {
            Box::pin(async move {
                *abi_result(results, 0)? =
                    Val::F64(number(primitive(caller, abi_arg(params, 0)?).await?)?.to_bits());
                Ok(())
            })
        },
    )?;
    register_host_fn_async(
        linker,
        MODULE_NAME,
        crate::mangle::prelude("__value_to_index"),
        FuncType::new(&engine, [value.clone()], [ValType::F64]),
        true,
        |caller, params, results| {
            Box::pin(async move {
                let key = primitive_with_hint(caller, abi_arg(params, 0)?, true).await?;
                let units = string(caller, key)?;
                let index = number(Primitive::String(units.clone()))?;
                if format_number_js(index).encode_utf16().collect::<Vec<_>>() != units {
                    return Err(range_error("Array index must be a canonical numeric key"));
                }
                *abi_result(results, 0)? = Val::F64(index.to_bits());
                Ok(())
            })
        },
    )?;
    register_host_fn_async(
        linker,
        MODULE_NAME,
        crate::mangle::prelude("__value_to_string"),
        FuncType::new(
            &engine,
            [value],
            [ValType::Ref(RefType::new(
                false,
                HeapType::ConcreteStruct(intr.string),
            ))],
        ),
        true,
        |caller, params, results| {
            Box::pin(async move {
                let value = primitive_with_hint(caller, abi_arg(params, 0)?, true).await?;
                let text = string(caller, value)?;
                *abi_result(results, 0)? = Val::AnyRef(Some(
                    write_submilli_string_struct_units(caller, &text)?.to_anyref(),
                ));
                Ok(())
            })
        },
    )?;
    Ok(())
}

pub(super) fn declare(defs: &mut PackageDeclaration) {
    for operation in ARITHMETIC.into_iter().chain(super::bitwise::OPERATIONS) {
        let name = format!("__value_{operation}");
        declare_method(
            defs,
            &name,
            crate::mangle::prelude(&name),
            vec![
                Param::new("lhs", Type::Unknown),
                Param::new("rhs", Type::Unknown),
            ],
            Type::Unknown,
        );
    }
    for operation in RELATIONAL {
        let name = format!("__value_{operation}");
        declare_method(
            defs,
            &name,
            crate::mangle::prelude(&name),
            vec![
                Param::new("lhs", Type::Unknown),
                Param::new("rhs", Type::Unknown),
            ],
            Type::Boolean,
        );
    }
    for operation in ["neg", "pos", "numeric", "inc", "dec", "bitnot"] {
        let name = format!("__value_{operation}");
        declare_method(
            defs,
            &name,
            crate::mangle::prelude(&name),
            vec![Param::new("value", Type::Unknown)],
            Type::Unknown,
        );
    }
    declare_method(
        defs,
        "__value_to_number",
        crate::mangle::prelude("__value_to_number"),
        vec![Param::new("value", Type::Unknown)],
        Type::Number,
    );
    declare_method(
        defs,
        "__value_to_index",
        crate::mangle::prelude("__value_to_index"),
        vec![Param::new("value", Type::Unknown)],
        Type::Number,
    );
    declare_method(
        defs,
        "__value_to_string",
        crate::mangle::prelude("__value_to_string"),
        vec![Param::new("value", Type::Unknown)],
        Type::String,
    );
}

#[derive(Clone)]
pub(super) enum Primitive {
    Null,
    Undefined,
    Number(f64),
    Boolean(bool),
    String(Vec<u16>),
    BigInt(num_bigint::BigInt),
}

fn compare(
    caller: &mut Caller<'_, StoreData>,
    lhs: &Primitive,
    rhs: &Primitive,
) -> wasmtime::Result<Option<Ordering>> {
    match (lhs, rhs) {
        (Primitive::String(lhs), Primitive::String(rhs)) => Ok(Some(lhs.cmp(rhs))),
        (Primitive::BigInt(lhs), Primitive::BigInt(rhs)) => Ok(Some(lhs.cmp(rhs))),
        (Primitive::BigInt(lhs), Primitive::String(rhs)) => {
            Ok(parse_bigint(caller, rhs)?.map(|rhs| lhs.cmp(&rhs)))
        }
        (Primitive::String(lhs), Primitive::BigInt(rhs)) => {
            Ok(parse_bigint(caller, lhs)?.map(|lhs| lhs.cmp(rhs)))
        }
        (Primitive::BigInt(lhs), rhs) => Ok(compare_bigint_number(lhs, number(rhs.clone())?)),
        (lhs, Primitive::BigInt(rhs)) => {
            Ok(compare_bigint_number(rhs, number(lhs.clone())?).map(Ordering::reverse))
        }
        _ => Ok(number(lhs.clone())?.partial_cmp(&number(rhs.clone())?)),
    }
}

fn parse_bigint(
    caller: &mut Caller<'_, StoreData>,
    units: &[u16],
) -> wasmtime::Result<Option<num_bigint::BigInt>> {
    if units.len() > super::bigint::ops::MAX_DECIMAL_INPUT_BYTES {
        return Err(range_error(
            "BigInt comparison: input exceeds 65536 code units; use a shorter value",
        ));
    }
    let Ok(text) = String::from_utf16(units) else {
        return Ok(None);
    };
    let text = text.trim_matches(crate::runtime::number::is_js_whitespace);
    if text.is_empty() {
        return Ok(Some(num_bigint::BigInt::zero()));
    }
    for (prefix, radix) in [
        ("0x", 16),
        ("0X", 16),
        ("0o", 8),
        ("0O", 8),
        ("0b", 2),
        ("0B", 2),
    ] {
        if let Some(digits) = text.strip_prefix(prefix) {
            return parse_bigint_digits(caller, digits, radix);
        }
    }
    let (negative, digits) = if let Some(digits) = text.strip_prefix('-') {
        (true, digits)
    } else {
        (false, text.strip_prefix('+').unwrap_or(text))
    };
    Ok(parse_bigint_digits(caller, digits, 10)?.map(|value| if negative { -value } else { value }))
}

fn parse_bigint_digits(
    caller: &mut Caller<'_, StoreData>,
    digits: &str,
    radix: u32,
) -> wasmtime::Result<Option<num_bigint::BigInt>> {
    crate::runtime::fuel::charge(
        &mut *caller,
        crate::runtime::fuel::SCAN,
        digits.len() as u64,
    )?;
    if digits.is_empty() || !digits.chars().all(|digit| digit.is_digit(radix)) {
        return Ok(None);
    }
    if radix == 10 {
        let n = digits.len() as u64;
        crate::runtime::fuel::charge(
            &mut *caller,
            crate::runtime::fuel::ELEM,
            n.saturating_mul(n.div_ceil(19)),
        )?;
    }
    Ok(num_bigint::BigInt::parse_bytes(digits.as_bytes(), radix))
}

fn compare_bigint_number(lhs: &num_bigint::BigInt, rhs: f64) -> Option<Ordering> {
    if rhs.is_nan() {
        return None;
    }
    if rhs == f64::INFINITY {
        return Some(Ordering::Less);
    }
    if rhs == f64::NEG_INFINITY {
        return Some(Ordering::Greater);
    }
    let integer = num_bigint::BigInt::from_f64(rhs.trunc())?;
    Some(match lhs.cmp(&integer) {
        Ordering::Equal if rhs.fract() > 0.0 => Ordering::Less,
        Ordering::Equal if rhs.fract() < 0.0 => Ordering::Greater,
        ordering => ordering,
    })
}

async fn primitive(caller: &mut Caller<'_, StoreData>, value: &Val) -> wasmtime::Result<Primitive> {
    primitive_with_hint(caller, value, false).await
}

pub(super) async fn primitive_with_hint(
    caller: &mut Caller<'_, StoreData>,
    value: &Val,
    prefers_string: bool,
) -> wasmtime::Result<Primitive> {
    if let Some(primitive) = read_primitive(caller, value)? {
        return Ok(primitive);
    }
    let methods = if prefers_string {
        ["toString", "valueOf"]
    } else {
        ["valueOf", "toString"]
    };
    for name in methods {
        let method = conversion_method(caller, value, name).await?;
        let converted = if let Some(method) = method {
            if !is_callable(caller, &method)? {
                continue;
            }
            super::closure::read(caller, &method, name)?
                .call_with_receiver(caller, *value, &[])
                .await?
        } else if name == "toString" {
            super::vtable::dispatch_vtable_slot(caller, value, 0, &[]).await?
        } else {
            // The inherited Object.valueOf returns the object itself.
            continue;
        };
        if let Some(primitive) = read_primitive(caller, &converted)? {
            return Ok(primitive);
        }
    }
    Err(type_error("Cannot convert object to primitive value"))
}

pub(super) async fn conversion_method(
    caller: &mut Caller<'_, StoreData>,
    value: &Val,
    name: &str,
) -> wasmtime::Result<Option<Val>> {
    // An optional conversion field left out holds `undefined` (or `null`), and
    // the value then converts as if it declared none.
    if let Some(method) = super::collection::object_field(caller, value, name)? {
        let declared = !super::undefined::is_nullish(caller, &method)?;
        return Ok(declared.then_some(method));
    }
    let Val::AnyRef(Some(reference)) = value else {
        return Ok(None);
    };
    let Some(object) = reference.as_struct(&mut *caller)? else {
        return Ok(None);
    };
    let Val::AnyRef(Some(vtable)) = object.field(&mut *caller, 0)? else {
        return Ok(None);
    };
    let Some(vtable) = vtable.as_struct(&mut *caller)? else {
        return Ok(None);
    };
    let intr = intrinsic_types(&mut *caller)?;
    if !vtable.matches_ty(&*caller, &intr.class_vtable)? {
        return Ok(None);
    }
    let getter_name = format!("get {name}");
    let Some(getter) = super::collection::object_accessor(caller, value, &getter_name)? else {
        return Ok(None);
    };
    if !is_callable(caller, &getter)? {
        return Ok(None);
    }
    // Callers bind a receiver to the method, which allocates, before they
    // call it.
    let kept_method = super::keep::KeptValue::new(caller)?;
    let method = super::closure::read(caller, &getter, &getter_name)?
        .call_with_receiver(caller, *value, &[])
        .await?;
    kept_method.set(caller, method)?;
    Ok(Some(method))
}

pub(super) fn is_callable(
    caller: &mut Caller<'_, StoreData>,
    value: &Val,
) -> wasmtime::Result<bool> {
    let Val::AnyRef(Some(reference)) = value else {
        return Ok(false);
    };
    let Some(object) = reference.as_struct(&mut *caller)? else {
        return Ok(false);
    };
    let closure = intrinsic_types(&mut *caller)?.closure.clone();
    object.matches_ty(&*caller, &closure)
}

pub(super) async fn to_number(
    caller: &mut Caller<'_, StoreData>,
    value: &Val,
) -> wasmtime::Result<f64> {
    number(primitive(caller, value).await?)
}

pub(super) fn truthy(caller: &mut Caller<'_, StoreData>, value: &Val) -> wasmtime::Result<bool> {
    Ok(match read_primitive(caller, value)? {
        Some(Primitive::Null | Primitive::Undefined) => false,
        Some(Primitive::Boolean(value)) => value,
        Some(Primitive::Number(value)) => value != 0.0 && !value.is_nan(),
        Some(Primitive::String(value)) => !value.is_empty(),
        Some(Primitive::BigInt(value)) => !value.is_zero(),
        None => true,
    })
}

fn read_primitive(
    caller: &mut Caller<'_, StoreData>,
    value: &Val,
) -> wasmtime::Result<Option<Primitive>> {
    let Val::AnyRef(Some(reference)) = value else {
        return Ok(matches!(value, Val::AnyRef(None)).then_some(Primitive::Null));
    };
    let Some(object) = reference.as_struct(&mut *caller)? else {
        return Ok(None);
    };
    let intr = intrinsic_types(&mut *caller)?;
    if object.matches_ty(&*caller, &intr.undefined)? {
        return Ok(Some(Primitive::Undefined));
    }
    if object.matches_ty(&*caller, &intr.boxed_number)? {
        let Val::F64(bits) = object.field(&mut *caller, 1)? else {
            return Err(crate::runtime::host::invariant_trap(
                "value: invalid boxed number payload",
            ));
        };
        return Ok(Some(Primitive::Number(f64::from_bits(bits))));
    }
    if object.matches_ty(&*caller, &intr.boxed_boolean)? {
        let Val::I32(value) = object.field(&mut *caller, 1)? else {
            return Err(crate::runtime::host::invariant_trap(
                "value: invalid boxed boolean payload",
            ));
        };
        return Ok(Some(Primitive::Boolean(value != 0)));
    }
    if object.matches_ty(&*caller, &intr.string)? {
        return Ok(Some(Primitive::String(super::vtable::read_string_units(
            caller,
            value,
            "string primitive",
        )?)));
    }
    if object.matches_ty(&*caller, &intr.bigint)? {
        let (sign, limbs) = read_bigint_struct(caller, value, "arithmetic")?;
        return Ok(Some(Primitive::BigInt(limbs_to_bigint(sign, &limbs)?)));
    }
    Ok(None)
}

fn arithmetic(
    caller: &mut Caller<'_, StoreData>,
    operation: &str,
    lhs: Primitive,
    rhs: Primitive,
) -> wasmtime::Result<Val> {
    if operation == "add"
        && (matches!(lhs, Primitive::String(_)) || matches!(rhs, Primitive::String(_)))
    {
        let mut text = string(caller, lhs)?;
        text.extend(string(caller, rhs)?);
        let value = write_submilli_string_struct_units(caller, &text)?;
        return Ok(Val::AnyRef(Some(value.to_anyref())));
    }
    if let (Primitive::BigInt(lhs), Primitive::BigInt(rhs)) = (&lhs, &rhs) {
        return bigint_arithmetic(caller, operation, lhs, rhs);
    }
    let lhs = number(lhs)?;
    let rhs = number(rhs)?;
    let result = match operation {
        "add" => lhs + rhs,
        "sub" => lhs - rhs,
        "mul" => lhs * rhs,
        "div" => lhs / rhs,
        "rem" => lhs % rhs,
        "pow" => crate::runtime::number::pow_js(lhs, rhs),
        name if super::bitwise::OPERATIONS.contains(&name) => {
            super::bitwise::number(name, lhs, rhs)?
        }
        _ => {
            return Err(crate::runtime::host::invariant_trap(
                "number arithmetic: unknown operation",
            ));
        }
    };
    Ok(Val::AnyRef(Some(
        write_boxed_number_struct(caller, result)?.to_anyref(),
    )))
}

fn number(value: Primitive) -> wasmtime::Result<f64> {
    Ok(match value {
        Primitive::Null => 0.0,
        Primitive::Undefined => f64::NAN,
        Primitive::Number(value) => value,
        Primitive::Boolean(value) => {
            if value {
                1.0
            } else {
                0.0
            }
        }
        // A lone surrogate cannot form a numeric literal. Reject it as NaN
        // without changing any string value or its UTF-16 representation.
        Primitive::String(units) => {
            String::from_utf16(&units).map_or(f64::NAN, |text| string_to_number_js(&text))
        }
        Primitive::BigInt(_) => return Err(type_error("Cannot mix BigInt and other types")),
    })
}

pub(super) fn string(
    caller: &mut Caller<'_, StoreData>,
    value: Primitive,
) -> wasmtime::Result<Vec<u16>> {
    let text = match value {
        Primitive::String(units) => return Ok(units),
        Primitive::Null => "null".to_owned(),
        Primitive::Undefined => "undefined".to_owned(),
        Primitive::Number(value) => format_number_js(value),
        Primitive::Boolean(value) => value.to_string(),
        Primitive::BigInt(value) => super::bigint::ops::format_bigint(caller, &value, 10)?,
    };
    Ok(text.encode_utf16().collect())
}

fn bigint_arithmetic(
    caller: &mut Caller<'_, StoreData>,
    operation: &str,
    lhs: &num_bigint::BigInt,
    rhs: &num_bigint::BigInt,
) -> wasmtime::Result<Val> {
    if matches!(operation, "div" | "rem") && rhs.is_zero() {
        return Err(range_error("Division by zero"));
    }
    let result = match operation {
        "add" => lhs + rhs,
        "sub" => lhs - rhs,
        "mul" => lhs * rhs,
        "div" => lhs / rhs,
        "rem" => lhs % rhs,
        "pow" => lhs.pow(
            rhs.to_u32()
                .ok_or_else(|| range_error("BigInt exponent must fit a non-negative u32"))?,
        ),
        name if super::bitwise::OPERATIONS.contains(&name) => {
            super::bitwise::bigint(caller, name, lhs, rhs)?
        }
        _ => {
            return Err(crate::runtime::host::invariant_trap(
                "bigint arithmetic: unknown operation",
            ));
        }
    };
    make_bigint_struct(caller, result)
}

fn unary(
    caller: &mut Caller<'_, StoreData>,
    operation: &str,
    value: Primitive,
) -> wasmtime::Result<Val> {
    Ok(
        if operation != "pos"
            && let Primitive::BigInt(value) = &value
        {
            let value = match operation {
                "neg" => -value,
                "bitnot" => super::bitwise::complement(caller, value)?,
                "inc" => value + 1,
                "dec" => value - 1,
                _ => value.clone(),
            };
            make_bigint_struct(caller, value)?
        } else {
            let number = number(value)?;
            let number = match operation {
                "neg" => -number,
                "bitnot" => f64::from(!crate::runtime::number::to_int32(number)),
                "inc" => number + 1.0,
                "dec" => number - 1.0,
                _ => number,
            };
            Val::AnyRef(Some(write_boxed_number_struct(caller, number)?.to_anyref()))
        },
    )
}

/// Inspect a number without invoking JavaScript coercion hooks.
pub(super) fn number_value(
    caller: &mut Caller<'_, StoreData>,
    value: &Val,
) -> wasmtime::Result<Option<f64>> {
    let Val::AnyRef(Some(reference)) = value else {
        return Ok(None);
    };
    let Some(object) = reference.as_struct(&mut *caller)? else {
        return Ok(None);
    };
    let intr = intrinsic_types(&mut *caller)?;
    if !object.matches_ty(&*caller, &intr.boxed_number)? {
        return Ok(None);
    }
    match object.field(&mut *caller, 1)? {
        Val::F64(bits) => Ok(Some(f64::from_bits(bits))),
        _ => Err(crate::runtime::host::fatal_host_error(
            "invalid boxed number payload",
        )),
    }
}

pub(super) async fn search_string(
    caller: &mut Caller<'_, StoreData>,
    value: &Val,
) -> wasmtime::Result<Vec<u16>> {
    let intr = intrinsic_types(&mut *caller)?;
    if let Val::AnyRef(Some(reference)) = value
        && let Some(object) = reference.as_struct(&mut *caller)?
        && object.matches_ty(&*caller, &intr.regex)?
    {
        return Err(type_error("String search argument must not be a RegExp"));
    }
    let value = primitive_with_hint(caller, value, true).await?;
    string(caller, value)
}