From 86953dd9d309db6a65c47aa0d4b2ee6c81ee9413 Mon Sep 17 00:00:00 2001 From: imxyy_soope_ Date: Mon, 26 Jan 2026 22:02:29 +0800 Subject: [PATCH] refactor: thunk scope --- nix-js/Cargo.toml | 2 +- .../{scc_optimization.rs => thunk_scope.rs} | 0 nix-js/src/codegen.rs | 105 +++---- nix-js/src/context.rs | 171 +++--------- nix-js/src/ir.rs | 19 +- nix-js/src/ir/downgrade.rs | 57 +--- nix-js/src/ir/utils.rs | 264 ++++++------------ .../{scc_optimization.rs => thunk_scope.rs} | 0 8 files changed, 180 insertions(+), 438 deletions(-) rename nix-js/benches/{scc_optimization.rs => thunk_scope.rs} (100%) rename nix-js/tests/{scc_optimization.rs => thunk_scope.rs} (100%) diff --git a/nix-js/Cargo.toml b/nix-js/Cargo.toml index 6f98400..c9f5f31 100644 --- a/nix-js/Cargo.toml +++ b/nix-js/Cargo.toml @@ -75,7 +75,7 @@ name = "builtins" harness = false [[bench]] -name = "scc_optimization" +name = "thunk_scope" harness = false [[bench]] diff --git a/nix-js/benches/scc_optimization.rs b/nix-js/benches/thunk_scope.rs similarity index 100% rename from nix-js/benches/scc_optimization.rs rename to nix-js/benches/thunk_scope.rs diff --git a/nix-js/src/codegen.rs b/nix-js/src/codegen.rs index e055611..cf2028e 100644 --- a/nix-js/src/codegen.rs +++ b/nix-js/src/codegen.rs @@ -107,23 +107,9 @@ impl Compile for Ir { Ir::List(x) => x.compile(ctx), Ir::Call(x) => x.compile(ctx), Ir::Arg(x) => format!("arg{}", x.inner.0), - Ir::Let(x) => x.compile(ctx), + Ir::TopLevel(x) => x.compile(ctx), Ir::Select(x) => x.compile(ctx), - &Ir::Thunk(Thunk { - inner: expr_id, - span, - }) => { - let inner = ctx.get_ir(expr_id).compile(ctx); - format!( - "Nix.createThunk(()=>({}),\"expr{} {}:{}:{}\")", - inner, - expr_id.0, - ctx.get_current_source().get_name(), - usize::from(span.start()), - usize::from(span.end()) - ) - } - &Ir::ExprRef(ExprRef { inner: expr_id, .. }) => { + &Ir::Thunk(Thunk { inner: expr_id, .. }) => { format!("expr{}", expr_id.0) } Ir::Builtins(_) => "Nix.builtins".to_string(), @@ -220,7 +206,13 @@ impl Compile for UnOp { impl Compile for Func { fn compile(&self, ctx: &Ctx) -> String { let id = ctx.get_ir(self.arg).as_ref().unwrap_arg().inner.0; - let body = ctx.get_ir(self.body).compile(ctx); + let thunk_defs = compile_thunks(&self.thunks, ctx); + let body_code = ctx.get_ir(self.body).compile(ctx); + let body = if thunk_defs.is_empty() { + body_code + } else { + format!("{{{}return {}}}", thunk_defs, body_code) + }; if let Some(Param { required, @@ -232,9 +224,17 @@ impl Compile for Func { let required = format!("[{}]", required.join(",")); let mut optional = optional.iter().map(|&sym| ctx.get_sym(sym).escape_quote()); let optional = format!("[{}]", optional.join(",")); - format!("Nix.mkFunction(arg{id}=>({body}),{required},{optional},{ellipsis})") + if thunk_defs.is_empty() { + format!("Nix.mkFunction(arg{id}=>({body}),{required},{optional},{ellipsis})") + } else { + format!("Nix.mkFunction(arg{id}=>{body},{required},{optional},{ellipsis})") + } } else { - format!("arg{id}=>({body})") + if thunk_defs.is_empty() { + format!("arg{id}=>({body})") + } else { + format!("arg{id}=>{body}") + } } } } @@ -248,51 +248,38 @@ impl Compile for Call { } } -/// Determines if a Thunk should be kept (not unwrapped) for non-recursive let bindings. -/// Returns true for complex expressions that should remain lazy to preserve Nix semantics. -fn should_keep_thunk(ir: &Ir) -> bool { - match ir { - // Simple literals can be evaluated eagerly - Ir::Int(_) | Ir::Float(_) | Ir::Bool(_) | Ir::Null(_) | Ir::Str(_) => false, - // Builtin references are safe to evaluate eagerly - Ir::Builtin(_) | Ir::Builtins(_) => false, - Ir::ExprRef(_) => true, - // Everything else should remain lazy: - _ => true, +fn compile_thunks(thunks: &[(ExprId, ExprId)], ctx: &Ctx) -> String { + if thunks.is_empty() { + return String::new(); } + thunks + .iter() + .map(|&(slot, inner)| { + let inner_code = ctx.get_ir(inner).compile(ctx); + let inner_span = ctx.get_ir(inner).span(); + format!( + "let expr{}=Nix.createThunk(()=>({}),\"expr{} {}:{}:{}\")", + slot.0, + inner_code, + slot.0, + ctx.get_current_source().get_name(), + usize::from(inner_span.start()), + usize::from(inner_span.end()) + ) + }) + .join(";") + + ";" } -fn unwrap_thunk(ir: &Ir, ctx: &impl CodegenContext) -> String { - if let Ir::Thunk(Thunk { inner, .. }) = ir { - let inner_ir = ctx.get_ir(*inner); - if should_keep_thunk(inner_ir) { - ir.compile(ctx) - } else { - inner_ir.compile(ctx) - } - } else { - ir.compile(ctx) - } -} - -impl Compile for Let { +impl Compile for TopLevel { fn compile(&self, ctx: &Ctx) -> String { - let info = &self.binding_sccs; - let mut js_statements = Vec::new(); - - for (scc_exprs, is_recursive) in info.sccs.iter() { - for &expr in scc_exprs { - let value = if *is_recursive { - ctx.get_ir(expr).compile(ctx) - } else { - unwrap_thunk(ctx.get_ir(expr), ctx) - }; - js_statements.push(format!("const expr{}={}", expr.0, value)); - } - } - + let thunk_defs = compile_thunks(&self.thunks, ctx); let body = ctx.get_ir(self.body).compile(ctx); - format!("(()=>{{{};return {}}})()", js_statements.join(";"), body) + if thunk_defs.is_empty() { + body + } else { + format!("(()=>{{{}return {}}})()", thunk_defs, body) + } } } diff --git a/nix-js/src/context.rs b/nix-js/src/context.rs index a285ec8..6bf4752 100644 --- a/nix-js/src/context.rs +++ b/nix-js/src/context.rs @@ -1,28 +1,21 @@ use std::path::Path; use std::ptr::NonNull; -use hashbrown::{HashMap, HashSet}; +use hashbrown::HashMap; use itertools::Itertools as _; -use petgraph::graphmap::DiGraphMap; use rnix::TextRange; use string_interner::DefaultStringInterner; use crate::codegen::{CodegenContext, compile}; use crate::error::{Error, Result, Source}; use crate::ir::{ - Arg, ArgId, Bool, Builtin, Downgrade as _, DowngradeContext, ExprId, ExprRef, Ir, Null, SymId, - Thunk, ToIr as _, synthetic_span, + Arg, ArgId, Bool, Builtin, Downgrade as _, DowngradeContext, ExprId, Ir, Null, SymId, Thunk, + ToIr as _, synthetic_span, }; use crate::runtime::{Runtime, RuntimeContext}; use crate::store::{Store, StoreBackend, StoreConfig}; use crate::value::Value; -#[derive(Debug)] -pub(crate) struct SccInfo { - /// list of SCCs (exprs, recursive) - pub(crate) sccs: Vec<(Vec, bool)>, -} - pub struct Context { ctx: Ctx, runtime: Runtime, @@ -256,14 +249,6 @@ impl RuntimeContext for Ctx { } } -struct DependencyTracker { - graph: DiGraphMap, - current_binding: Option, - let_scope_exprs: HashSet, - // The outer binding that owns this tracker (for nested let scopes in function params) - owner_binding: Option, -} - enum Scope<'ctx> { Global(&'ctx HashMap), Let(HashMap), @@ -292,7 +277,7 @@ pub struct DowngradeCtx<'ctx> { irs: Vec>, scopes: Vec>, arg_id: usize, - dep_tracker_stack: Vec, + thunk_scopes: Vec>, } impl<'ctx> DowngradeCtx<'ctx> { @@ -301,7 +286,7 @@ impl<'ctx> DowngradeCtx<'ctx> { scopes: vec![Scope::Global(global)], irs: vec![], arg_id: 0, - dep_tracker_stack: Vec::new(), + thunk_scopes: vec![Vec::new()], ctx, } } @@ -346,14 +331,15 @@ impl DowngradeContext for DowngradeCtx<'_> { | Ir::Float(_) | Ir::Bool(_) | Ir::Null(_) - | Ir::Str(_) => id, - _ => self.new_expr( - Thunk { - inner: id, - span: ir.span(), - } - .to_ir(), - ), + | Ir::Str(_) + | Ir::Thunk(_) => id, + _ => { + let span = ir.span(); + let slot = self.reserve_slots(1).next().expect("reserve_slots failed"); + self.replace_ir(slot, Thunk { inner: slot, span }.to_ir()); + self.register_thunk(slot, id); + slot + } } } @@ -375,45 +361,7 @@ impl DowngradeContext for DowngradeCtx<'_> { } Scope::Let(let_scope) => { if let Some(&expr) = let_scope.get(&sym) { - // Find which tracker contains this expression - let expr_tracker_idx = self - .dep_tracker_stack - .iter() - .position(|t| t.let_scope_exprs.contains(&expr)); - - // Find the innermost tracker with a current_binding - let current_tracker_idx = self - .dep_tracker_stack - .iter() - .rposition(|t| t.current_binding.is_some()); - - // Record dependency if both exist - if let (Some(expr_idx), Some(curr_idx)) = - (expr_tracker_idx, current_tracker_idx) - { - let current_binding = self.dep_tracker_stack[curr_idx] - .current_binding - .expect("current_binding not set"); - let owner_binding = self.dep_tracker_stack[curr_idx].owner_binding; - - // If referencing from inner scope to outer scope - if curr_idx >= expr_idx { - let tracker = &mut self.dep_tracker_stack[expr_idx]; - let from_node = current_binding; - let to_node = expr; - if curr_idx > expr_idx { - // Cross-scope reference: use owner_binding if available - if let Some(owner) = owner_binding { - tracker.graph.add_edge(owner, expr, ()); - } - } else { - // Same-level reference: record directly - tracker.graph.add_edge(from_node, to_node, ()); - } - } - } - - return Ok(self.new_expr(ExprRef { inner: expr, span }.to_ir())); + return Ok(self.new_expr(Thunk { inner: expr, span }.to_ir())); } } &Scope::Param(param_sym, expr) => { @@ -486,11 +434,15 @@ impl DowngradeContext for DowngradeCtx<'_> { } fn downgrade(mut self, root: rnix::ast::Expr) -> Result { - let root = root.downgrade(&mut self)?; + use crate::ir::TopLevel; + let body = root.downgrade(&mut self)?; + let thunks = self.pop_thunk_scope(); + let span = self.get_ir(body).span(); + let top_level = self.new_expr(TopLevel { body, thunks, span }.to_ir()); self.ctx .irs .extend(self.irs.into_iter().map(Option::unwrap)); - Ok(root) + Ok(top_level) } fn with_let_scope(&mut self, bindings: HashMap, f: F) -> R @@ -515,83 +467,26 @@ impl DowngradeContext for DowngradeCtx<'_> { where F: FnOnce(&mut Self) -> R, { + let namespace = self.maybe_thunk(namespace); self.scopes.push(Scope::With(namespace)); let mut guard = ScopeGuard { ctx: self }; f(guard.as_ctx()) } - fn push_dep_tracker(&mut self, slots: &[ExprId]) { - let mut graph = DiGraphMap::new(); - let mut let_scope_exprs = HashSet::new(); - - for &expr in slots.iter() { - graph.add_node(expr); - let_scope_exprs.insert(expr); - } - - self.dep_tracker_stack.push(DependencyTracker { - graph, - current_binding: None, - let_scope_exprs, - owner_binding: None, - }); + fn push_thunk_scope(&mut self) { + self.thunk_scopes.push(Vec::new()); } - fn push_dep_tracker_with_owner(&mut self, slots: &[ExprId], owner: ExprId) { - let mut graph = DiGraphMap::new(); - let mut let_scope_exprs = HashSet::new(); - - for &expr in slots.iter() { - graph.add_node(expr); - let_scope_exprs.insert(expr); - } - - self.dep_tracker_stack.push(DependencyTracker { - graph, - current_binding: None, - let_scope_exprs, - owner_binding: Some(owner), - }); - } - - fn get_current_binding(&self) -> Option { - self.dep_tracker_stack - .last() - .and_then(|t| t.current_binding) - } - - fn set_current_binding(&mut self, expr: Option) { - if let Some(tracker) = self.dep_tracker_stack.last_mut() { - tracker.current_binding = expr; - } - } - - fn pop_dep_tracker(&mut self) -> Result { - let tracker = self - .dep_tracker_stack + fn pop_thunk_scope(&mut self) -> Vec<(ExprId, ExprId)> { + self.thunk_scopes .pop() - .expect("pop_dep_tracker without active tracker"); + .expect("pop_thunk_scope without active scope") + } - use petgraph::algo::kosaraju_scc; - let sccs = kosaraju_scc(&tracker.graph); - - let mut sccs_topo = Vec::new(); - - for scc_nodes in sccs.iter() { - let mut scc_exprs = Vec::new(); - let mut is_recursive = scc_nodes.len() > 1; - - for &expr in scc_nodes { - scc_exprs.push(expr); - - if !is_recursive && tracker.graph.contains_edge(expr, expr) { - is_recursive = true; - } - } - - sccs_topo.push((scc_exprs, is_recursive)); - } - - Ok(SccInfo { sccs: sccs_topo }) + fn register_thunk(&mut self, slot: ExprId, inner: ExprId) { + self.thunk_scopes + .last_mut() + .expect("register_thunk without active scope") + .push((slot, inner)); } } diff --git a/nix-js/src/ir.rs b/nix-js/src/ir.rs index 064a3a1..ae9e0b0 100644 --- a/nix-js/src/ir.rs +++ b/nix-js/src/ir.rs @@ -3,7 +3,6 @@ use hashbrown::HashMap; use rnix::{TextRange, ast}; use string_interner::symbol::SymbolU32; -use crate::context::SccInfo; use crate::error::{Error, Result, Source}; use crate::value::format_symbol; use nix_js_macros::ir; @@ -44,11 +43,9 @@ pub trait DowngradeContext { where F: FnOnce(&mut Self) -> R; - fn push_dep_tracker(&mut self, slots: &[ExprId]); - fn push_dep_tracker_with_owner(&mut self, slots: &[ExprId], owner: ExprId); - fn get_current_binding(&self) -> Option; - fn set_current_binding(&mut self, expr: Option); - fn pop_dep_tracker(&mut self) -> Result; + fn push_thunk_scope(&mut self); + fn pop_thunk_scope(&mut self) -> Vec<(ExprId, ExprId)>; + fn register_thunk(&mut self, slot: ExprId, inner: ExprId); } ir! { @@ -71,10 +68,9 @@ ir! { Assert { pub assertion: ExprId, pub expr: ExprId, pub assertion_raw: String }, ConcatStrings { pub parts: Vec }, Path { pub expr: ExprId }, - Func { pub body: ExprId, pub param: Option, pub arg: ExprId }, - Let { pub binding_sccs: SccInfo, pub body: ExprId }, + Func { pub body: ExprId, pub param: Option, pub arg: ExprId, pub thunks: Vec<(ExprId, ExprId)> }, + TopLevel { pub body: ExprId, pub thunks: Vec<(ExprId, ExprId)> }, Arg(ArgId), - ExprRef(ExprId), Thunk(ExprId), Builtins, Builtin(SymId), @@ -101,10 +97,9 @@ impl Ir { Ir::ConcatStrings(c) => c.span, Ir::Path(p) => p.span, Ir::Func(f) => f.span, - Ir::Let(l) => l.span, + Ir::TopLevel(t) => t.span, Ir::Arg(a) => a.span, - Ir::ExprRef(e) => e.span, - Ir::Thunk(t) => t.span, + Ir::Thunk(e) => e.span, Ir::Builtins(b) => b.span, Ir::Builtin(b) => b.span, Ir::CurPos(c) => c.span, diff --git a/nix-js/src/ir/downgrade.rs b/nix-js/src/ir/downgrade.rs index 2c27d5b..bf84981 100644 --- a/nix-js/src/ir/downgrade.rs +++ b/nix-js/src/ir/downgrade.rs @@ -158,7 +158,7 @@ impl Downgrade for ast::Str { ast::InterpolPart::Literal(lit) => Ok(ctx.new_expr(Str { val: lit, span }.to_ir())), ast::InterpolPart::Interpolation(interpol) => { let inner = interpol.expr().unwrap().downgrade(ctx)?; - Ok(ctx.new_expr(Thunk { inner, span }.to_ir())) + Ok(ctx.maybe_thunk(inner)) } }) .collect::>>()?; @@ -220,7 +220,7 @@ impl Downgrade for ast::AttrSet { // rec { a = 1; b = a; } => let a = 1; b = a; in { inherit a b; } let entries: Vec<_> = self.entries().collect(); - let (binding_sccs, body) = downgrade_let_bindings(entries, ctx, |ctx, binding_keys| { + downgrade_let_bindings(entries, ctx, span, |ctx, binding_keys| { // Create plain attrset as body with inherit let mut attrs = AttrSet { stcs: HashMap::new(), @@ -229,22 +229,12 @@ impl Downgrade for ast::AttrSet { }; for sym in binding_keys { - // FIXME: span let expr = ctx.lookup(*sym, synthetic_span())?; attrs.stcs.insert(*sym, (expr, synthetic_span())); } Ok(ctx.new_expr(attrs.to_ir())) - })?; - - Ok(ctx.new_expr( - Let { - body, - binding_sccs, - span, - } - .to_ir(), - )) + }) } } @@ -308,17 +298,8 @@ impl Downgrade for ast::Select { let expr = self.expr().unwrap().downgrade(ctx)?; let attrpath = downgrade_attrpath(self.attrpath().unwrap(), ctx)?; let default = if let Some(default) = self.default_expr() { - let span = default.syntax().text_range(); let default_expr = default.downgrade(ctx)?; - Some( - ctx.new_expr( - Thunk { - inner: default_expr, - span, - } - .to_ir(), - ), - ) + Some(ctx.maybe_thunk(default_expr)) } else { None }; @@ -378,17 +359,9 @@ impl Downgrade for ast::LetIn { let body_expr = self.body().unwrap(); let span = self.syntax().text_range(); - let (binding_sccs, body) = - downgrade_let_bindings(entries, ctx, |ctx, _binding_keys| body_expr.downgrade(ctx))?; - - Ok(ctx.new_expr( - Let { - body, - binding_sccs, - span, - } - .to_ir(), - )) + downgrade_let_bindings(entries, ctx, span, |ctx, _binding_keys| { + body_expr.downgrade(ctx) + }) } } @@ -412,9 +385,10 @@ impl Downgrade for ast::Lambda { let raw_param = self.param().unwrap(); let arg = ctx.new_arg(raw_param.syntax().text_range()); + ctx.push_thunk_scope(); + let param; let body; - let span = self.body().unwrap().syntax().text_range(); match raw_param { ast::Param::IdentParam(id) => { @@ -436,7 +410,6 @@ impl Downgrade for ast::Lambda { let PatternBindings { body: inner_body, - scc_info, required, optional, } = downgrade_pattern_bindings(pat_entries, alias, arg, ctx, |ctx, _| { @@ -449,24 +422,18 @@ impl Downgrade for ast::Lambda { ellipsis, }); - body = ctx.new_expr( - Let { - body: inner_body, - binding_sccs: scc_info, - span, - } - .to_ir(), - ); + body = inner_body; } } + let thunks = ctx.pop_thunk_scope(); let span = self.syntax().text_range(); - // The function's body and parameters are now stored directly in the `Func` node. Ok(ctx.new_expr( Func { body, param, arg, + thunks, span, } .to_ir(), diff --git a/nix-js/src/ir/utils.rs b/nix-js/src/ir/utils.rs index 0462f81..701a2f4 100644 --- a/nix-js/src/ir/utils.rs +++ b/nix-js/src/ir/utils.rs @@ -5,6 +5,7 @@ use hashbrown::hash_map::Entry; use hashbrown::{HashMap, HashSet}; use itertools::Itertools as _; use rnix::ast; +use rnix::TextRange; use rowan::ast::AstNode; use crate::error::{Error, Result}; @@ -99,13 +100,7 @@ pub fn downgrade_inherit( } .to_ir(), ); - ctx.new_expr( - Thunk { - inner: select_expr, - span, - } - .to_ir(), - ) + ctx.maybe_thunk(select_expr) } else { ctx.lookup(ident, span)? }; @@ -221,20 +216,12 @@ pub fn downgrade_static_attrpathvalue( pub struct PatternBindings { pub body: ExprId, - pub scc_info: SccInfo, pub required: Vec, pub optional: Vec, } -/// Helper function for Lambda pattern parameters with SCC analysis. -/// Processes pattern entries like `{ a, b ? 2, ... }@alias` and creates optimized bindings. -/// -/// # Parameters -/// - `pat_entries`: Iterator over pattern entries from the AST -/// - `alias`: Optional alias symbol (from @alias syntax) -/// - `arg`: The argument expression to extract from -/// -/// Returns a tuple of (binding slots, body, SCC info, required params, allowed params) +/// Helper function for Lambda pattern parameters. +/// Processes pattern entries like `{ a, b ? 2, ... }@alias` and creates bindings. pub fn downgrade_pattern_bindings( pat_entries: impl Iterator, alias: Option, @@ -294,98 +281,6 @@ where } }); - // Get the owner from outer tracker's current_binding - let owner = ctx.get_current_binding(); - - let (scc_info, body) = downgrade_bindings_generic_with_owner( - ctx, - binding_keys, - |ctx, sym_to_slot| { - let mut bindings = HashMap::new(); - - for Param { - sym, - sym_span, - default, - span, - } in params - { - let slot = *sym_to_slot.get(&sym).unwrap(); - ctx.set_current_binding(Some(slot)); - - let default = if let Some(default) = default { - let default = default.clone().downgrade(ctx)?; - Some(ctx.maybe_thunk(default)) - } else { - None - }; - - let select_expr = ctx.new_expr( - Select { - expr: arg, - attrpath: vec![Attr::Str(sym, sym_span)], - default, - span, - } - .to_ir(), - ); - bindings.insert(sym, select_expr); - ctx.set_current_binding(None); - } - - if let Some(alias_sym) = alias { - bindings.insert(alias_sym, arg); - } - - Ok(bindings) - }, - body_fn, - owner, // Pass the owner to track cross-scope dependencies - )?; - - Ok(PatternBindings { - body, - scc_info, - required, - optional, - }) -} - -/// Generic helper function to downgrade bindings with SCC analysis. -/// This is the core logic for let bindings, extracted for reuse. -/// -/// # Parameters -/// - `binding_keys`: The symbols for all bindings -/// - `compute_bindings_fn`: Called in let scope with sym_to_slot mapping to compute binding values -/// - `body_fn`: Called in let scope to compute the body expression -/// -/// Returns a tuple of (binding slots, body result, SCC info) -pub fn downgrade_bindings_generic( - ctx: &mut Ctx, - binding_keys: Vec, - compute_bindings_fn: B, - body_fn: F, -) -> Result<(SccInfo, ExprId)> -where - Ctx: DowngradeContext, - B: FnOnce(&mut Ctx, &HashMap) -> Result>, - F: FnOnce(&mut Ctx, &[SymId]) -> Result, -{ - downgrade_bindings_generic_with_owner(ctx, binding_keys, compute_bindings_fn, body_fn, None) -} - -pub fn downgrade_bindings_generic_with_owner( - ctx: &mut Ctx, - binding_keys: Vec, - compute_bindings_fn: B, - body_fn: F, - owner: Option, -) -> Result<(SccInfo, ExprId)> -where - Ctx: DowngradeContext, - B: FnOnce(&mut Ctx, &HashMap) -> Result>, - F: FnOnce(&mut Ctx, &[SymId]) -> Result, -{ let slots: Vec<_> = ctx.reserve_slots(binding_keys.len()).collect(); let let_bindings: HashMap<_, _> = binding_keys .iter() @@ -393,53 +288,63 @@ where .zip(slots.iter().copied()) .collect(); - if let Some(owner_binding) = owner { - ctx.push_dep_tracker_with_owner(&slots, owner_binding); - } else { - ctx.push_dep_tracker(&slots); + for &slot in &slots { + let span = synthetic_span(); + ctx.replace_ir(slot, Thunk { inner: slot, span }.to_ir()); } ctx.with_let_scope(let_bindings.clone(), |ctx| { - let bindings = compute_bindings_fn(ctx, &let_bindings)?; + for Param { + sym, + sym_span, + default, + span, + } in params + { + let slot = *let_bindings.get(&sym).unwrap(); - let scc_info = ctx.pop_dep_tracker()?; - - for (sym, slot) in binding_keys.iter().copied().zip(slots.iter()) { - if let Some(&expr) = bindings.get(&sym) { - ctx.replace_ir( - *slot, - Thunk { - inner: expr, - span: ctx.get_ir(expr).span(), - } - .to_ir(), - ); + let default = if let Some(default) = default { + let default = default.clone().downgrade(ctx)?; + Some(ctx.maybe_thunk(default)) } else { - return Err(Error::internal(format!( - "binding '{}' not found", - format_symbol(ctx.get_sym(sym)) - ))); - } + None + }; + + let select_expr = ctx.new_expr( + Select { + expr: arg, + attrpath: vec![Attr::Str(sym, sym_span)], + default, + span, + } + .to_ir(), + ); + ctx.register_thunk(slot, select_expr); + } + + if let Some(alias_sym) = alias { + let slot = *let_bindings.get(&alias_sym).unwrap(); + ctx.register_thunk(slot, arg); } let body = body_fn(ctx, &binding_keys)?; - Ok((scc_info, body)) + Ok(PatternBindings { + body, + required, + optional, + }) }) } /// Helper function to downgrade entries with let bindings semantics. /// This extracts common logic for both `rec` attribute sets and `let...in` expressions. -/// -/// Returns a tuple of (binding slots, body result, SCC info) where: -/// - binding slots: pre-allocated expression slots for the bindings -/// - body result: the result of calling `body_fn` in the let scope -/// - SCC info: strongly connected components information for optimization pub fn downgrade_let_bindings( entries: Vec, ctx: &mut Ctx, + _span: TextRange, body_fn: F, -) -> Result<(SccInfo, ExprId)> +) -> Result where Ctx: DowngradeContext, F: FnOnce(&mut Ctx, &[SymId]) -> Result, @@ -469,8 +374,6 @@ where let attrpath = value.attrpath().unwrap(); let attrs_vec: Vec<_> = attrpath.attrs().collect(); - // Only check for duplicate definitions if this is a top-level binding (path length == 1) - // For nested paths (e.g., types.a, types.b), they will be merged into the same attrset if attrs_vec.len() == 1 { if let Some(ast::Attr::Ident(ident)) = attrs_vec.first() { let sym = ctx.new_sym(ident.to_string()); @@ -486,7 +389,6 @@ where } } } else if attrs_vec.len() > 1 { - // For nested paths, just record the first-level name without checking duplicates if let Some(ast::Attr::Ident(ident)) = attrs_vec.first() { let sym = ctx.new_sym(ident.to_string()); binding_syms.insert(sym); @@ -497,51 +399,47 @@ where } let binding_keys: Vec<_> = binding_syms.into_iter().collect(); + let slots: Vec<_> = ctx.reserve_slots(binding_keys.len()).collect(); + let let_bindings: HashMap<_, _> = binding_keys + .iter() + .copied() + .zip(slots.iter().copied()) + .collect(); - downgrade_bindings_generic( - ctx, - binding_keys, - |ctx, sym_to_slot| { - let mut temp_attrs = AttrSet { - stcs: HashMap::new(), - dyns: Vec::new(), - span: synthetic_span(), - }; + for &slot in &slots { + let span = synthetic_span(); + ctx.replace_ir(slot, Thunk { inner: slot, span }.to_ir()); + } - for entry in entries { - match entry { - ast::Entry::Inherit(inherit) => { - for attr in inherit.attrs() { - if let ast::Attr::Ident(ident) = attr { - let sym = ctx.new_sym(ident.to_string()); - let slot = *sym_to_slot.get(&sym).unwrap(); - ctx.set_current_binding(Some(slot)); - } - } - downgrade_inherit(inherit, &mut temp_attrs.stcs, ctx)?; - ctx.set_current_binding(None); - } - ast::Entry::AttrpathValue(value) => { - let attrpath = value.attrpath().unwrap(); - if let Some(first_attr) = attrpath.attrs().next() - && let ast::Attr::Ident(ident) = first_attr - { - let sym = ctx.new_sym(ident.to_string()); - let slot = *sym_to_slot.get(&sym).unwrap(); - ctx.set_current_binding(Some(slot)); - } - downgrade_static_attrpathvalue(value, &mut temp_attrs, ctx)?; - ctx.set_current_binding(None); - } + ctx.with_let_scope(let_bindings.clone(), |ctx| { + let mut temp_attrs = AttrSet { + stcs: HashMap::new(), + dyns: Vec::new(), + span: synthetic_span(), + }; + + for entry in entries { + match entry { + ast::Entry::Inherit(inherit) => { + downgrade_inherit(inherit, &mut temp_attrs.stcs, ctx)?; + } + ast::Entry::AttrpathValue(value) => { + downgrade_static_attrpathvalue(value, &mut temp_attrs, ctx)?; } } + } - Ok(temp_attrs - .stcs - .into_iter() - .map(|(k, (v, _))| (k, v)) - .collect()) - }, - body_fn, - ) + for (sym, slot) in binding_keys.iter().copied().zip(slots.iter()) { + if let Some(&(expr, _)) = temp_attrs.stcs.get(&sym) { + ctx.register_thunk(*slot, expr); + } else { + return Err(Error::internal(format!( + "binding '{}' not found", + format_symbol(ctx.get_sym(sym)) + ))); + } + } + + body_fn(ctx, &binding_keys) + }) } diff --git a/nix-js/tests/scc_optimization.rs b/nix-js/tests/thunk_scope.rs similarity index 100% rename from nix-js/tests/scc_optimization.rs rename to nix-js/tests/thunk_scope.rs