use std::collections::{BTreeMap, BTreeSet};
use super::branch_value_folding::{
fold_branch_value_locals_in_block, matches_local_lvalue,
};
use super::temp_touch::{
expr_touches_any_temp, stmt_consumes_temps_only_in_control_head,
stmt_contains_nested_nonlocal_control, stmt_touches_any_temp, stmts_touch_temp,
};
use crate::hir::common::{
HirAssign, HirBlock, HirCallExpr, HirExpr, HirLValue, HirLocalDecl,
HirProto, HirStmt, HirTableConstructor, HirTableField, HirTableKey, LocalId, TempId,
};
use crate::hir::promotion::ProtoPromotionFacts;
#[cfg_attr(not(test), allow(dead_code))]
pub(super) fn promote_temps_to_locals_in_proto(proto: &mut HirProto) -> bool {
promote_temps_to_locals_in_proto_with_facts(proto, &ProtoPromotionFacts::default())
}
pub(super) fn promote_temps_to_locals_in_proto_with_facts(
proto: &mut HirProto,
facts: &ProtoPromotionFacts,
) -> bool {
let mut next_local_index = proto.locals.len();
let mut new_locals = Vec::new();
let mut new_local_debug_hints = Vec::new();
let mut ctx = PromotionCtx {
facts,
temp_debug_locals: &proto.temp_debug_locals,
next_local_index: &mut next_local_index,
new_locals: &mut new_locals,
new_local_debug_hints: &mut new_local_debug_hints,
};
let result = promote_block(
&mut ctx,
&mut proto.body,
&BTreeMap::new(),
&BTreeMap::new(),
);
proto.locals.extend(new_locals);
proto.local_debug_hints.extend(new_local_debug_hints);
result.changed
}
#[derive(Debug, Clone)]
struct PromotionPlan {
decl_index: usize,
local: LocalId,
home_slot: Option<usize>,
temps: BTreeSet<TempId>,
removable_aliases: BTreeSet<usize>,
init: PromotionInit,
action: PromotionAction,
}
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
enum PromotionInit {
FromAssign,
Empty,
}
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
enum PromotionAction {
AllocateLocal,
ReuseExistingLocal,
}
#[derive(Debug, Clone, Default)]
struct FallthroughSummary {
falls_through: bool,
assigned_temps: BTreeSet<TempId>,
}
struct PromotionResult {
changed: bool,
trailing_mapping: BTreeMap<TempId, LocalId>,
}
struct PromotionCtx<'a> {
facts: &'a ProtoPromotionFacts,
temp_debug_locals: &'a [Option<String>],
next_local_index: &'a mut usize,
new_locals: &'a mut Vec<LocalId>,
new_local_debug_hints: &'a mut Vec<Option<String>>,
}
struct PlanAllocator<'a> {
temp_debug_locals: &'a [Option<String>],
plans: &'a mut Vec<PromotionPlan>,
reserved_temps: &'a mut BTreeSet<TempId>,
reserved_alias_indices: &'a mut BTreeSet<usize>,
next_local_index: &'a mut usize,
new_locals: &'a mut Vec<LocalId>,
new_local_debug_hints: &'a mut Vec<Option<String>>,
}
impl PlanAllocator<'_> {
fn allocate_local(
&mut self,
decl_index: usize,
home_slot: Option<usize>,
temps: BTreeSet<TempId>,
removable_aliases: BTreeSet<usize>,
init: PromotionInit,
) {
let local = LocalId(*self.next_local_index);
*self.next_local_index += 1;
self.new_locals.push(local);
self.new_local_debug_hints
.push(debug_hint_for_temp_group(self.temp_debug_locals, &temps));
self.reserved_temps.extend(temps.iter().copied());
self.reserved_alias_indices
.extend(removable_aliases.iter().copied());
self.plans.push(PromotionPlan {
decl_index,
local,
home_slot,
temps,
removable_aliases,
init,
action: PromotionAction::AllocateLocal,
});
}
fn reuse_existing_local(
&mut self,
decl_index: usize,
local: LocalId,
home_slot: Option<usize>,
temps: BTreeSet<TempId>,
removable_aliases: BTreeSet<usize>,
init: PromotionInit,
) {
self.reserved_temps.extend(temps.iter().copied());
self.reserved_alias_indices
.extend(removable_aliases.iter().copied());
self.plans.push(PromotionPlan {
decl_index,
local,
home_slot,
temps,
removable_aliases,
init,
action: PromotionAction::ReuseExistingLocal,
});
}
}
fn promote_block(
ctx: &mut PromotionCtx<'_>,
block: &mut HirBlock,
inherited: &BTreeMap<TempId, LocalId>,
inherited_sticky_slots: &BTreeMap<usize, LocalId>,
) -> PromotionResult {
let plans = collect_plans(ctx, block, inherited, inherited_sticky_slots);
let plan_by_decl = plans.iter().fold(
BTreeMap::<usize, Vec<&PromotionPlan>>::new(),
|mut grouped, plan| {
grouped.entry(plan.decl_index).or_default().push(plan);
grouped
},
);
let removable = plans
.iter()
.flat_map(|plan| plan.removable_aliases.iter().copied())
.collect::<BTreeSet<_>>();
let mut changed = !plans.is_empty();
let mut mapping = inherited.clone();
let mut slot_candidates = inherited_sticky_slots.clone();
let mut active_sticky_slots = inherited_sticky_slots.clone();
let original_stmts = std::mem::take(&mut block.stmts);
let mut rewritten = Vec::with_capacity(original_stmts.len());
for (index, mut stmt) in original_stmts.into_iter().enumerate() {
let mut replaced_stmt = false;
if let Some(plans) = plan_by_decl.get(&index) {
let mapping_before_decl = mapping.clone();
for plan in plans {
if let Some(anchor_stmt) =
rewrite_plan_anchor_stmt(&stmt, plan, &mapping_before_decl)
{
rewritten.push(anchor_stmt);
}
for temp in &plan.temps {
mapping.insert(*temp, plan.local);
}
if let Some(slot) = plan.home_slot
&& matches!(plan.action, PromotionAction::AllocateLocal)
{
slot_candidates.entry(slot).or_insert(plan.local);
}
replaced_stmt |= plan_replaces_original_stmt(plan);
}
}
activate_captured_slots_in_stmt(
&stmt,
ctx.facts,
&slot_candidates,
&mut active_sticky_slots,
);
if replaced_stmt {
continue;
}
if removable.contains(&index) {
continue;
}
let stmt_changed = rewrite_stmt(ctx, &mut stmt, &mapping, &active_sticky_slots);
changed |= stmt_changed;
rewritten.push(stmt);
}
block.stmts = rewritten;
changed |= fold_branch_value_locals_in_block(&mut block.stmts);
changed |= merge_adjacent_local_assigns_in_block(&mut block.stmts);
PromotionResult {
changed,
trailing_mapping: mapping,
}
}
fn merge_adjacent_local_assigns_in_block(stmts: &mut Vec<HirStmt>) -> bool {
let mut changed = false;
let mut index = 0;
while index + 1 < stmts.len() {
let Some(merged) = try_merge_empty_local_with_assign(&stmts[index], &stmts[index + 1])
else {
index += 1;
continue;
};
stmts[index] = HirStmt::LocalDecl(Box::new(merged));
stmts.remove(index + 1);
changed = true;
}
changed
}
fn try_merge_empty_local_with_assign(
decl_stmt: &HirStmt,
assign_stmt: &HirStmt,
) -> Option<HirLocalDecl> {
let HirStmt::LocalDecl(local_decl) = decl_stmt else {
return None;
};
let HirStmt::Assign(assign) = assign_stmt else {
return None;
};
if !local_decl.values.is_empty() || local_decl.bindings.is_empty() {
return None;
}
if local_decl.bindings.len() != assign.targets.len() || assign.values.is_empty() {
return None;
}
if !local_decl
.bindings
.iter()
.zip(assign.targets.iter())
.all(|(binding, target)| matches_local_lvalue(target, *binding))
{
return None;
}
Some(HirLocalDecl {
bindings: local_decl.bindings.clone(),
values: assign.values.clone(),
})
}
fn collect_plans(
ctx: &mut PromotionCtx<'_>,
block: &HirBlock,
inherited: &BTreeMap<TempId, LocalId>,
inherited_sticky_slots: &BTreeMap<usize, LocalId>,
) -> Vec<PromotionPlan> {
if block.stmts.iter().any(|stmt| {
matches!(
stmt,
HirStmt::Continue | HirStmt::Goto(_) | HirStmt::Label(_) | HirStmt::Unstructured(_)
)
}) {
return Vec::new();
}
let facts = ctx.facts;
let temp_debug_locals = ctx.temp_debug_locals;
let mut plans = Vec::new();
let mut reserved_temps = inherited.keys().copied().collect::<BTreeSet<_>>();
let mut reserved_alias_indices = BTreeSet::new();
let mut slot_candidates = inherited_sticky_slots.clone();
let mut sticky_slots = inherited_sticky_slots.clone();
for (decl_index, stmt) in block.stmts.iter().enumerate() {
if reserved_alias_indices.contains(&decl_index) {
activate_captured_slots_in_stmt(stmt, facts, &slot_candidates, &mut sticky_slots);
continue;
}
let Some(root_temp) = simple_temp_assign_target(stmt) else {
activate_captured_slots_in_stmt(stmt, facts, &slot_candidates, &mut sticky_slots);
continue;
};
if reserved_temps.contains(&root_temp) {
activate_captured_slots_in_stmt(stmt, facts, &slot_candidates, &mut sticky_slots);
continue;
}
if stmt_self_updates_temp(stmt, root_temp) {
activate_captured_slots_in_stmt(stmt, facts, &slot_candidates, &mut sticky_slots);
continue;
}
let mut group = BTreeSet::from([root_temp]);
let mut removable_aliases = BTreeSet::new();
let mut has_future_touch = false;
for future_index in decl_index + 1..block.stmts.len() {
if removable_aliases.contains(&future_index) {
continue;
}
let future_stmt = &block.stmts[future_index];
if let Some(alias_temp) = alias_temp_for_group(future_stmt, &group)
&& !reserved_temps.contains(&alias_temp)
&& !group.contains(&alias_temp)
&& !stmts_touch_temp(&block.stmts[decl_index + 1..future_index], alias_temp)
{
group.insert(alias_temp);
removable_aliases.insert(future_index);
continue;
}
if stmt_touches_any_temp(future_stmt, &group) {
has_future_touch = true;
}
}
let sticky_local = facts
.home_slot(root_temp)
.and_then(|slot| sticky_slots.get(&slot).copied());
if sticky_local.is_none() && !has_future_touch {
activate_captured_slots_in_stmt(stmt, facts, &slot_candidates, &mut sticky_slots);
continue;
}
if sticky_local.is_none() {
let touching_stmt_indices = (decl_index + 1..block.stmts.len())
.filter(|future_index| !removable_aliases.contains(future_index))
.filter(|future_index| stmt_touches_any_temp(&block.stmts[*future_index], &group))
.collect::<Vec<_>>();
if touching_stmt_indices.len() == 1
&& stmt_consumes_temps_only_in_control_head(
&block.stmts[touching_stmt_indices[0]],
&group,
)
{
activate_captured_slots_in_stmt(stmt, facts, &slot_candidates, &mut sticky_slots);
continue;
}
if touching_stmt_indices
.iter()
.copied()
.any(|stmt_index| stmt_contains_nested_nonlocal_control(&block.stmts[stmt_index]))
{
activate_captured_slots_in_stmt(stmt, facts, &slot_candidates, &mut sticky_slots);
continue;
}
}
let mut allocator = PlanAllocator {
temp_debug_locals,
plans: &mut plans,
reserved_temps: &mut reserved_temps,
reserved_alias_indices: &mut reserved_alias_indices,
next_local_index: ctx.next_local_index,
new_locals: ctx.new_locals,
new_local_debug_hints: ctx.new_local_debug_hints,
};
if let Some(local) = sticky_local {
allocator.reuse_existing_local(
decl_index,
local,
facts.home_slot(root_temp),
group.clone(),
removable_aliases,
PromotionInit::FromAssign,
);
} else {
let slot = facts.home_slot(root_temp);
allocator.allocate_local(
decl_index,
slot,
group.clone(),
removable_aliases,
PromotionInit::FromAssign,
);
if let Some(slot) = slot
&& let Some(local) = allocator.plans.last().map(|plan| plan.local)
{
slot_candidates.entry(slot).or_insert(local);
}
}
activate_captured_slots_in_stmt(stmt, facts, &slot_candidates, &mut sticky_slots);
}
let mut sticky_slots = inherited_sticky_slots.clone();
for (decl_index, stmt) in block.stmts.iter().enumerate() {
let merge_temps = if_merge_candidate_temps(
stmt,
&block.stmts[..decl_index],
&block.stmts[decl_index + 1..],
&reserved_temps,
);
for temp in merge_temps {
let mut allocator = PlanAllocator {
temp_debug_locals,
plans: &mut plans,
reserved_temps: &mut reserved_temps,
reserved_alias_indices: &mut reserved_alias_indices,
next_local_index: ctx.next_local_index,
new_locals: ctx.new_locals,
new_local_debug_hints: ctx.new_local_debug_hints,
};
if let Some(local) = facts
.home_slot(temp)
.and_then(|slot| sticky_slots.get(&slot).copied())
{
allocator.reuse_existing_local(
decl_index,
local,
facts.home_slot(temp),
BTreeSet::from([temp]),
BTreeSet::new(),
PromotionInit::Empty,
);
} else {
let slot = facts.home_slot(temp);
allocator.allocate_local(
decl_index,
slot,
BTreeSet::from([temp]),
BTreeSet::new(),
PromotionInit::Empty,
);
if let Some(slot) = slot
&& let Some(local) = allocator.plans.last().map(|plan| plan.local)
{
slot_candidates.entry(slot).or_insert(local);
}
}
}
activate_captured_slots_in_stmt(stmt, facts, &slot_candidates, &mut sticky_slots);
}
plans
}
fn activate_captured_slots_in_stmt(
stmt: &HirStmt,
facts: &ProtoPromotionFacts,
slot_candidates: &BTreeMap<usize, LocalId>,
sticky_slots: &mut BTreeMap<usize, LocalId>,
) {
let mut captured_slots = BTreeSet::new();
facts.collect_captured_home_slots_in_stmt(stmt, &mut captured_slots);
for slot in captured_slots {
if let Some(local) = slot_candidates.get(&slot).copied() {
sticky_slots.insert(slot, local);
}
}
}
fn simple_temp_assign_target(stmt: &HirStmt) -> Option<TempId> {
let HirStmt::Assign(assign) = stmt else {
return None;
};
let [HirLValue::Temp(temp)] = assign.targets.as_slice() else {
return None;
};
let [_value] = assign.values.as_slice() else {
return None;
};
Some(*temp)
}
fn alias_temp_for_group(stmt: &HirStmt, group: &BTreeSet<TempId>) -> Option<TempId> {
let HirStmt::Assign(assign) = stmt else {
return None;
};
let [HirLValue::Temp(alias)] = assign.targets.as_slice() else {
return None;
};
let [HirExpr::TempRef(source)] = assign.values.as_slice() else {
return None;
};
group.contains(source).then_some(*alias)
}
fn stmt_self_updates_temp(stmt: &HirStmt, temp: TempId) -> bool {
let HirStmt::Assign(assign) = stmt else {
return false;
};
matches!(assign.targets.as_slice(), [HirLValue::Temp(id)] if *id == temp)
&& assign
.values
.iter()
.any(|value| expr_touches_any_temp(value, &BTreeSet::from([temp])))
}
fn if_merge_candidate_temps(
stmt: &HirStmt,
prior_stmts: &[HirStmt],
future_stmts: &[HirStmt],
reserved_temps: &BTreeSet<TempId>,
) -> Vec<TempId> {
let HirStmt::If(if_stmt) = stmt else {
return Vec::new();
};
let Some(else_block) = &if_stmt.else_block else {
return Vec::new();
};
let then_summary = summarize_block_fallthrough_assignments(&if_stmt.then_block);
let else_summary = summarize_block_fallthrough_assignments(else_block);
let Some(common_temps) =
intersect_fallthrough_assignment_sets([then_summary.as_ref(), else_summary.as_ref()])
else {
return Vec::new();
};
common_temps
.into_iter()
.filter(|temp| !reserved_temps.contains(temp))
.filter(|temp| !stmts_touch_temp(prior_stmts, *temp))
.filter(|temp| stmts_touch_temp(future_stmts, *temp))
.collect()
}
fn summarize_block_fallthrough_assignments(block: &HirBlock) -> Option<FallthroughSummary> {
let mut assigned_temps = BTreeSet::new();
let mut falls_through = true;
for stmt in &block.stmts {
if !falls_through {
break;
}
let stmt_summary = summarize_stmt_fallthrough_assignments(stmt)?;
if stmt_summary.falls_through {
assigned_temps.extend(stmt_summary.assigned_temps);
} else {
falls_through = false;
}
}
Some(FallthroughSummary {
falls_through,
assigned_temps,
})
}
fn summarize_stmt_fallthrough_assignments(stmt: &HirStmt) -> Option<FallthroughSummary> {
match stmt {
HirStmt::LocalDecl(_)
| HirStmt::ErrNil(_)
| HirStmt::ToBeClosed(_)
| HirStmt::Close(_)
| HirStmt::CallStmt(_)
| HirStmt::Label(_) => Some(FallthroughSummary {
falls_through: true,
assigned_temps: BTreeSet::new(),
}),
HirStmt::Assign(assign) => Some(FallthroughSummary {
falls_through: true,
assigned_temps: assign
.targets
.iter()
.filter_map(|target| match target {
HirLValue::Temp(temp) => Some(*temp),
HirLValue::Local(_)
| HirLValue::Upvalue(_)
| HirLValue::Global(_)
| HirLValue::TableAccess(_) => None,
})
.collect(),
}),
HirStmt::TableSetList(_) => None,
HirStmt::Return(_) | HirStmt::Goto(_) | HirStmt::Break | HirStmt::Continue => {
Some(FallthroughSummary {
falls_through: false,
assigned_temps: BTreeSet::new(),
})
}
HirStmt::If(if_stmt) => {
let else_block = if_stmt.else_block.as_ref()?;
let then_summary = summarize_block_fallthrough_assignments(&if_stmt.then_block)?;
let else_summary = summarize_block_fallthrough_assignments(else_block)?;
let assigned_temps =
intersect_fallthrough_assignment_sets([Some(&then_summary), Some(&else_summary)])
.unwrap_or_default();
Some(FallthroughSummary {
falls_through: then_summary.falls_through || else_summary.falls_through,
assigned_temps,
})
}
HirStmt::Block(block) => summarize_block_fallthrough_assignments(block),
HirStmt::While(_)
| HirStmt::Repeat(_)
| HirStmt::NumericFor(_)
| HirStmt::GenericFor(_)
| HirStmt::Unstructured(_) => None,
}
}
fn intersect_fallthrough_assignment_sets<'a>(
summaries: impl IntoIterator<Item = Option<&'a FallthroughSummary>>,
) -> Option<BTreeSet<TempId>> {
let mut fallthrough_sets = summaries
.into_iter()
.flatten()
.filter(|summary| summary.falls_through)
.map(|summary| summary.assigned_temps.clone());
let mut intersection = fallthrough_sets.next()?;
for set in fallthrough_sets {
intersection = intersection
.intersection(&set)
.copied()
.collect::<BTreeSet<_>>();
}
Some(intersection)
}
fn rewrite_plan_anchor_stmt(
stmt: &HirStmt,
plan: &PromotionPlan,
mapping: &BTreeMap<TempId, LocalId>,
) -> Option<HirStmt> {
let values = match plan.init {
PromotionInit::FromAssign => {
let HirStmt::Assign(assign) = stmt else {
return None;
};
let [HirLValue::Temp(_temp)] = assign.targets.as_slice() else {
return None;
};
assign
.values
.iter()
.cloned()
.map(|mut expr| {
rewrite_expr(&mut expr, mapping);
expr
})
.collect::<Vec<_>>()
}
PromotionInit::Empty => Vec::new(),
};
match (plan.action, plan.init) {
(PromotionAction::AllocateLocal, _) => Some(HirStmt::LocalDecl(Box::new(HirLocalDecl {
bindings: vec![plan.local],
values,
}))),
(PromotionAction::ReuseExistingLocal, PromotionInit::FromAssign) => {
Some(HirStmt::Assign(Box::new(HirAssign {
targets: vec![HirLValue::Local(plan.local)],
values,
})))
}
(PromotionAction::ReuseExistingLocal, PromotionInit::Empty) => None,
}
}
fn plan_replaces_original_stmt(plan: &PromotionPlan) -> bool {
matches!(plan.init, PromotionInit::FromAssign)
}
fn rewrite_stmt(
ctx: &mut PromotionCtx<'_>,
stmt: &mut HirStmt,
mapping: &BTreeMap<TempId, LocalId>,
sticky_slots: &BTreeMap<usize, LocalId>,
) -> bool {
match stmt {
HirStmt::LocalDecl(local_decl) => {
local_decl.values.iter_mut().fold(false, |changed, expr| {
rewrite_expr(expr, mapping) || changed
})
}
HirStmt::Assign(assign) => {
let targets_changed = assign.targets.iter_mut().fold(false, |changed, target| {
rewrite_lvalue(target, mapping) || changed
});
let values_changed = assign.values.iter_mut().fold(false, |changed, expr| {
rewrite_expr(expr, mapping) || changed
});
targets_changed || values_changed
}
HirStmt::TableSetList(set_list) => {
let base_changed = rewrite_expr(&mut set_list.base, mapping);
let values_changed = set_list.values.iter_mut().fold(false, |changed, expr| {
rewrite_expr(expr, mapping) || changed
});
let trailing_changed = set_list
.trailing_multivalue
.as_mut()
.is_some_and(|expr| rewrite_expr(expr, mapping));
base_changed || values_changed || trailing_changed
}
HirStmt::ErrNil(err_nil) => rewrite_expr(&mut err_nil.value, mapping),
HirStmt::ToBeClosed(to_be_closed) => rewrite_expr(&mut to_be_closed.value, mapping),
HirStmt::CallStmt(call_stmt) => rewrite_call_expr(&mut call_stmt.call, mapping),
HirStmt::Return(ret) => ret.values.iter_mut().fold(false, |changed, expr| {
rewrite_expr(expr, mapping) || changed
}),
HirStmt::If(if_stmt) => {
let cond_changed = rewrite_expr(&mut if_stmt.cond, mapping);
let then_changed =
promote_block(ctx, &mut if_stmt.then_block, mapping, sticky_slots).changed;
let else_changed = if_stmt.else_block.as_mut().is_some_and(|else_block| {
promote_block(ctx, else_block, mapping, sticky_slots).changed
});
cond_changed || then_changed || else_changed
}
HirStmt::While(while_stmt) => {
let cond_changed = rewrite_expr(&mut while_stmt.cond, mapping);
let body_changed =
promote_block(ctx, &mut while_stmt.body, mapping, sticky_slots).changed;
cond_changed || body_changed
}
HirStmt::Repeat(repeat_stmt) => {
let body_result = promote_block(ctx, &mut repeat_stmt.body, mapping, sticky_slots);
let cond_changed = rewrite_expr(&mut repeat_stmt.cond, &body_result.trailing_mapping);
body_result.changed || cond_changed
}
HirStmt::NumericFor(numeric_for) => {
let start_changed = rewrite_expr(&mut numeric_for.start, mapping);
let limit_changed = rewrite_expr(&mut numeric_for.limit, mapping);
let step_changed = rewrite_expr(&mut numeric_for.step, mapping);
let body_changed =
promote_block(ctx, &mut numeric_for.body, mapping, sticky_slots).changed;
start_changed || limit_changed || step_changed || body_changed
}
HirStmt::GenericFor(generic_for) => {
let iterator_changed = generic_for
.iterator
.iter_mut()
.fold(false, |changed, expr| {
rewrite_expr(expr, mapping) || changed
});
let body_changed =
promote_block(ctx, &mut generic_for.body, mapping, sticky_slots).changed;
iterator_changed || body_changed
}
HirStmt::Block(block) => promote_block(ctx, block, mapping, sticky_slots).changed,
HirStmt::Unstructured(unstructured) => {
promote_block(ctx, &mut unstructured.body, mapping, sticky_slots).changed
}
HirStmt::Break
| HirStmt::Close(_)
| HirStmt::Continue
| HirStmt::Goto(_)
| HirStmt::Label(_) => false,
}
}
fn debug_hint_for_temp_group(
temp_debug_locals: &[Option<String>],
temps: &BTreeSet<TempId>,
) -> Option<String> {
temps
.iter()
.find_map(|temp| temp_debug_locals.get(temp.index()).cloned().flatten())
}
fn rewrite_call_expr(call: &mut HirCallExpr, mapping: &BTreeMap<TempId, LocalId>) -> bool {
let callee_changed = rewrite_expr(&mut call.callee, mapping);
let args_changed = call
.args
.iter_mut()
.fold(false, |changed, arg| rewrite_expr(arg, mapping) || changed);
callee_changed || args_changed
}
fn rewrite_expr(expr: &mut HirExpr, mapping: &BTreeMap<TempId, LocalId>) -> bool {
match expr {
HirExpr::TempRef(temp) => {
if let Some(local) = mapping.get(temp) {
*expr = HirExpr::LocalRef(*local);
true
} else {
false
}
}
HirExpr::TableAccess(access) => {
let base_changed = rewrite_expr(&mut access.base, mapping);
let key_changed = rewrite_expr(&mut access.key, mapping);
base_changed || key_changed
}
HirExpr::Unary(unary) => rewrite_expr(&mut unary.expr, mapping),
HirExpr::Binary(binary) => {
let lhs_changed = rewrite_expr(&mut binary.lhs, mapping);
let rhs_changed = rewrite_expr(&mut binary.rhs, mapping);
lhs_changed || rhs_changed
}
HirExpr::LogicalAnd(logical) | HirExpr::LogicalOr(logical) => {
let lhs_changed = rewrite_expr(&mut logical.lhs, mapping);
let rhs_changed = rewrite_expr(&mut logical.rhs, mapping);
lhs_changed || rhs_changed
}
HirExpr::Decision(decision) => decision.nodes.iter_mut().fold(false, |changed, node| {
let test_changed = rewrite_expr(&mut node.test, mapping);
let truthy_changed = rewrite_decision_target(&mut node.truthy, mapping);
let falsy_changed = rewrite_decision_target(&mut node.falsy, mapping);
changed || test_changed || truthy_changed || falsy_changed
}),
HirExpr::Call(call) => rewrite_call_expr(call, mapping),
HirExpr::TableConstructor(table) => rewrite_table_constructor(table, mapping),
HirExpr::Closure(closure) => closure.captures.iter_mut().fold(false, |changed, capture| {
rewrite_expr(&mut capture.value, mapping) || changed
}),
HirExpr::Nil
| HirExpr::Boolean(_)
| HirExpr::Integer(_)
| HirExpr::Number(_)
| HirExpr::String(_)
| HirExpr::Int64(_)
| HirExpr::UInt64(_)
| HirExpr::Complex { .. }
| HirExpr::ParamRef(_)
| HirExpr::LocalRef(_)
| HirExpr::UpvalueRef(_)
| HirExpr::GlobalRef(_)
| HirExpr::VarArg
| HirExpr::Unresolved(_) => false,
}
}
fn rewrite_decision_target(
target: &mut crate::hir::common::HirDecisionTarget,
mapping: &BTreeMap<TempId, LocalId>,
) -> bool {
match target {
crate::hir::common::HirDecisionTarget::Expr(expr) => rewrite_expr(expr, mapping),
crate::hir::common::HirDecisionTarget::Node(_)
| crate::hir::common::HirDecisionTarget::CurrentValue => false,
}
}
fn rewrite_table_constructor(
table: &mut HirTableConstructor,
mapping: &BTreeMap<TempId, LocalId>,
) -> bool {
let fields_changed = table.fields.iter_mut().fold(false, |changed, field| {
let field_changed = match field {
HirTableField::Array(expr) => rewrite_expr(expr, mapping),
HirTableField::Record(field) => {
let key_changed = match &mut field.key {
HirTableKey::Name(_) => false,
HirTableKey::Expr(expr) => rewrite_expr(expr, mapping),
};
let value_changed = rewrite_expr(&mut field.value, mapping);
key_changed || value_changed
}
};
changed || field_changed
});
let trailing_changed = table
.trailing_multivalue
.as_mut()
.is_some_and(|expr| rewrite_expr(expr, mapping));
fields_changed || trailing_changed
}
fn rewrite_lvalue(lvalue: &mut HirLValue, mapping: &BTreeMap<TempId, LocalId>) -> bool {
match lvalue {
HirLValue::Temp(temp) => {
if let Some(local) = mapping.get(temp) {
*lvalue = HirLValue::Local(*local);
true
} else {
false
}
}
HirLValue::TableAccess(access) => {
let base_changed = rewrite_expr(&mut access.base, mapping);
let key_changed = rewrite_expr(&mut access.key, mapping);
base_changed || key_changed
}
HirLValue::Local(_) | HirLValue::Upvalue(_) | HirLValue::Global(_) => false,
}
}
#[cfg(test)]
mod tests;