profiler: Fill in the profiler state functions (#1226)

* profiler: Fill in the profiler state functions

This patch is a follow-up to #1213

It fills in the profiler state implementation, so that it can be
usable from the whamm monitor.

The most notable change that is not described in the original issue,
is the profile result encoding format, described in the `format.rs` module.

* Review comments

Author: saulecabrera

crates/profiler-lib/src/format.rs

@@ -0,0 +1,127 @@
+//! Binary encoding of the profile report.
+//!
+//! A report is a magic + version header followed by fixed-width,
+//! little-endian records, in the following format:
+//!
+//! ```text
+//! magic   : b"JPRF"   (4 bytes)
+//! version : u8        (== VERSION)
+//! records : u32       (number of records that follow)
+//! record[] :          (16 bytes each, ordered by (func_addr, target))
+//!   func_addr : u32
+//!   target    : u32
+//!   count     : u64
+//! ```
+
+use anyhow::{Result, ensure};
+
+const MAGIC: &[u8; 4] = b"JPRF";
+const VERSION: u8 = 1;
+const HEADER_LEN: usize = MAGIC.len() + size_of::<u8>() + size_of::<u32>();
+const RECORD_LEN: usize = size_of::<u32>() + size_of::<u32>() + size_of::<u64>();
+
+/// A single `(func_addr, target) -> count` report entry.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct Record {
+    /// QuickJS bytecode buffer start address, identifying the JS function.
+    pub func_addr: u32,
+    /// The `br_table` target, i.e. the QuickJS opcode.
+    pub target: u32,
+    /// Total countable Wasm instructions attributed to that opcode.
+    pub count: u64,
+}
+
+/// Serialize `records` into the binary report format.
+pub fn write<I>(records: I) -> Vec<u8>
+where
+    I: ExactSizeIterator<Item = Record>,
+{
+    let mut out = Vec::with_capacity(HEADER_LEN + records.len() * RECORD_LEN);
+    out.extend_from_slice(MAGIC);
+    out.push(VERSION);
+    out.extend_from_slice(&(records.len() as u32).to_le_bytes());
+    for r in records {
+        out.extend_from_slice(&r.func_addr.to_le_bytes());
+        out.extend_from_slice(&r.target.to_le_bytes());
+        out.extend_from_slice(&r.count.to_le_bytes());
+    }
+    out
+}
+
+/// Parse a binary report, validating the magic and version.
+pub fn read(bytes: &[u8]) -> Result<Vec<Record>> {
+    ensure!(bytes.len() >= HEADER_LEN, "report shorter than header");
+    ensure!(&bytes[..MAGIC.len()] == MAGIC, "bad report magic");
+
+    let version = bytes[MAGIC.len()];
+    ensure!(version == VERSION, "unsupported report version {version}");
+
+    let count_off = MAGIC.len() + size_of::<u8>();
+    let count =
+        u32::from_le_bytes(bytes[count_off..count_off + size_of::<u32>()].try_into()?) as usize;
+
+    let body = &bytes[HEADER_LEN..];
+    ensure!(
+        body.len() == count * RECORD_LEN,
+        "report body is {} bytes, expected {} for {count} records",
+        body.len(),
+        count * RECORD_LEN
+    );
+
+    let records = body
+        .chunks_exact(RECORD_LEN)
+        .map(|c| Record {
+            func_addr: u32::from_le_bytes(c[0..4].try_into().unwrap()),
+            target: u32::from_le_bytes(c[4..8].try_into().unwrap()),
+            count: u64::from_le_bytes(c[8..16].try_into().unwrap()),
+        })
+        .collect();
+
+    Ok(records)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn roundtrips() {
+        let records = vec![
+            Record {
+                func_addr: 0x1000,
+                target: 5,
+                count: 3,
+            },
+            Record {
+                func_addr: 0x1000,
+                target: 7,
+                count: 5,
+            },
+        ];
+        let bytes = write(records.clone().into_iter());
+        assert_eq!(read(&bytes).unwrap(), records);
+    }
+
+    #[test]
+    fn empty_roundtrips() {
+        let bytes = write(std::iter::empty());
+        assert!(read(&bytes).unwrap().is_empty());
+    }
+
+    #[test]
+    fn rejects_bad_magic() {
+        let mut bytes = write(std::iter::empty());
+        bytes[0] = b'X';
+        assert!(read(&bytes).is_err());
+    }
+
+    #[test]
+    fn rejects_truncated_body() {
+        let bytes = write(std::iter::once(Record {
+            func_addr: 1,
+            target: 2,
+            count: 3,
+        }));
+        assert!(read(&bytes[..bytes.len() - 1]).is_err());
+    }
+}

crates/profiler-lib/src/interpreter.rs

@@ -186,6 +186,32 @@ pub fn is_byte_load(load: &Load) -> bool {
     matches!(load.kind, LoadKind::I32_8 { .. } | LoadKind::I64_8 { .. })
 }
 
+/// Collect the byte offsets of "countable" opcodes.
+pub(crate) fn countable_opcodes(func: &LocalFunction) -> BTreeSet<u32> {
+    #[derive(Default)]
+    struct Collect {
+        countable: BTreeSet<u32>,
+    }
+    impl<'instr> Visitor<'instr> for Collect {
+        fn visit_instr(&mut self, instr: &'instr Instr, loc: &'instr InstrLocId) {
+            let excluded = matches!(
+                instr,
+                Instr::Block(_)
+                    | Instr::Loop(_)
+                    | Instr::Drop(_)
+                    | Instr::Return(_)
+                    | Instr::Unreachable(_)
+            );
+            if !excluded {
+                self.countable.insert(loc.data());
+            }
+        }
+    }
+    let mut v = Collect::default();
+    dfs_in_order(&mut v, func, func.entry_block());
+    v.countable
+}
+
 impl<'f, 'instr> Visitor<'instr> for AbstractInterp<'f> {
     fn visit_instr(&mut self, _: &'instr Instr, loc: &'instr InstrLocId) {
         // Save the program counter before visiting each operator.

crates/profiler-lib/src/lib.rs

@@ -1,18 +1,30 @@
+pub mod format;
 mod interpreter;
 mod state;
 
-use state::State;
+use state::{Profiler, State};
+use std::cell::RefCell;
 use std::io::Read;
 use std::sync::OnceLock;
 
 static STATE: OnceLock<State> = OnceLock::new();
 
+thread_local! {
+    /// Runtime profiling state.
+    static PROFILER: RefCell<Profiler> = RefCell::new(Profiler::new());
+}
+
 fn state() -> &'static State {
     STATE
         .get()
         .expect("STATE must be initialized via `wizer.initialize`")
 }
 
+/// Run `f` with mutable access to the runtime [`Profiler`].
+fn with_profiler<R>(f: impl FnOnce(&mut Profiler) -> R) -> R {
+    PROFILER.with_borrow_mut(f)
+}
+
 /// Use Wizer to pre-initialize the user library module that will be
 /// passed to Whamm. Note that this is temporary workaround until
 /// Whamm grows such capability (see
@@ -63,32 +75,55 @@ pub extern "C" fn is_dispatch_load(fid: u32, pc: u32) -> bool {
 /// Start a new JS function frame.
 #[unsafe(no_mangle)]
 pub extern "C" fn start_func() {
-    todo!()
+    with_profiler(|p| p.start_func());
 }
 
-/// Exit the top most JS function frame.
+/// Pop the top most JS function frame. The outermost activation closes
+/// out the final opcode with the instruction count at this point.
 #[unsafe(no_mangle)]
-pub extern "C" fn exit_func() {
-    todo!()
+pub extern "C" fn exit_func(instruction_count: i64) {
+    with_profiler(|p| p.exit_func(instruction_count as u64));
 }
 
-/// Set current dispatch function target. i.e., the `br_table` target.
+/// Switch to the dispatch `br_table` target, closing out the previous
+/// opcode with the instruction count at this point.
 #[unsafe(no_mangle)]
-pub extern "C" fn set_dispatch_target(_target: u32) {
-    todo!()
+pub extern "C" fn set_dispatch_target(target: u32, instruction_count: i64) {
+    with_profiler(|p| p.set_dispatch_target(target, instruction_count as u64));
 }
 
 /// Set the effective address (i.e., the start address) of the current
 /// function which uniquely identifies it.
 #[unsafe(no_mangle)]
-pub extern "C" fn set_func_addr(_addr: u32) {
-    todo!()
+pub extern "C" fn set_func_addr(addr: u32) {
+    with_profiler(|p| p.set_func_addr(addr));
+}
+
+/// Whether the opcode at offset `pc` in function `fid` is counted by
+/// the profiler.
+#[unsafe(no_mangle)]
+pub extern "C" fn is_countable_opcode(fid: u32, pc: u32) -> bool {
+    state().is_countable_opcode(fid, pc)
+}
+
+/// Flush the profiler results into a buffer.
+#[unsafe(no_mangle)]
+pub extern "C" fn report() {
+    with_profiler(|p| p.report());
+}
+
+// TODO: Validate that `report` must run exactly once per profile
+// invocation.
+/// Linear-memory offset of the serialized report buffer.
+#[unsafe(no_mangle)]
+pub extern "C" fn report_ptr() -> u32 {
+    with_profiler(|p| p.report_bytes().as_ptr() as u32)
 }
 
-/// Handle the execution of the given opcode.
+/// Length in bytes of the serialized report buffer.
 #[unsafe(no_mangle)]
-pub extern "C" fn handle_opcode(_pc: u32) {
-    todo!()
+pub extern "C" fn report_len() -> u32 {
+    with_profiler(|p| p.report_bytes().len() as u32)
 }
 
 #[cfg(test)]

crates/profiler-lib/src/monitor.mm

@@ -1,57 +1,46 @@
 use {{LIBRARY_NAME}};
 
-// Ensure that we only read the first address of the function
-// bytecode, which serves as the function identifier.
-var expect_first_load: bool;
+// Global count of executed countable Wasm instructions, across every
+// function.
+var instruction_count: i64;
 
-// True only while executing a handler body, that is the window
-// between a dispatch `br_table` and the next dispatch load. Provides
-// more accurate counting of Wasm-instructions-per-JS-opcode.
-var counting: bool;
+// Set to true on the dispatch function entry to capture the function's
+// address from its first dispatch load.
+var expect_first_load: bool;
 
 wasm:func:entry / @static {{LIBRARY_NAME}}.is_dispatch_func(fid as i32) as bool / {
-    // When entering the dispatch function, track a new JS function frame.
+    // Push a new JS function frame.
     {{LIBRARY_NAME}}.start_func();
-    // When entering the dispatch function, we expect the first load later on.
     expect_first_load = true;
 }
 
 wasm:func:exit / @static {{LIBRARY_NAME}}.is_dispatch_func(fid as i32) as bool / {
-    // When exiting the dispatch function, pop the current JS function frame.
-    {{LIBRARY_NAME}}.exit_func();
+    // Pop the topmost JS function frame. The outermost activation closes
+    // out the final opcode with the current instruction count.
+    {{LIBRARY_NAME}}.exit_func(instruction_count);
 }
 
+// Switching the dispatch target closes out the opcode that just ran
+// and begins the new one.
+// Declared before the increment probe so this `br_table` counts toward
+// current JS opcode it dispatches to, not the previous one.
 wasm:opcode:br_table:before / @static {{LIBRARY_NAME}}.is_dispatch_func(fid as i32) as bool / {
-    // We are in the dispatch loop, so we track the current dispatch target.
-    {{LIBRARY_NAME}}.set_dispatch_target(target as i32);
-    // Entering a handler body for `target`, so we begin counting.
-    counting = true;
+    {{LIBRARY_NAME}}.set_dispatch_target(target as i32, instruction_count);
 }
 
-wasm:opcode:*load*:before / @static {{LIBRARY_NAME}}.is_dispatch_load(fid as i32, pc as i32) as bool / {
-    // When hitting the dispatch load, stop counting, the current and
-    // following instructions do not belong to any JS opcode handler.
-    counting = false;
+// Increment the instruction count, iff the opcode is countable.
+wasm:opcode:*:before / @static {{LIBRARY_NAME}}.is_countable_opcode(fid as i32, pc as i32) as bool / {
+    instruction_count = instruction_count + 1;
 }
 
 wasm:opcode:*load*:before / expect_first_load && @static {{LIBRARY_NAME}}.is_dispatch_load(fid as i32, pc as i32) as bool / {
-    // Set the current function identifier, which is the effective address of the dispatch load.
-    // We no longer care about `expect_first_load`, this probe should
-    // only fire once per JS function invocation.
+    // First load of the frame: its effective address identifies the JS
+    // function.
     {{LIBRARY_NAME}}.set_func_addr(effective_addr as i32);
     expect_first_load = false;
 }
 
-wasm:opcode:*:before / @static {{LIBRARY_NAME}}.is_dispatch_func(fid as i32) as bool
-    && counting
-    && opname != "nop"
-    && opname != "drop"
-    && opname != "block"
-    && opname != "loop"
-    && opname != "unreachable"
-    && opname != "return"
-    && opname != "else"
-    && opname != "end" / {
-
-    {{LIBRARY_NAME}}.handle_opcode(pc as i32);
+wasm:func:exit / fname == "_start" / {
+    // At program exit emit the single, whole-execution report.
+    {{LIBRARY_NAME}}.report();
 }