186. PGO in Rust

The same three commands

Chapter 109 profiled this program with the TypeScript CLI. Its 20 elements are pushed inside a function, so the analysis can’t count them:

const ps = [];
function spawn(x) {
    ps.push({ x, vx: 1 });
}
for (let i = 0; i < 20; i++) spawn(i);
for (let f = 0; f < 100; f++) {
    for (let i = 0; i < ps.length; i++) {
        ps[i].x += ps[i].vx;
    }
}
console.log(ps[0].x);

The Rust binary takes the same commands. A real session (release build, rustc 1.97, Node 22.23; long lines wrapped):

$ viscosity explain spawn.js
ps (line 1) {x, vx} → struct of arrays
  · …
  · size unknown statically
  · 2 of 3 element accesses proven in
    bounds
$ viscosity profile spawn.js -o p.json
$ cat p.json
{
  "ps": {
    "maxLength": 20
  }
}
$ viscosity explain spawn.js \
    --profile p.json
ps (line 1) {x, vx} → kept as array of
  objects
  · about 20 elements (measured): below
    the 64-element threshold, where the
    benefit is too small

The decision flips, as it should, and every line matches the TypeScript CLI’s (chapter 182 tells how the last one once didn’t). The JSON file is byte-identical to the one the TypeScript CLI writes, too. The reason is worth a look.

Writing a program that counts

From crates/cli/src/profile.rs, adapted:

pub fn instrument(
    source: &str,
) -> Result<String, String> {
    let lowered = JsPlugin
        .lower(source)
        .map_err(|e| e.to_string())?;
    let a = analyze(&lowered.module);
    let mut edits: Vec<(u32, u32, &str)> =
        a.collections.iter().map(|c| {
            let m = &lowered.module;
            let src = m.inst(c.init).src
                .expect("has a source");
            (src.start, src.end, &*c.name)
        }).collect();
    edits.sort_by_key(|e| Reverse(e.0));
    let mut out = source.to_string();
    for (start, end, name) in edits {
        let (lo, hi) = (start as usize,
            end as usize);
        let q = serde_json::to_string(name)
            .expect("strings serialize");
        let w = format!(
            "__vsc_track({q}, {})",
            &out[lo..hi]);
        out.replace_range(lo..hi, &w);
    }
    Ok(format!("{PRELUDE}{out}"))
}

It finds collections exactly as the compiler does (the plugin’s lower, then analyze), takes each initializer’s span, and wraps it. The edits are sorted by start, last first (Reverse flips the order), so replacing one span never moves the byte offsets of the spans still to come. replace_range edits the String in place, and the offsets are UTF-8 bytes from oxc, which is what String is indexed by. The name is quoted by serde_json, so it is a valid JavaScript string literal.

PRELUDE is a Rust &str holding JavaScript, the TypeScript CLI’s prelude character for character: __vsc_track, and an exit hook that writes JSON.stringify(p, null, 2). So in both engines the JSON is written by Node, with the same code. That’s why the files are identical. Rust never formats it.

Running node from Rust

let status = Command::new("node")
    .arg(&inst)
    .env("VSC_PROFILE_OUT", &out)
    .stdin(Stdio::null())
    .stdout(Stdio::null())
    .status()
    .map_err(|e| format!(
        "could not run node (profiling \
         needs Node.js on the PATH): \
         {e}"))?;
if !status.success() {
    return Err(match status.code() {
        Some(c) => format!(
            "the profiled program \
             failed (exit code {c})"),
        None => "the profiled program \
            failed (killed by a signal)"
            .to_string(),
    });
}

std::process::Command builds a child process, .status() runs it and waits, and the program’s output is thrown away (stderr still shows). The instrumented file is instrumented.cjs in a temporary directory, owned by a small guard value whose Drop deletes the directory on every way out of the function, error or not. The TypeScript CLI runs process.execPath, the very Node running it. The Rust binary has no Node of its own, so it looks node up on the PATH, and says so if that fails:

$ PATH=/nowhere viscosity profile spawn.js
viscosity: could not run node (profiling
needs Node.js on the PATH): No such file
or directory (os error 2)

Reading it back

read_profile in main.rs parses the file with serde_json into a Value, requires an object, and requires every entry to be an object with a finite, non-negative maxLength (adapted; bad stands for the error message):

obj.iter().map(|(name, v)| {
    let n = v.as_object()
        .and_then(|e| e.get("maxLength"))
        .and_then(Value::as_f64);
    match n {
        Some(n) if n.is_finite()
            && n >= 0.0 =>
            Ok((name.clone(), n)),
        _ => Err(bad(name)),
    }
}).collect()

Collecting an iterator of Results into one Result stops at the first Err. The planner then looks up the collection’s name, and a hit wins over the estimate (chapter 182). When this chapter was first written, the CLIs disagreed on hand-edited files: given {"ps": {"maxLength": "20"}}, Rust skipped the entry while TypeScript compared "20" < 64 as JavaScript does and kept the array. Now both apply this one rule, with the same words: the profile entry "ps" must be an object whose maxLength is a non-negative number.

What a profile can’t say

The code is short enough to read its limits straight off it.

• Keyed by name. Give two functions each their own ps, one with 1000 elements and one with 10, and the profile has one entry, "ps": { "maxLength": 1000 }. With it, the Rust CLI reports “about 1000 elements (measured by profiling)” for both.
• The length at exit. The hook reads a.length once, at exit, and takes the maximum over every array tracked under that name. For a collection Viscosity can transform that is its peak, since it only grows. Every tracked array also stays alive until exit, because __vsc_seen holds it.
• One run, no input. The program gets no arguments and an empty stdin.
• Scripts only. The file is run as CommonJS, so a program with import fails (“Cannot use import statement outside a module”). A program that throws is an error too, “the profiled program failed (exit code 1)”, even though its exit hook wrote the JSON: a partial run’s sizes could mislead the planner.
• JavaScript only. viscosity profile x.py refuses, with the reference’s words: “profile runs JavaScript programs only; Python profiling is not supported yet”.

Your turn

// Wraps each span src[start..end] (byte
// offsets) in __vsc_track("name", ...).
fn instrument(
    src: &str,
    spans: &[(usize, usize, &str)],
) -> String {
    todo!()
}

What you learned

• instrument finds collections the way the compiler does and wraps each initializer in __vsc_track(...), applying edits last first so the byte offsets stay valid.
• The JSON is written by Node, from a prelude copied from the TypeScript CLI, so it is byte-identical. Rust only starts node with std::process::Command.
• --profile becomes a HashMap<String, f64>, under one rule both CLIs share: every entry needs a non-negative maxLength.
• Profiles are keyed by name, record the length at exit, and need a CommonJS script that exits cleanly. Python isn’t profiled.