176. Dominators and SSA in Rust

Same algorithms, new containers

Chapters 76 to 83 built dominators, SSA and natural loops in TypeScript. engine/rust/crates/core/src/cfg.rs and builder.rs port them almost line for line, and both name their sources in their first comment: Cooper, Harvey and Kennedy’s “A Simple, Fast Dominance Algorithm”, and Braun et al.’s on-the-fly SSA construction^[1].

What changes is where things live. The TS code keeps Maps keyed by block numbers, and instructions are objects that hold arrays of other instructions’ ids. In Rust, blocks and values sit in the module’s arenas (chapter 170), so a map from blocks becomes a Vec indexed by BlockId, and “absent” becomes None.

Reverse postorder and dominators

cfg.ts computes reverse postorder with a recursive visit. cfg.rs uses an explicit stack of (block, successors, next) that tries the successors in the same order. The result is identical, and a long chain of blocks can’t overflow the call stack.

The dominator loop, adapted:

let mut idom: Vec<Option<BlockId>> =
    vec![None; f.blocks.len()];
idom[f.entry.index()] = Some(f.entry);
let intersect =
    |idom: &[Option<BlockId>], a, b| {
        // climb a and b by rpo index
    };
// ...then, for each block in rpo:
let mut d = first;
for &p in rest {
    d = intersect(&idom, p, d);
}
idom[b.index()] = Some(d);

Why does the closure take idom as a parameter, when it could just use the variable? Try it. A closure that captures idom holds a shared borrow for as long as it is used, and the loop assigns to idom in between calls. rustc 1.97, on a small copy of this loop (trimmed and rewrapped):

error[E0502]: cannot borrow `idom` as
  mutable because it is also borrowed
  as immutable
  | let up = |mut a: usize, ...| {
  |          --- immutable borrow
  |              occurs here
  ...
  | idom[b] = Some(d);
  | ^^^^ mutable borrow occurs here

Passing &idom at each call makes each borrow last one call. The TypeScript closure captures freely; nothing stops it from reading a map it is updating.

Natural loops

natural_loops finds each back edge t → h (where h dominates t) and walks predecessors from t back to h. Loops are kept in an IndexMap<BlockId, BTreeSet<BlockId>>: the IndexMap lists headers in the order they are first found, as the TS Map does, and the BTreeSet keeps each body sorted. One quirk is kept for parity, and pinned by a unit test: an unreachable predecessor of a loop block is swept into the body, as in the reference.

use std::collections::BTreeSet;

// Does a dominate b? Climb b's idoms.
fn dominates(
    idom: &[usize],
    a: usize,
    mut b: usize,
) -> bool {
    loop {
        if b == a {
            return true;
        }
        if idom[b] == b {
            return false; // the entry
        }
        b = idom[b];
    }
}

// The body of the loop of back edge
// t -> h: h, plus everything that
// reaches t without passing h.
fn body(
    preds: &[Vec<usize>],
    h: usize,
    t: usize,
) -> BTreeSet<usize> {
    let mut body = BTreeSet::from([h]);
    let mut stack = vec![t];
    while let Some(x) = stack.pop() {
        if body.insert(x) {
            stack.extend(&preds[x]);
        }
    }
    body
}

fn main() {
    // 0 entry, 1 head, 2 body,
    // 3 throw, 4 exit, 5 latch
    let succ = vec![
        vec![1],
        vec![2, 4],
        vec![3, 5],
        vec![],
        vec![],
        vec![1],
    ];
    let mut preds = vec![vec![]; 6];
    for b in 0..6 {
        for &s in &succ[b] {
            preds[s].push(b);
        }
    }
    // From Cooper-Harvey-Kennedy.
    let idom = [0, 0, 1, 2, 1, 2];
    // A back edge goes to a dominator.
    let mut back = vec![];
    for t in 0..6 {
        for &h in &succ[t] {
            if dominates(&idom, h, t) {
                back.push((t, h));
            }
        }
    }
    for (t, h) in back {
        let b = body(&preds, h, t);
        println!("{t} -> {h}: {b:?}");
    }
    let d = dominates(&idom, 2, 3);
    println!("2 dominates 3: {d}");
}

The SSA builder

FunctionBuilder<K> is generic over the variable key K, and the JS plugin’s key is an enum: a resolved binding, or the hidden index of a for...of. It doesn’t own the Module: every method takes &mut Module. A front-end lowering a nested function keeps two builders alive, the outer and the inner, and both write into the same module. If each held its own &mut Module, the borrow checker would reject the second. Reading a variable with no definition in the block, trimmed:

let val = if !self.sealed.contains(&block) {
    let phi = self.new_phi(m, block);
    self.incomplete.entry(block)
        .or_default()
        .insert(v.clone(), phi);
    phi
} else if preds.len() == 1 {
    self.read_variable(m, v, preds[0])
} else {
    let phi = self.new_phi(m, block);
    // record it first: breaks cycles
    self.write_variable(
        v.clone(), block, phi);
    self.add_phi_operands(m, v, phi)
};

(A third branch, with no predecessors, makes an undefined constant.) incomplete is an IndexMap, because sealing fills phis in that order, which decides value numbers. new_phi inserts at the front of the block, so the newest phi comes first. In the VIR of let s = 0; for (let i = 0; i < 4; i++) s = s + i; the loop header reads:

  b1:  ; preds b0 b3
    %5 = phi %0, %6
    %2 = phi %1, %9

%2 (for i) was made when the header was lowered. %5 (for s) was made later, from the body, and went above it.

Trivial phis and the throw bug

A phi is trivial if, ignoring itself, it merges just one value. Then try_remove_trivial_phi replaces it with that value. In the instructions that’s a loop over operands. Terminators need their own match (T is Terminator):

match &mut blocks[bi].term {
    T::CondBr { cond, .. }
        if *cond == phi => *cond = same,
    T::Return { value: Some(v), .. }
        if *v == phi => *v = same,
    T::Throw { value, .. }
        if *value == phi => *value = same,
    _ => {}
}

The TypeScript version had the first two cases only. The third line went into builder.ts in commit 16c7767, “fix four bugs found by the Rust port’s parity work”. Here is what it changes, on const e = ps[0]; for (...) { if (i === 1) throw e; }. Real listings from the builder, the second with that line deleted:

b9:  ; preds b6        b9:  ; preds b6
  throw %14              throw %21

%21 was the header’s phi for e. It was removed, and the throw still names it. Step through:

Could Rust’s exhaustive match have caught it? Not as written: the _ => {} arm covers whatever isn’t listed. A match that names every kind, with Br and Unreachable doing nothing, would have made the missing case visible, and turns a future terminator kind into a compile error at each such site.

Your turn

struct Doms {
    // ix[b]: b's position in rpo.
    ix: Vec<usize>,
    // idom[b]: None until known.
    idom: Vec<Option<usize>>,
}

// preds[b]: predecessors of block b.
// rpo: reachable blocks, entry first.
fn idoms(
    preds: &[Vec<usize>],
    rpo: &[usize],
) -> Vec<Option<usize>> {
    let n = preds.len();
    let mut d = Doms {
        ix: vec![usize::MAX; n],
        idom: vec![None; n],
    };
    for i in 0..rpo.len() {
        d.ix[rpo[i]] = i;
    }
    d.idom[rpo[0]] = Some(rpo[0]);
    while pass(&mut d, preds, rpo) {}
    d.idom
}

impl Doms {
    // Climb from a and b, always moving
    // the one later in rpo, until they
    // meet.
    fn intersect(
        &self,
        mut a: usize,
        mut b: usize,
    ) -> usize {
        todo!()
    }
}

// One pass over rpo; true if any idom
// changed.
fn pass(
    d: &mut Doms,
    preds: &[Vec<usize>],
    rpo: &[usize],
) -> bool {
    todo!()
}

What you learned

• cfg.rs and builder.rs are the TS algorithms with Rust containers: Vec<Option<BlockId>> for idoms, an explicit DFS stack, IndexMap where order is observable.
• A closure that captures state borrows it. intersect takes &idom per call so that the loop can update idom (E0502 otherwise).
• FunctionBuilder<K> takes &mut Module in every method, so nested functions’ builders can share one module. SSA over arena IR means rewriting ids in place.
• Phis go at the front of a block. A trivial phi is replaced everywhere, including in terminators. Missing throw there was a soundness bug in the reference, found by the port.