use std::collections::{HashMap, HashSet};
use wasm_encoder::ValType;
use crate::ast::{Literal, Spanned};
use crate::ir::hir::{
ResolvedCallee, ResolvedExpr, ResolvedFnBody, ResolvedFnDef, ResolvedPattern, ResolvedStmt,
ResolvedStrPart,
};
use crate::types::Type;
use super::super::WasmGcError;
use super::super::types::{TypeRegistry, aver_to_wasm};
use super::FnMap;
use super::infer::{
arm_is_option_pattern_resolved, arm_is_result_pattern_resolved, aver_type_str_of,
};
pub(super) struct SlotTable {
pub(super) by_slot: Vec<ValType>,
pub(super) subject_scratch: Option<u32>,
pub(super) args_get_scratch: Option<[u32; 4]>,
pub(super) vector_set_scratch: HashMap<String, u32>,
pub(super) console_print_wasip2_scratch: Option<[u32; 2]>,
pub(super) args_get_wasip2_retptr_scratch: Option<u32>,
pub(super) env_get_wasip2_scratch: Option<[u32; 2]>,
pub(super) random_int_wasip2_min_scratch: Option<u32>,
}
impl SlotTable {
pub(super) fn build_for_fn(
fd: &ResolvedFnDef,
registry: &TypeRegistry,
_fn_map: &FnMap,
) -> Result<Self, WasmGcError> {
let mut by_slot: Vec<ValType> = Vec::new();
if let Some(resolution) = fd.resolution.as_ref() {
for ty in resolution.local_slot_types.iter() {
let aver_str = ty.display();
let v = aver_to_wasm(&aver_str, Some(registry))
.unwrap_or(None)
.unwrap_or(ValType::I32);
by_slot.push(v);
}
} else {
for (_, ty) in &fd.params {
let v = aver_to_wasm(&ty.display(), Some(registry))
.unwrap_or(None)
.unwrap_or(ValType::I32);
by_slot.push(v);
}
}
let needs_scratch = fn_needs_subject_scratch(fd, registry);
let subject_scratch = if needs_scratch {
let scratch_ty = ValType::Ref(wasm_encoder::RefType {
nullable: true,
heap_type: wasm_encoder::HeapType::Abstract {
shared: false,
ty: wasm_encoder::AbstractHeapType::Eq,
},
});
let idx = by_slot.len() as u32;
by_slot.push(scratch_ty);
Some(idx)
} else {
None
};
let args_get_scratch = if fn_needs_args_get_scratch(fd) {
let i64_ty = ValType::I64;
let list_ref = registry.list_type_idx("List<String>").map(|idx| {
ValType::Ref(wasm_encoder::RefType {
nullable: true,
heap_type: wasm_encoder::HeapType::Concrete(idx),
})
});
let str_ref = registry.string_array_type_idx.map(|idx| {
ValType::Ref(wasm_encoder::RefType {
nullable: true,
heap_type: wasm_encoder::HeapType::Concrete(idx),
})
});
match (list_ref, str_ref) {
(Some(list_ty), Some(s_ty)) => {
let i_idx = by_slot.len() as u32;
by_slot.push(i64_ty);
let len_idx = by_slot.len() as u32;
by_slot.push(i64_ty);
let acc_idx = by_slot.len() as u32;
by_slot.push(list_ty);
let s_idx = by_slot.len() as u32;
by_slot.push(s_ty);
Some([i_idx, len_idx, acc_idx, s_idx])
}
_ => {
return Err(WasmGcError::Validation(
"Args.get() requires List<String> and String slots in registry — \
pre-register them by ensuring the program reaches a List<String> \
literal or String value first"
.into(),
));
}
}
} else {
None
};
let mut vector_set_canonicals: HashSet<String> = HashSet::new();
collect_vector_set_canonicals(fd, &mut vector_set_canonicals);
let mut vector_set_scratch: HashMap<String, u32> = HashMap::new();
let mut sorted: Vec<String> = vector_set_canonicals.into_iter().collect();
sorted.sort(); for canonical in sorted {
if let Some(vec_idx) = registry.vector_type_idx(&canonical) {
let ty = ValType::Ref(wasm_encoder::RefType {
nullable: true,
heap_type: wasm_encoder::HeapType::Concrete(vec_idx),
});
let local_idx = by_slot.len() as u32;
by_slot.push(ty);
vector_set_scratch.insert(canonical, local_idx);
}
}
let console_print_wasip2_scratch = if fn_needs_console_print_wasip2_scratch(fd) {
let len_idx = by_slot.len() as u32;
by_slot.push(ValType::I32);
let off_idx = by_slot.len() as u32;
by_slot.push(ValType::I32);
Some([len_idx, off_idx])
} else {
None
};
let args_get_wasip2_retptr_scratch = if fn_needs_args_get_scratch(fd) {
let idx = by_slot.len() as u32;
by_slot.push(ValType::I32);
Some(idx)
} else {
None
};
let env_get_wasip2_scratch = if fn_needs_env_get_wasip2_scratch(fd) {
let retptr = by_slot.len() as u32;
by_slot.push(ValType::I32);
let key_len = by_slot.len() as u32;
by_slot.push(ValType::I32);
Some([retptr, key_len])
} else {
None
};
let random_int_wasip2_min_scratch = if fn_needs_random_int_wasip2_scratch(fd) {
let idx = by_slot.len() as u32;
by_slot.push(ValType::I64);
Some(idx)
} else {
None
};
Ok(Self {
by_slot,
subject_scratch,
args_get_scratch,
vector_set_scratch,
console_print_wasip2_scratch,
args_get_wasip2_retptr_scratch,
env_get_wasip2_scratch,
random_int_wasip2_min_scratch,
})
}
pub(super) fn extra_locals(&self, params_count: usize) -> Vec<ValType> {
self.by_slot.iter().skip(params_count).copied().collect()
}
}
pub(super) fn fn_needs_args_get_scratch(fd: &ResolvedFnDef) -> bool {
let ResolvedFnBody::Block(stmts) = fd.body.as_ref();
stmts.iter().any(stmt_reaches_args_get_no_args)
}
fn stmt_reaches_args_get_no_args(stmt: &ResolvedStmt) -> bool {
match stmt {
ResolvedStmt::Binding { value, .. } | ResolvedStmt::Expr(value) => {
expr_reaches_args_get_no_args(&value.node)
}
}
}
fn expr_reaches_builtin_call<F>(expr: &ResolvedExpr, matches: &F) -> bool
where
F: Fn(&str, &[Spanned<ResolvedExpr>]) -> bool,
{
match expr {
ResolvedExpr::Call(callee, args) => {
if let ResolvedCallee::Builtin(name) = callee
&& matches(name.as_str(), args)
{
return true;
}
args.iter()
.any(|a| expr_reaches_builtin_call(&a.node, matches))
}
ResolvedExpr::BinOp(_, l, r) => {
expr_reaches_builtin_call(&l.node, matches)
|| expr_reaches_builtin_call(&r.node, matches)
}
ResolvedExpr::Neg(inner) => expr_reaches_builtin_call(&inner.node, matches),
ResolvedExpr::Match { subject, arms } => {
expr_reaches_builtin_call(&subject.node, matches)
|| arms
.iter()
.any(|a| expr_reaches_builtin_call(&a.body.node, matches))
}
ResolvedExpr::TailCall { args, .. } => args
.iter()
.any(|a| expr_reaches_builtin_call(&a.node, matches)),
ResolvedExpr::Attr(obj, _) => expr_reaches_builtin_call(&obj.node, matches),
ResolvedExpr::ErrorProp(inner) => expr_reaches_builtin_call(&inner.node, matches),
ResolvedExpr::Ctor(_, args) => args
.iter()
.any(|a| expr_reaches_builtin_call(&a.node, matches)),
ResolvedExpr::RecordCreate { fields, .. } => fields
.iter()
.any(|(_, e)| expr_reaches_builtin_call(&e.node, matches)),
ResolvedExpr::RecordUpdate { base, updates, .. } => {
expr_reaches_builtin_call(&base.node, matches)
|| updates
.iter()
.any(|(_, e)| expr_reaches_builtin_call(&e.node, matches))
}
ResolvedExpr::List(items)
| ResolvedExpr::Tuple(items)
| ResolvedExpr::IndependentProduct(items, _) => items
.iter()
.any(|e| expr_reaches_builtin_call(&e.node, matches)),
ResolvedExpr::MapLiteral(entries) => entries.iter().any(|(k, v)| {
expr_reaches_builtin_call(&k.node, matches)
|| expr_reaches_builtin_call(&v.node, matches)
}),
ResolvedExpr::InterpolatedStr(parts) => parts.iter().any(|p| {
if let ResolvedStrPart::Parsed(inner) = p {
expr_reaches_builtin_call(&inner.node, matches)
} else {
false
}
}),
_ => false,
}
}
fn expr_reaches_args_get_no_args(expr: &ResolvedExpr) -> bool {
expr_reaches_builtin_call(expr, &|name, args| name == "Args.get" && args.is_empty())
}
pub(super) fn fn_needs_console_print_wasip2_scratch(fd: &ResolvedFnDef) -> bool {
let ResolvedFnBody::Block(stmts) = fd.body.as_ref();
stmts.iter().any(stmt_reaches_console_print)
}
fn stmt_reaches_console_print(stmt: &ResolvedStmt) -> bool {
match stmt {
ResolvedStmt::Binding { value, .. } | ResolvedStmt::Expr(value) => {
expr_reaches_console_print(&value.node)
}
}
}
fn expr_reaches_console_print(expr: &ResolvedExpr) -> bool {
expr_reaches_builtin_call(expr, &|name, _| {
matches!(name, "Console.print" | "Console.error" | "Console.warn")
})
}
pub(super) fn fn_needs_env_get_wasip2_scratch(fd: &ResolvedFnDef) -> bool {
let ResolvedFnBody::Block(stmts) = fd.body.as_ref();
stmts.iter().any(stmt_reaches_env_get)
}
fn stmt_reaches_env_get(stmt: &ResolvedStmt) -> bool {
match stmt {
ResolvedStmt::Binding { value, .. } | ResolvedStmt::Expr(value) => {
expr_reaches_env_get(&value.node)
}
}
}
fn expr_reaches_env_get(expr: &ResolvedExpr) -> bool {
expr_reaches_builtin_call(expr, &|name, _| name == "Env.get")
}
pub(super) fn fn_needs_random_int_wasip2_scratch(fd: &ResolvedFnDef) -> bool {
let ResolvedFnBody::Block(stmts) = fd.body.as_ref();
stmts.iter().any(stmt_reaches_random_int)
}
fn stmt_reaches_random_int(stmt: &ResolvedStmt) -> bool {
match stmt {
ResolvedStmt::Binding { value, .. } | ResolvedStmt::Expr(value) => {
expr_reaches_random_int(&value.node)
}
}
}
fn expr_reaches_random_int(expr: &ResolvedExpr) -> bool {
expr_reaches_builtin_call(expr, &|name, _| name == "Random.int")
}
pub(super) fn fn_needs_subject_scratch(fd: &ResolvedFnDef, registry: &TypeRegistry) -> bool {
let ResolvedFnBody::Block(stmts) = fd.body.as_ref();
stmts.iter().any(|s| stmt_needs_scratch(s, registry))
}
pub(super) fn stmt_needs_scratch(stmt: &ResolvedStmt, registry: &TypeRegistry) -> bool {
match stmt {
ResolvedStmt::Binding { value, .. } | ResolvedStmt::Expr(value) => {
expr_needs_scratch(&value.node, registry)
}
}
}
#[allow(clippy::only_used_in_recursion)]
pub(super) fn expr_needs_scratch(expr: &ResolvedExpr, registry: &TypeRegistry) -> bool {
match expr {
ResolvedExpr::Match { subject, arms } => {
if expr_needs_scratch(&subject.node, registry) {
return true;
}
if arms.iter().any(arm_is_option_pattern_resolved) {
return true;
}
if arms.iter().any(arm_is_result_pattern_resolved) {
return true;
}
if arms.iter().any(|a| {
matches!(
&a.pattern,
ResolvedPattern::EmptyList | ResolvedPattern::Cons(_, _)
)
}) {
return true;
}
if arms
.iter()
.any(|a| matches!(&a.pattern, ResolvedPattern::Tuple(_)))
{
return true;
}
if arms
.iter()
.any(|a| matches!(&a.pattern, ResolvedPattern::Literal(Literal::Str(_))))
{
return true;
}
if arms
.iter()
.any(|a| matches!(&a.pattern, ResolvedPattern::Ctor(_, _)))
{
return true;
}
arms.iter()
.any(|a| expr_needs_scratch(&a.body.node, registry))
}
ResolvedExpr::BinOp(_, l, r) => {
expr_needs_scratch(&l.node, registry) || expr_needs_scratch(&r.node, registry)
}
ResolvedExpr::Neg(inner) => expr_needs_scratch(&inner.node, registry),
ResolvedExpr::Call(callee, args) => {
if let ResolvedCallee::Builtin(name) = callee
&& matches!(
name.as_str(),
"Option.withDefault" | "Result.withDefault" | "Option.toResult"
)
{
return true;
}
args.iter().any(|a| expr_needs_scratch(&a.node, registry))
}
ResolvedExpr::TailCall { args, .. } => {
args.iter().any(|a| expr_needs_scratch(&a.node, registry))
}
ResolvedExpr::Attr(obj, _) => expr_needs_scratch(&obj.node, registry),
ResolvedExpr::ErrorProp(_) => true,
ResolvedExpr::Ctor(_, args) => args.iter().any(|a| expr_needs_scratch(&a.node, registry)),
ResolvedExpr::RecordCreate { fields, .. } => fields
.iter()
.any(|(_, e)| expr_needs_scratch(&e.node, registry)),
ResolvedExpr::List(items) if !items.is_empty() => true,
ResolvedExpr::List(items) => items.iter().any(|e| expr_needs_scratch(&e.node, registry)),
ResolvedExpr::MapLiteral(entries) => entries.iter().any(|(k, v)| {
expr_needs_scratch(&k.node, registry) || expr_needs_scratch(&v.node, registry)
}),
ResolvedExpr::IndependentProduct(items, unwrap) => {
*unwrap || items.iter().any(|e| expr_needs_scratch(&e.node, registry))
}
ResolvedExpr::Tuple(items) => items.iter().any(|e| expr_needs_scratch(&e.node, registry)),
ResolvedExpr::RecordUpdate { base, updates, .. } => {
expr_needs_scratch(&base.node, registry)
|| updates
.iter()
.any(|(_, e)| expr_needs_scratch(&e.node, registry))
}
ResolvedExpr::InterpolatedStr(parts) => parts.iter().any(|p| {
if let ResolvedStrPart::Parsed(inner) = p {
expr_needs_scratch(&inner.node, registry)
} else {
false
}
}),
_ => false,
}
}
pub(super) fn count_value_params(params: &[(String, Type)]) -> usize {
params
.iter()
.filter(|(_, ty)| ty.display().trim() != "Unit")
.count()
}
fn collect_vector_set_canonicals(fd: &ResolvedFnDef, out: &mut HashSet<String>) {
let ResolvedFnBody::Block(stmts) = fd.body.as_ref();
for stmt in stmts {
match stmt {
ResolvedStmt::Binding { value, .. } | ResolvedStmt::Expr(value) => {
walk_expr_for_vector_set(value, out)
}
}
}
}
fn walk_expr_for_vector_set(expr: &Spanned<ResolvedExpr>, out: &mut HashSet<String>) {
match &expr.node {
ResolvedExpr::Call(callee, args) => {
if let ResolvedCallee::Builtin(name) = callee
&& name == "Vector.set"
&& args.len() == 3
{
let vec_aver = aver_type_str_of(&args[0]);
let canonical: String = vec_aver.chars().filter(|c| !c.is_whitespace()).collect();
if canonical.starts_with("Vector<") {
out.insert(canonical);
}
}
for a in args {
walk_expr_for_vector_set(a, out);
}
}
ResolvedExpr::BinOp(_, l, r) => {
walk_expr_for_vector_set(l, out);
walk_expr_for_vector_set(r, out);
}
ResolvedExpr::Neg(inner) => walk_expr_for_vector_set(inner, out),
ResolvedExpr::Match { subject, arms } => {
walk_expr_for_vector_set(subject, out);
for arm in arms {
walk_expr_for_vector_set(&arm.body, out);
}
}
ResolvedExpr::TailCall { args, .. } => {
for a in args {
walk_expr_for_vector_set(a, out);
}
}
ResolvedExpr::Attr(obj, _) => walk_expr_for_vector_set(obj, out),
ResolvedExpr::ErrorProp(inner) => walk_expr_for_vector_set(inner, out),
ResolvedExpr::Ctor(_, args) => {
for a in args {
walk_expr_for_vector_set(a, out);
}
}
ResolvedExpr::RecordCreate { fields, .. } => {
for (_, e) in fields {
walk_expr_for_vector_set(e, out);
}
}
ResolvedExpr::RecordUpdate { base, updates, .. } => {
walk_expr_for_vector_set(base, out);
for (_, e) in updates {
walk_expr_for_vector_set(e, out);
}
}
ResolvedExpr::List(items)
| ResolvedExpr::Tuple(items)
| ResolvedExpr::IndependentProduct(items, _) => {
for e in items {
walk_expr_for_vector_set(e, out);
}
}
ResolvedExpr::MapLiteral(entries) => {
for (k, v) in entries {
walk_expr_for_vector_set(k, out);
walk_expr_for_vector_set(v, out);
}
}
ResolvedExpr::InterpolatedStr(parts) => {
for p in parts {
if let ResolvedStrPart::Parsed(inner) = p {
walk_expr_for_vector_set(inner, out);
}
}
}
_ => {}
}
}