178. Collection discovery in Rust

Same pass, new home

Chapter 96 read the first two steps of analysis.ts: find candidate bindings, then follow every load of each one through its uses. The TypeScript keeps everything in local variables of one big analyze function, and the steps are closures that share them.

analysis.rs produces the same collections, the same bailouts and the same messages in the same order. What changes is where the state lives, and how one collection refers to another thing.

The Analyzer struct

The TS closures become methods of one struct, trimmed:

struct Analyzer<'m> {
    m: &'m Module,
    uses: Uses,
    stores: HashMap<BindingId,
        Vec<ValueId>>,
    loads: HashMap<BindingId,
        Vec<ValueId>>,
    call_sites: HashMap<FuncId,
        Vec<ValueId>>,
    collections: Vec<Collection>,
    collection_of_load:
        HashMap<ValueId, usize>,
    element_of: HashMap<ValueId, usize>,
    // ...numeric state, a cfg cache
}

HashMap here is FxHashMap from rustc-hash, renamed on import: a faster hasher for small keys^[1]. Its iteration order is arbitrary, which is fine, because the analyzer only looks things up in these maps. Anything whose order reaches explain is a Vec or an IndexMap^[2].

The analyzer borrows the module for 'm, and inst hands out instructions for that lifetime, not for the borrow of self:

fn inst(&self, v: ValueId) -> &'m Inst {
    self.m.inst(v)
}

So you can hold an instruction and still change self.collections. Chapter 179 shows the error you get without it.

Collections by number

In TypeScript, collectionOfLoad maps a value to the Collection object, and refsOf is keyed by it. Rust’s collection_of_load maps a value to a usize, the collection’s index in collections. A &mut Collection in the map would keep the Vec borrowed for as long as the map exists, and the pass pushes to that Vec and changes its entries all the time. An index borrows nothing.

Finding candidates is otherwise a line-for-line port. Where TS filters every store and tests the length, Rust stops at the first array store, with let-else:

let st = self.stores_of(binding.id);
let Some(&first_array_store) =
    st.iter().find(|&&x| matches!(
        self.inst(self.inst(x).args[0]).op,
        Op::ArrayEmpty | Op::ArrayNew))
else {
    continue;
};

Some details exist only to match the reference’s text. When an array literal holds something other than a record, the TS message interpolates init.src?.line, which prints undefined if there’s no position. The port spells that out: init.src.map_or("undefined".... And every bailout goes through one helper, which keeps the first source position seen for each message, for --frames:

fn bail(&mut self, msg: String,
        at: Option<SourceRef>) {
    self.bailout_sites
        .entry(msg.clone())
        .or_insert(at);
    self.bailouts.push(msg);
}

The worklist and the borrow checker

Step 2, from shapes_and_uses, trimmed:

let mut work: Vec<ValueId> = self
    .loads.get(&binding).cloned()
    .unwrap_or_default();
let mut seen_param = HashSet::default();
while let Some(l) = work.pop() {
    for u in self.uses_of(l) {
        if let Some(next) = self
            .use_of_collection(ci, l, u,
                &mut seen_param)
        {
            work.push(next);
        }
    }
}

Two differences from TS. First, the TS closure pushes a parameter onto work itself; use_of_collection returns it, because it takes &mut self and the stack is a local of the caller. Second, uses_of(l) returns a copy of the use list. Write self.uses.of(l) in the for instead, and the loop holds a shared borrow of self while its body calls a &mut self method. rustc 1.97 on a small copy of this loop (trimmed and rewrapped):

error[E0502]: cannot borrow `*self` as
  mutable because it is also borrowed
  as immutable
  | for &u in self.uses.of(l) {
  |           ---------------
  |           immutable borrow occurs
  |           here
  |     self.check(u);
  |     ^^^^^^^^^^^^^ mutable borrow
  |                   occurs here

The copy costs a few words per use. It is the simplest thing that lets the method mutate freely. Try the error yourself:

struct Uses(Vec<Vec<usize>>);

impl Uses {
    fn of(&self, v: usize) -> &[usize] {
        &self.0[v]
    }
}

struct A {
    uses: Uses,
    bails: Vec<String>,
}

impl A {
    // Like Analyzer::uses_of: a copy,
    // so self isn't borrowed after.
    fn uses_of(
        &self,
        v: usize,
    ) -> Vec<usize> {
        self.uses.of(v).to_vec()
    }
    fn check(&mut self, u: usize) {
        let m = format!("use {u}");
        self.bails.push(m);
    }
    fn walk(&mut self, l: usize) {
        // Try self.uses.of(l) here.
        for u in self.uses_of(l) {
            self.check(u);
        }
    }
}

fn main() {
    let mut a = A {
        uses: Uses(vec![vec![1, 2]]),
        bails: vec![],
    };
    a.walk(0);
    println!("{:?}", a.bails);
}

Order is output

work.pop() takes the last load first, as work.pop() does in TypeScript. That order is visible: it decides which push sets the shape and the order of the messages. Here it is on real VIR, printed from the port for a program that pushes into ps and passes it to a function step:

<module>:
  %23 = load ps
  %25 = record {x, y} %20, %24
  %26 = push %23, %25
  %30 = load step
  %31 = load ps
  %32 = call %30, %31
  %35 = load ps
  %37 = index.get %35, %36
step(%0):
  %4 = field.get %0 .length
  %6 = index.get %0, %2
  %8 = index.get %0, %2

seen_param stops the walk visiting a parameter twice when two call sites pass the collection. Note what it does not do: a function that passes the parameter on to itself, walk(qs, k + 1), isn’t followed, because the recursive call site passes qs, which isn’t yet known to be ps when the sites are checked. Both engines report it the same way (real output, wrapped):

ps (line 1) {x} → kept as array of
  objects
  · unsafe: line 10: passed to walk(),
    which is also called with other
    values

That refusal is safe, just less precise than it could be.

Your turn

use std::collections::HashSet;

// One use of a value that refers to
// the collection.
pub enum Use {
    // .length
    Len,
    // .push(r): the record r
    Push(usize),
    // [i]: the element value
    Index(usize),
    // passed to known function f as
    // argument k (every call site of f
    // passes this collection there)
    Call(usize, usize),
    // anything else, and why
    Other(&'static str),
}

pub struct Prog {
    // uses[v]: the uses of value v
    pub uses: Vec<Vec<Use>>,
    // params[f][k]: the value of f's
    // k-th parameter
    pub params: Vec<Vec<usize>>,
}

#[derive(Debug, Default, PartialEq)]
pub struct Found {
    pub pushes: Vec<usize>,
    pub elements: Vec<usize>,
    pub bails: Vec<String>,
}

fn one(
    u: &Use,
    p: &Prog,
    out: &mut Found,
) -> Option<usize> {
    todo!()
}

pub fn discover(
    p: &Prog,
    loads: &[usize],
) -> Found {
    todo!()
}

What you learned

• The collection pass in Rust keeps the TS closures’ state in Analyzer<'m>, with inst returning &'m Inst so instructions don’t borrow the analyzer.
• Collections are referred to by index (usize), never by reference, so the Vec of collections can keep changing.
• uses_of copies the use list so that the loop body can call &mut self methods (E0502 otherwise).
• The worklist is a stack, as in the reference: the last load first. Known functions are followed by returning the parameter; recursion through the parameter isn’t.