add copy and patch jit infrastructure with a stencil compilation pipeline

Author: friendlymatthew

Cargo.lock

@@ -1,5 +1,5 @@
 [workspace]
-members = ["gabagool", "gabagool-debug-adapter"]
+members = ["gabagool", "gabagool-debug-adapter", "gabagool-stencils"]
 default-members = ["gabagool"]
 resolver = "2"
 

Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "gabagool-stencils"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+license.workspace = true
+
+[dependencies]

gabagool-stencils/Cargo.toml

@@ -0,0 +1 @@
+

gabagool-stencils/src/lib.rs

@@ -0,0 +1,36 @@
+
+#ifndef STENCIL_CONTEXT_H
+#define STENCIL_CONTEXT_H
+
+#include <stdint.h>
+
+typedef struct StencilContext StencilContext;
+
+// every stencil must be a funciton with this signature
+typedef void (*StencilFn)(StencilContext *ctx);
+
+typedef struct {
+  uint64_t imm0;
+  uint64_t imm1;
+} OpImmediate;
+
+// this is the entire execution state, passed to every stencil
+struct StencilContext {
+  uint64_t *stack;
+  uint64_t stack_pointer;
+  uint64_t *locals;
+  uint64_t *mem_base;
+  uint64_t mem_len;
+  uint64_t *globals;
+  const OpImmediate *imm_table;
+  const StencilFn *fn_table;
+  uint32_t pc;
+  uint8_t snapshot_flag;
+  uint8_t exit_reason;
+  uint32_t exit_value;
+};
+
+#define EXIT_SNAPSHOT 0
+#define EXIT_RETURN 1
+
+#endif

gabagool-stencils/src/stencil_context.h

@@ -0,0 +1,29 @@
+
+#include "stencil_context.h"
+
+void nop(StencilContext *ctx) {
+  ctx->pc += 1;
+
+  if (ctx->snapshot_flag) {
+    ctx->exit_reason = EXIT_SNAPSHOT;
+    return;
+  }
+
+  __attribute__((musttail)) return ctx->fn_table[ctx->pc](ctx);
+}
+
+void i32_const(StencilContext *ctx) {
+  ctx->stack[ctx->stack_pointer] = (uint32_t)ctx->imm_table[ctx->pc].imm0;
+
+  ctx->stack_pointer += 1;
+  ctx->pc += 1;
+
+  if (ctx->snapshot_flag) {
+    ctx->exit_reason = EXIT_SNAPSHOT;
+    return;
+  }
+
+  __attribute__((musttail)) return ctx->fn_table[ctx->pc](ctx);
+}
+
+void return_(StencilContext *ctx) { ctx->exit_reason = EXIT_RETURN; }

gabagool-stencils/src/stencils.c

@@ -10,11 +10,19 @@ license.workspace = true
 
 [dev-dependencies]
 insta = "1"
+wat = "1"
+
+[dependencies]
+libc = { version = "0.2.183", optional = true }
 
 [build-dependencies]
 wast = { version = "245.0.1", optional = true }
+cc = { version = "1.2.57", optional = true }
+object = { version = "0.38.1", optional = true }
+
 
 [features]
 core-tests = ["wast"]
 component-tests = []
 debugger = []
+jit = ["cc", "libc", "object"]

gabagool/Cargo.toml

@@ -24,6 +24,19 @@ fn main() {
         )
         .unwrap();
     }
+
+    #[cfg(feature = "jit")]
+    jit::generate();
+
+    #[cfg(not(feature = "jit"))]
+    {
+        let out_dir = std::env::var("OUT_DIR").unwrap();
+        std::fs::write(
+            std::path::Path::new(&out_dir).join("stencils_generated.rs"),
+            "",
+        )
+        .unwrap();
+    }
 }
 
 #[cfg(feature = "component-tests")]
@@ -607,3 +620,79 @@ mod spec_tests {
         rendered.map(|v| v.join(", "))
     }
 }
+
+#[cfg(feature = "jit")]
+mod jit {
+    use std::{env, fs, path::Path};
+
+    use object::{Object, ObjectSection, ObjectSymbol, SymbolKind};
+
+    const STENCIL_NAMES: &[&str] = &["nop", "i32_const", "return_"];
+
+    pub fn generate() {
+        println!("cargo::rerun-if-changed=../gabagool-stencils/src/stencils.c");
+        println!("cargo::rerun-if-changed=../gabagool-stencils/src/stencil_context.h");
+
+        let out_dir = env::var("OUT_DIR").unwrap();
+
+        let objects = cc::Build::new()
+            .file("../gabagool-stencils/src/stencils.c")
+            .include("../gabagool-stencils/src")
+            .opt_level(3)
+            .flag("-fno-stack-protector")
+            .flag("-fno-asynchronous-unwind-tables")
+            .flag("-fno-exceptions")
+            .cargo_metadata(false)
+            .compile_intermediates();
+
+        let objects = objects.first().expect("expect .o file");
+
+        let obj_data = fs::read(objects).expect("should exist");
+        let obj_file = object::File::parse(&*obj_data).expect("should parse");
+
+        let text_section = obj_file
+            .sections()
+            .find(|s| s.name() == Ok("__text") || s.name() == Ok(".text"))
+            .expect("text section should exist");
+
+        let text_data = text_section.data().unwrap();
+        let text_addr = text_section.address();
+
+        let mut sym_addrs = obj_file
+            .symbols()
+            .filter(|s| s.kind() == SymbolKind::Text && s.section_index().is_some())
+            .filter_map(|s| Some((s.name().ok()?, s.address())))
+            .collect::<Vec<_>>();
+
+        sym_addrs.sort_by_key(|&(_, a)| a);
+
+        let mut generated = String::new();
+
+        for stencil in STENCIL_NAMES {
+            let prefixed = format!("_{}", stencil);
+            let idx = sym_addrs
+                .iter()
+                .position(|&(name, _addr)| name == *stencil || name == prefixed)
+                .unwrap_or_else(|| panic!("symbol {stencil} not found"));
+
+            let addr = sym_addrs[idx].1;
+            let next_addr = sym_addrs
+                .get(idx + 1)
+                .map(|s| s.1)
+                .unwrap_or(text_addr + text_data.len() as u64);
+
+            let offset = (addr - text_addr) as usize;
+            let size = (next_addr - addr) as usize;
+
+            let bs = text_data.get(offset..offset + size).expect("valid bytes");
+
+            generated.push_str(&format!(
+                "pub const STENCIL_{}: &[u8] = &{:?};\n",
+                stencil.to_uppercase(),
+                bs,
+            ));
+        }
+
+        fs::write(Path::new(&out_dir).join("stencils_generated.rs"), generated).unwrap();
+    }
+}

gabagool/build.rs

@@ -320,7 +320,7 @@ impl Debugger {
     }
 }
 
-#[cfg(test)]
+#[cfg(all(test, not(feature = "jit")))]
 mod tests {
     use super::*;
     use crate::Module;