From 55662d88bb1e91f40f77b52adae7b6841a9eec44 Mon Sep 17 00:00:00 2001 From: kgtkr Date: Thu, 8 Jan 2026 15:40:47 +0900 Subject: [PATCH 01/15] =?UTF-8?q?Closure=E3=81=AB=E5=AE=9A=E6=95=B0?= =?UTF-8?q?=E6=83=85=E5=A0=B1=E3=82=92=E3=81=A4=E3=81=91=E3=82=8B=E3=82=88?= =?UTF-8?q?=E3=81=86=E3=81=AB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- webschembly-compiler-crates/ir/src/id.rs | 68 +++++++++++++++++++ webschembly-compiler-crates/ir/src/typ.rs | 28 +++++++- webschembly-compiler/src/compiler.rs | 16 ++--- .../src/ir_generator/module_generator.rs | 14 ++-- .../src/ir_processor/desugar.rs | 2 +- .../src/jit/bb_index_manager.rs | 3 - .../src/jit/closure_global_layout.rs | 15 ++-- .../src/jit/env_index_manager.rs | 15 ++-- webschembly-compiler/src/jit/jit_ctx.rs | 8 +-- webschembly-compiler/src/jit/jit_func.rs | 35 +++++----- webschembly-compiler/src/jit/jit_module.rs | 20 +++--- webschembly-compiler/src/jit/mod.rs | 19 +++--- .../src/wasm_generator/module_generator.rs | 10 +-- 13 files changed, 169 insertions(+), 84 deletions(-) diff --git a/webschembly-compiler-crates/ir/src/id.rs b/webschembly-compiler-crates/ir/src/id.rs index 277f2549..d36715a8 100644 --- a/webschembly-compiler-crates/ir/src/id.rs +++ b/webschembly-compiler-crates/ir/src/id.rs @@ -161,3 +161,71 @@ impl fmt::Display for Display<'_, JitBasicBlockId> { Ok(()) } } + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, From, Into)] +pub struct ClosureEnvIndex(pub usize); + +impl ClosureEnvIndex { + pub fn display<'a>(&self, meta: &'a Meta) -> Display<'a, ClosureEnvIndex> { + Display { value: *self, meta } + } +} + +impl fmt::Display for Display<'_, ClosureEnvIndex> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "closure_env{}", self.value.0) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, From, Into)] +pub struct ClosureArgIndex(pub usize); + +impl ClosureArgIndex { + pub fn display<'a>(&self, meta: &'a Meta) -> Display<'a, ClosureArgIndex> { + Display { value: *self, meta } + } +} + +impl fmt::Display for Display<'_, ClosureArgIndex> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "closure_arg{}", self.value.0) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, From, Into)] +pub struct BBIndex(pub usize); + +impl BBIndex { + pub fn display<'a>(&self, meta: &'a Meta) -> Display<'a, BBIndex> { + Display { value: *self, meta } + } +} + +impl fmt::Display for Display<'_, BBIndex> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "bb_index{}", self.value.0) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct ConstantClosure { + pub func_id: JitFuncId, + pub env_index: ClosureEnvIndex, +} + +impl ConstantClosure { + pub fn display<'a>(&self, meta: &'a Meta) -> Display<'a, ConstantClosure> { + Display { value: *self, meta } + } +} + +impl fmt::Display for Display<'_, ConstantClosure> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "constant_closure({}, {})", + self.value.func_id.display(self.meta), + self.value.env_index.display(self.meta) + ) + } +} diff --git a/webschembly-compiler-crates/ir/src/typ.rs b/webschembly-compiler-crates/ir/src/typ.rs index 7535a73f..9091b111 100644 --- a/webschembly-compiler-crates/ir/src/typ.rs +++ b/webschembly-compiler-crates/ir/src/typ.rs @@ -3,6 +3,8 @@ TypeもしくはmutableなTypeを表す Type自体にRefを含めて再帰的にしてしまうと無限種類の型を作れるようになってしまうので、IRではそれを避けるためこのような構造になっている TODO: LocalTypeという名前は適切ではない */ +use crate::id::ConstantClosure; + #[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, derive_more::Display)] pub enum LocalType { #[display("ref<{}>", _0)] @@ -38,6 +40,14 @@ impl LocalType { _ => None, } } + + pub fn remove_constant(self) -> Self { + match self { + LocalType::Ref(typ) => LocalType::Ref(typ.remove_constant()), + LocalType::Type(typ) => LocalType::Type(typ.remove_constant()), + _ => self, + } + } } #[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, derive_more::Display)] @@ -55,6 +65,13 @@ impl Type { Type::Obj => None, } } + + pub fn remove_constant(self) -> Self { + match self { + Type::Val(val_type) => Type::Val(val_type.remove_constant()), + Type::Obj => Type::Obj, + } + } } impl From for Type { @@ -88,7 +105,16 @@ pub enum ValType { #[display("uvector<{0}>", _0)] UVector(UVectorKind), #[display("closure")] - Closure, + Closure(Option), +} + +impl ValType { + pub fn remove_constant(self) -> Self { + match self { + ValType::Closure(_) => ValType::Closure(None), + _ => self, + } + } } #[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, derive_more::Display)] diff --git a/webschembly-compiler/src/compiler.rs b/webschembly-compiler/src/compiler.rs index f7b94e7f..937f908a 100644 --- a/webschembly-compiler/src/compiler.rs +++ b/webschembly-compiler/src/compiler.rs @@ -133,8 +133,8 @@ impl Compiler { &mut self.global_manager, module_id, func_id, - crate::jit::env_index_manager::EnvIndex(env_index), - crate::jit::closure_global_layout::ClosureIndex(func_index), + ir::ClosureEnvIndex(env_index), + ir::ClosureArgIndex(func_index), ); preprocess_module(&mut module); @@ -167,10 +167,10 @@ impl Compiler { let mut module = jit.instantiate_bb( module_id, func_id, - crate::jit::env_index_manager::EnvIndex(env_index), - crate::jit::closure_global_layout::ClosureIndex(func_index), + ir::ClosureEnvIndex(env_index), + ir::ClosureArgIndex(func_index), bb_id, - crate::jit::bb_index_manager::BBIndex(index), + ir::BBIndex(index), &mut self.global_manager, ); preprocess_module(&mut module); @@ -211,12 +211,12 @@ impl Compiler { &mut self.global_manager, module_id, func_id, - crate::jit::env_index_manager::EnvIndex(env_index), - crate::jit::closure_global_layout::ClosureIndex(func_index), + ir::ClosureEnvIndex(env_index), + ir::ClosureArgIndex(func_index), bb_id, kind, ir::BasicBlockId::from(source_bb_id), - crate::jit::bb_index_manager::BBIndex(source_index), + ir::BBIndex(source_index), ) .map(|mut module| { preprocess_module(&mut module); diff --git a/webschembly-compiler/src/ir_generator/module_generator.rs b/webschembly-compiler/src/ir_generator/module_generator.rs index f01608e2..b07b5e0c 100644 --- a/webschembly-compiler/src/ir_generator/module_generator.rs +++ b/webschembly-compiler/src/ir_generator/module_generator.rs @@ -266,7 +266,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { let bb_entry = self.builder.bbs.allocate_key(); self.builder.current_bb_id = Some(bb_entry); - let self_closure = self.builder.local(Type::Val(ValType::Closure)); + let self_closure = self.builder.local(Type::Val(ValType::Closure(None))); let args = self.builder.local(LocalType::VariadicArgs); let args_len_local = self.builder.local(Type::Val(ValType::Int)); let expected_args_len_local = self.builder.local(Type::Val(ValType::Int)); @@ -559,7 +559,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { ast::Expr::Lambda(x, lambda) => { let func_id = self.module_generator.gen_func(x, ast.span, lambda); let func_local = self.builder.local(LocalType::FuncRef); - let val_type_local = self.builder.local(Type::Val(ValType::Closure)); + let val_type_local = self.builder.local(Type::Val(ValType::Closure(None))); self.builder.exprs.push(Instr { local: Some(func_local), kind: InstrKind::FuncRef(func_id), @@ -617,7 +617,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { }); self.builder.exprs.push(Instr { local: result, - kind: InstrKind::ToObj(ValType::Closure, val_type_local), + kind: InstrKind::ToObj(ValType::Closure(None), val_type_local), }); } ast::Expr::If(_, ast::If { cond, then, els }) => { @@ -887,7 +887,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { let is_closure_local = self.builder.local(Type::Val(ValType::Bool)); self.builder.exprs.push(Instr { local: Some(is_closure_local), - kind: InstrKind::Is(ValType::Closure, obj_func_local), + kind: InstrKind::Is(ValType::Closure(None), obj_func_local), }); let then_bb_id = self.builder.bbs.allocate_key(); @@ -912,10 +912,10 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { self.builder.current_bb_id = Some(then_bb_id); - let closure_local = self.builder.local(ValType::Closure); + let closure_local = self.builder.local(ValType::Closure(None)); self.builder.exprs.push(Instr { local: Some(closure_local), - kind: InstrKind::FromObj(ValType::Closure, obj_func_local), + kind: InstrKind::FromObj(ValType::Closure(None), obj_func_local), }); let args_local = self.builder.local(LocalType::VariadicArgs); @@ -1449,7 +1449,7 @@ impl BuiltinConversionRule { ir_gen: |ctx, arg1| { ctx.builder.exprs.push(Instr { local: Some(ctx.dest), - kind: InstrKind::Is(ValType::Closure, arg1), + kind: InstrKind::Is(ValType::Closure(None), arg1), }); }, }], diff --git a/webschembly-compiler/src/ir_processor/desugar.rs b/webschembly-compiler/src/ir_processor/desugar.rs index 9ea73f51..2219bd54 100644 --- a/webschembly-compiler/src/ir_processor/desugar.rs +++ b/webschembly-compiler/src/ir_processor/desugar.rs @@ -92,7 +92,7 @@ fn desugar_call_closure( func_type: FuncType { args: { let mut args = Vec::new(); - args.push(ValType::Closure.into()); + args.push(ValType::Closure(None).into()); args.extend(call_closure.arg_types); args }, diff --git a/webschembly-compiler/src/jit/bb_index_manager.rs b/webschembly-compiler/src/jit/bb_index_manager.rs index 4b75d2e6..88ebd683 100644 --- a/webschembly-compiler/src/jit/bb_index_manager.rs +++ b/webschembly-compiler/src/jit/bb_index_manager.rs @@ -10,9 +10,6 @@ use super::index_flag::IndexFlag; pub const BB_LAYOUT_MAX_SIZE: usize = 32; pub const BB_LAYOUT_DEFAULT_INDEX: BBIndex = BBIndex(0); -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct BBIndex(pub usize); - #[derive(Debug)] pub struct BBIndexManager { type_params_to_index: FxBiHashMap, BBIndex>, diff --git a/webschembly-compiler/src/jit/closure_global_layout.rs b/webschembly-compiler/src/jit/closure_global_layout.rs index 759d9727..d8cfea6c 100644 --- a/webschembly-compiler/src/jit/closure_global_layout.rs +++ b/webschembly-compiler/src/jit/closure_global_layout.rs @@ -6,10 +6,7 @@ use crate::fxbihashmap::FxBiHashMap; use super::index_flag::IndexFlag; pub const CLOSURE_LAYOUT_MAX_SIZE: usize = 32; -pub const CLOSURE_LAYOUT_DEFAULT_INDEX: ClosureIndex = ClosureIndex(0); - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct ClosureIndex(pub usize); +pub const CLOSURE_LAYOUT_DEFAULT_INDEX: ClosureArgIndex = ClosureArgIndex(0); #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum ClosureArgs { @@ -19,8 +16,8 @@ pub enum ClosureArgs { #[derive(Debug)] pub struct ClosureGlobalLayout { - args_to_index: FxBiHashMap, - instantiated_idx: FxHashSet, + args_to_index: FxBiHashMap, + instantiated_idx: FxHashSet, } impl Default for ClosureGlobalLayout { @@ -39,12 +36,12 @@ impl ClosureGlobalLayout { } } - pub fn idx(&mut self, args: &ClosureArgs) -> Option<(ClosureIndex, IndexFlag)> { + pub fn idx(&mut self, args: &ClosureArgs) -> Option<(ClosureArgIndex, IndexFlag)> { // TODO: argsの長さに上限を設定 let index = if let Some(&index) = self.args_to_index.get_by_left(args) { index } else if self.args_to_index.len() < CLOSURE_LAYOUT_MAX_SIZE { - let index = ClosureIndex(self.args_to_index.len()); + let index = ClosureArgIndex(self.args_to_index.len()); self.args_to_index.insert(args.clone(), index); index } else { @@ -58,7 +55,7 @@ impl ClosureGlobalLayout { Some((index, flag)) } - pub fn arg_types(&self, index: ClosureIndex) -> &ClosureArgs { + pub fn arg_types(&self, index: ClosureArgIndex) -> &ClosureArgs { self.args_to_index.get_by_right(&index).unwrap() } } diff --git a/webschembly-compiler/src/jit/env_index_manager.rs b/webschembly-compiler/src/jit/env_index_manager.rs index 9d39e31d..8e924aca 100644 --- a/webschembly-compiler/src/jit/env_index_manager.rs +++ b/webschembly-compiler/src/jit/env_index_manager.rs @@ -8,15 +8,12 @@ use crate::ir_generator::GlobalManager; use super::index_flag::IndexFlag; pub const ENV_LAYOUT_MAX_SIZE: usize = 32; -pub const ENV_LAYOUT_DEFAULT_INDEX: EnvIndex = EnvIndex(0); - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct EnvIndex(pub usize); +pub const ENV_LAYOUT_DEFAULT_INDEX: ClosureEnvIndex = ClosureEnvIndex(0); #[derive(Debug)] pub struct EnvIndexManager { - env_types_to_index: FxBiHashMap, EnvIndex>, - index_to_table_global: FxHashMap, + env_types_to_index: FxBiHashMap, ClosureEnvIndex>, + index_to_table_global: FxHashMap, } impl Default for EnvIndexManager { @@ -40,7 +37,7 @@ impl EnvIndexManager { &mut self, env_types: &VecMap, global_manager: &mut GlobalManager, - ) -> Option<(Option, EnvIndex, IndexFlag)> { + ) -> Option<(Option, ClosureEnvIndex, IndexFlag)> { if let Some(&index) = self .env_types_to_index .get_by_left(VecMapEq::from_ref(env_types)) @@ -49,7 +46,7 @@ impl EnvIndexManager { debug_assert!(index == ENV_LAYOUT_DEFAULT_INDEX || global.is_some()); Some((global.copied(), index, IndexFlag::ExistingInstance)) } else if self.env_types_to_index.len() < ENV_LAYOUT_MAX_SIZE { - let index = EnvIndex(self.env_types_to_index.len()); + let index = ClosureEnvIndex(self.env_types_to_index.len()); self.env_types_to_index .insert(env_types.clone().into(), index); let global = global_manager.gen_global(LocalType::EntrypointTable); @@ -60,7 +57,7 @@ impl EnvIndexManager { } } - pub fn env_types(&self, index: EnvIndex) -> (&VecMap, Option) { + pub fn env_types(&self, index: ClosureEnvIndex) -> (&VecMap, Option) { debug_assert!( index == ENV_LAYOUT_DEFAULT_INDEX || self.index_to_table_global.contains_key(&index) ); diff --git a/webschembly-compiler/src/jit/jit_ctx.rs b/webschembly-compiler/src/jit/jit_ctx.rs index 37f12cba..f556bfd9 100644 --- a/webschembly-compiler/src/jit/jit_ctx.rs +++ b/webschembly-compiler/src/jit/jit_ctx.rs @@ -1,6 +1,6 @@ use rustc_hash::FxHashMap; -use super::closure_global_layout::{ClosureGlobalLayout, ClosureIndex}; +use super::closure_global_layout::ClosureGlobalLayout; use super::jit_config::JitConfig; use webschembly_compiler_ir::*; @@ -12,7 +12,7 @@ pub struct JitCtx { is_instantiated: bool, // 0..CLOSURE_LAYOUT_MAX_SIZEまでのindexに対応する関数のスタブが入ったMutFuncRef // func_indexがインスタンス化されるときにMutFuncRefにFuncRefがセットされる - stub_globals: FxHashMap, + stub_globals: FxHashMap, // instantiate_funcの結果を保存するグローバル instantiate_func_global: Option, } @@ -32,7 +32,7 @@ impl JitCtx { self.config } - pub fn stub_global(&self, index: ClosureIndex) -> Global { + pub fn stub_global(&self, index: ClosureArgIndex) -> Global { debug_assert!(self.is_instantiated); self.stub_globals[&index] } @@ -48,7 +48,7 @@ impl JitCtx { pub fn init_instantiated( &mut self, - stub_globals: FxHashMap, + stub_globals: FxHashMap, instantiate_func_global: Global, ) { debug_assert!(!self.is_instantiated); diff --git a/webschembly-compiler/src/jit/jit_func.rs b/webschembly-compiler/src/jit/jit_func.rs index a985c3ec..93419c2f 100644 --- a/webschembly-compiler/src/jit/jit_func.rs +++ b/webschembly-compiler/src/jit/jit_func.rs @@ -1,11 +1,10 @@ use rustc_hash::{FxHashMap, FxHashSet}; -use super::bb_index_manager::{BB_LAYOUT_DEFAULT_INDEX, BBIndex, BBIndexManager}; +use super::bb_index_manager::{BB_LAYOUT_DEFAULT_INDEX, BBIndexManager}; use super::closure_global_layout::{ CLOSURE_LAYOUT_DEFAULT_INDEX, CLOSURE_LAYOUT_MAX_SIZE, ClosureArgs, ClosureGlobalLayout, - ClosureIndex, }; -use super::env_index_manager::{ENV_LAYOUT_DEFAULT_INDEX, EnvIndex, EnvIndexManager}; +use super::env_index_manager::{ENV_LAYOUT_DEFAULT_INDEX, EnvIndexManager}; use super::index_flag::IndexFlag; use super::jit_ctx::JitCtx; use crate::fxbihashmap::FxBiHashMap; @@ -20,7 +19,7 @@ use webschembly_compiler_ir::*; #[derive(Debug)] pub struct JitFunc { - pub jit_specialized_env_funcs: FxHashMap, + pub jit_specialized_env_funcs: FxHashMap, } impl JitFunc { @@ -49,7 +48,7 @@ impl JitFunc { #[derive(Debug)] pub struct JitSpecializedEnvFunc { - pub jit_specialized_arg_funcs: FxHashMap, + pub jit_specialized_arg_funcs: FxHashMap, pub func: Func, } @@ -58,7 +57,7 @@ impl JitSpecializedEnvFunc { module_id: JitModuleId, global_manager: &mut GlobalManager, func: &Func, - env_index: EnvIndex, + env_index: ClosureEnvIndex, env_index_manager: &EnvIndexManager, jit_ctx: &mut JitCtx, ) -> Self { @@ -84,8 +83,8 @@ impl JitSpecializedEnvFunc { #[derive(Debug)] pub struct JitSpecializedArgFunc { module_id: JitModuleId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, func: Func, jit_bbs: VecMap, } @@ -95,8 +94,8 @@ impl JitSpecializedArgFunc { module_id: JitModuleId, global_manager: &mut GlobalManager, func: &Func, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, jit_ctx: &mut JitCtx, ) -> Self { let mut func = func.clone(); @@ -722,7 +721,7 @@ impl JitSpecializedArgFunc { instrs.push(Instr { local: Some(stub), kind: InstrKind::GlobalGet( - jit_ctx.stub_global(ClosureIndex(index)).id, + jit_ctx.stub_global(ClosureArgIndex(index)).id, ), }); locals.push(stub); @@ -916,7 +915,7 @@ impl JitSpecializedArgFunc { let mut args = Vec::new(); let closure_local = locals.push_with(|id| Local { id, - typ: ValType::Closure.into(), + typ: ValType::Closure(None).into(), }); let mut arg_locals = Vec::new(); args.push(closure_local); @@ -1074,7 +1073,7 @@ impl JitSpecializedArgFunc { }); instrs.push(Instr { local: Some(stub), - kind: InstrKind::GlobalGet(jit_ctx.stub_global(ClosureIndex(index)).id), + kind: InstrKind::GlobalGet(jit_ctx.stub_global(ClosureArgIndex(index)).id), }); entrypoint_table_locals.push(stub); } @@ -1199,7 +1198,7 @@ fn specialize_call_closure( def_use_chain: &DefUseChain, bbs: &VecMap, closure_global_layout: &mut ClosureGlobalLayout, - required_closure_idx: &mut Vec, + required_closure_idx: &mut Vec, ) -> Option { if call_closure.func_index != CLOSURE_LAYOUT_DEFAULT_INDEX.0 { return None; @@ -1433,7 +1432,7 @@ fn calculate_args_to_pass( fn closure_func_assign_env_types( func: &mut Func, - env_index: EnvIndex, + env_index: ClosureEnvIndex, env_index_manager: &EnvIndexManager, ) { if env_index == ENV_LAYOUT_DEFAULT_INDEX { @@ -1463,7 +1462,7 @@ fn closure_func_assign_env_types( } let new_closure_arg = func.locals.push_with(|id| Local { id, - typ: LocalType::Type(Type::Val(ValType::Closure)), + typ: LocalType::Type(Type::Val(ValType::Closure(None))), }); let prev_closure_arg = func.args[0]; func.args[0] = new_closure_arg; @@ -1527,7 +1526,7 @@ fn closure_func_assign_env_types( fn closure_func_assign_types( func: &mut Func, - func_index: ClosureIndex, + func_index: ClosureArgIndex, closure_global_layout: &ClosureGlobalLayout, ) { let ClosureArgs::Specified(args_type) = closure_global_layout.arg_types(func_index) else { @@ -1541,7 +1540,7 @@ fn closure_func_assign_types( .map(|&arg| func.locals[arg].typ) .collect::>(), vec![ - LocalType::Type(Type::Val(ValType::Closure)), + LocalType::Type(Type::Val(ValType::Closure(None))), LocalType::VariadicArgs ] ); diff --git a/webschembly-compiler/src/jit/jit_module.rs b/webschembly-compiler/src/jit/jit_module.rs index ded0d1f8..312bc5b5 100644 --- a/webschembly-compiler/src/jit/jit_module.rs +++ b/webschembly-compiler/src/jit/jit_module.rs @@ -1,13 +1,13 @@ use rustc_hash::FxHashMap; -use super::bb_index_manager::BBIndex; -use super::closure_global_layout::{CLOSURE_LAYOUT_MAX_SIZE, ClosureIndex}; -use super::env_index_manager::{EnvIndex, EnvIndexManager}; +use super::closure_global_layout::CLOSURE_LAYOUT_MAX_SIZE; +use super::env_index_manager::EnvIndexManager; use super::jit_ctx::JitCtx; use super::jit_func::{JitFunc, JitSpecializedArgFunc}; use crate::ir_generator::GlobalManager; use crate::jit::jit_func::JitSpecializedEnvFunc; use vec_map::{HasId, VecMap}; +use webschembly_compiler_ir::{BBIndex, ClosureArgIndex, ClosureEnvIndex}; use webschembly_compiler_ir::*; #[derive(Debug)] pub struct JitModule { @@ -170,7 +170,7 @@ impl JitModule { if !jit_ctx.is_instantiated() { let mut stub_globals = FxHashMap::default(); for func_index in 0..CLOSURE_LAYOUT_MAX_SIZE { - let func_index = ClosureIndex(func_index); + let func_index = ClosureArgIndex(func_index); let stub_global = global_manager.gen_global(LocalType::MutFuncRef); stub_globals.insert(func_index, stub_global); let stub_local = entry_func.locals.push_with(|id| Local { @@ -217,8 +217,8 @@ impl JitModule { &mut self, global_manager: &mut GlobalManager, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, jit_ctx: &mut JitCtx, ) -> Module { let jit_func_entry = self.jit_funcs.get_mut(&func_id).unwrap(); @@ -264,8 +264,8 @@ impl JitModule { pub fn instantiate_bb( &mut self, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, bb_id: BasicBlockId, index: BBIndex, global_manager: &mut GlobalManager, @@ -297,8 +297,8 @@ impl JitModule { global_manager: &mut GlobalManager, jit_ctx: &mut JitCtx, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, bb_id: BasicBlockId, kind: BranchKind, source_bb_id: BasicBlockId, diff --git a/webschembly-compiler/src/jit/mod.rs b/webschembly-compiler/src/jit/mod.rs index a0879b2a..816a32f3 100644 --- a/webschembly-compiler/src/jit/mod.rs +++ b/webschembly-compiler/src/jit/mod.rs @@ -1,17 +1,17 @@ use crate::ir_generator::GlobalManager; use vec_map::VecMap; -use bb_index_manager::BBIndex; -use closure_global_layout::ClosureIndex; -use env_index_manager::EnvIndex; + mod jit_config; mod jit_module; pub use jit_config::JitConfig; use jit_module::JitModule; mod jit_ctx; use jit_ctx::JitCtx; +use webschembly_compiler_ir::{BBIndex, ClosureArgIndex, ClosureEnvIndex}; use webschembly_compiler_ir::*; pub mod bb_index_manager; + pub mod closure_global_layout; pub mod env_index_manager; pub mod index_flag; @@ -51,12 +51,13 @@ impl Jit { global_manager: &mut GlobalManager, module_id: JitModuleId, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, ) -> Module { self.jit_module[module_id].instantiate_func( global_manager, func_id, + env_index, func_index, &mut self.ctx, @@ -67,8 +68,8 @@ impl Jit { &mut self, module_id: JitModuleId, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, bb_id: BasicBlockId, index: BBIndex, global_manager: &mut GlobalManager, @@ -89,8 +90,8 @@ impl Jit { global_manager: &mut GlobalManager, module_id: JitModuleId, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, bb_id: BasicBlockId, kind: BranchKind, source_bb_id: BasicBlockId, diff --git a/webschembly-compiler/src/wasm_generator/module_generator.rs b/webschembly-compiler/src/wasm_generator/module_generator.rs index a91cb29b..97659278 100644 --- a/webschembly-compiler/src/wasm_generator/module_generator.rs +++ b/webschembly-compiler/src/wasm_generator/module_generator.rs @@ -931,7 +931,7 @@ impl<'a> ModuleGenerator<'a> { nullable: true, heap_type: HeapType::Concrete(self.cons_type), }), - ir::ValType::Closure => ValType::Ref(RefType { + ir::ValType::Closure(_) => ValType::Ref(RefType { nullable: true, heap_type: HeapType::Concrete(self.closure_type), }), @@ -979,7 +979,7 @@ impl<'a> ModuleGenerator<'a> { ir::LocalType::Type(ir::Type::Val(ir::ValType::Cons)) => { Instruction::RefNull(HeapType::Concrete(self.cons_type)) } - ir::LocalType::Type(ir::Type::Val(ir::ValType::Closure)) => { + ir::LocalType::Type(ir::Type::Val(ir::ValType::Closure(_))) => { Instruction::RefNull(HeapType::Concrete(self.closure_type)) } ir::LocalType::Type(ir::Type::Val(ir::ValType::Vector)) => { @@ -1488,7 +1488,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { self.module_generator.cons_type, ))); } - ir::ValType::Closure => { + ir::ValType::Closure(_) => { function.instruction(&Instruction::LocalGet(self.local_id_to_idx(*val))); function.instruction(&Instruction::RefCastNonNull(HeapType::Concrete( self.module_generator.closure_type, @@ -1551,7 +1551,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { ir::ValType::Cons => { function.instruction(&Instruction::LocalGet(self.local_id_to_idx(*val))); } - ir::ValType::Closure => { + ir::ValType::Closure(_) => { function.instruction(&Instruction::LocalGet(self.local_id_to_idx(*val))); } ir::ValType::Vector => { @@ -1699,7 +1699,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { ir::ValType::Symbol => self.module_generator.symbol_type, ir::ValType::Nil => self.module_generator.nil_type, ir::ValType::Cons => self.module_generator.cons_type, - ir::ValType::Closure => self.module_generator.closure_type, + ir::ValType::Closure(_) => self.module_generator.closure_type, ir::ValType::Vector => self.module_generator.vector_type, ir::ValType::UVector(kind) => { self.module_generator.uvector_kind_to_type_idx(*kind) From a35a4e73c68fc8716c8f62090e437540c5e6513b Mon Sep 17 00:00:00 2001 From: kgtkr Date: Thu, 8 Jan 2026 15:52:53 +0900 Subject: [PATCH 02/15] remove_constant --- webschembly-compiler/src/compiler.rs | 2 ++ .../src/ir_generator/module_generator.rs | 15 ++++++-- webschembly-compiler/src/ir_processor/mod.rs | 1 + .../src/ir_processor/remove_constant.rs | 35 +++++++++++++++++++ 4 files changed, 51 insertions(+), 2 deletions(-) create mode 100644 webschembly-compiler/src/ir_processor/remove_constant.rs diff --git a/webschembly-compiler/src/compiler.rs b/webschembly-compiler/src/compiler.rs index 937f908a..5e34a288 100644 --- a/webschembly-compiler/src/compiler.rs +++ b/webschembly-compiler/src/compiler.rs @@ -273,6 +273,8 @@ fn postprocess(module: &mut ir::Module, global_manager: &mut GlobalManager) { register_allocation(func); remove_unused_local(func); + + crate::ir_processor::remove_constant::remove_constant(func); } // モジュールごとにグローバルを真面目に管理するのは大変なのでここで計算 diff --git a/webschembly-compiler/src/ir_generator/module_generator.rs b/webschembly-compiler/src/ir_generator/module_generator.rs index b07b5e0c..7caeef8a 100644 --- a/webschembly-compiler/src/ir_generator/module_generator.rs +++ b/webschembly-compiler/src/ir_generator/module_generator.rs @@ -559,7 +559,12 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { ast::Expr::Lambda(x, lambda) => { let func_id = self.module_generator.gen_func(x, ast.span, lambda); let func_local = self.builder.local(LocalType::FuncRef); - let val_type_local = self.builder.local(Type::Val(ValType::Closure(None))); + let val_type_local = self.builder.local(Type::Val(ValType::Closure(Some( + ConstantClosure { + func_id: JitFuncId::from(func_id), + env_index: ClosureEnvIndex(0), + }, + )))); self.builder.exprs.push(Instr { local: Some(func_local), kind: InstrKind::FuncRef(func_id), @@ -617,7 +622,13 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { }); self.builder.exprs.push(Instr { local: result, - kind: InstrKind::ToObj(ValType::Closure(None), val_type_local), + kind: InstrKind::ToObj( + ValType::Closure(Some(ConstantClosure { + func_id: JitFuncId::from(func_id), + env_index: ClosureEnvIndex(0), + })), + val_type_local, + ), }); } ast::Expr::If(_, ast::If { cond, then, els }) => { diff --git a/webschembly-compiler/src/ir_processor/mod.rs b/webschembly-compiler/src/ir_processor/mod.rs index c9a2d875..01a0953d 100644 --- a/webschembly-compiler/src/ir_processor/mod.rs +++ b/webschembly-compiler/src/ir_processor/mod.rs @@ -5,3 +5,4 @@ pub mod optimizer; pub mod register_allocation; pub mod ssa; pub mod ssa_optimizer; +pub mod remove_constant; diff --git a/webschembly-compiler/src/ir_processor/remove_constant.rs b/webschembly-compiler/src/ir_processor/remove_constant.rs new file mode 100644 index 00000000..519ba6c6 --- /dev/null +++ b/webschembly-compiler/src/ir_processor/remove_constant.rs @@ -0,0 +1,35 @@ +use webschembly_compiler_ir::Func; +use webschembly_compiler_ir::InstrKind; + +pub fn remove_constant(func: &mut Func) { + for local in func.locals.values_mut() { + local.typ = local.typ.remove_constant(); + } + + func.ret_type = func.ret_type.remove_constant(); + + if let Some(meta) = &mut func.closure_meta { + for env_type in &mut meta.env_types { + *env_type = env_type.remove_constant(); + } + } + + for bb in func.bbs.values_mut() { + for instr in &mut bb.instrs { + match &mut instr.kind { + InstrKind::CreateRef(typ) | InstrKind::DerefRef(typ, _) | InstrKind::SetRef(typ, _, _) => { + *typ = typ.remove_constant(); + } + InstrKind::ToObj(val_type, _) | InstrKind::FromObj(val_type, _) | InstrKind::Is(val_type, _) => { + *val_type = val_type.remove_constant(); + } + InstrKind::Closure { env_types, .. } | InstrKind::ClosureEnv(env_types, _, _) | InstrKind::ClosureSetEnv(env_types, _, _, _) => { + for env_type in env_types { + *env_type = env_type.remove_constant(); + } + } + _ => {} + } + } + } +} From c90e91b60522a63b9406c84b283a1a75b30ca4a6 Mon Sep 17 00:00:00 2001 From: kgtkr Date: Thu, 8 Jan 2026 17:43:01 +0900 Subject: [PATCH 03/15] fix bug --- .../src/ir_generator/module_generator.rs | 5 +--- .../src/ir_processor/remove_constant.rs | 25 ------------------- 2 files changed, 1 insertion(+), 29 deletions(-) diff --git a/webschembly-compiler/src/ir_generator/module_generator.rs b/webschembly-compiler/src/ir_generator/module_generator.rs index 7caeef8a..4d7c6e91 100644 --- a/webschembly-compiler/src/ir_generator/module_generator.rs +++ b/webschembly-compiler/src/ir_generator/module_generator.rs @@ -623,10 +623,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { self.builder.exprs.push(Instr { local: result, kind: InstrKind::ToObj( - ValType::Closure(Some(ConstantClosure { - func_id: JitFuncId::from(func_id), - env_index: ClosureEnvIndex(0), - })), + ValType::Closure(None), val_type_local, ), }); diff --git a/webschembly-compiler/src/ir_processor/remove_constant.rs b/webschembly-compiler/src/ir_processor/remove_constant.rs index 519ba6c6..4d0d8788 100644 --- a/webschembly-compiler/src/ir_processor/remove_constant.rs +++ b/webschembly-compiler/src/ir_processor/remove_constant.rs @@ -7,29 +7,4 @@ pub fn remove_constant(func: &mut Func) { } func.ret_type = func.ret_type.remove_constant(); - - if let Some(meta) = &mut func.closure_meta { - for env_type in &mut meta.env_types { - *env_type = env_type.remove_constant(); - } - } - - for bb in func.bbs.values_mut() { - for instr in &mut bb.instrs { - match &mut instr.kind { - InstrKind::CreateRef(typ) | InstrKind::DerefRef(typ, _) | InstrKind::SetRef(typ, _, _) => { - *typ = typ.remove_constant(); - } - InstrKind::ToObj(val_type, _) | InstrKind::FromObj(val_type, _) | InstrKind::Is(val_type, _) => { - *val_type = val_type.remove_constant(); - } - InstrKind::Closure { env_types, .. } | InstrKind::ClosureEnv(env_types, _, _) | InstrKind::ClosureSetEnv(env_types, _, _, _) => { - for env_type in env_types { - *env_type = env_type.remove_constant(); - } - } - _ => {} - } - } - } } From 026f6a3678ec4a4c00d8026d18e47cea9d983bb7 Mon Sep 17 00:00:00 2001 From: kgtkr Date: Thu, 8 Jan 2026 18:24:30 +0900 Subject: [PATCH 04/15] propagate_types --- webschembly-compiler/src/ir_processor/mod.rs | 1 + .../src/ir_processor/propagate_types.rs | 464 ++++++++++++++++++ .../src/ir_processor/remove_constant.rs | 1 - .../src/ir_processor/ssa_optimizer.rs | 5 +- webschembly-compiler/src/jit/jit_func.rs | 12 +- 5 files changed, 480 insertions(+), 3 deletions(-) create mode 100644 webschembly-compiler/src/ir_processor/propagate_types.rs diff --git a/webschembly-compiler/src/ir_processor/mod.rs b/webschembly-compiler/src/ir_processor/mod.rs index 01a0953d..9f1b1e5a 100644 --- a/webschembly-compiler/src/ir_processor/mod.rs +++ b/webschembly-compiler/src/ir_processor/mod.rs @@ -6,3 +6,4 @@ pub mod register_allocation; pub mod ssa; pub mod ssa_optimizer; pub mod remove_constant; +pub mod propagate_types; diff --git a/webschembly-compiler/src/ir_processor/propagate_types.rs b/webschembly-compiler/src/ir_processor/propagate_types.rs new file mode 100644 index 00000000..ae2028aa --- /dev/null +++ b/webschembly-compiler/src/ir_processor/propagate_types.rs @@ -0,0 +1,464 @@ +use crate::ir_processor::cfg_analyzer::calculate_rpo; +use vec_map::VecMap; +use webschembly_compiler_ir::*; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LatticeValue { + Top, + Constant(ConstantClosure), + Bottom, +} + +impl LatticeValue { + fn meet(self, other: LatticeValue) -> LatticeValue { + match (self, other) { + (LatticeValue::Top, x) | (x, LatticeValue::Top) => x, + (LatticeValue::Constant(a), LatticeValue::Constant(b)) => { + if a == b { + LatticeValue::Constant(a) + } else { + LatticeValue::Bottom + } + } + (LatticeValue::Bottom, _) | (_, LatticeValue::Bottom) => LatticeValue::Bottom, + } + } +} + +pub fn propagate_types(func: &mut Func) { + let mut lattice = VecMap::new(); + // Initialize args to Bottom (we don't know anything about them) + // Other locals are implicitly Top (not in map or handled as default) + for &arg in &func.args { + lattice.insert(arg, LatticeValue::Bottom); + } + + // Initialize all other locals to Top + for local_id in func.locals.keys() { + if !lattice.contains_key(local_id) { + lattice.insert(local_id, LatticeValue::Top); + } + } + + let rpo = calculate_rpo(&func.bbs, func.bb_entry); + + // Convert RPO map to a list of BBs sorted by RPO index + let mut sorted_bbs: Vec = rpo.keys().cloned().collect(); + sorted_bbs.sort_by_key(|bb| rpo.get(bb).unwrap()); + + let mut changed = true; + while changed { + changed = false; + + for &bb_id in &sorted_bbs { + let bb = &func.bbs[bb_id]; + for instr in &bb.instrs { + if let Some(dest) = instr.local { + let old_val = *lattice.get(dest).unwrap_or(&LatticeValue::Top); + let new_val = match &instr.kind { + InstrKind::Closure { func_id, env_index, .. } => { + LatticeValue::Constant(ConstantClosure { + func_id: *func_id, + env_index: ClosureEnvIndex(*env_index), + }) + } + InstrKind::Move(src) => { + *lattice.get(*src).unwrap_or(&LatticeValue::Top) + } + InstrKind::Phi { incomings, non_exhaustive } => { + if *non_exhaustive { + LatticeValue::Bottom + } else { + let mut val = LatticeValue::Top; + for incoming in incomings { + let incoming_val = *lattice.get(incoming.local).unwrap_or(&LatticeValue::Top); + val = val.meet(incoming_val); + } + val + } + } + // If we have a direct reference to a known constant logic in other instructions + // we might want to propagate that, but for now we only care about propagating definitions. + // Any other instruction that defines a value produces Bottom for that value (unknown type/value) + _ => LatticeValue::Bottom, + }; + + if old_val != new_val { + lattice.insert(dest, new_val); + changed = true; + } + } + } + } + } + + // Update types based on lattice results + for (local_id, val) in lattice { + match val { + LatticeValue::Constant(constant) => { + let local = &mut func.locals[local_id]; + match local.typ { + LocalType::Type(Type::Val(ValType::Closure(_))) => { + local.typ = LocalType::Type(Type::Val(ValType::Closure(Some(constant)))); + } + // If it was Type::Obj (upcasted), we can refine it to Closure + LocalType::Type(Type::Obj) => { + local.typ = LocalType::Type(Type::Val(ValType::Closure(Some(constant)))); + } + _ => {} + } + } + // If Bottom, we might need to revert to generic Closure(None) if it was previously set to something specific + // but here we are primarily taking generic Code (Closure(None)) and refining it. + // If the code was *already* specialized/refined, we should be careful not to overwrite it with less specific info + // unless we are sure. But this pass is usually run to refine. + LatticeValue::Bottom => { + // Ensure if we failed to prove constant, it remains as generic closure if it was a closure + /* + Note: We don't generally downgrade types here explicitly because the lattice implies + we only UPGRADE to Constant or stay Bottom. + However, if this pass is run multiple times or if we have partial information, + we might want to be safe. But currently the IR starts with Closure(None) usually. + */ + } + LatticeValue::Top => { + // Dead code or uninitialized. Leave as is or set to Bottom equivalent? + // Usually safe to ignore. + } + } + } +} + + +#[cfg(test)] +mod tests { + use super::*; + use webschembly_compiler_ir::{Func, BasicBlock, Instr, InstrKind, Local, Type, ValType, TerminatorInstr, ExitInstr}; + + fn create_dummy_func() -> Func { + Func { + id: FuncId::from(0), + args: vec![], + locals: VecMap::new(), + bbs: VecMap::new(), + bb_entry: BasicBlockId::from(0), + ret_type: LocalType::Type(Type::Val(ValType::Nil)), + closure_meta: None, + } + } + + #[test] + fn test_propagate_simple_move() { + let mut func = create_dummy_func(); + let l0 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l1 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + + let bb0 = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![ + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // Target ID + entrypoint_table: l1, // Dummy + original_entrypoint_table: l1, // Dummy + }, + }, + Instr { + local: Some(l1), + kind: InstrKind::Move(l0), + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l1)).into()) + } + ], + }); + func.bb_entry = bb0; + + propagate_types(&mut func); + + let t1 = func.locals[l1].typ; + if let LocalType::Type(Type::Val(ValType::Closure(Some(c)))) = t1 { + assert_eq!(c.func_id, JitFuncId::from(10)); + assert_eq!(c.env_index, ClosureEnvIndex(0)); + } else { + panic!("Expected ConstantClosure, got {:?}", t1); + } + } + + #[test] + fn test_propagate_phi_same() { + let mut func = create_dummy_func(); + // l0 and l1 are same closure constant + let l0 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l1 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l2 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); // Phi result + + /* + bb0: + l0 = closure(10, 0) + jump bb1 + bb1: + l1 = closure(10, 0) + jump bb2 + bb2: + l2 = phi(l0: bb0, l1: bb1) + */ + + let bb0 = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![], // Fill later + }); + let bb1 = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![], // Fill later + }); + let bb2 = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![], // Fill later + }); + + func.bbs[bb0].instrs = vec![ + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + } + ]; + + func.bbs[bb1].instrs = vec![ + Instr { + local: Some(l1), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + } + ]; + + func.bbs[bb2].instrs = vec![ + Instr { + local: Some(l2), + kind: InstrKind::Phi { + incomings: vec![ + PhiIncomingValue { bb: bb0, local: l0 }, + PhiIncomingValue { bb: bb1, local: l1 }, + ], + non_exhaustive: false, + } + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))), + } + ]; + + // entry -> bb0 (cond jump is hard so just make one path reachable for simple test structure or use Jump from entry) + // Let's make entry jump to bb0 and bb1? No, just linear flow is fine? + // Wait, for Phi to work we need valid preds. + // Let's make entry conditional jump to bb0 and bb1. + let l_cond = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Bool)) }); + func.bb_entry = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![ + Instr { + local: Some(l_cond), + kind: InstrKind::Bool(true), + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)) + } + ] + }); + + propagate_types(&mut func); + + let t2 = func.locals[l2].typ; + if let LocalType::Type(Type::Val(ValType::Closure(Some(c)))) = t2 { + assert_eq!(c.func_id, JitFuncId::from(10)); + } else { + panic!("Expected ConstantClosure for l2, got {:?}", t2); + } + } + + #[test] + fn test_propagate_phi_different() { + let mut func = create_dummy_func(); + let l0 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l1 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l2 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + + let bb0 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + let bb1 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + let bb2 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + + func.bbs[bb0].instrs = vec![ + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // ID 10 + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)) } + ]; + + func.bbs[bb1].instrs = vec![ + Instr { + local: Some(l1), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(11), // ID 11 (Different) + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)) } + ]; + + func.bbs[bb2].instrs = vec![ + Instr { + local: Some(l2), + kind: InstrKind::Phi { + incomings: vec![ + PhiIncomingValue { bb: bb0, local: l0 }, + PhiIncomingValue { bb: bb1, local: l1 }, + ], + non_exhaustive: false, + } + }, + Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))) } + ]; + + let l_cond = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Bool)) }); + func.bb_entry = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![ + Instr { local: Some(l_cond), kind: InstrKind::Bool(true) }, + Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)) } + ] + }); + + propagate_types(&mut func); + + let t2 = func.locals[l2].typ; + // Should remain generic Closure(None) because input constants are different + if let LocalType::Type(Type::Val(ValType::Closure(None))) = t2 { + // OK + } else { + panic!("Expected generic Closure(None) (Bottom) for l2, got {:?}", t2); + } + } + + #[test] + fn test_propagate_phi_non_exhaustive() { + let mut func = create_dummy_func(); + let l0 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l1 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l2 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + + let bb0 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + let bb1 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + let bb2 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + + func.bbs[bb0].instrs = vec![ + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // ID 10 + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)) } + ]; + + func.bbs[bb1].instrs = vec![ + Instr { + local: Some(l1), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // ID 10 (Same as l0) + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)) } + ]; + + func.bbs[bb2].instrs = vec![ + Instr { + local: Some(l2), + kind: InstrKind::Phi { + incomings: vec![ + PhiIncomingValue { bb: bb0, local: l0 }, + PhiIncomingValue { bb: bb1, local: l1 }, + ], + // Even though l0 and l1 are same constant, non_exhaustive should make result Bottom + non_exhaustive: true, + } + }, + Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))) } + ]; + + let l_cond = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Bool)) }); + func.bb_entry = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![ + Instr { local: Some(l_cond), kind: InstrKind::Bool(true) }, + Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)) } + ] + }); + + propagate_types(&mut func); + + let t2 = func.locals[l2].typ; + // Should remain generic Closure(None) because non_exhaustive is true + if let LocalType::Type(Type::Val(ValType::Closure(None))) = t2 { + // OK + } else { + panic!("Expected generic Closure(None) (Bottom) for l2 due to non_exhaustive, got {:?}", t2); + } + } +} diff --git a/webschembly-compiler/src/ir_processor/remove_constant.rs b/webschembly-compiler/src/ir_processor/remove_constant.rs index 4d0d8788..584a5483 100644 --- a/webschembly-compiler/src/ir_processor/remove_constant.rs +++ b/webschembly-compiler/src/ir_processor/remove_constant.rs @@ -1,5 +1,4 @@ use webschembly_compiler_ir::Func; -use webschembly_compiler_ir::InstrKind; pub fn remove_constant(func: &mut Func) { for local in func.locals.values_mut() { diff --git a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs index 0b6f2996..44e1b6a2 100644 --- a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs +++ b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs @@ -3,7 +3,8 @@ use rustc_hash::{FxHashMap, FxHashSet}; use crate::ir_processor::{ cfg_analyzer::{DomTreeNode, build_dom_tree, calc_doms, calc_predecessors, calculate_rpo}, optimizer::remove_unreachable_bb, - ssa::{DefUseChain, debug_assert_ssa}, + propagate_types::propagate_types, + ssa::{debug_assert_ssa, DefUseChain}, }; use vec_map::VecMap; use webschembly_compiler_ir::*; @@ -462,6 +463,8 @@ pub fn ssa_optimize(func: &mut Func, config: SsaOptimizerConfig) { let doms = calc_doms(&func.bbs, &rpo, func.bb_entry, &predecessors); let dom_tree = build_dom_tree(&func.bbs, &rpo, func.bb_entry, &doms); + propagate_types(func); + for _ in 0..config.iterations { debug_assert_ssa(func); copy_propagation(func, &rpo); diff --git a/webschembly-compiler/src/jit/jit_func.rs b/webschembly-compiler/src/jit/jit_func.rs index 93419c2f..73a572d3 100644 --- a/webschembly-compiler/src/jit/jit_func.rs +++ b/webschembly-compiler/src/jit/jit_func.rs @@ -868,6 +868,14 @@ impl JitSpecializedArgFunc { local: Some(entrypoint_table_local), kind: InstrKind::GlobalGet(entrypoint_table_global.id), }); + if let Some(local) = body_func.bbs[bb_id].instrs[instr_idx].local { + body_func.locals[local].typ = + (ValType::Closure(Some(ConstantClosure { + func_id: *func_id, + env_index: ClosureEnvIndex(env_index.0), + }))) + .into(); + } body_func.bbs[bb_id].instrs[instr_idx].kind = InstrKind::Closure { envs: new_envs, env_types: new_env_types.clone(), @@ -1073,7 +1081,9 @@ impl JitSpecializedArgFunc { }); instrs.push(Instr { local: Some(stub), - kind: InstrKind::GlobalGet(jit_ctx.stub_global(ClosureArgIndex(index)).id), + kind: InstrKind::GlobalGet( + jit_ctx.stub_global(ClosureArgIndex(index)).id, + ), }); entrypoint_table_locals.push(stub); } From a2656a8ae2cb5d70c780bd0fec81d581125f2786 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 8 Jan 2026 07:56:15 +0000 Subject: [PATCH 05/15] chore: auto-fix format and lint --- .../src/ir_generator/module_generator.rs | 12 ++++++------ webschembly-compiler/src/ir_processor/mod.rs | 4 ++-- webschembly-compiler/src/jit/jit_module.rs | 2 +- webschembly-compiler/src/jit/mod.rs | 4 +--- 4 files changed, 10 insertions(+), 12 deletions(-) diff --git a/webschembly-compiler/src/ir_generator/module_generator.rs b/webschembly-compiler/src/ir_generator/module_generator.rs index 4d7c6e91..0d78f17a 100644 --- a/webschembly-compiler/src/ir_generator/module_generator.rs +++ b/webschembly-compiler/src/ir_generator/module_generator.rs @@ -559,12 +559,12 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { ast::Expr::Lambda(x, lambda) => { let func_id = self.module_generator.gen_func(x, ast.span, lambda); let func_local = self.builder.local(LocalType::FuncRef); - let val_type_local = self.builder.local(Type::Val(ValType::Closure(Some( - ConstantClosure { - func_id: JitFuncId::from(func_id), - env_index: ClosureEnvIndex(0), - }, - )))); + let val_type_local = + self.builder + .local(Type::Val(ValType::Closure(Some(ConstantClosure { + func_id: JitFuncId::from(func_id), + env_index: ClosureEnvIndex(0), + })))); self.builder.exprs.push(Instr { local: Some(func_local), kind: InstrKind::FuncRef(func_id), diff --git a/webschembly-compiler/src/ir_processor/mod.rs b/webschembly-compiler/src/ir_processor/mod.rs index 9f1b1e5a..2cf14d5f 100644 --- a/webschembly-compiler/src/ir_processor/mod.rs +++ b/webschembly-compiler/src/ir_processor/mod.rs @@ -2,8 +2,8 @@ pub mod cfg_analyzer; pub mod dataflow; pub mod desugar; pub mod optimizer; +pub mod propagate_types; pub mod register_allocation; +pub mod remove_constant; pub mod ssa; pub mod ssa_optimizer; -pub mod remove_constant; -pub mod propagate_types; diff --git a/webschembly-compiler/src/jit/jit_module.rs b/webschembly-compiler/src/jit/jit_module.rs index 312bc5b5..c32c7dee 100644 --- a/webschembly-compiler/src/jit/jit_module.rs +++ b/webschembly-compiler/src/jit/jit_module.rs @@ -7,8 +7,8 @@ use super::jit_func::{JitFunc, JitSpecializedArgFunc}; use crate::ir_generator::GlobalManager; use crate::jit::jit_func::JitSpecializedEnvFunc; use vec_map::{HasId, VecMap}; -use webschembly_compiler_ir::{BBIndex, ClosureArgIndex, ClosureEnvIndex}; use webschembly_compiler_ir::*; +use webschembly_compiler_ir::{BBIndex, ClosureArgIndex, ClosureEnvIndex}; #[derive(Debug)] pub struct JitModule { module_id: JitModuleId, diff --git a/webschembly-compiler/src/jit/mod.rs b/webschembly-compiler/src/jit/mod.rs index 816a32f3..e445b693 100644 --- a/webschembly-compiler/src/jit/mod.rs +++ b/webschembly-compiler/src/jit/mod.rs @@ -1,15 +1,14 @@ use crate::ir_generator::GlobalManager; use vec_map::VecMap; - mod jit_config; mod jit_module; pub use jit_config::JitConfig; use jit_module::JitModule; mod jit_ctx; use jit_ctx::JitCtx; -use webschembly_compiler_ir::{BBIndex, ClosureArgIndex, ClosureEnvIndex}; use webschembly_compiler_ir::*; +use webschembly_compiler_ir::{BBIndex, ClosureArgIndex, ClosureEnvIndex}; pub mod bb_index_manager; pub mod closure_global_layout; @@ -57,7 +56,6 @@ impl Jit { self.jit_module[module_id].instantiate_func( global_manager, func_id, - env_index, func_index, &mut self.ctx, From 7e96a04d926298cbc09a7fd4dc09bd406fe398c1 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 8 Jan 2026 09:27:54 +0000 Subject: [PATCH 06/15] chore: auto-fix format and lint --- .../src/ir_generator/module_generator.rs | 5 +- .../src/ir_processor/propagate_types.rs | 475 ++++++++++-------- .../src/ir_processor/ssa_optimizer.rs | 2 +- 3 files changed, 281 insertions(+), 201 deletions(-) diff --git a/webschembly-compiler/src/ir_generator/module_generator.rs b/webschembly-compiler/src/ir_generator/module_generator.rs index 0d78f17a..35fade19 100644 --- a/webschembly-compiler/src/ir_generator/module_generator.rs +++ b/webschembly-compiler/src/ir_generator/module_generator.rs @@ -622,10 +622,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { }); self.builder.exprs.push(Instr { local: result, - kind: InstrKind::ToObj( - ValType::Closure(None), - val_type_local, - ), + kind: InstrKind::ToObj(ValType::Closure(None), val_type_local), }); } ast::Expr::If(_, ast::If { cond, then, els }) => { diff --git a/webschembly-compiler/src/ir_processor/propagate_types.rs b/webschembly-compiler/src/ir_processor/propagate_types.rs index ae2028aa..cbf14617 100644 --- a/webschembly-compiler/src/ir_processor/propagate_types.rs +++ b/webschembly-compiler/src/ir_processor/propagate_types.rs @@ -41,7 +41,7 @@ pub fn propagate_types(func: &mut Func) { } let rpo = calculate_rpo(&func.bbs, func.bb_entry); - + // Convert RPO map to a list of BBs sorted by RPO index let mut sorted_bbs: Vec = rpo.keys().cloned().collect(); sorted_bbs.sort_by_key(|bb| rpo.get(bb).unwrap()); @@ -56,30 +56,32 @@ pub fn propagate_types(func: &mut Func) { if let Some(dest) = instr.local { let old_val = *lattice.get(dest).unwrap_or(&LatticeValue::Top); let new_val = match &instr.kind { - InstrKind::Closure { func_id, env_index, .. } => { - LatticeValue::Constant(ConstantClosure { - func_id: *func_id, - env_index: ClosureEnvIndex(*env_index), - }) - } - InstrKind::Move(src) => { - *lattice.get(*src).unwrap_or(&LatticeValue::Top) - } - InstrKind::Phi { incomings, non_exhaustive } => { + InstrKind::Closure { + func_id, env_index, .. + } => LatticeValue::Constant(ConstantClosure { + func_id: *func_id, + env_index: ClosureEnvIndex(*env_index), + }), + InstrKind::Move(src) => *lattice.get(*src).unwrap_or(&LatticeValue::Top), + InstrKind::Phi { + incomings, + non_exhaustive, + } => { if *non_exhaustive { LatticeValue::Bottom } else { let mut val = LatticeValue::Top; for incoming in incomings { - let incoming_val = *lattice.get(incoming.local).unwrap_or(&LatticeValue::Top); + let incoming_val = + *lattice.get(incoming.local).unwrap_or(&LatticeValue::Top); val = val.meet(incoming_val); } val } } - // If we have a direct reference to a known constant logic in other instructions - // we might want to propagate that, but for now we only care about propagating definitions. - // Any other instruction that defines a value produces Bottom for that value (unknown type/value) + // If we have a direct reference to a known constant logic in other instructions + // we might want to propagate that, but for now we only care about propagating definitions. + // Any other instruction that defines a value produces Bottom for that value (unknown type/value) _ => LatticeValue::Bottom, }; @@ -101,25 +103,25 @@ pub fn propagate_types(func: &mut Func) { LocalType::Type(Type::Val(ValType::Closure(_))) => { local.typ = LocalType::Type(Type::Val(ValType::Closure(Some(constant)))); } - // If it was Type::Obj (upcasted), we can refine it to Closure + // If it was Type::Obj (upcasted), we can refine it to Closure LocalType::Type(Type::Obj) => { local.typ = LocalType::Type(Type::Val(ValType::Closure(Some(constant)))); } _ => {} } } - // If Bottom, we might need to revert to generic Closure(None) if it was previously set to something specific - // but here we are primarily taking generic Code (Closure(None)) and refining it. - // If the code was *already* specialized/refined, we should be careful not to overwrite it with less specific info - // unless we are sure. But this pass is usually run to refine. + // If Bottom, we might need to revert to generic Closure(None) if it was previously set to something specific + // but here we are primarily taking generic Code (Closure(None)) and refining it. + // If the code was *already* specialized/refined, we should be careful not to overwrite it with less specific info + // unless we are sure. But this pass is usually run to refine. LatticeValue::Bottom => { - // Ensure if we failed to prove constant, it remains as generic closure if it was a closure - /* - Note: We don't generally downgrade types here explicitly because the lattice implies - we only UPGRADE to Constant or stay Bottom. - However, if this pass is run multiple times or if we have partial information, - we might want to be safe. But currently the IR starts with Closure(None) usually. - */ + // Ensure if we failed to prove constant, it remains as generic closure if it was a closure + /* + Note: We don't generally downgrade types here explicitly because the lattice implies + we only UPGRADE to Constant or stay Bottom. + However, if this pass is run multiple times or if we have partial information, + we might want to be safe. But currently the IR starts with Closure(None) usually. + */ } LatticeValue::Top => { // Dead code or uninitialized. Leave as is or set to Bottom equivalent? @@ -129,11 +131,12 @@ pub fn propagate_types(func: &mut Func) { } } - #[cfg(test)] mod tests { use super::*; - use webschembly_compiler_ir::{Func, BasicBlock, Instr, InstrKind, Local, Type, ValType, TerminatorInstr, ExitInstr}; + use webschembly_compiler_ir::{ + BasicBlock, ExitInstr, Func, Instr, InstrKind, Local, TerminatorInstr, Type, ValType, + }; fn create_dummy_func() -> Func { Func { @@ -150,8 +153,14 @@ mod tests { #[test] fn test_propagate_simple_move() { let mut func = create_dummy_func(); - let l0 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); - let l1 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l0 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l1 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); let bb0 = func.bbs.push_with(|id| BasicBlock { id, @@ -163,8 +172,8 @@ mod tests { env_types: vec![], env_index: 0, module_id: JitModuleId::from(0), - func_id: JitFuncId::from(10), // Target ID - entrypoint_table: l1, // Dummy + func_id: JitFuncId::from(10), // Target ID + entrypoint_table: l1, // Dummy original_entrypoint_table: l1, // Dummy }, }, @@ -174,8 +183,10 @@ mod tests { }, Instr { local: None, - kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l1)).into()) - } + kind: InstrKind::Terminator( + TerminatorInstr::Exit(ExitInstr::Return(l1)), + ), + }, ], }); func.bb_entry = bb0; @@ -195,9 +206,18 @@ mod tests { fn test_propagate_phi_same() { let mut func = create_dummy_func(); // l0 and l1 are same closure constant - let l0 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); - let l1 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); - let l2 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); // Phi result + let l0 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l1 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l2 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); // Phi result /* bb0: @@ -224,77 +244,80 @@ mod tests { }); func.bbs[bb0].instrs = vec![ - Instr { - local: Some(l0), - kind: InstrKind::Closure { - envs: vec![], - env_types: vec![], - env_index: 0, - module_id: JitModuleId::from(0), - func_id: JitFuncId::from(10), - entrypoint_table: l0, - original_entrypoint_table: l0, - }, + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), + entrypoint_table: l0, + original_entrypoint_table: l0, }, - Instr { - local: None, - kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), - } + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, ]; - - func.bbs[bb1].instrs = vec![ - Instr { - local: Some(l1), - kind: InstrKind::Closure { - envs: vec![], - env_types: vec![], - env_index: 0, - module_id: JitModuleId::from(0), - func_id: JitFuncId::from(10), - entrypoint_table: l0, - original_entrypoint_table: l0, - }, + + func.bbs[bb1].instrs = vec![ + Instr { + local: Some(l1), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), + entrypoint_table: l0, + original_entrypoint_table: l0, }, - Instr { - local: None, - kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), - } + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, ]; - func.bbs[bb2].instrs = vec![ - Instr { - local: Some(l2), - kind: InstrKind::Phi { - incomings: vec![ - PhiIncomingValue { bb: bb0, local: l0 }, - PhiIncomingValue { bb: bb1, local: l1 }, - ], - non_exhaustive: false, - } - }, - Instr { - local: None, - kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))), - } + func.bbs[bb2].instrs = vec![ + Instr { + local: Some(l2), + kind: InstrKind::Phi { + incomings: vec![ + PhiIncomingValue { bb: bb0, local: l0 }, + PhiIncomingValue { bb: bb1, local: l1 }, + ], + non_exhaustive: false, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))), + }, ]; - + // entry -> bb0 (cond jump is hard so just make one path reachable for simple test structure or use Jump from entry) - // Let's make entry jump to bb0 and bb1? No, just linear flow is fine? + // Let's make entry jump to bb0 and bb1? No, just linear flow is fine? // Wait, for Phi to work we need valid preds. // Let's make entry conditional jump to bb0 and bb1. - let l_cond = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Bool)) }); + let l_cond = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Bool)), + }); func.bb_entry = func.bbs.push_with(|id| BasicBlock { - id, - instrs: vec![ - Instr { - local: Some(l_cond), - kind: InstrKind::Bool(true), - }, - Instr { - local: None, - kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)) - } - ] + id, + instrs: vec![ + Instr { + local: Some(l_cond), + kind: InstrKind::Bool(true), + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)), + }, + ], }); propagate_types(&mut func); @@ -310,67 +333,94 @@ mod tests { #[test] fn test_propagate_phi_different() { let mut func = create_dummy_func(); - let l0 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); - let l1 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); - let l2 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l0 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l1 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l2 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); let bb0 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); let bb1 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); let bb2 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); func.bbs[bb0].instrs = vec![ - Instr { - local: Some(l0), - kind: InstrKind::Closure { - envs: vec![], - env_types: vec![], - env_index: 0, - module_id: JitModuleId::from(0), - func_id: JitFuncId::from(10), // ID 10 - entrypoint_table: l0, - original_entrypoint_table: l0, - }, + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // ID 10 + entrypoint_table: l0, + original_entrypoint_table: l0, }, - Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)) } + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, ]; - - func.bbs[bb1].instrs = vec![ - Instr { - local: Some(l1), - kind: InstrKind::Closure { - envs: vec![], - env_types: vec![], - env_index: 0, - module_id: JitModuleId::from(0), - func_id: JitFuncId::from(11), // ID 11 (Different) - entrypoint_table: l0, - original_entrypoint_table: l0, - }, + + func.bbs[bb1].instrs = vec![ + Instr { + local: Some(l1), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(11), // ID 11 (Different) + entrypoint_table: l0, + original_entrypoint_table: l0, }, - Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)) } + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, ]; - func.bbs[bb2].instrs = vec![ - Instr { - local: Some(l2), - kind: InstrKind::Phi { - incomings: vec![ - PhiIncomingValue { bb: bb0, local: l0 }, - PhiIncomingValue { bb: bb1, local: l1 }, - ], - non_exhaustive: false, - } - }, - Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))) } + func.bbs[bb2].instrs = vec![ + Instr { + local: Some(l2), + kind: InstrKind::Phi { + incomings: vec![ + PhiIncomingValue { bb: bb0, local: l0 }, + PhiIncomingValue { bb: bb1, local: l1 }, + ], + non_exhaustive: false, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))), + }, ]; - let l_cond = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Bool)) }); + let l_cond = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Bool)), + }); func.bb_entry = func.bbs.push_with(|id| BasicBlock { - id, - instrs: vec![ - Instr { local: Some(l_cond), kind: InstrKind::Bool(true) }, - Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)) } - ] + id, + instrs: vec![ + Instr { + local: Some(l_cond), + kind: InstrKind::Bool(true), + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)), + }, + ], }); propagate_types(&mut func); @@ -380,75 +430,105 @@ mod tests { if let LocalType::Type(Type::Val(ValType::Closure(None))) = t2 { // OK } else { - panic!("Expected generic Closure(None) (Bottom) for l2, got {:?}", t2); + panic!( + "Expected generic Closure(None) (Bottom) for l2, got {:?}", + t2 + ); } } #[test] fn test_propagate_phi_non_exhaustive() { let mut func = create_dummy_func(); - let l0 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); - let l1 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); - let l2 = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Closure(None))) }); + let l0 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l1 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l2 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); let bb0 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); let bb1 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); let bb2 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); func.bbs[bb0].instrs = vec![ - Instr { - local: Some(l0), - kind: InstrKind::Closure { - envs: vec![], - env_types: vec![], - env_index: 0, - module_id: JitModuleId::from(0), - func_id: JitFuncId::from(10), // ID 10 - entrypoint_table: l0, - original_entrypoint_table: l0, - }, + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // ID 10 + entrypoint_table: l0, + original_entrypoint_table: l0, }, - Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)) } + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, ]; - - func.bbs[bb1].instrs = vec![ - Instr { - local: Some(l1), - kind: InstrKind::Closure { - envs: vec![], - env_types: vec![], - env_index: 0, - module_id: JitModuleId::from(0), - func_id: JitFuncId::from(10), // ID 10 (Same as l0) - entrypoint_table: l0, - original_entrypoint_table: l0, - }, + + func.bbs[bb1].instrs = vec![ + Instr { + local: Some(l1), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // ID 10 (Same as l0) + entrypoint_table: l0, + original_entrypoint_table: l0, }, - Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)) } + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, ]; - func.bbs[bb2].instrs = vec![ - Instr { - local: Some(l2), - kind: InstrKind::Phi { - incomings: vec![ - PhiIncomingValue { bb: bb0, local: l0 }, - PhiIncomingValue { bb: bb1, local: l1 }, - ], - // Even though l0 and l1 are same constant, non_exhaustive should make result Bottom - non_exhaustive: true, - } - }, - Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))) } + func.bbs[bb2].instrs = vec![ + Instr { + local: Some(l2), + kind: InstrKind::Phi { + incomings: vec![ + PhiIncomingValue { bb: bb0, local: l0 }, + PhiIncomingValue { bb: bb1, local: l1 }, + ], + // Even though l0 and l1 are same constant, non_exhaustive should make result Bottom + non_exhaustive: true, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))), + }, ]; - let l_cond = func.locals.push_with(|id| Local { id, typ: LocalType::Type(Type::Val(ValType::Bool)) }); + let l_cond = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Bool)), + }); func.bb_entry = func.bbs.push_with(|id| BasicBlock { - id, - instrs: vec![ - Instr { local: Some(l_cond), kind: InstrKind::Bool(true) }, - Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)) } - ] + id, + instrs: vec![ + Instr { + local: Some(l_cond), + kind: InstrKind::Bool(true), + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)), + }, + ], }); propagate_types(&mut func); @@ -458,7 +538,10 @@ mod tests { if let LocalType::Type(Type::Val(ValType::Closure(None))) = t2 { // OK } else { - panic!("Expected generic Closure(None) (Bottom) for l2 due to non_exhaustive, got {:?}", t2); + panic!( + "Expected generic Closure(None) (Bottom) for l2 due to non_exhaustive, got {:?}", + t2 + ); } } } diff --git a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs index 44e1b6a2..0aafd2e3 100644 --- a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs +++ b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs @@ -4,7 +4,7 @@ use crate::ir_processor::{ cfg_analyzer::{DomTreeNode, build_dom_tree, calc_doms, calc_predecessors, calculate_rpo}, optimizer::remove_unreachable_bb, propagate_types::propagate_types, - ssa::{debug_assert_ssa, DefUseChain}, + ssa::{DefUseChain, debug_assert_ssa}, }; use vec_map::VecMap; use webschembly_compiler_ir::*; From f4adb9117ba12fde84bafa79f24dd8f613504e35 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 8 Jan 2026 09:29:48 +0000 Subject: [PATCH 07/15] chore: auto-fix format and lint --- webschembly-compiler/src/ir_processor/propagate_types.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/webschembly-compiler/src/ir_processor/propagate_types.rs b/webschembly-compiler/src/ir_processor/propagate_types.rs index cbf14617..86657f5c 100644 --- a/webschembly-compiler/src/ir_processor/propagate_types.rs +++ b/webschembly-compiler/src/ir_processor/propagate_types.rs @@ -183,9 +183,7 @@ mod tests { }, Instr { local: None, - kind: InstrKind::Terminator( - TerminatorInstr::Exit(ExitInstr::Return(l1)), - ), + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l1))), }, ], }); From bfa245d0d99791bc81b1e729dd5e80b39cfadf1c Mon Sep 17 00:00:00 2001 From: kgtkr Date: Thu, 8 Jan 2026 19:47:01 +0900 Subject: [PATCH 08/15] =?UTF-8?q?=E7=89=B9=E6=AE=8A=E5=8C=96=E3=81=A7?= =?UTF-8?q?=E5=AE=9A=E6=95=B0=E3=82=92=E8=80=83=E6=85=AE=E3=81=99=E3=82=8B?= =?UTF-8?q?=E3=82=88=E3=81=86=E3=81=AB=E3=81=97=E3=81=9F=E3=80=82=E3=81=BE?= =?UTF-8?q?=E3=81=9F=E5=AE=9A=E6=95=B0=E4=BC=9D=E6=92=AD=E3=81=AE=E3=83=90?= =?UTF-8?q?=E3=82=B0=E4=BF=AE=E6=AD=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- DEBUGGING.md | 48 +++++++++++++++++++ .../src/ir_processor/ssa_optimizer.rs | 3 +- webschembly-compiler/src/jit/jit_func.rs | 17 +++++++ 3 files changed, 67 insertions(+), 1 deletion(-) create mode 100644 DEBUGGING.md diff --git a/DEBUGGING.md b/DEBUGGING.md new file mode 100644 index 00000000..eb8eab4a --- /dev/null +++ b/DEBUGGING.md @@ -0,0 +1,48 @@ +# Debugging Guide + +This document outlines the workflow and tools for debugging the Webschembly compiler and JIT. + +## Running Tests with Logs + +To debug specific Scheme files or tests, use `just run` within the `webschembly-js` directory. + +### Environment Variables + +- `LOG_STDOUT=1`: Enables `log::debug!` output from the runtime and compiler (if configured). Use this to see runtime execution trace, JIT instantiation events, and error messages. +- `LOG=1`: Dumps the generated Intermediate Representation (IR) and Wasm binaries to the `webschembly-js/log/` directory. + +### Command Examples + +```bash +# Run a specific fixture with runtime logs +cd webschembly-js +just LOG_STDOUT=1 run ./fixtures/rec.scm + +# Run with IR dumping +just LOG=1 run ./fixtures/rec.scm +``` + +## Analyzing Generated IR + +When `LOG=1` is used, the `webschembly-js/log/` directory will contain files named with the timestamp and content description. + +- `checks-TIMESTAMP-filename-0.ir`: Typically the main module IR (initial compilation). +- `checks-TIMESTAMP-filename-N.ir`: IR for JIT-compiled functions or stubs. `instantiate_func` or `instantiate_bb` calls in the runtime logs (seen with `LOG_STDOUT=1`) will reference `module_id` and `func_id`, which correlate to these files (though the mapping requires checking the `instantiate` log id vs file index). + +**Tip**: Look for "instantiate: id:X" in the runtime logs. The corresponding IR file is often suffix `-X.ir`. + +## Common Issues & Fixes + +### "call target is not a closure" + +This error occurs when the compiled code attempts to invoke a value as a closure, but the compile-time or run-time check fails. + +- **Runtime**: The value on the stack is not a closure object (e.g., encoded as `val_type` mismatch). +- **Compile-time (Optimization)**: If the error appears unconditionally in the IR (e.g., `error "call target is not a closure"`), it means `constant_folding` or another pass determined the check `Is(Closure(None), target)` is false. + - *Watch out for*: Type mismatches in `InstrKind::Is`. Ensure `Closure(None)` (generic check) correctly matches specialized types like `Closure(Some(C))` (constant closure). Using `.remove_constant()` on types before comparison is crucial in `ssa_optimizer.rs`. + +## JIT Optimization Logic + +- **`propagate_types` Pass**: Analyzes dataflow to identify constant closures (`Closure(Some(C))`). Logic in `src/ir_processor/propagate_types.rs`. +- **Specialization**: `jit_func.rs` uses available type info (from `locals`) to unbox closure entries or arguments. +- **Constant Folding**: `ssa_optimizer.rs` folds constants. Be careful with strict equality checks on specialized types. diff --git a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs index 0aafd2e3..f1ca5265 100644 --- a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs +++ b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs @@ -371,7 +371,8 @@ pub fn constant_folding( if let Some(&InstrKind::ToObj(typ2, _)) = def_use.get_def_non_move_expr(&func.bbs, src) => { - func.bbs[*bb_id].instrs[expr_idx].kind = InstrKind::Bool(typ1 == typ2); + func.bbs[*bb_id].instrs[expr_idx].kind = + InstrKind::Bool(typ1.remove_constant() == typ2.remove_constant()); } InstrKind::ClosureEnv(_, closure, index) if let Some(InstrKind::Closure { envs, .. }) = diff --git a/webschembly-compiler/src/jit/jit_func.rs b/webschembly-compiler/src/jit/jit_func.rs index 73a572d3..34012137 100644 --- a/webschembly-compiler/src/jit/jit_func.rs +++ b/webschembly-compiler/src/jit/jit_func.rs @@ -612,6 +612,10 @@ impl JitSpecializedArgFunc { if let Some(&InstrKind::ToObj(typ, val_local)) = def_use_chain.get_def_non_move_expr(&body_func.bbs, obj_local) { + let mut typ = typ; + if let LocalType::Type(Type::Val(t)) = body_func.locals[val_local].typ { + typ = t; + } Some(TypedObj { typ, val_local }) } else { typed_objs.get(&obj_local).copied() @@ -762,6 +766,7 @@ impl JitSpecializedArgFunc { call_closure, &def_use_chain, &body_func.bbs, + &body_func.locals, jit_ctx.closure_global_layout(), &mut required_closure_idx, ) => @@ -779,6 +784,7 @@ impl JitSpecializedArgFunc { call_closure, &def_use_chain, &body_func.bbs, + &body_func.locals, jit_ctx.closure_global_layout(), &mut required_closure_idx, ) => @@ -822,6 +828,12 @@ impl JitSpecializedArgFunc { def_use_chain .get_def_non_move_expr(&body_func.bbs, value_local) { + let mut typ = typ; + if let LocalType::Type(Type::Val(t)) = + body_func.locals[val_local].typ + { + typ = t; + } Some(TypedObj { typ, val_local }) } else { None @@ -1207,6 +1219,7 @@ fn specialize_call_closure( call_closure: &InstrCallClosure, def_use_chain: &DefUseChain, bbs: &VecMap, + locals: &VecMap, closure_global_layout: &mut ClosureGlobalLayout, required_closure_idx: &mut Vec, ) -> Option { @@ -1228,6 +1241,10 @@ fn specialize_call_closure( def_use_chain.get_def_non_move_expr(bbs, obj_arg) { fixed_args.push(val_local); + let mut typ = typ; + if let LocalType::Type(Type::Val(t)) = locals[val_local].typ { + typ = t; + } fixed_arg_types.push(Type::from(typ)); } else { fixed_args.push(obj_arg); From 083a3edafda1a8b7d80d646408710c23cfb45db8 Mon Sep 17 00:00:00 2001 From: kgtkr Date: Thu, 8 Jan 2026 21:00:17 +0900 Subject: [PATCH 09/15] =?UTF-8?q?WIP:=20inline=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- webschembly-compiler/src/compiler.rs | 14 +- .../src/ir_processor/inline.rs | 450 ++++++++++++++++++ webschembly-compiler/src/ir_processor/mod.rs | 1 + .../src/ir_processor/ssa_optimizer.rs | 2 +- 4 files changed, 457 insertions(+), 10 deletions(-) create mode 100644 webschembly-compiler/src/ir_processor/inline.rs diff --git a/webschembly-compiler/src/compiler.rs b/webschembly-compiler/src/compiler.rs index 5e34a288..1121fc2e 100644 --- a/webschembly-compiler/src/compiler.rs +++ b/webschembly-compiler/src/compiler.rs @@ -8,9 +8,7 @@ use crate::ir_processor::optimizer::remove_unused_local; use crate::ir_processor::register_allocation::register_allocation; use crate::ir_processor::ssa::split_critical_edges; use crate::ir_processor::ssa::{debug_assert_ssa, remove_phi}; -use crate::ir_processor::ssa_optimizer::ModuleInliner; use crate::ir_processor::ssa_optimizer::SsaOptimizerConfig; -use crate::ir_processor::ssa_optimizer::inlining; use crate::ir_processor::ssa_optimizer::ssa_optimize; use crate::jit::{Jit, JitConfig}; use crate::lexer; @@ -243,13 +241,8 @@ fn preprocess_module(module: &mut ir::Module) { } fn optimize_module(module: &mut ir::Module, config: SsaOptimizerConfig) { - let mut module_inliner = ModuleInliner::new(module); - let n = 5; - for i in 0..n { - if config.enable_inlining { - // inliningはInstrKind::Closureのfunc_idに依存しているので、JIT後のモジュールには使えない - inlining(module, &mut module_inliner, i == n - 1); - } + let n = 10; + for _ in 0..n { for func in module.funcs.values_mut() { ssa_optimize( func, @@ -259,6 +252,9 @@ fn optimize_module(module: &mut ir::Module, config: SsaOptimizerConfig) { }, ); } + if config.enable_inlining { + crate::ir_processor::inline::inline_module(module); + } } } diff --git a/webschembly-compiler/src/ir_processor/inline.rs b/webschembly-compiler/src/ir_processor/inline.rs new file mode 100644 index 00000000..0cbdc65f --- /dev/null +++ b/webschembly-compiler/src/ir_processor/inline.rs @@ -0,0 +1,450 @@ +use rustc_hash::FxHashMap; +use vec_map::VecMap; +use webschembly_compiler_ir::*; + +use crate::ir_processor::optimizer::remove_unreachable_bb; + +const LIMIT_BB_COUNT: usize = 100; + +pub fn inline_module(module: &mut Module) { + let mut new_funcs = VecMap::new(); + + for (func_id_usize, func) in module.funcs.iter() { + let _func_id = FuncId::from(func_id_usize); + let mut new_func = func.clone(); + run_inlining(&mut new_func, module); + new_funcs.insert(func_id_usize, new_func); + } + + module.funcs = new_funcs; +} + +struct InlineContext<'a> { + module: &'a Module, + tail_instances: FxHashMap, +} + +#[derive(Debug, Clone)] +struct TailCallInfo { + entry_bb: BasicBlockId, + arg_phis: Vec, +} + +fn run_inlining(func: &mut Func, module: &Module) { + let mut ctx = InlineContext { + module, + tail_instances: FxHashMap::default(), + }; + + let mut worklist: Vec = func.bbs.keys().map(BasicBlockId::from).collect(); + + while let Some(bb_id) = worklist.pop() { + if func.bbs.iter().count() > LIMIT_BB_COUNT { + log::debug!("BB limit reached for func {:?}", func.id); + break; + } + if !func.bbs.contains_key(bb_id.into()) { + continue; + } + + // 1. Check Non-Tail Calls (Instrs) + let mut call_found = None; + { + let bb = &func.bbs[bb_id.into()]; + for (idx, instr) in bb.instrs.iter().enumerate() { + if let InstrKind::CallClosure(call) = &instr.kind { + let closure_local = call.closure; + if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = + func.locals[closure_local.into()].typ + { + call_found = Some((idx, constant, call.clone())); + break; + } + } + } + } + + if let Some((idx, constant, call)) = call_found { + log::debug!( + "Non-tail call candidate found in BB {:?} to func {:?}", + bb_id, + constant.func_id + ); + let result_local = func.bbs[bb_id.into()].instrs[idx].local; + let continuation_bb_id = inline_non_tail( + func, + &mut ctx, + &mut worklist, + bb_id, + idx, + constant, + &call, + result_local, + ); + worklist.push(continuation_bb_id); + continue; + } + + // 2. Check Tail Call (Terminator) + let mut tail_call_found = None; + { + let bb = &func.bbs[bb_id.into()]; + if let TerminatorInstr::Exit(ExitInstr::TailCallClosure(call)) = bb.terminator() { + let closure_local = call.closure; + if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = + func.locals[closure_local.into()].typ + { + tail_call_found = Some((constant, call.clone())); + } + } + } + + if let Some((constant, call)) = tail_call_found { + log::debug!( + "Tail call candidate found in BB {:?} to func {:?}", + bb_id, + constant.func_id + ); + inline_tail(func, &mut ctx, &mut worklist, bb_id, constant, &call); + continue; + } + } + + log::debug!( + "Inlining finished for func {:?}. Final BB count: {}", + func.id, + func.bbs.iter().count() + ); + remove_unreachable_bb(func); +} + +fn inline_non_tail( + func: &mut Func, + ctx: &mut InlineContext, + worklist: &mut Vec, + caller_bb_id: BasicBlockId, + instr_idx: usize, + constant: ConstantClosure, + call: &InstrCallClosure, + result_local: Option, +) -> BasicBlockId { + let continuation_bb_id = func.bbs.allocate_key(); + let continuation_bb_id = BasicBlockId::from(continuation_bb_id); + let original_terminator = func.bbs[caller_bb_id.into()].terminator().clone(); + + let instrs_after = func.bbs[caller_bb_id.into()] + .instrs + .split_off(instr_idx + 1); + func.bbs[caller_bb_id.into()].instrs.pop(); + + func.bbs.insert_node(BasicBlock { + id: continuation_bb_id, + instrs: instrs_after, + }); + *func.bbs[continuation_bb_id.into()].terminator_mut() = original_terminator; + + let callee_id = FuncId::from(constant.func_id); + let callee = &ctx.module.funcs[callee_id.into()]; + + let mut local_map = FxHashMap::default(); + let mut bb_map = FxHashMap::default(); + + for old_bb_id_usize in callee.bbs.keys() { + let old_bb_id = BasicBlockId::from(old_bb_id_usize); + let new_bb_id_usize = func.bbs.allocate_key(); + let new_bb_id = BasicBlockId::from(new_bb_id_usize); + bb_map.insert(old_bb_id, new_bb_id); + worklist.push(new_bb_id); + } + + for (i, &arg_local) in callee.args.iter().enumerate() { + let replacement = if i == 0 { + call.closure + } else if i - 1 < call.args.len() { + call.args[i - 1] + } else { + call.args[call.args.len() - 1] + }; + local_map.insert(arg_local, replacement); + } + + for (local_id_usize, local) in callee.locals.iter() { + let local_id = LocalId::from(local_id_usize); + if !local_map.contains_key(&local_id) { + let new_id_usize = func.locals.push_with(|id| Local { + id, + typ: local.typ, + ..*local + }); + let new_id = LocalId::from(new_id_usize); + local_map.insert(local_id, new_id); + } + } + + let mut phi_incomings = Vec::new(); + + for (old_bb_id_usize, old_bb) in callee.bbs.iter() { + let old_bb_id = BasicBlockId::from(old_bb_id_usize); + let new_bb_id = bb_map[&old_bb_id]; + let mut new_instrs = Vec::new(); + + for instr in &old_bb.instrs { + let mut new_instr = instr.clone(); + if let Some(local) = new_instr.local { + if let Some(&mapped) = local_map.get(&local) { + new_instr.local = Some(mapped); + } + } + rewrite_usages(&mut new_instr.kind, &local_map); + for bb_ref in new_instr.kind.bb_ids_mut() { + if let Some(&mapped) = bb_map.get(bb_ref) { + *bb_ref = mapped; + } + } + new_instrs.push(new_instr); + } + + new_instrs.pop(); + + let mut new_terminator = old_bb.terminator().clone(); + rewrite_terminator_usages(&mut new_terminator, &local_map); + for bb_ref in new_terminator.bb_ids_mut() { + if let Some(&mapped) = bb_map.get(bb_ref) { + *bb_ref = mapped; + } + } + + if let TerminatorInstr::Exit(ExitInstr::Return(val)) = &new_terminator { + // Note: ExitInstr::Return contains LocalId (not Option) + phi_incomings.push(PhiIncomingValue { + local: *val, + bb: new_bb_id, + }); + new_terminator = TerminatorInstr::Jump(continuation_bb_id); + } else if let TerminatorInstr::Exit(ExitInstr::TailCallClosure(call)) = &new_terminator { + if let Some(dst) = result_local { + let dst_typ = func.locals[dst.into()].typ.clone(); + let temp_local_usize = func.locals.push_with(|id| Local { id, typ: dst_typ }); + let temp_local = LocalId::from(temp_local_usize); + + new_instrs.push(Instr { + local: Some(temp_local), + kind: InstrKind::CallClosure(call.clone()), + }); + phi_incomings.push(PhiIncomingValue { + local: temp_local, + bb: new_bb_id, + }); + } else { + new_instrs.push(Instr { + local: None, + kind: InstrKind::CallClosure(call.clone()), + }); + } + new_terminator = TerminatorInstr::Jump(continuation_bb_id); + } + + new_instrs.push(Instr { + local: None, + kind: InstrKind::Terminator(new_terminator), + }); + + func.bbs.insert_node(BasicBlock { + id: new_bb_id, + instrs: new_instrs, + }); + } + + if let Some(dst) = result_local { + if !phi_incomings.is_empty() { + func.bbs[continuation_bb_id.into()].instrs.insert( + 0, + Instr { + local: Some(dst), + kind: InstrKind::Phi { + incomings: phi_incomings, + non_exhaustive: false, + }, + }, + ); + } + } + + let new_entry_id = bb_map[&callee.bb_entry]; + log::debug!( + "Inlining non-tail: Caller {:?} -> Cont {:?}, Entry {:?}", + caller_bb_id, + continuation_bb_id, + new_entry_id + ); + + func.bbs[caller_bb_id.into()].instrs.push(Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(new_entry_id)), + }); + + continuation_bb_id +} + +fn inline_tail( + func: &mut Func, + ctx: &mut InlineContext, + worklist: &mut Vec, + caller_bb_id: BasicBlockId, + constant: ConstantClosure, + call: &InstrCallClosure, +) { + if let Some(info) = ctx.tail_instances.get(&constant) { + log::debug!("Reusing tail instance at entry BB {:?}", info.entry_bb); + for (i, &phi_local) in info.arg_phis.iter().enumerate() { + let val = if i == 0 { + call.closure + } else { + call.args[i - 1] + }; + let entry_bb = &mut func.bbs[info.entry_bb.into()]; + for instr in &mut entry_bb.instrs { + if let InstrKind::Phi { incomings, .. } = &mut instr.kind { + if instr.local == Some(phi_local) { + incomings.push(PhiIncomingValue { + local: val, + bb: caller_bb_id, + }); + break; + } + } + } + } + *func.bbs[caller_bb_id.into()].terminator_mut() = TerminatorInstr::Jump(info.entry_bb); + return; + } + + let callee_id = FuncId::from(constant.func_id); + let callee = &ctx.module.funcs[callee_id.into()]; + + let mut local_map = FxHashMap::default(); + let mut bb_map = FxHashMap::default(); + + for old_bb_id_usize in callee.bbs.keys() { + let old_bb_id = BasicBlockId::from(old_bb_id_usize); + let new_bb_id_usize = func.bbs.allocate_key(); + let new_bb_id = BasicBlockId::from(new_bb_id_usize); + bb_map.insert(old_bb_id, new_bb_id); + worklist.push(new_bb_id); + } + + let new_entry_id = bb_map[&callee.bb_entry]; + let mut arg_phis = Vec::new(); + let mut entry_phi_instrs = Vec::new(); + + for (i, &arg_local) in callee.args.iter().enumerate() { + let new_arg_local_usize = func.locals.push_with(|id| Local { + id, + typ: callee.locals[arg_local.into()].typ, + ..callee.locals[arg_local.into()] + }); + let new_arg_local = LocalId::from(new_arg_local_usize); + local_map.insert(arg_local, new_arg_local); + arg_phis.push(new_arg_local); + + let val = if i == 0 { + call.closure + } else { + call.args[i - 1] + }; + entry_phi_instrs.push(Instr { + local: Some(new_arg_local), + kind: InstrKind::Phi { + incomings: vec![PhiIncomingValue { + local: val, + bb: caller_bb_id, + }], + non_exhaustive: false, + }, + }); + } + + for (local_id_usize, local) in callee.locals.iter() { + let local_id = LocalId::from(local_id_usize); + if !local_map.contains_key(&local_id) { + let new_id_usize = func.locals.push_with(|id| Local { + id, + typ: local.typ, + ..*local + }); + let new_id = LocalId::from(new_id_usize); + local_map.insert(local_id, new_id); + } + } + + for (old_bb_id_usize, old_bb) in callee.bbs.iter() { + let old_bb_id = BasicBlockId::from(old_bb_id_usize); + let new_bb_id = bb_map[&old_bb_id]; + let mut new_instrs = Vec::new(); + + if old_bb_id == callee.bb_entry { + new_instrs.extend(entry_phi_instrs.clone()); + } + + for instr in &old_bb.instrs { + let mut new_instr = instr.clone(); + if let Some(local) = new_instr.local { + if let Some(&mapped) = local_map.get(&local) { + new_instr.local = Some(mapped); + } + } + rewrite_usages(&mut new_instr.kind, &local_map); + for bb_ref in new_instr.kind.bb_ids_mut() { + if let Some(&mapped) = bb_map.get(bb_ref) { + *bb_ref = mapped; + } + } + new_instrs.push(new_instr); + } + + let mut new_terminator = old_bb.terminator().clone(); + rewrite_terminator_usages(&mut new_terminator, &local_map); + for bb_ref in new_terminator.bb_ids_mut() { + if let Some(&mapped) = bb_map.get(bb_ref) { + *bb_ref = mapped; + } + } + + func.bbs.insert_node(BasicBlock { + id: new_bb_id, + instrs: new_instrs, + }); + *func.bbs[new_bb_id.into()].terminator_mut() = new_terminator; + } + + *func.bbs[caller_bb_id.into()].terminator_mut() = TerminatorInstr::Jump(new_entry_id); + + log::debug!( + "Created new tail instance for func {:?} at entry BB {:?}", + constant.func_id, + new_entry_id + ); + ctx.tail_instances.insert( + constant, + TailCallInfo { + entry_bb: new_entry_id, + arg_phis, + }, + ); +} + +fn rewrite_usages(kind: &mut InstrKind, map: &FxHashMap) { + for (local, _) in kind.local_usages_mut() { + if let Some(&mapped) = map.get(local) { + *local = mapped; + } + } +} + +fn rewrite_terminator_usages(term: &mut TerminatorInstr, map: &FxHashMap) { + for local in term.local_ids_mut() { + if let Some(&mapped) = map.get(local) { + *local = mapped; + } + } +} diff --git a/webschembly-compiler/src/ir_processor/mod.rs b/webschembly-compiler/src/ir_processor/mod.rs index 2cf14d5f..7d9e70c3 100644 --- a/webschembly-compiler/src/ir_processor/mod.rs +++ b/webschembly-compiler/src/ir_processor/mod.rs @@ -1,6 +1,7 @@ pub mod cfg_analyzer; pub mod dataflow; pub mod desugar; +pub mod inline; pub mod optimizer; pub mod propagate_types; pub mod register_allocation; diff --git a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs index f1ca5265..a9bfd4d8 100644 --- a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs +++ b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs @@ -449,7 +449,7 @@ impl Default for SsaOptimizerConfig { SsaOptimizerConfig { enable_cse: true, enable_dce: true, - enable_inlining: false, // true, + enable_inlining: true, iterations: 5, } } From 78187abdf6152e5a8d90b1174a1887015e8bdc2a Mon Sep 17 00:00:00 2001 From: kgtkr Date: Thu, 8 Jan 2026 21:23:18 +0900 Subject: [PATCH 10/15] =?UTF-8?q?debbug=E4=B8=AD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/ir_processor/inline.rs | 31 +++++++++++++++++++ webschembly-compiler/src/ir_processor/ssa.rs | 5 ++- .../src/wasm_generator/module_generator.rs | 1 + 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/webschembly-compiler/src/ir_processor/inline.rs b/webschembly-compiler/src/ir_processor/inline.rs index 0cbdc65f..b0fb02ed 100644 --- a/webschembly-compiler/src/ir_processor/inline.rs +++ b/webschembly-compiler/src/ir_processor/inline.rs @@ -116,6 +116,23 @@ fn run_inlining(func: &mut Func, module: &Module) { func.bbs.iter().count() ); remove_unreachable_bb(func); + + // Verify Phi position + for (id, bb) in func.bbs.iter() { + let mut phi_mode = true; + for (idx, instr) in bb.instrs.iter().enumerate() { + if let InstrKind::Phi { .. } = instr.kind { + if !phi_mode { + panic!( + "INLINE_MODULE_END: Phi at non-start! BB: {:?}, Index: {}, Instrs: {:#?}", + id, idx, bb.instrs + ); + } + } else if !matches!(instr.kind, InstrKind::Nop) { + phi_mode = false; + } + } + } } fn inline_non_tail( @@ -257,6 +274,13 @@ fn inline_non_tail( if let Some(dst) = result_local { if !phi_incomings.is_empty() { + let id_val: usize = continuation_bb_id.into(); + if id_val == 69 { + panic!( + "DEBUG_INLINE_HIT_69: BEFORE INSERT: {:#?}", + func.bbs[continuation_bb_id.into()].instrs + ); + } func.bbs[continuation_bb_id.into()].instrs.insert( 0, Instr { @@ -267,10 +291,17 @@ fn inline_non_tail( }, }, ); + if id_val == 69 { + panic!( + "DEBUG_INLINE_HIT_69: AFTER INSERT: {:#?}", + func.bbs[continuation_bb_id.into()].instrs + ); + } } } let new_entry_id = bb_map[&callee.bb_entry]; + log::debug!( "Inlining non-tail: Caller {:?} -> Cont {:?}, Entry {:?}", caller_bb_id, diff --git a/webschembly-compiler/src/ir_processor/ssa.rs b/webschembly-compiler/src/ir_processor/ssa.rs index cf8767e4..7c449137 100644 --- a/webschembly-compiler/src/ir_processor/ssa.rs +++ b/webschembly-compiler/src/ir_processor/ssa.rs @@ -210,7 +210,10 @@ fn assert_ssa(func: &Func) { phi_area = false; } } else if let InstrKind::Phi { .. } = expr.kind { - panic!("phi instruction must be at the beginning of a basic block"); + panic!( + "phi instruction must be at the beginning of a basic block. BB: {:?}, Instrs: {:#?}", + bb.id, bb.instrs + ); } } } diff --git a/webschembly-compiler/src/wasm_generator/module_generator.rs b/webschembly-compiler/src/wasm_generator/module_generator.rs index 97659278..f3be567a 100644 --- a/webschembly-compiler/src/wasm_generator/module_generator.rs +++ b/webschembly-compiler/src/wasm_generator/module_generator.rs @@ -1081,6 +1081,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { if let ir::InstrKind::Terminator(..) = expr.kind { debug_assert!(expr.local.is_none()); debug_assert_eq!(i, bb.instrs.len() - 1); + // debug_assert_eq!(i, bb.instrs.len() - 1); } else { self.gen_assign(function, expr); } From 77b25f7f3a23f145dbbff9776aa3f5753586a7b6 Mon Sep 17 00:00:00 2001 From: kgtkr Date: Thu, 8 Jan 2026 21:52:59 +0900 Subject: [PATCH 11/15] wip --- .../src/ir_processor/inline.rs | 82 +++++++++++++++---- 1 file changed, 67 insertions(+), 15 deletions(-) diff --git a/webschembly-compiler/src/ir_processor/inline.rs b/webschembly-compiler/src/ir_processor/inline.rs index b0fb02ed..4ff34b38 100644 --- a/webschembly-compiler/src/ir_processor/inline.rs +++ b/webschembly-compiler/src/ir_processor/inline.rs @@ -7,12 +7,27 @@ use crate::ir_processor::optimizer::remove_unreachable_bb; const LIMIT_BB_COUNT: usize = 100; pub fn inline_module(module: &mut Module) { + let mut global_map = FxHashMap::default(); + let entry_func = &module.funcs[module.entry]; + // Scan entry function for GlobalSet with ConstantClosure + for bb in entry_func.bbs.values() { + for instr in &bb.instrs { + if let InstrKind::GlobalSet(global_id, val_local) = instr.kind { + if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = + entry_func.locals[val_local.into()].typ + { + global_map.insert(global_id, constant); + } + } + } + } + let mut new_funcs = VecMap::new(); for (func_id_usize, func) in module.funcs.iter() { let _func_id = FuncId::from(func_id_usize); let mut new_func = func.clone(); - run_inlining(&mut new_func, module); + run_inlining(&mut new_func, module, &global_map); new_funcs.insert(func_id_usize, new_func); } @@ -30,7 +45,11 @@ struct TailCallInfo { arg_phis: Vec, } -fn run_inlining(func: &mut Func, module: &Module) { +fn run_inlining( + func: &mut Func, + module: &Module, + global_map: &FxHashMap, +) { let mut ctx = InlineContext { module, tail_instances: FxHashMap::default(), @@ -54,9 +73,21 @@ fn run_inlining(func: &mut Func, module: &Module) { for (idx, instr) in bb.instrs.iter().enumerate() { if let InstrKind::CallClosure(call) = &instr.kind { let closure_local = call.closure; + let mut constant_opt = None; + if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = func.locals[closure_local.into()].typ { + constant_opt = Some(constant); + } else if let Some(def_instr) = find_local_def(func, closure_local) { + if let InstrKind::GlobalGet(global_id) = def_instr.kind { + if let Some(&constant) = global_map.get(&global_id) { + constant_opt = Some(constant); + } + } + } + + if let Some(constant) = constant_opt { call_found = Some((idx, constant, call.clone())); break; } @@ -91,9 +122,21 @@ fn run_inlining(func: &mut Func, module: &Module) { let bb = &func.bbs[bb_id.into()]; if let TerminatorInstr::Exit(ExitInstr::TailCallClosure(call)) = bb.terminator() { let closure_local = call.closure; + let mut constant_opt = None; + if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = func.locals[closure_local.into()].typ { + constant_opt = Some(constant); + } else if let Some(def_instr) = find_local_def(func, closure_local) { + if let InstrKind::GlobalGet(global_id) = def_instr.kind { + if let Some(&constant) = global_map.get(&global_id) { + constant_opt = Some(constant); + } + } + } + + if let Some(constant) = constant_opt { tail_call_found = Some((constant, call.clone())); } } @@ -135,6 +178,28 @@ fn run_inlining(func: &mut Func, module: &Module) { } } +fn find_local_def(func: &Func, local: LocalId) -> Option<&Instr> { + // Simple scan for single definition (SSA-like), but inefficient. + // However, for GlobalGet, it's usually near the top. + // Optimally, use DefUseChain or similar, but we don't have it here easily. + // We scan all blocks? No, that's too slow. + // But typically instructions are defined before use in the same block or dominator. + // For now, scan all instructions in all blocks (very slow!). + // BETTER: Build a def map at start of inlining? + // OR: Just scan the current block backwards? GlobalGet is usually in the same block for simple code. + // Let's scan ALL blocks for now as a quick fix, optimizing later if needed. + // Actually, `webschembly_compiler_ir` might have a helper? + // I'll implement a simple full scan. + for bb in func.bbs.values() { + for instr in &bb.instrs { + if instr.local == Some(local) { + return Some(instr); + } + } + } + None +} + fn inline_non_tail( func: &mut Func, ctx: &mut InlineContext, @@ -274,13 +339,6 @@ fn inline_non_tail( if let Some(dst) = result_local { if !phi_incomings.is_empty() { - let id_val: usize = continuation_bb_id.into(); - if id_val == 69 { - panic!( - "DEBUG_INLINE_HIT_69: BEFORE INSERT: {:#?}", - func.bbs[continuation_bb_id.into()].instrs - ); - } func.bbs[continuation_bb_id.into()].instrs.insert( 0, Instr { @@ -291,12 +349,6 @@ fn inline_non_tail( }, }, ); - if id_val == 69 { - panic!( - "DEBUG_INLINE_HIT_69: AFTER INSERT: {:#?}", - func.bbs[continuation_bb_id.into()].instrs - ); - } } } From f29668638a5a90261c33477b32cfddabeb609d2c Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 8 Jan 2026 10:48:27 +0000 Subject: [PATCH 12/15] chore: auto-fix format and lint --- DEBUGGING.md | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/DEBUGGING.md b/DEBUGGING.md index eb8eab4a..aae4556b 100644 --- a/DEBUGGING.md +++ b/DEBUGGING.md @@ -8,8 +8,8 @@ To debug specific Scheme files or tests, use `just run` within the `webschembly- ### Environment Variables -- `LOG_STDOUT=1`: Enables `log::debug!` output from the runtime and compiler (if configured). Use this to see runtime execution trace, JIT instantiation events, and error messages. -- `LOG=1`: Dumps the generated Intermediate Representation (IR) and Wasm binaries to the `webschembly-js/log/` directory. +- `LOG_STDOUT=1`: Enables `log::debug!` output from the runtime and compiler (if configured). Use this to see runtime execution trace, JIT instantiation events, and error messages. +- `LOG=1`: Dumps the generated Intermediate Representation (IR) and Wasm binaries to the `webschembly-js/log/` directory. ### Command Examples @@ -26,8 +26,8 @@ just LOG=1 run ./fixtures/rec.scm When `LOG=1` is used, the `webschembly-js/log/` directory will contain files named with the timestamp and content description. -- `checks-TIMESTAMP-filename-0.ir`: Typically the main module IR (initial compilation). -- `checks-TIMESTAMP-filename-N.ir`: IR for JIT-compiled functions or stubs. `instantiate_func` or `instantiate_bb` calls in the runtime logs (seen with `LOG_STDOUT=1`) will reference `module_id` and `func_id`, which correlate to these files (though the mapping requires checking the `instantiate` log id vs file index). +- `checks-TIMESTAMP-filename-0.ir`: Typically the main module IR (initial compilation). +- `checks-TIMESTAMP-filename-N.ir`: IR for JIT-compiled functions or stubs. `instantiate_func` or `instantiate_bb` calls in the runtime logs (seen with `LOG_STDOUT=1`) will reference `module_id` and `func_id`, which correlate to these files (though the mapping requires checking the `instantiate` log id vs file index). **Tip**: Look for "instantiate: id:X" in the runtime logs. The corresponding IR file is often suffix `-X.ir`. @@ -37,12 +37,12 @@ When `LOG=1` is used, the `webschembly-js/log/` directory will contain files nam This error occurs when the compiled code attempts to invoke a value as a closure, but the compile-time or run-time check fails. -- **Runtime**: The value on the stack is not a closure object (e.g., encoded as `val_type` mismatch). -- **Compile-time (Optimization)**: If the error appears unconditionally in the IR (e.g., `error "call target is not a closure"`), it means `constant_folding` or another pass determined the check `Is(Closure(None), target)` is false. - - *Watch out for*: Type mismatches in `InstrKind::Is`. Ensure `Closure(None)` (generic check) correctly matches specialized types like `Closure(Some(C))` (constant closure). Using `.remove_constant()` on types before comparison is crucial in `ssa_optimizer.rs`. +- **Runtime**: The value on the stack is not a closure object (e.g., encoded as `val_type` mismatch). +- **Compile-time (Optimization)**: If the error appears unconditionally in the IR (e.g., `error "call target is not a closure"`), it means `constant_folding` or another pass determined the check `Is(Closure(None), target)` is false. + - _Watch out for_: Type mismatches in `InstrKind::Is`. Ensure `Closure(None)` (generic check) correctly matches specialized types like `Closure(Some(C))` (constant closure). Using `.remove_constant()` on types before comparison is crucial in `ssa_optimizer.rs`. ## JIT Optimization Logic -- **`propagate_types` Pass**: Analyzes dataflow to identify constant closures (`Closure(Some(C))`). Logic in `src/ir_processor/propagate_types.rs`. -- **Specialization**: `jit_func.rs` uses available type info (from `locals`) to unbox closure entries or arguments. -- **Constant Folding**: `ssa_optimizer.rs` folds constants. Be careful with strict equality checks on specialized types. +- **`propagate_types` Pass**: Analyzes dataflow to identify constant closures (`Closure(Some(C))`). Logic in `src/ir_processor/propagate_types.rs`. +- **Specialization**: `jit_func.rs` uses available type info (from `locals`) to unbox closure entries or arguments. +- **Constant Folding**: `ssa_optimizer.rs` folds constants. Be careful with strict equality checks on specialized types. From f9b661c5c2782372049968d91008bfd8ab4858b2 Mon Sep 17 00:00:00 2001 From: kgtkr Date: Fri, 9 Jan 2026 01:34:40 +0900 Subject: [PATCH 13/15] =?UTF-8?q?ModuleId=E3=82=92=E5=90=AB=E3=82=81?= =?UTF-8?q?=E3=82=8B=E3=82=88=E3=81=86=E3=81=AB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- webschembly-compiler-crates/ir/src/id.rs | 4 +++- webschembly-compiler/src/ir_generator/module_generator.rs | 1 + webschembly-compiler/src/ir_processor/propagate_types.rs | 2 ++ webschembly-compiler/src/jit/jit_func.rs | 1 + 4 files changed, 7 insertions(+), 1 deletion(-) diff --git a/webschembly-compiler-crates/ir/src/id.rs b/webschembly-compiler-crates/ir/src/id.rs index d36715a8..0c5c771e 100644 --- a/webschembly-compiler-crates/ir/src/id.rs +++ b/webschembly-compiler-crates/ir/src/id.rs @@ -209,6 +209,7 @@ impl fmt::Display for Display<'_, BBIndex> { #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct ConstantClosure { + pub module_id: JitModuleId, pub func_id: JitFuncId, pub env_index: ClosureEnvIndex, } @@ -223,7 +224,8 @@ impl fmt::Display for Display<'_, ConstantClosure> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( f, - "constant_closure({}, {})", + "constant_closure({}, {}, {})", + self.value.module_id.display(self.meta), self.value.func_id.display(self.meta), self.value.env_index.display(self.meta) ) diff --git a/webschembly-compiler/src/ir_generator/module_generator.rs b/webschembly-compiler/src/ir_generator/module_generator.rs index 35fade19..a745c7d2 100644 --- a/webschembly-compiler/src/ir_generator/module_generator.rs +++ b/webschembly-compiler/src/ir_generator/module_generator.rs @@ -562,6 +562,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { let val_type_local = self.builder .local(Type::Val(ValType::Closure(Some(ConstantClosure { + module_id: self.module_generator.id, func_id: JitFuncId::from(func_id), env_index: ClosureEnvIndex(0), })))); diff --git a/webschembly-compiler/src/ir_processor/propagate_types.rs b/webschembly-compiler/src/ir_processor/propagate_types.rs index 86657f5c..61d597dc 100644 --- a/webschembly-compiler/src/ir_processor/propagate_types.rs +++ b/webschembly-compiler/src/ir_processor/propagate_types.rs @@ -56,12 +56,14 @@ pub fn propagate_types(func: &mut Func) { if let Some(dest) = instr.local { let old_val = *lattice.get(dest).unwrap_or(&LatticeValue::Top); let new_val = match &instr.kind { + /* InstrKind::Closure { func_id, env_index, .. } => LatticeValue::Constant(ConstantClosure { func_id: *func_id, env_index: ClosureEnvIndex(*env_index), }), + */ InstrKind::Move(src) => *lattice.get(*src).unwrap_or(&LatticeValue::Top), InstrKind::Phi { incomings, diff --git a/webschembly-compiler/src/jit/jit_func.rs b/webschembly-compiler/src/jit/jit_func.rs index 34012137..70e769a6 100644 --- a/webschembly-compiler/src/jit/jit_func.rs +++ b/webschembly-compiler/src/jit/jit_func.rs @@ -883,6 +883,7 @@ impl JitSpecializedArgFunc { if let Some(local) = body_func.bbs[bb_id].instrs[instr_idx].local { body_func.locals[local].typ = (ValType::Closure(Some(ConstantClosure { + module_id: *module_id, func_id: *func_id, env_index: ClosureEnvIndex(env_index.0), }))) From 113d95d52723aea8d71b3e0bf70d283aab62749f Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 8 Jan 2026 13:07:03 +0000 Subject: [PATCH 14/15] chore: auto-fix format and lint --- .../src/ir_processor/inline.rs | 123 ++++++++---------- 1 file changed, 57 insertions(+), 66 deletions(-) diff --git a/webschembly-compiler/src/ir_processor/inline.rs b/webschembly-compiler/src/ir_processor/inline.rs index 4ff34b38..f0f55ec5 100644 --- a/webschembly-compiler/src/ir_processor/inline.rs +++ b/webschembly-compiler/src/ir_processor/inline.rs @@ -12,20 +12,19 @@ pub fn inline_module(module: &mut Module) { // Scan entry function for GlobalSet with ConstantClosure for bb in entry_func.bbs.values() { for instr in &bb.instrs { - if let InstrKind::GlobalSet(global_id, val_local) = instr.kind { - if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = - entry_func.locals[val_local.into()].typ + if let InstrKind::GlobalSet(global_id, val_local) = instr.kind + && let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = + entry_func.locals[val_local].typ { global_map.insert(global_id, constant); } - } } } let mut new_funcs = VecMap::new(); for (func_id_usize, func) in module.funcs.iter() { - let _func_id = FuncId::from(func_id_usize); + let _func_id = func_id_usize; let mut new_func = func.clone(); run_inlining(&mut new_func, module, &global_map); new_funcs.insert(func_id_usize, new_func); @@ -55,37 +54,35 @@ fn run_inlining( tail_instances: FxHashMap::default(), }; - let mut worklist: Vec = func.bbs.keys().map(BasicBlockId::from).collect(); + let mut worklist: Vec = func.bbs.keys().collect(); while let Some(bb_id) = worklist.pop() { if func.bbs.iter().count() > LIMIT_BB_COUNT { log::debug!("BB limit reached for func {:?}", func.id); break; } - if !func.bbs.contains_key(bb_id.into()) { + if !func.bbs.contains_key(bb_id) { continue; } // 1. Check Non-Tail Calls (Instrs) let mut call_found = None; { - let bb = &func.bbs[bb_id.into()]; + let bb = &func.bbs[bb_id]; for (idx, instr) in bb.instrs.iter().enumerate() { if let InstrKind::CallClosure(call) = &instr.kind { let closure_local = call.closure; let mut constant_opt = None; if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = - func.locals[closure_local.into()].typ + func.locals[closure_local].typ { constant_opt = Some(constant); - } else if let Some(def_instr) = find_local_def(func, closure_local) { - if let InstrKind::GlobalGet(global_id) = def_instr.kind { - if let Some(&constant) = global_map.get(&global_id) { + } else if let Some(def_instr) = find_local_def(func, closure_local) + && let InstrKind::GlobalGet(global_id) = def_instr.kind + && let Some(&constant) = global_map.get(&global_id) { constant_opt = Some(constant); } - } - } if let Some(constant) = constant_opt { call_found = Some((idx, constant, call.clone())); @@ -101,7 +98,7 @@ fn run_inlining( bb_id, constant.func_id ); - let result_local = func.bbs[bb_id.into()].instrs[idx].local; + let result_local = func.bbs[bb_id].instrs[idx].local; let continuation_bb_id = inline_non_tail( func, &mut ctx, @@ -119,22 +116,20 @@ fn run_inlining( // 2. Check Tail Call (Terminator) let mut tail_call_found = None; { - let bb = &func.bbs[bb_id.into()]; + let bb = &func.bbs[bb_id]; if let TerminatorInstr::Exit(ExitInstr::TailCallClosure(call)) = bb.terminator() { let closure_local = call.closure; let mut constant_opt = None; if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = - func.locals[closure_local.into()].typ + func.locals[closure_local].typ { constant_opt = Some(constant); - } else if let Some(def_instr) = find_local_def(func, closure_local) { - if let InstrKind::GlobalGet(global_id) = def_instr.kind { - if let Some(&constant) = global_map.get(&global_id) { + } else if let Some(def_instr) = find_local_def(func, closure_local) + && let InstrKind::GlobalGet(global_id) = def_instr.kind + && let Some(&constant) = global_map.get(&global_id) { constant_opt = Some(constant); } - } - } if let Some(constant) = constant_opt { tail_call_found = Some((constant, call.clone())); @@ -211,30 +206,30 @@ fn inline_non_tail( result_local: Option, ) -> BasicBlockId { let continuation_bb_id = func.bbs.allocate_key(); - let continuation_bb_id = BasicBlockId::from(continuation_bb_id); - let original_terminator = func.bbs[caller_bb_id.into()].terminator().clone(); + let continuation_bb_id = continuation_bb_id; + let original_terminator = func.bbs[caller_bb_id].terminator().clone(); - let instrs_after = func.bbs[caller_bb_id.into()] + let instrs_after = func.bbs[caller_bb_id] .instrs .split_off(instr_idx + 1); - func.bbs[caller_bb_id.into()].instrs.pop(); + func.bbs[caller_bb_id].instrs.pop(); func.bbs.insert_node(BasicBlock { id: continuation_bb_id, instrs: instrs_after, }); - *func.bbs[continuation_bb_id.into()].terminator_mut() = original_terminator; + *func.bbs[continuation_bb_id].terminator_mut() = original_terminator; let callee_id = FuncId::from(constant.func_id); - let callee = &ctx.module.funcs[callee_id.into()]; + let callee = &ctx.module.funcs[callee_id]; let mut local_map = FxHashMap::default(); let mut bb_map = FxHashMap::default(); for old_bb_id_usize in callee.bbs.keys() { - let old_bb_id = BasicBlockId::from(old_bb_id_usize); + let old_bb_id = old_bb_id_usize; let new_bb_id_usize = func.bbs.allocate_key(); - let new_bb_id = BasicBlockId::from(new_bb_id_usize); + let new_bb_id = new_bb_id_usize; bb_map.insert(old_bb_id, new_bb_id); worklist.push(new_bb_id); } @@ -251,32 +246,31 @@ fn inline_non_tail( } for (local_id_usize, local) in callee.locals.iter() { - let local_id = LocalId::from(local_id_usize); - if !local_map.contains_key(&local_id) { + let local_id = local_id_usize; + if let std::collections::hash_map::Entry::Vacant(e) = local_map.entry(local_id) { let new_id_usize = func.locals.push_with(|id| Local { id, typ: local.typ, ..*local }); - let new_id = LocalId::from(new_id_usize); - local_map.insert(local_id, new_id); + let new_id = new_id_usize; + e.insert(new_id); } } let mut phi_incomings = Vec::new(); for (old_bb_id_usize, old_bb) in callee.bbs.iter() { - let old_bb_id = BasicBlockId::from(old_bb_id_usize); + let old_bb_id = old_bb_id_usize; let new_bb_id = bb_map[&old_bb_id]; let mut new_instrs = Vec::new(); for instr in &old_bb.instrs { let mut new_instr = instr.clone(); - if let Some(local) = new_instr.local { - if let Some(&mapped) = local_map.get(&local) { + if let Some(local) = new_instr.local + && let Some(&mapped) = local_map.get(&local) { new_instr.local = Some(mapped); } - } rewrite_usages(&mut new_instr.kind, &local_map); for bb_ref in new_instr.kind.bb_ids_mut() { if let Some(&mapped) = bb_map.get(bb_ref) { @@ -305,9 +299,9 @@ fn inline_non_tail( new_terminator = TerminatorInstr::Jump(continuation_bb_id); } else if let TerminatorInstr::Exit(ExitInstr::TailCallClosure(call)) = &new_terminator { if let Some(dst) = result_local { - let dst_typ = func.locals[dst.into()].typ.clone(); + let dst_typ = func.locals[dst].typ; let temp_local_usize = func.locals.push_with(|id| Local { id, typ: dst_typ }); - let temp_local = LocalId::from(temp_local_usize); + let temp_local = temp_local_usize; new_instrs.push(Instr { local: Some(temp_local), @@ -337,9 +331,9 @@ fn inline_non_tail( }); } - if let Some(dst) = result_local { - if !phi_incomings.is_empty() { - func.bbs[continuation_bb_id.into()].instrs.insert( + if let Some(dst) = result_local + && !phi_incomings.is_empty() { + func.bbs[continuation_bb_id].instrs.insert( 0, Instr { local: Some(dst), @@ -350,7 +344,6 @@ fn inline_non_tail( }, ); } - } let new_entry_id = bb_map[&callee.bb_entry]; @@ -361,7 +354,7 @@ fn inline_non_tail( new_entry_id ); - func.bbs[caller_bb_id.into()].instrs.push(Instr { + func.bbs[caller_bb_id].instrs.push(Instr { local: None, kind: InstrKind::Terminator(TerminatorInstr::Jump(new_entry_id)), }); @@ -385,33 +378,32 @@ fn inline_tail( } else { call.args[i - 1] }; - let entry_bb = &mut func.bbs[info.entry_bb.into()]; + let entry_bb = &mut func.bbs[info.entry_bb]; for instr in &mut entry_bb.instrs { - if let InstrKind::Phi { incomings, .. } = &mut instr.kind { - if instr.local == Some(phi_local) { + if let InstrKind::Phi { incomings, .. } = &mut instr.kind + && instr.local == Some(phi_local) { incomings.push(PhiIncomingValue { local: val, bb: caller_bb_id, }); break; } - } } } - *func.bbs[caller_bb_id.into()].terminator_mut() = TerminatorInstr::Jump(info.entry_bb); + *func.bbs[caller_bb_id].terminator_mut() = TerminatorInstr::Jump(info.entry_bb); return; } let callee_id = FuncId::from(constant.func_id); - let callee = &ctx.module.funcs[callee_id.into()]; + let callee = &ctx.module.funcs[callee_id]; let mut local_map = FxHashMap::default(); let mut bb_map = FxHashMap::default(); for old_bb_id_usize in callee.bbs.keys() { - let old_bb_id = BasicBlockId::from(old_bb_id_usize); + let old_bb_id = old_bb_id_usize; let new_bb_id_usize = func.bbs.allocate_key(); - let new_bb_id = BasicBlockId::from(new_bb_id_usize); + let new_bb_id = new_bb_id_usize; bb_map.insert(old_bb_id, new_bb_id); worklist.push(new_bb_id); } @@ -423,10 +415,10 @@ fn inline_tail( for (i, &arg_local) in callee.args.iter().enumerate() { let new_arg_local_usize = func.locals.push_with(|id| Local { id, - typ: callee.locals[arg_local.into()].typ, - ..callee.locals[arg_local.into()] + typ: callee.locals[arg_local].typ, + ..callee.locals[arg_local] }); - let new_arg_local = LocalId::from(new_arg_local_usize); + let new_arg_local = new_arg_local_usize; local_map.insert(arg_local, new_arg_local); arg_phis.push(new_arg_local); @@ -448,20 +440,20 @@ fn inline_tail( } for (local_id_usize, local) in callee.locals.iter() { - let local_id = LocalId::from(local_id_usize); - if !local_map.contains_key(&local_id) { + let local_id = local_id_usize; + if let std::collections::hash_map::Entry::Vacant(e) = local_map.entry(local_id) { let new_id_usize = func.locals.push_with(|id| Local { id, typ: local.typ, ..*local }); - let new_id = LocalId::from(new_id_usize); - local_map.insert(local_id, new_id); + let new_id = new_id_usize; + e.insert(new_id); } } for (old_bb_id_usize, old_bb) in callee.bbs.iter() { - let old_bb_id = BasicBlockId::from(old_bb_id_usize); + let old_bb_id = old_bb_id_usize; let new_bb_id = bb_map[&old_bb_id]; let mut new_instrs = Vec::new(); @@ -471,11 +463,10 @@ fn inline_tail( for instr in &old_bb.instrs { let mut new_instr = instr.clone(); - if let Some(local) = new_instr.local { - if let Some(&mapped) = local_map.get(&local) { + if let Some(local) = new_instr.local + && let Some(&mapped) = local_map.get(&local) { new_instr.local = Some(mapped); } - } rewrite_usages(&mut new_instr.kind, &local_map); for bb_ref in new_instr.kind.bb_ids_mut() { if let Some(&mapped) = bb_map.get(bb_ref) { @@ -497,10 +488,10 @@ fn inline_tail( id: new_bb_id, instrs: new_instrs, }); - *func.bbs[new_bb_id.into()].terminator_mut() = new_terminator; + *func.bbs[new_bb_id].terminator_mut() = new_terminator; } - *func.bbs[caller_bb_id.into()].terminator_mut() = TerminatorInstr::Jump(new_entry_id); + *func.bbs[caller_bb_id].terminator_mut() = TerminatorInstr::Jump(new_entry_id); log::debug!( "Created new tail instance for func {:?} at entry BB {:?}", From 3707d670ff95b7ca8fa0c3e56ac571ec68684225 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 8 Jan 2026 13:10:01 +0000 Subject: [PATCH 15/15] chore: auto-fix format and lint --- .../src/ir_processor/inline.rs | 76 ++++++++++--------- 1 file changed, 40 insertions(+), 36 deletions(-) diff --git a/webschembly-compiler/src/ir_processor/inline.rs b/webschembly-compiler/src/ir_processor/inline.rs index f0f55ec5..a422bd85 100644 --- a/webschembly-compiler/src/ir_processor/inline.rs +++ b/webschembly-compiler/src/ir_processor/inline.rs @@ -15,9 +15,9 @@ pub fn inline_module(module: &mut Module) { if let InstrKind::GlobalSet(global_id, val_local) = instr.kind && let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = entry_func.locals[val_local].typ - { - global_map.insert(global_id, constant); - } + { + global_map.insert(global_id, constant); + } } } @@ -80,9 +80,10 @@ fn run_inlining( constant_opt = Some(constant); } else if let Some(def_instr) = find_local_def(func, closure_local) && let InstrKind::GlobalGet(global_id) = def_instr.kind - && let Some(&constant) = global_map.get(&global_id) { - constant_opt = Some(constant); - } + && let Some(&constant) = global_map.get(&global_id) + { + constant_opt = Some(constant); + } if let Some(constant) = constant_opt { call_found = Some((idx, constant, call.clone())); @@ -127,9 +128,10 @@ fn run_inlining( constant_opt = Some(constant); } else if let Some(def_instr) = find_local_def(func, closure_local) && let InstrKind::GlobalGet(global_id) = def_instr.kind - && let Some(&constant) = global_map.get(&global_id) { - constant_opt = Some(constant); - } + && let Some(&constant) = global_map.get(&global_id) + { + constant_opt = Some(constant); + } if let Some(constant) = constant_opt { tail_call_found = Some((constant, call.clone())); @@ -209,9 +211,7 @@ fn inline_non_tail( let continuation_bb_id = continuation_bb_id; let original_terminator = func.bbs[caller_bb_id].terminator().clone(); - let instrs_after = func.bbs[caller_bb_id] - .instrs - .split_off(instr_idx + 1); + let instrs_after = func.bbs[caller_bb_id].instrs.split_off(instr_idx + 1); func.bbs[caller_bb_id].instrs.pop(); func.bbs.insert_node(BasicBlock { @@ -268,9 +268,10 @@ fn inline_non_tail( for instr in &old_bb.instrs { let mut new_instr = instr.clone(); if let Some(local) = new_instr.local - && let Some(&mapped) = local_map.get(&local) { - new_instr.local = Some(mapped); - } + && let Some(&mapped) = local_map.get(&local) + { + new_instr.local = Some(mapped); + } rewrite_usages(&mut new_instr.kind, &local_map); for bb_ref in new_instr.kind.bb_ids_mut() { if let Some(&mapped) = bb_map.get(bb_ref) { @@ -332,18 +333,19 @@ fn inline_non_tail( } if let Some(dst) = result_local - && !phi_incomings.is_empty() { - func.bbs[continuation_bb_id].instrs.insert( - 0, - Instr { - local: Some(dst), - kind: InstrKind::Phi { - incomings: phi_incomings, - non_exhaustive: false, - }, + && !phi_incomings.is_empty() + { + func.bbs[continuation_bb_id].instrs.insert( + 0, + Instr { + local: Some(dst), + kind: InstrKind::Phi { + incomings: phi_incomings, + non_exhaustive: false, }, - ); - } + }, + ); + } let new_entry_id = bb_map[&callee.bb_entry]; @@ -381,13 +383,14 @@ fn inline_tail( let entry_bb = &mut func.bbs[info.entry_bb]; for instr in &mut entry_bb.instrs { if let InstrKind::Phi { incomings, .. } = &mut instr.kind - && instr.local == Some(phi_local) { - incomings.push(PhiIncomingValue { - local: val, - bb: caller_bb_id, - }); - break; - } + && instr.local == Some(phi_local) + { + incomings.push(PhiIncomingValue { + local: val, + bb: caller_bb_id, + }); + break; + } } } *func.bbs[caller_bb_id].terminator_mut() = TerminatorInstr::Jump(info.entry_bb); @@ -464,9 +467,10 @@ fn inline_tail( for instr in &old_bb.instrs { let mut new_instr = instr.clone(); if let Some(local) = new_instr.local - && let Some(&mapped) = local_map.get(&local) { - new_instr.local = Some(mapped); - } + && let Some(&mapped) = local_map.get(&local) + { + new_instr.local = Some(mapped); + } rewrite_usages(&mut new_instr.kind, &local_map); for bb_ref in new_instr.kind.bb_ids_mut() { if let Some(&mapped) = bb_map.get(bb_ref) {