Refactor AST substitution logic
This commit refactors the AST substitution logic to correctly handle nested expansions and improve the inlining of lambda expressions. The changes include: - Extracting the core value from potentially expanded AST nodes before checking their kind. - Ensuring that lambda expressions with empty upvalues are correctly substituted. - Adding support for inlining global get expressions as AST substitutions. - Improving the `UsageInfo` collection to correctly track assigned values.
This commit is contained in:
@@ -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)
|
||||
{
|
||||
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);
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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, .. } => {
|
||||
|
||||
Reference in New Issue
Block a user