From d664d433dc234c1a45984b08836878b5416533e7 Mon Sep 17 00:00:00 2001 From: imxyy_soope_ Date: Sun, 11 May 2025 14:57:05 +0800 Subject: [PATCH] feat: JIT (WIP) --- flake.nix | 1 + src/bin/repl.rs | 7 +++++-- src/builtins/mod.rs | 2 +- src/ty/internal/attrset.rs | 5 +++-- src/ty/internal/func.rs | 7 +++++-- src/ty/internal/list.rs | 2 +- src/ty/internal/mod.rs | 8 ++++---- src/ty/internal/primop.rs | 8 ++++---- src/vm/mod.rs | 14 +++++++++----- src/vm/test.rs | 6 +++++- 10 files changed, 38 insertions(+), 22 deletions(-) diff --git a/flake.nix b/flake.nix index 7689e82..4233d1e 100644 --- a/flake.nix +++ b/flake.nix @@ -36,6 +36,7 @@ libffi libxml2 ncurses + libz ]; in builtins.concatStringsSep ":" (map (lib: "${lib}/lib") libs) diff --git a/src/bin/repl.rs b/src/bin/repl.rs index 5e0027e..3a11149 100644 --- a/src/bin/repl.rs +++ b/src/bin/repl.rs @@ -1,9 +1,10 @@ +use inkwell::context::Context; use rustyline::error::ReadlineError; use rustyline::{DefaultEditor, Result}; use nixjit::compile::compile; use nixjit::ir::downgrade; -use nixjit::vm::run; +use nixjit::vm::{run, JITContext}; macro_rules! unwrap { ($e:expr) => { @@ -31,7 +32,9 @@ fn main() -> Result<()> { rnix::Root::parse(expr.as_str()).tree().expr().unwrap() )); let prog = compile(downgraded); - println!("{}", unwrap!(run(prog))); + let ctx = Context::create(); + let jit = JITContext::new(&ctx); + println!("{}", unwrap!(run(prog, jit))); } Err(ReadlineError::Interrupted) => { println!("CTRL-C"); diff --git a/src/builtins/mod.rs b/src/builtins/mod.rs index bf80df7..c530ec6 100644 --- a/src/builtins/mod.rs +++ b/src/builtins/mod.rs @@ -1,4 +1,4 @@ -use std::cell::{OnceCell, RefCell}; +use std::cell::RefCell; use std::rc::Rc; use crate::ty::common::Symbol; diff --git a/src/ty/internal/attrset.rs b/src/ty/internal/attrset.rs index e491250..61b0a57 100644 --- a/src/ty/internal/attrset.rs +++ b/src/ty/internal/attrset.rs @@ -10,6 +10,7 @@ use super::super::common::Symbol; use super::super::public as p; use super::{ToPublic, Value}; +#[repr(C)] #[derive(Debug, Constructor, Clone, PartialEq)] pub struct AttrSet<'vm> { data: HashTrieMap>, @@ -82,7 +83,7 @@ impl<'vm> AttrSet<'vm> { &self.data } - pub fn force_deep(&mut self, vm: &VM<'vm>) -> Result<()> { + pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<()> { let mut map: Vec<_> = self .data .into_iter() @@ -169,7 +170,7 @@ impl<'vm> RecAttrSet<'vm> { RecAttrSet { data } } - pub fn force_deep(&mut self, vm: &VM<'vm>) -> Result<()> { + pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<()> { let mut map: Vec<_> = self .data .map diff --git a/src/ty/internal/func.rs b/src/ty/internal/func.rs index a3fda30..fc6ff2f 100644 --- a/src/ty/internal/func.rs +++ b/src/ty/internal/func.rs @@ -8,7 +8,7 @@ use crate::bytecode::OpCodes; use crate::error::Result; use crate::ir; use crate::ty::internal::Value; -use crate::vm::{CapturedEnv, VM}; +use crate::vm::{CapturedEnv, Env, VM}; #[derive(Debug, Clone)] pub enum Param { @@ -40,15 +40,18 @@ impl From for Param { } } +pub type JITFunc<'vm> = unsafe extern "C" fn(vm: *mut VM<'vm, '_>, *mut Env<'vm>, *mut Value<'vm>) -> Value<'vm>; + #[derive(Debug, Clone)] pub struct Func<'vm> { pub env: OnceCell>, pub param: Param, pub opcodes: OpCodes, + pub compiled: Option> } impl<'vm> Func<'vm> { - pub fn call(&'vm self, vm: &VM<'vm>, arg: Value<'vm>) -> Result> { + pub fn call(&'vm self, vm: &VM<'vm, '_>, arg: Value<'vm>) -> Result> { use Param::*; let env = self.env.get().unwrap().clone().released(); diff --git a/src/ty/internal/list.rs b/src/ty/internal/list.rs index 47ab8d2..a703d8b 100644 --- a/src/ty/internal/list.rs +++ b/src/ty/internal/list.rs @@ -30,7 +30,7 @@ impl<'vm> List<'vm> { self } - pub fn force_deep(&mut self, vm: &VM<'vm>) -> Result<()> { + pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<()> { let mut vec: Vec<_> = self.data.iter().cloned().collect(); vec.iter_mut() .map(|v| v.force_deep(vm).map(|_| ())) diff --git a/src/ty/internal/mod.rs b/src/ty/internal/mod.rs index 6046921..f39939b 100644 --- a/src/ty/internal/mod.rs +++ b/src/ty/internal/mod.rs @@ -127,7 +127,7 @@ impl<'vm> Value<'vm> { } } - pub fn call(self, vm: &VM<'vm>, args: Vec>) -> Result> { + pub fn call(self, vm: &VM<'vm, '_>, args: Vec>) -> Result> { use Value::*; Ok(match self { PrimOp(func) => func.call(vm, args), @@ -365,7 +365,7 @@ impl<'vm> Value<'vm> { self } - pub fn force(&mut self, vm: &VM<'vm>) -> Result<&mut Self> { + pub fn force(&mut self, vm: &VM<'vm, '_>) -> Result<&mut Self> { if let Value::Thunk(thunk) = self { let value = thunk.force(vm)?; *self = value @@ -376,7 +376,7 @@ impl<'vm> Value<'vm> { Ok(self) } - pub fn force_deep(&mut self, vm: &VM<'vm>) -> Result<&mut Self> { + pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<&mut Self> { match self { Value::Thunk(thunk) => { let mut value = thunk.force(vm)?; @@ -447,7 +447,7 @@ impl<'vm> Thunk<'vm> { *self.env.borrow_mut() = Some(env); } - pub fn force(&self, vm: &VM<'vm>) -> Result> { + pub fn force(&self, vm: &VM<'vm, '_>) -> Result> { match &*self.thunk.borrow() { _Thunk::Value(value) => return Ok(value.as_ref().clone()), _Thunk::SuspendedFrom(from) => { diff --git a/src/ty/internal/primop.rs b/src/ty/internal/primop.rs index d48943a..3626cf8 100644 --- a/src/ty/internal/primop.rs +++ b/src/ty/internal/primop.rs @@ -8,7 +8,7 @@ use super::Value; pub struct PrimOp<'vm> { pub name: &'static str, arity: u8, - func: fn(&VM<'vm>, Vec>) -> Value<'vm>, + func: fn(&VM<'vm, '_>, Vec>) -> Value<'vm>, } impl PartialEq for PrimOp<'_> { @@ -18,7 +18,7 @@ impl PartialEq for PrimOp<'_> { } impl<'vm> PrimOp<'vm> { - pub fn call(self, vm: &VM<'vm>, args: Vec>) -> Value<'vm> { + pub fn call(self, vm: &VM<'vm, '_>, args: Vec>) -> Value<'vm> { if (args.len() as u8) < self.arity { Value::PartialPrimOp(PartialPrimOp { name: self.name, @@ -39,7 +39,7 @@ pub struct PartialPrimOp<'vm> { pub name: &'static str, arity: u8, args: Vec>, - func: fn(&VM<'vm>, Vec>) -> Value<'vm>, + func: fn(&VM<'vm, '_>, Vec>) -> Value<'vm>, } impl PartialEq for PartialPrimOp<'_> { @@ -49,7 +49,7 @@ impl PartialEq for PartialPrimOp<'_> { } impl<'vm> PartialPrimOp<'vm> { - pub fn call(mut self, vm: &VM<'vm>, args: Vec>) -> Value<'vm> { + pub fn call(mut self, vm: &VM<'vm, '_>, args: Vec>) -> Value<'vm> { let len = args.len() as u8; self.args.extend(args); if len < self.arity { diff --git a/src/vm/mod.rs b/src/vm/mod.rs index b029b22..9cdda9c 100644 --- a/src/vm/mod.rs +++ b/src/vm/mod.rs @@ -11,6 +11,7 @@ use crate::ty::public as p; use stack::{STACK_SIZE, Stack}; pub use env::{CapturedEnv, Env}; +pub use jit::JITContext; mod env; mod jit; @@ -19,7 +20,7 @@ mod stack; #[cfg(test)] mod test; -pub fn run(prog: Program) -> Result { +pub fn run(prog: Program, jit: JITContext<'_>) -> Result { let vm = VM::new( prog.thunks, prog.funcs @@ -28,8 +29,10 @@ pub fn run(prog: Program) -> Result { env: OnceCell::new(), param: f.param, opcodes: f.opcodes, + compiled: None, }) .collect(), + jit ); let env = env(); let temp = vm.eval(prog.top_level, env)?; @@ -37,18 +40,19 @@ pub fn run(prog: Program) -> Result { Ok(temp) } -pub struct VM<'vm> { +pub struct VM<'vm, 'jit> { thunks: Box<[Thunk<'vm>]>, funcs: Box<[Func<'vm>]>, + jit: JITContext<'jit>, } -impl<'vm> VM<'vm> { - fn new(thunks: Box<[OpCodes]>, funcs: Box<[Func<'vm>]>) -> Self { +impl<'vm, 'jit> VM<'vm, 'jit> { + fn new(thunks: Box<[OpCodes]>, funcs: Box<[Func<'vm>]>, jit: JITContext<'jit>) -> Self { let thunks = thunks .into_iter() .map(|opcodes| Thunk::new(opcodes)) .collect(); - VM { thunks, funcs } + VM { thunks, funcs, jit } } pub fn get_thunk(&self, idx: usize) -> &'vm Thunk<'vm> { diff --git a/src/vm/test.rs b/src/vm/test.rs index bb17c1b..6b0d309 100644 --- a/src/vm/test.rs +++ b/src/vm/test.rs @@ -1,5 +1,6 @@ extern crate test; +use inkwell::context::Context; use test::{Bencher, black_box}; use ecow::EcoString; @@ -9,6 +10,7 @@ use crate::compile::compile; use crate::ir::downgrade; use crate::ty::common::Symbol; use crate::ty::public::*; +use crate::vm::JITContext; use super::run; @@ -17,7 +19,9 @@ fn test_expr(expr: &str, expected: Value) { let downgraded = downgrade(rnix::Root::parse(expr).tree().expr().unwrap()).unwrap(); let prog = compile(downgraded); dbg!(&prog); - assert_eq!(run(prog).unwrap(), expected); + let ctx = Context::create(); + let jit = JITContext::new(&ctx); + assert_eq!(run(prog, jit).unwrap(), expected); } macro_rules! thunk {