use crate::runtime::host::{abi_arg, abi_result};
use std::collections::BTreeMap;
use wasmtime::{Caller, FuncType, Global, GlobalType, Linker, Mutability, Store, Val, ValType};
use crate::runtime::StoreData;
use crate::runtime::host::register_host_fn;
use crate::runtime::intrinsic_types::intrinsic_types;
use crate::runtime::number::{to_int32, to_uint32};
use crate::runtime::prelude::MODULE_NAME;
use crate::{
MangledName, NamespaceSymbol, PackageDeclaration, Param, Span, Type, ValueKind, ValueSymbol,
};
type UnaryFn = fn(f64) -> f64;
type BinaryFn = fn(f64, f64) -> f64;
const CONSTANTS: &[(&str, f64)] = &[
("PI", std::f64::consts::PI),
("E", std::f64::consts::E),
("LN2", std::f64::consts::LN_2),
("LN10", std::f64::consts::LN_10),
("LOG2E", std::f64::consts::LOG2_E),
("LOG10E", std::f64::consts::LOG10_E),
("SQRT2", std::f64::consts::SQRT_2),
("SQRT1_2", std::f64::consts::FRAC_1_SQRT_2),
];
const UNARY: &[(&str, UnaryFn)] = &[
("abs", f64::abs),
("ceil", f64::ceil),
("floor", f64::floor),
("trunc", f64::trunc),
("sqrt", f64::sqrt),
("round", math_round),
("sign", math_sign),
("fround", math_fround),
("clz32", math_clz32),
("exp", f64::exp),
("expm1", f64::exp_m1),
("log", f64::ln),
("log1p", f64::ln_1p),
("log2", f64::log2),
("log10", f64::log10),
("cbrt", f64::cbrt),
("sin", f64::sin),
("cos", f64::cos),
("tan", f64::tan),
("asin", f64::asin),
("acos", f64::acos),
("atan", f64::atan),
("sinh", f64::sinh),
("cosh", f64::cosh),
("tanh", f64::tanh),
("asinh", f64::asinh),
("acosh", f64::acosh),
("atanh", f64::atanh),
];
const BINARY: &[(&str, BinaryFn)] = &[
("imul", math_imul),
("pow", crate::runtime::number::pow_js),
("atan2", f64::atan2),
];
const VARIADIC: &[&str] = &["min", "max", "hypot"];
pub(crate) fn math_key(name: &str) -> MangledName {
crate::mangle::extend(&crate::mangle::prelude("Math"), name)
}
pub(crate) fn namespace_symbol() -> NamespaceSymbol {
let math_prefix = crate::mangle::prelude("Math");
let mut math = NamespaceSymbol {
name: "Math".to_string(),
mangled_prefix: math_prefix.clone(),
declaration_span: Span::at(crate::FileId::MATH),
values: BTreeMap::new(),
types: BTreeMap::new(),
namespaces: BTreeMap::new(),
doc: None,
};
for (name, _) in CONSTANTS {
math.values.insert(
(*name).to_string(),
ValueSymbol {
name: (*name).to_string(),
mangled_name: math_key(name),
declaration_span: Span::at(crate::FileId::MATH),
kind: ValueKind::Const {
ty: Type::Number,
doc: None,
},
},
);
}
for (name, _) in UNARY {
math.values.insert(
(*name).to_string(),
ValueSymbol {
name: (*name).to_string(),
mangled_name: math_key(name),
declaration_span: Span::at(crate::FileId::MATH),
kind: ValueKind::Function {
generics: Vec::new(),
params: vec![Param::new("x", Type::Number)],
ret: Type::Number,
type_predicate: None,
doc: None,
},
},
);
}
math.values.insert(
"random".to_string(),
ValueSymbol {
name: "random".to_string(),
mangled_name: math_key("random"),
declaration_span: Span::at(crate::FileId::MATH),
kind: ValueKind::Function {
generics: Vec::new(),
params: Vec::new(),
ret: Type::Number,
type_predicate: None,
doc: crate::doc(
crate::FileId::MATH,
"/** Pseudo-random number in `[0, 1)`. Not suitable for cryptographic use — reach for `submilli:crypto.randomBytes` instead. */",
),
},
},
);
for (name, _) in BINARY {
let params = match *name {
"pow" => vec![
Param::new("base", Type::Number),
Param::new("exponent", Type::Number),
],
"atan2" => vec![Param::new("y", Type::Number), Param::new("x", Type::Number)],
"imul" => vec![Param::new("a", Type::Number), Param::new("b", Type::Number)],
_ => vec![Param::new("a", Type::Number), Param::new("b", Type::Number)],
};
math.values.insert(
(*name).to_string(),
ValueSymbol {
name: (*name).to_string(),
mangled_name: math_key(name),
declaration_span: Span::at(crate::FileId::MATH),
kind: ValueKind::Function {
generics: Vec::new(),
params,
ret: Type::Number,
type_predicate: None,
doc: None,
},
},
);
}
for name in VARIADIC {
math.values.insert(
(*name).to_string(),
ValueSymbol {
name: (*name).to_string(),
mangled_name: math_key(name),
declaration_span: Span::at(crate::FileId::MATH),
kind: ValueKind::Function {
generics: Vec::new(),
params: vec![Param::rest("values", Type::Array(Box::new(Type::Number)))],
ret: Type::Number,
type_predicate: None,
doc: None,
},
},
);
}
math
}
pub fn declare(defs: &mut PackageDeclaration) {
defs.namespaces
.insert("Math".to_string(), namespace_symbol());
}
pub(crate) fn install_constants(
linker: &mut Linker<StoreData>,
store: &mut Store<StoreData>,
) -> wasmtime::Result<()> {
for (name, value) in CONSTANTS {
let gty = GlobalType::new(ValType::F64, Mutability::Var);
let global = Global::new(&mut *store, gty, Val::F64(value.to_bits()))?;
let mangled = math_key(name);
linker.define(&mut *store, MODULE_NAME, mangled.as_str(), global)?;
}
Ok(())
}
pub fn install(linker: &mut Linker<StoreData>) -> wasmtime::Result<()> {
let engine = linker.engine().clone();
let unary_ty = FuncType::new(&engine, [ValType::F64], [ValType::F64]);
let binary_ty = FuncType::new(&engine, [ValType::F64, ValType::F64], [ValType::F64]);
let array_ref = ValType::Ref(wasmtime::RefType::new(
false,
wasmtime::HeapType::ConcreteStruct(
crate::runtime::intrinsic_types::build_intrinsic_types(&engine)?.array,
),
));
let variadic_ty = FuncType::new(&engine, [array_ref], [ValType::F64]);
for (name, op) in UNARY {
let op = *op;
register_host_fn(
linker,
MODULE_NAME,
math_key(name),
unary_ty.clone(),
true,
move |_caller, params, results| -> wasmtime::Result<()> {
let x = read_f64(abi_arg(params, 0)?, "Math unary")?;
*abi_result(results, 0)? = Val::F64(op(x).to_bits());
Ok(())
},
)?;
}
for (name, op) in BINARY {
let op = *op;
register_host_fn(
linker,
MODULE_NAME,
math_key(name),
binary_ty.clone(),
true,
move |_caller, params, results| -> wasmtime::Result<()> {
let a = read_f64(abi_arg(params, 0)?, "Math binary")?;
let b = read_f64(abi_arg(params, 1)?, "Math binary")?;
*abi_result(results, 0)? = Val::F64(op(a, b).to_bits());
Ok(())
},
)?;
}
for name in VARIADIC {
let op = *name;
register_host_fn(
linker,
MODULE_NAME,
math_key(name),
variadic_ty.clone(),
true,
move |caller, params, results| -> wasmtime::Result<()> {
let values = read_variadic_numbers(caller, abi_arg(params, 0)?, op)?;
let result = match op {
"min" => math_min(&values),
"max" => math_max(&values),
"hypot" => math_hypot(&values),
_ => {
return Err(crate::runtime::host::invariant_trap(
"Math variadic: unknown operation",
));
}
};
*abi_result(results, 0)? = Val::F64(result.to_bits());
Ok(())
},
)?;
}
register_host_fn(
linker,
MODULE_NAME,
math_key("random"),
FuncType::new(&engine, [], [ValType::F64]),
false,
|_caller, _params, results| -> wasmtime::Result<()> {
let mut bytes = [0u8; 8];
getrandom::getrandom(&mut bytes)
.map_err(|e| wasmtime::Error::msg(format!("Math.random getrandom failed: {e}")))?;
let raw = u64::from_le_bytes(bytes);
let mantissa = raw & ((1u64 << 52) - 1);
let bits = (0x3FFu64 << 52) | mantissa;
*abi_result(results, 0)? = Val::F64((f64::from_bits(bits) - 1.0).to_bits());
Ok(())
},
)?;
Ok(())
}
fn read_f64(val: &Val, name: &str) -> wasmtime::Result<f64> {
match val {
Val::F64(bits) => Ok(f64::from_bits(*bits)),
other => Err(crate::runtime::host::type_error(format!(
"{name} expects f64, got {other:?}"
))),
}
}
fn read_variadic_numbers(
caller: &mut Caller<'_, StoreData>,
val: &Val,
name: &str,
) -> wasmtime::Result<Vec<f64>> {
let elements = crate::runtime::prelude::collection::read_array_vals(caller, val)?;
let boxed_number = intrinsic_types(&mut *caller)?.boxed_number.clone();
let mut out = Vec::with_capacity(elements.len());
for element in elements {
let Val::AnyRef(Some(any)) = element else {
return Err(crate::runtime::host::type_error(format!(
"Math.{name} expects boxed number arguments"
)));
};
let st = any.as_struct(&mut *caller)?.ok_or_else(|| {
crate::runtime::host::type_error(format!("Math.{name}: expected boxed number"))
})?;
if !st.matches_ty(&*caller, &boxed_number)? {
return Err(crate::runtime::host::type_error(format!(
"Math.{name}: expected boxed number"
)));
}
out.push(read_f64(&st.field(&mut *caller, 1)?, name)?);
}
Ok(out)
}
fn math_round(x: f64) -> f64 {
if x.is_nan() || x == 0.0 || x.is_infinite() {
return x;
}
if (-0.5..0.0).contains(&x) || x == -0.5 {
return -0.0;
}
(x + 0.5).floor()
}
fn math_sign(x: f64) -> f64 {
if x > 0.0 {
1.0
} else if x < 0.0 {
-1.0
} else {
x
}
}
fn math_fround(x: f64) -> f64 {
(x as f32) as f64
}
fn math_clz32(x: f64) -> f64 {
f64::from(to_uint32(x).leading_zeros())
}
fn math_imul(a: f64, b: f64) -> f64 {
let a = to_int32(a);
let b = to_int32(b);
f64::from(a.wrapping_mul(b))
}
fn math_min(values: &[f64]) -> f64 {
let mut out = f64::INFINITY;
for &value in values {
if value.is_nan() {
return f64::NAN;
}
if value < out || (value == 0.0 && out == 0.0 && value.is_sign_negative()) {
out = value;
}
}
out
}
fn math_max(values: &[f64]) -> f64 {
let mut out = f64::NEG_INFINITY;
for &value in values {
if value.is_nan() {
return f64::NAN;
}
if value > out || (value == 0.0 && out == 0.0 && value.is_sign_positive()) {
out = value;
}
}
out
}
fn math_hypot(values: &[f64]) -> f64 {
if values.iter().any(|v| v.is_infinite()) {
return f64::INFINITY;
}
if values.iter().any(|v| v.is_nan()) {
return f64::NAN;
}
values.iter().fold(0.0, |acc, value| acc.hypot(*value))
}