diff --git a/src/ast/compiler/optimizer/engine.rs b/src/ast/compiler/optimizer/engine.rs index c5e9bbe..8a2e522 100644 --- a/src/ast/compiler/optimizer/engine.rs +++ b/src/ast/compiler/optimizer/engine.rs @@ -504,17 +504,29 @@ impl Optimizer { } } + let unmapped_addr = if let BoundKind::Define { addr, .. } = &e.kind { + Some(*addr) + } else { + None + }; + let opt = self.visit_node(e.clone(), sub, path); - if let BoundKind::Define { addr, value, .. } = &opt.kind - && !info.assigned.contains(addr) + if let BoundKind::Define { value, .. } = &opt.kind + && let Some(orig_addr) = unmapped_addr + && !info.assigned.contains(&orig_addr) { - if let BoundKind::Constant(val) = &value.kind { - sub.add_value(*addr, val.clone()); - } else if let BoundKind::Lambda { upvalues, .. } = &value.kind + let mut core_value = value.as_ref(); + while let BoundKind::Expansion { bound_expanded, .. } = &core_value.kind { + core_value = bound_expanded.as_ref(); + } + + if let BoundKind::Constant(val) = &core_value.kind { + sub.add_value(orig_addr, val.clone()); + } else if let BoundKind::Lambda { upvalues, .. } = &core_value.kind && upvalues.is_empty() { - sub.add_ast_substitution(*addr, (**value).clone()); + sub.add_ast_substitution(orig_addr, core_value.clone()); } } @@ -560,10 +572,13 @@ impl Optimizer { for (old_idx, capture_addr) in upvalues.iter().enumerate() { let mut inlined_val = None; - if !sub.assigned.contains(capture_addr) - && let Some(val) = sub.get_value(capture_addr) - { - inlined_val = Some(val.clone()); + let mut inlined_ast = None; + if !sub.assigned.contains(capture_addr) { + if let Some(val) = sub.get_value(capture_addr) { + inlined_val = Some(val.clone()); + } else if let Some(ast) = sub.ast_substitutions.get(capture_addr) { + inlined_ast = Some(Rc::clone(ast)); + } } if let Address::Local(slot) = capture_addr { @@ -575,6 +590,11 @@ impl Optimizer { .add_value(Address::Upvalue(UpvalueIdx(old_idx as u32)), val); mapping.push(None); upvalues_changed = true; + } else if let Some(ast) = inlined_ast { + next_inner_subs + .add_ast_substitution(Address::Upvalue(UpvalueIdx(old_idx as u32)), (*ast).clone()); + mapping.push(None); + upvalues_changed = true; } else { mapping.push(Some(new_upvalues.len() as u32)); let mapped = sub.map_address(*capture_addr); diff --git a/src/ast/compiler/optimizer/inliner.rs b/src/ast/compiler/optimizer/inliner.rs index fbc1a26..0169c61 100644 --- a/src/ast/compiler/optimizer/inliner.rs +++ b/src/ast/compiler/optimizer/inliner.rs @@ -112,18 +112,23 @@ impl<'a> Inliner<'a> { if let Some(arg) = args.get(*offset) && !body_usage.is_assigned(addr) { - if let BoundKind::Constant(val) = &arg.kind { + let mut core_arg = arg.as_ref(); + while let BoundKind::Expansion { bound_expanded, .. } = &core_arg.kind { + core_arg = bound_expanded.as_ref(); + } + + if let BoundKind::Constant(val) = &core_arg.kind { sub.add_value(*addr, val.clone()); - } else if let BoundKind::Lambda { upvalues, .. } = &arg.kind + } else if let BoundKind::Lambda { upvalues, .. } = &core_arg.kind && upvalues.is_empty() { - sub.add_ast_substitution(*addr, (**arg).clone()); + sub.add_ast_substitution(*addr, core_arg.clone()); } else if let BoundKind::Get { addr: Address::Global(_), .. - } = &arg.kind + } = &core_arg.kind { - sub.add_ast_substitution(*addr, (**arg).clone()); + sub.add_ast_substitution(*addr, core_arg.clone()); } } diff --git a/src/ast/compiler/optimizer/utils.rs b/src/ast/compiler/optimizer/utils.rs index 470284d..4bd204b 100644 --- a/src/ast/compiler/optimizer/utils.rs +++ b/src/ast/compiler/optimizer/utils.rs @@ -143,8 +143,7 @@ impl UsageInfo { self.collect(v); } } - BoundKind::Define { addr, value, .. } => { - self.assigned.insert(*addr); + BoundKind::Define { value, .. } => { self.collect(value); } BoundKind::Expansion { bound_expanded, .. } => {