diff --git a/sway-core/src/asm_generation/fuel/register_allocator.rs b/sway-core/src/asm_generation/fuel/register_allocator.rs index f00231ceda5..81e9069da08 100644 --- a/sway-core/src/asm_generation/fuel/register_allocator.rs +++ b/sway-core/src/asm_generation/fuel/register_allocator.rs @@ -1,12 +1,14 @@ use crate::{ asm_generation::fuel::compiler_constants, - asm_lang::{allocated_ops::AllocatedRegister, virtual_register::*, Op, VirtualOp}, + asm_lang::{ + allocated_ops::AllocatedRegister, virtual_register::*, ControlFlowOp, Label, Op, VirtualOp, + }, }; -use std::collections::{BTreeSet, HashMap}; - use either::Either; use petgraph::graph::{node_index, NodeIndex}; +use rustc_hash::FxHashSet; +use std::collections::{BTreeSet, HashMap}; pub type InterferenceGraph = petgraph::stable_graph::StableGraph, (), petgraph::Undirected>; @@ -127,13 +129,19 @@ impl RegisterPool { /// This function finally returns `live_out` because it has all the liveness information needed. /// `live_in` is computed because it is needed to compute `live_out` iteratively. /// -pub(crate) fn liveness_analysis(ops: &[Op]) -> HashMap> { - // Hash maps that will reprsent the live_in and live_out tables. The key of each hash map is - // simply the index of each instruction in the `ops` vector. - let mut live_in: HashMap> = - HashMap::from_iter((0..ops.len()).into_iter().map(|idx| (idx, BTreeSet::new()))); - let mut live_out: HashMap> = - HashMap::from_iter((0..ops.len()).into_iter().map(|idx| (idx, BTreeSet::new()))); +pub(crate) fn liveness_analysis(ops: &[Op]) -> Vec> { + // Vectors representing maps that will reprsent the live_in and live_out tables. Each entry + // corresponds to an instruction in `ops`. + let mut live_in: Vec> = vec![FxHashSet::default(); ops.len()]; + let mut live_out: Vec> = vec![FxHashSet::default(); ops.len()]; + let mut label_to_index: HashMap = HashMap::new(); + + // Keep track of an map between jump labels and op indices. Useful to compute op successors. + for (idx, op) in ops.iter().enumerate() { + if let Either::Right(ControlFlowOp::Label(op_label)) = op.opcode { + label_to_index.insert(op_label, idx); + } + } let mut modified = true; while modified { @@ -141,6 +149,7 @@ pub(crate) fn liveness_analysis(ops: &[Op]) -> HashMap HashMap HashMap>, + live_out: &[FxHashSet], ) -> (InterferenceGraph, HashMap) { let mut interference_graph = InterferenceGraph::with_capacity(0, 0); @@ -225,8 +226,8 @@ pub(crate) fn create_interference_graph( reg_to_node_map.insert(reg.clone(), interference_graph.add_node(Some(reg.clone()))); }); - for (ix, regs) in live_out { - match &ops[*ix].opcode { + for (ix, regs) in live_out.iter().enumerate() { + match &ops[ix].opcode { Either::Left(VirtualOp::MOVE(v, c)) => { if let Some(ix1) = reg_to_node_map.get(v) { for b in regs.iter() { @@ -242,7 +243,7 @@ pub(crate) fn create_interference_graph( } } _ => { - for v in &ops[*ix].def_registers() { + for v in &ops[ix].def_registers() { if let Some(ix1) = reg_to_node_map.get(v) { for b in regs.iter() { if let Some(ix2) = reg_to_node_map.get(b) { diff --git a/sway-core/src/asm_lang/mod.rs b/sway-core/src/asm_lang/mod.rs index c572f4d6b8f..0c4490cbaab 100644 --- a/sway-core/src/asm_lang/mod.rs +++ b/sway-core/src/asm_lang/mod.rs @@ -1042,10 +1042,15 @@ impl Op { } } - pub(crate) fn successors(&self, index: usize, ops: &[Op]) -> Vec { + pub(crate) fn successors( + &self, + index: usize, + ops: &[Op], + label_to_index: &HashMap, + ) -> Vec { match &self.opcode { Either::Left(virt_op) => virt_op.successors(index, ops), - Either::Right(org_op) => org_op.successors(index, ops), + Either::Right(org_op) => org_op.successors(index, ops, label_to_index), } } @@ -1667,7 +1672,12 @@ impl ControlFlowOp { } } - pub(crate) fn successors(&self, index: usize, ops: &[Op]) -> Vec { + pub(crate) fn successors( + &self, + index: usize, + ops: &[Op], + label_to_index: &HashMap, + ) -> Vec { use ControlFlowOp::*; let mut next_ops = Vec::new(); @@ -1687,15 +1697,7 @@ impl ControlFlowOp { | PopAll(_) => (), Jump(jump_label) | JumpIfNotEq(_, _, jump_label) | JumpIfNotZero(_, jump_label) => { - // Find the label in the ops list. - for (idx, op) in ops.iter().enumerate() { - if let Either::Right(ControlFlowOp::Label(op_label)) = op.opcode { - if op_label == *jump_label { - next_ops.push(idx); - break; - } - } - } + next_ops.push(label_to_index[jump_label]); } };