refactor(cfg): use IndexVec for storing basic blocks (#6323)

Use an `IndexVec` when storing basic blocks. This makes the link between nodes in `.graph` and elements of `.basic_blocks` more clear. I had to rename `BasicBlockId` to `BlockNodeId` to avoid a name collision. I wasn't sure what else to name the `Idx` type for the basic blocks vec.

crates/oxc_cfg/Cargo.toml

@@ -21,9 +21,11 @@ test = false
 doctest = false
 
 [dependencies]
+oxc_index = { workspace = true }
 oxc_syntax = { workspace = true }
 
 bitflags = { workspace = true }
 itertools = { workspace = true }
+nonmax = { workspace = true }
 petgraph = { workspace = true }
 rustc-hash = { workspace = true }

crates/oxc_cfg/src/block.rs

@@ -1,7 +1,4 @@
 use oxc_syntax::node::NodeId;
-use petgraph::stable_graph::NodeIndex;
-
-pub type BasicBlockId = NodeIndex;
 
 #[derive(Debug, Clone)]
 pub struct BasicBlock {

crates/oxc_cfg/src/builder/context.rs

@@ -1,5 +1,5 @@
 use super::ControlFlowGraphBuilder;
-use crate::{BasicBlockId, EdgeType};
+use crate::{BlockNodeId, EdgeType};
 
 bitflags::bitflags! {
     #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -15,9 +15,9 @@ bitflags::bitflags! {
 pub(super) struct Ctx<'a> {
     flags: CtxFlags,
     label: Option<&'a str>,
-    entries: Vec<(CtxFlags, BasicBlockId)>,
-    break_jmp: Option<BasicBlockId>,
-    continue_jmp: Option<BasicBlockId>,
+    entries: Vec<(CtxFlags, BlockNodeId)>,
+    break_jmp: Option<BlockNodeId>,
+    continue_jmp: Option<BlockNodeId>,
 }
 
 impl<'a> Ctx<'a> {
@@ -29,54 +29,54 @@ impl<'a> Ctx<'a> {
         self.label.as_ref().is_some_and(|it| *it == label)
     }
 
-    fn r#break(&mut self, entry: BasicBlockId) {
+    fn r#break(&mut self, entry: BlockNodeId) {
         self.entries.push((CtxFlags::BREAK, entry));
     }
 
-    fn r#continue(&mut self, entry: BasicBlockId) {
+    fn r#continue(&mut self, entry: BlockNodeId) {
         self.entries.push((CtxFlags::CONTINUE, entry));
     }
 }
 
 pub trait CtxCursor {
     #![allow(clippy::return_self_not_must_use)]
     /// Marks the break jump position in the current context.
-    fn mark_break(self, jmp_pos: BasicBlockId) -> Self;
+    fn mark_break(self, jmp_pos: BlockNodeId) -> Self;
     /// Marks the continue jump position in the current context.
-    fn mark_continue(self, jmp_pos: BasicBlockId) -> Self;
+    fn mark_continue(self, jmp_pos: BlockNodeId) -> Self;
     /// Creates a break entry in the current context.
-    fn r#break(self, bb: BasicBlockId) -> Self;
+    fn r#break(self, bb: BlockNodeId) -> Self;
     /// Creates a continue entry in the current context.
-    fn r#continue(self, bb: BasicBlockId) -> Self;
+    fn r#continue(self, bb: BlockNodeId) -> Self;
 }
 
 pub struct QueryCtx<'a, 'c>(&'c mut ControlFlowGraphBuilder<'a>, /* label */ Option<&'a str>);
 
 impl<'a, 'c> CtxCursor for QueryCtx<'a, 'c> {
-    fn mark_break(self, jmp_pos: BasicBlockId) -> Self {
+    fn mark_break(self, jmp_pos: BlockNodeId) -> Self {
         self.0.in_break_context(self.1, |ctx| {
             debug_assert!(ctx.break_jmp.is_none());
             ctx.break_jmp = Some(jmp_pos);
         });
         self
     }
 
-    fn mark_continue(self, jmp_pos: BasicBlockId) -> Self {
+    fn mark_continue(self, jmp_pos: BlockNodeId) -> Self {
         self.0.in_continue_context(self.1, |ctx| {
             debug_assert!(ctx.continue_jmp.is_none());
             ctx.continue_jmp = Some(jmp_pos);
         });
         self
     }
 
-    fn r#break(self, bb: BasicBlockId) -> Self {
+    fn r#break(self, bb: BlockNodeId) -> Self {
         self.0.in_break_context(self.1, |ctx| {
             ctx.r#break(bb);
         });
         self
     }
 
-    fn r#continue(self, bb: BasicBlockId) -> Self {
+    fn r#continue(self, bb: BlockNodeId) -> Self {
         self.0.in_continue_context(self.1, |ctx| {
             ctx.r#continue(bb);
         });
@@ -192,24 +192,24 @@ impl<'a, 'c> RefCtxCursor<'a, 'c> {
 }
 
 impl<'a, 'c> CtxCursor for RefCtxCursor<'a, 'c> {
-    fn mark_break(self, jmp_pos: BasicBlockId) -> Self {
+    fn mark_break(self, jmp_pos: BlockNodeId) -> Self {
         debug_assert!(self.0.break_jmp.is_none());
         self.0.break_jmp = Some(jmp_pos);
         self
     }
 
-    fn mark_continue(self, jmp_pos: BasicBlockId) -> Self {
+    fn mark_continue(self, jmp_pos: BlockNodeId) -> Self {
         debug_assert!(self.0.continue_jmp.is_none());
         self.0.continue_jmp = Some(jmp_pos);
         self
     }
 
-    fn r#break(self, bb: BasicBlockId) -> Self {
+    fn r#break(self, bb: BlockNodeId) -> Self {
         self.0.r#break(bb);
         self
     }
 
-    fn r#continue(self, bb: BasicBlockId) -> Self {
+    fn r#continue(self, bb: BlockNodeId) -> Self {
         self.0.r#continue(bb);
         self
     }

crates/oxc_cfg/src/builder/mod.rs

@@ -2,28 +2,29 @@ mod context;
 
 use context::Ctx;
 pub use context::{CtxCursor, CtxFlags};
+use oxc_index::IndexVec;
 use oxc_syntax::node::NodeId;
 use petgraph::Direction;
 
 use super::{
-    BasicBlock, BasicBlockId, ControlFlowGraph, EdgeType, ErrorEdgeKind, Graph, Instruction,
+    BasicBlock, BlockNodeId, ControlFlowGraph, EdgeType, ErrorEdgeKind, Graph, Instruction,
     InstructionKind, IterationInstructionKind, LabeledInstruction,
 };
-use crate::ReturnInstructionKind;
+use crate::{BasicBlockId, ReturnInstructionKind};
 
 #[derive(Debug, Default)]
-struct ErrorHarness(ErrorEdgeKind, BasicBlockId);
+struct ErrorHarness(ErrorEdgeKind, BlockNodeId);
 
 #[derive(Debug, Default)]
 pub struct ControlFlowGraphBuilder<'a> {
     pub graph: Graph,
-    pub basic_blocks: Vec<BasicBlock>,
-    pub current_node_ix: BasicBlockId,
+    pub basic_blocks: IndexVec<BasicBlockId, BasicBlock>,
+    pub current_node_ix: BlockNodeId,
     ctx_stack: Vec<Ctx<'a>>,
     /// Contains the error unwinding path represented as a stack of `ErrorHarness`es
     error_path: Vec<ErrorHarness>,
     /// Stack of finalizers, the top most element is always the appropriate one for current node.
-    finalizers: Vec<Option<BasicBlockId>>,
+    finalizers: Vec<Option<BlockNodeId>>,
 }
 
 impl<'a> ControlFlowGraphBuilder<'a> {
@@ -36,7 +37,7 @@ impl<'a> ControlFlowGraphBuilder<'a> {
     }
 
     /// # Panics
-    pub fn basic_block(&self, basic_block: BasicBlockId) -> &BasicBlock {
+    pub fn basic_block(&self, basic_block: BlockNodeId) -> &BasicBlock {
         let idx = *self
             .graph
             .node_weight(basic_block)
@@ -47,7 +48,7 @@ impl<'a> ControlFlowGraphBuilder<'a> {
     }
 
     /// # Panics
-    pub fn basic_block_mut(&mut self, basic_block: BasicBlockId) -> &mut BasicBlock {
+    pub fn basic_block_mut(&mut self, basic_block: BlockNodeId) -> &mut BasicBlock {
         let idx = *self
             .graph
             .node_weight(basic_block)
@@ -57,23 +58,22 @@ impl<'a> ControlFlowGraphBuilder<'a> {
             .expect("expected `self.current_node_ix` to be a valid node index in self.graph")
     }
 
-    pub(self) fn new_basic_block(&mut self) -> BasicBlockId {
+    pub(self) fn new_basic_block(&mut self) -> BlockNodeId {
         // current length would be the index of block we are adding on the next line.
-        let basic_block_ix = self.basic_blocks.len();
-        self.basic_blocks.push(BasicBlock::new());
+        let basic_block_ix = self.basic_blocks.push(BasicBlock::new());
         self.graph.add_node(basic_block_ix)
     }
 
     #[must_use]
-    pub fn new_basic_block_function(&mut self) -> BasicBlockId {
+    pub fn new_basic_block_function(&mut self) -> BlockNodeId {
         // we might want to differentiate between function blocks and normal blocks down the road.
         self.new_basic_block_normal()
     }
 
     /// # Panics
     /// if there is no error harness to attach to.
     #[must_use]
-    pub fn new_basic_block_normal(&mut self) -> BasicBlockId {
+    pub fn new_basic_block_normal(&mut self) -> BlockNodeId {
         let graph_ix = self.new_basic_block();
         self.current_node_ix = graph_ix;
 
@@ -89,7 +89,7 @@ impl<'a> ControlFlowGraphBuilder<'a> {
         graph_ix
     }
 
-    pub fn add_edge(&mut self, a: BasicBlockId, b: BasicBlockId, weight: EdgeType) {
+    pub fn add_edge(&mut self, a: BlockNodeId, b: BlockNodeId, weight: EdgeType) {
         if matches!(weight, EdgeType::NewFunction) {
             self.basic_block_mut(b).mark_as_reachable();
         } else if matches!(weight, EdgeType::Unreachable) || self.basic_block(a).is_unreachable() {
@@ -117,7 +117,7 @@ impl<'a> ControlFlowGraphBuilder<'a> {
 
     /// Creates and push a new `BasicBlockId` onto `self.error_path` stack.
     /// Returns the `BasicBlockId` of the created error harness block.
-    pub fn attach_error_harness(&mut self, kind: ErrorEdgeKind) -> BasicBlockId {
+    pub fn attach_error_harness(&mut self, kind: ErrorEdgeKind) -> BlockNodeId {
         let graph_ix = self.new_basic_block();
         self.error_path.push(ErrorHarness(kind, graph_ix));
         graph_ix
@@ -126,7 +126,7 @@ impl<'a> ControlFlowGraphBuilder<'a> {
     /// # Panics
     /// if there is no error harness pushed onto the stack,
     /// Or last harness doesn't match the expected `BasicBlockId`.
-    pub fn release_error_harness(&mut self, expect: BasicBlockId) {
+    pub fn release_error_harness(&mut self, expect: BlockNodeId) {
         let harness = self
             .error_path
             .pop()
@@ -139,7 +139,7 @@ impl<'a> ControlFlowGraphBuilder<'a> {
 
     /// Creates and push a new `BasicBlockId` onto `self.finalizers` stack.
     /// Returns the `BasicBlockId` of the created finalizer block.
-    pub fn attach_finalizer(&mut self) -> BasicBlockId {
+    pub fn attach_finalizer(&mut self) -> BlockNodeId {
         let graph_ix = self.new_basic_block();
         self.finalizers.push(Some(graph_ix));
         graph_ix
@@ -156,7 +156,7 @@ impl<'a> ControlFlowGraphBuilder<'a> {
 
     /// # Panics
     /// if last finalizer doesn't match the expected `BasicBlockId`.
-    pub fn release_finalizer(&mut self, expect: BasicBlockId) {
+    pub fn release_finalizer(&mut self, expect: BlockNodeId) {
         // return early if there is no finalizer.
         let Some(finalizer) = self.finalizers.pop() else { return };
         assert_eq!(
@@ -166,7 +166,7 @@ impl<'a> ControlFlowGraphBuilder<'a> {
         );
     }
 
-    pub fn append_condition_to(&mut self, block: BasicBlockId, node: Option<NodeId>) {
+    pub fn append_condition_to(&mut self, block: BlockNodeId, node: Option<NodeId>) {
         self.push_instruction_to(block, InstructionKind::Condition, node);
     }
 
@@ -228,7 +228,7 @@ impl<'a> ControlFlowGraphBuilder<'a> {
     #[inline]
     pub(self) fn push_instruction_to(
         &mut self,
-        block: BasicBlockId,
+        block: BlockNodeId,
         kind: InstructionKind,
         node_id: Option<NodeId>,
     ) {

crates/oxc_cfg/src/lib.rs

@@ -3,7 +3,11 @@ mod builder;
 pub mod dot;
 pub mod visit;
 
+use std::fmt;
+
 use itertools::Itertools;
+use nonmax::NonMaxU32;
+use oxc_index::{Idx, IndexVec};
 use petgraph::{
     visit::{Control, DfsEvent, EdgeRef},
     Direction,
@@ -23,7 +27,30 @@ pub use builder::{ControlFlowGraphBuilder, CtxCursor, CtxFlags};
 pub use dot::DisplayDot;
 use visit::set_depth_first_search;
 
-pub type Graph = petgraph::graph::DiGraph<usize, EdgeType>;
+pub type BlockNodeId = petgraph::stable_graph::NodeIndex;
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct BasicBlockId(NonMaxU32);
+
+impl Idx for BasicBlockId {
+    #[allow(clippy::cast_possible_truncation)]
+    fn from_usize(idx: usize) -> Self {
+        assert!(idx < u32::MAX as usize);
+        // SAFETY: We just checked `idx` is valid for `NonMaxU32`
+        Self(unsafe { NonMaxU32::new_unchecked(idx as u32) })
+    }
+
+    fn index(self) -> usize {
+        self.0.get() as usize
+    }
+}
+
+impl fmt::Display for BasicBlockId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+pub type Graph = petgraph::graph::DiGraph<BasicBlockId, EdgeType>;
 
 #[derive(Debug, Clone)]
 pub enum EdgeType {
@@ -65,7 +92,7 @@ pub enum EvalConstConditionResult {
 #[derive(Debug)]
 pub struct ControlFlowGraph {
     pub graph: Graph,
-    pub basic_blocks: Vec<BasicBlock>,
+    pub basic_blocks: IndexVec<BasicBlockId, BasicBlock>,
 }
 
 impl ControlFlowGraph {
@@ -74,25 +101,25 @@ impl ControlFlowGraph {
     }
 
     /// # Panics
-    pub fn basic_block(&self, id: BasicBlockId) -> &BasicBlock {
+    pub fn basic_block(&self, id: BlockNodeId) -> &BasicBlock {
         let ix = *self.graph.node_weight(id).expect("expected a valid node id in self.graph");
         self.basic_blocks.get(ix).expect("expected a valid node id in self.basic_blocks")
     }
 
     /// # Panics
-    pub fn basic_block_mut(&mut self, id: BasicBlockId) -> &mut BasicBlock {
+    pub fn basic_block_mut(&mut self, id: BlockNodeId) -> &mut BasicBlock {
         let ix = *self.graph.node_weight(id).expect("expected a valid node id in self.graph");
         self.basic_blocks.get_mut(ix).expect("expected a valid node id in self.basic_blocks")
     }
 
-    pub fn is_reachable(&self, from: BasicBlockId, to: BasicBlockId) -> bool {
+    pub fn is_reachable(&self, from: BlockNodeId, to: BlockNodeId) -> bool {
         self.is_reachable_filtered(from, to, |_| Control::Continue)
     }
 
-    pub fn is_reachable_filtered<F: Fn(BasicBlockId) -> Control<bool>>(
+    pub fn is_reachable_filtered<F: Fn(BlockNodeId) -> Control<bool>>(
         &self,
-        from: BasicBlockId,
-        to: BasicBlockId,
+        from: BlockNodeId,
+        to: BlockNodeId,
         filter: F,
     ) -> bool {
         if from == to {
@@ -127,13 +154,13 @@ impl ControlFlowGraph {
     /// Otherwise returns `Some(loop_start, loop_end)`.
     pub fn is_infinite_loop_start<F>(
         &self,
-        node: BasicBlockId,
+        node: BlockNodeId,
         try_eval_const_condition: F,
-    ) -> Option<(BasicBlockId, BasicBlockId)>
+    ) -> Option<(BlockNodeId, BlockNodeId)>
     where
         F: Fn(&Instruction) -> EvalConstConditionResult,
     {
-        fn get_jump_target(graph: &Graph, node: BasicBlockId) -> Option<BasicBlockId> {
+        fn get_jump_target(graph: &Graph, node: BlockNodeId) -> Option<BlockNodeId> {
             graph
                 .edges_directed(node, Direction::Outgoing)
                 .find_or_first(|e| matches!(e.weight(), EdgeType::Jump))
@@ -186,7 +213,7 @@ impl ControlFlowGraph {
         }
     }
 
-    pub fn is_cyclic(&self, node: BasicBlockId) -> bool {
+    pub fn is_cyclic(&self, node: BlockNodeId) -> bool {
         set_depth_first_search(&self.graph, Some(node), |event| match event {
             DfsEvent::BackEdge(_, id) if id == node => Err(()),
             _ => Ok(()),