refactor: split VmContext

fix/src/lib.rs

@@ -8,33 +8,33 @@ use fix_common::{StringId, Symbol};
 use fix_error::{Error, Result, Source};
 use fix_ir::downgrade::{Downgrade as _, DowngradeContext};
 use fix_ir::{Ir, IrRef, MaybeThunk, RawIrRef, ThunkId};
-use fix_vm::{ForceMode, StaticValue, Vm, VmContext};
+use fix_vm::{ForceMode, StaticValue, Vm, VmCode, VmContext, VmRuntimeCtx};
 use ghost_cell::{GhostCell, GhostToken};
 use hashbrown::{HashMap, HashSet};
 use string_interner::{DefaultStringInterner, Symbol as _};
 
-// mod fetcher;
-// mod nar;
-// mod nix_utils;
-// mod store;
-// mod string_context;
 mod derivation;
 pub mod logging;
 
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
+pub struct RuntimeState {
+    pub strings: DefaultStringInterner,
+    pub constants: Constants,
+}
+
+pub struct CodeState {
+    pub bytecode: Vec<u8>,
+    pub sources: Vec<Source>,
+    pub spans: Vec<(usize, rnix::TextRange)>,
+    pub thunk_count: usize,
+    pub global_env: HashMap<StringId, Ir<'static, RawIrRef<'static>>>,
+}
+
 pub struct Evaluator {
-    bytecode: Vec<u8>,
-    constants: Constants,
-    strings: DefaultStringInterner,
-
-    sources: Vec<Source>,
-    spans: Vec<(usize, rnix::TextRange)>,
-    // FIXME: remove?
-    thunk_count: usize,
-
-    global_env: HashMap<StringId, Ir<'static, RawIrRef<'static>>>,
+    pub runtime: RuntimeState,
+    pub code: CodeState,
 }
 
 impl Default for Evaluator {
@@ -48,14 +48,17 @@ impl Evaluator {
         let mut strings = DefaultStringInterner::new();
         let global_env = fix_ir::new_global_env(&mut strings);
         Self {
-            sources: Vec::new(),
-            spans: Vec::new(),
-            strings,
-            thunk_count: 0,
-            bytecode: Vec::new(),
-            constants: Constants::default(),
-
-            global_env,
+            runtime: RuntimeState {
+                strings,
+                constants: Constants::default(),
+            },
+            code: CodeState {
+                sources: Vec::new(),
+                spans: Vec::new(),
+                thunk_count: 0,
+                bytecode: Vec::new(),
+                global_env,
+            },
         }
     }
 
@@ -85,8 +88,13 @@ impl Evaluator {
         extra_scope: Option<Scope<'ctx>>,
         force_mode: ForceMode,
     ) -> Result<fix_common::Value> {
-        let root = self.downgrade(source, extra_scope)?;
-        let ip = fix_codegen::compile_bytecode(root.as_ref(), self);
+        let ip = {
+            let mut compiler = CompilerCtx {
+                code: &mut self.code,
+                runtime: &mut self.runtime,
+            };
+            compiler.compile_bytecode(source, extra_scope)?
+        };
         Vm::run(self, ip, force_mode)
     }
 
@@ -100,55 +108,81 @@ impl Evaluator {
     }
 
     pub fn compile_bytecode(&mut self, source: Source) -> Result<InstructionPtr> {
-        let root = self.downgrade(source, None)?;
-        let ip = fix_codegen::compile_bytecode(root.as_ref(), self);
-        Ok(ip)
+        let mut compiler = CompilerCtx {
+            code: &mut self.code,
+            runtime: &mut self.runtime,
+        };
+        compiler.compile_bytecode(source, None)
     }
 
     pub fn disassemble_colored(&self, ip: InstructionPtr) -> String {
         Disassembler::new(ip, self).disassemble_colored()
     }
+}
 
-    fn downgrade_ctx<'a, 'bump, 'id>(
-        &'a mut self,
-        bump: &'bump Bump,
-        token: GhostToken<'id>,
-        extra_scope: Option<Scope<'a>>,
-    ) -> DowngradeCtx<'a, 'id, 'bump> {
-        let Self {
-            global_env,
-            sources,
-            thunk_count,
-            strings,
-            ..
-        } = self;
-        DowngradeCtx {
-            bump,
-            token,
-            strings,
-            source: sources.last().expect("no current source").clone(),
-            scopes: [Scope::Global(global_env)]
-                .into_iter()
-                .chain(extra_scope)
-                .collect(),
-            with_scope_count: 0,
-            arg_count: 0,
-            thunk_count,
-            thunk_scopes: vec![ThunkScope::new_in(bump)],
-        }
+impl VmRuntimeCtx for RuntimeState {
+    fn intern_string(&mut self, s: impl AsRef<str>) -> StringId {
+        StringId(self.strings.get_or_intern(s))
+    }
+    fn resolve_string(&self, id: StringId) -> &str {
+        #[allow(clippy::unwrap_used)]
+        self.strings.resolve(id.0).unwrap()
+    }
+    fn get_const(&self, id: u32) -> StaticValue {
+        #[allow(clippy::unwrap_used)]
+        self.constants.get(id).unwrap()
+    }
+    fn add_const(&mut self, val: StaticValue) -> u32 {
+        self.constants.insert(val)
+    }
+}
+
+impl VmCode for CodeState {
+    fn bytecode(&self) -> &[u8] {
+        &self.bytecode
+    }
+    fn compile(
+        &mut self,
+        source: Source,
+        runtime: &mut impl VmRuntimeCtx,
+    ) -> Result<InstructionPtr> {
+        let mut compiler = CompilerCtx {
+            code: self,
+            runtime,
+        };
+        compiler.compile_bytecode(source, None)
+    }
+}
+
+impl VmContext for Evaluator {
+    fn split(&mut self) -> (&mut impl VmCode, &mut impl VmRuntimeCtx) {
+        (&mut self.code, &mut self.runtime)
+    }
+}
+
+struct CompilerCtx<'a, R: VmRuntimeCtx> {
+    code: &'a mut CodeState,
+    runtime: &'a mut R,
+}
+
+impl<'a, R: VmRuntimeCtx> CompilerCtx<'a, R> {
+    fn compile_bytecode(
+        &mut self,
+        source: Source,
+        extra_scope: Option<Scope>,
+    ) -> Result<InstructionPtr> {
+        let root = self.downgrade(source, extra_scope)?;
+        let ip = fix_codegen::compile_bytecode(root.as_ref(), self);
+        Ok(ip)
     }
 
-    fn downgrade<'a>(
-        &'a mut self,
-        source: Source,
-        extra_scope: Option<Scope<'a>>,
-    ) -> Result<OwnedIr> {
+    fn downgrade(&mut self, source: Source, extra_scope: Option<Scope>) -> Result<OwnedIr> {
         tracing::debug!("Parsing Nix expression");
 
-        self.sources.push(source.clone());
+        self.code.sources.push(source.clone());
 
         let root = rnix::Root::parse(&source.src);
-        handle_parse_error(root.errors(), source).map_or(Ok(()), Err)?;
+        handle_parse_error(root.errors(), source.clone()).map_or(Ok(()), Err)?;
 
         tracing::debug!("Downgrading Nix expression");
         let expr = root
@@ -157,38 +191,63 @@ impl Evaluator {
             .ok_or_else(|| Error::parse_error("unexpected EOF".into()))?;
         let bump = Bump::new();
         GhostToken::new(|token| {
-            let ir = self
-                .downgrade_ctx(&bump, token, extra_scope)
-                .downgrade_toplevel(expr)?;
+            let downgrade_ctx = DowngradeCtx::new(
+                &bump,
+                token,
+                self.runtime,
+                &self.code.global_env,
+                extra_scope,
+                &mut self.code.thunk_count,
+                source,
+            );
+            let ir = downgrade_ctx.downgrade_toplevel(expr)?;
             let ir = unsafe { std::mem::transmute::<RawIrRef<'_>, RawIrRef<'static>>(ir) };
             Ok(OwnedIr { _bump: bump, ir })
         })
     }
 }
 
-impl VmContext for Evaluator {
-    fn intern_string(&mut self, s: impl AsRef<str>) -> StringId {
-        StringId(self.strings.get_or_intern(s))
-    }
-    fn resolve_string(&self, id: StringId) -> &str {
-        #[allow(clippy::unwrap_used)]
-        self.strings.resolve(id.0).unwrap()
-    }
-    fn bytecode(&self) -> &[u8] {
-        &self.bytecode
-    }
-    fn get_const(&self, id: u32) -> StaticValue {
-        #[allow(clippy::unwrap_used)]
-        self.constants.get(id).unwrap()
+impl<'a, R: VmRuntimeCtx> BytecodeContext for CompilerCtx<'a, R> {
+    fn intern_string(&mut self, s: &str) -> StringId {
+        self.runtime.intern_string(s)
     }
 
-    fn compile(&mut self, _source: Source) {
-        todo!();
+    fn register_span(&mut self, range: rnix::TextRange) -> u32 {
+        let id = self.code.spans.len();
+        let source_id = self
+            .code
+            .sources
+            .len()
+            .checked_sub(1)
+            .expect("current_source not set");
+        self.code.spans.push((source_id, range));
+        id as u32
+    }
+
+    fn get_code(&self) -> &[u8] {
+        &self.code.bytecode
+    }
+
+    fn get_code_mut(&mut self) -> &mut Vec<u8> {
+        &mut self.code.bytecode
+    }
+
+    fn add_constant(&mut self, val: fix_codegen::Const) -> u32 {
+        use fix_codegen::Const::*;
+        let val = match val {
+            Smi(x) => StaticValue::new_inline(x),
+            Float(x) => StaticValue::new_float(x),
+            Bool(x) => StaticValue::new_inline(x),
+            String(x) => StaticValue::new_inline(x),
+            PrimOp { id, arity } => StaticValue::new_primop(id, arity),
+            Null => StaticValue::default(),
+        };
+        self.runtime.add_const(val)
     }
 }
 
 #[derive(Default)]
-struct Constants {
+pub struct Constants {
     data: Vec<StaticValue>,
     dedup: HashMap<u64, u32>,
 }
@@ -236,10 +295,10 @@ fn handle_parse_error<'a>(
     None
 }
 
-struct DowngradeCtx<'ctx, 'id, 'ir> {
+struct DowngradeCtx<'ctx, 'id, 'ir, R: VmRuntimeCtx> {
     bump: &'ir Bump,
     token: GhostToken<'id>,
-    strings: &'ctx mut DefaultStringInterner,
+    runtime: &'ctx mut R,
     source: Source,
     scopes: Vec<Scope<'ctx>>,
     with_scope_count: u32,
@@ -262,11 +321,11 @@ fn should_thunk<'id>(ir: IrRef<'id, '_>, token: &GhostToken<'id>) -> bool {
     )
 }
 
-impl<'ctx, 'id, 'ir> DowngradeCtx<'ctx, 'id, 'ir> {
+impl<'ctx, 'id, 'ir, R: VmRuntimeCtx> DowngradeCtx<'ctx, 'id, 'ir, R> {
     fn new(
         bump: &'ir Bump,
         token: GhostToken<'id>,
-        symbols: &'ctx mut DefaultStringInterner,
+        runtime: &'ctx mut R,
         global: &'ctx HashMap<StringId, Ir<'static, RawIrRef<'static>>>,
         extra_scope: Option<Scope<'ctx>>,
         thunk_count: &'ctx mut usize,
@@ -275,7 +334,7 @@ impl<'ctx, 'id, 'ir> DowngradeCtx<'ctx, 'id, 'ir> {
         Self {
             bump,
             token,
-            strings: symbols,
+            runtime,
             source,
             scopes: std::iter::once(Scope::Global(global))
                 .chain(extra_scope)
@@ -288,7 +347,9 @@ impl<'ctx, 'id, 'ir> DowngradeCtx<'ctx, 'id, 'ir> {
     }
 }
 
-impl<'ctx: 'ir, 'id, 'ir> DowngradeContext<'id, 'ir> for DowngradeCtx<'ctx, 'id, 'ir> {
+impl<'ctx: 'ir, 'id, 'ir, R: VmRuntimeCtx> DowngradeContext<'id, 'ir>
+    for DowngradeCtx<'ctx, 'id, 'ir, R>
+{
     fn new_expr(&self, expr: Ir<'ir, IrRef<'id, 'ir>>) -> IrRef<'id, 'ir> {
         IrRef::new(self.bump.alloc(GhostCell::new(expr)))
     }
@@ -317,11 +378,11 @@ impl<'ctx: 'ir, 'id, 'ir> DowngradeContext<'id, 'ir> for DowngradeCtx<'ctx, 'id,
     }
 
     fn intern_string(&mut self, sym: impl AsRef<str>) -> StringId {
-        StringId(self.strings.get_or_intern(sym))
+        self.runtime.intern_string(sym)
     }
 
     fn resolve_sym(&self, id: StringId) -> Symbol<'_> {
-        self.strings.resolve(id.0).expect("no symbol found").into()
+        self.runtime.resolve_string(id).into()
     }
 
     fn lookup(&self, sym: StringId, span: rnix::TextRange) -> Result<MaybeThunk> {
@@ -383,9 +444,9 @@ impl<'ctx: 'ir, 'id, 'ir> DowngradeContext<'id, 'ir> for DowngradeCtx<'ctx, 'id,
         self.source.clone()
     }
 
-    fn with_let_scope<F, R>(&mut self, keys: &[StringId], f: F) -> Result<R>
+    fn with_let_scope<F, Ret>(&mut self, keys: &[StringId], f: F) -> Result<Ret>
     where
-        F: FnOnce(&mut Self) -> Result<(bumpalo::collections::Vec<'ir, IrRef<'id, 'ir>>, R)>,
+        F: FnOnce(&mut Self) -> Result<(bumpalo::collections::Vec<'ir, IrRef<'id, 'ir>>, Ret)>,
     {
         let base = *self.thunk_count;
         *self.thunk_count = self
@@ -407,9 +468,9 @@ impl<'ctx: 'ir, 'id, 'ir> DowngradeContext<'id, 'ir> for DowngradeCtx<'ctx, 'id,
         Ok(ret)
     }
 
-    fn with_param_scope<F, R>(&mut self, sym: StringId, f: F) -> R
+    fn with_param_scope<F, Ret>(&mut self, sym: StringId, f: F) -> Ret
     where
-        F: FnOnce(&mut Self) -> R,
+        F: FnOnce(&mut Self) -> Ret,
     {
         self.scopes.push(Scope::Param {
             sym,
@@ -419,9 +480,9 @@ impl<'ctx: 'ir, 'id, 'ir> DowngradeContext<'id, 'ir> for DowngradeCtx<'ctx, 'id,
         f(guard.as_ctx())
     }
 
-    fn with_with_scope<F, R>(&mut self, f: F) -> R
+    fn with_with_scope<F, Ret>(&mut self, f: F) -> Ret
     where
-        F: FnOnce(&mut Self) -> R,
+        F: FnOnce(&mut Self) -> Ret,
     {
         self.with_scope_count += 1;
         let ret = f(self);
@@ -429,15 +490,15 @@ impl<'ctx: 'ir, 'id, 'ir> DowngradeContext<'id, 'ir> for DowngradeCtx<'ctx, 'id,
         ret
     }
 
-    fn with_thunk_scope<F, R>(
+    fn with_thunk_scope<F, Ret>(
         &mut self,
         f: F,
     ) -> (
-        R,
+        Ret,
         bumpalo::collections::Vec<'ir, (ThunkId, IrRef<'id, 'ir>)>,
     )
     where
-        F: FnOnce(&mut Self) -> R,
+        F: FnOnce(&mut Self) -> Ret,
     {
         self.thunk_scopes.push(ThunkScope::new_in(self.bump));
         let ret = f(self);
@@ -455,7 +516,7 @@ impl<'ctx: 'ir, 'id, 'ir> DowngradeContext<'id, 'ir> for DowngradeCtx<'ctx, 'id,
     }
 }
 
-impl<'id, 'ir, 'ctx: 'ir> DowngradeCtx<'ctx, 'id, 'ir> {
+impl<'id, 'ir, 'ctx: 'ir, R: VmRuntimeCtx> DowngradeCtx<'ctx, 'id, 'ir, R> {
     fn downgrade_toplevel(mut self, root: rnix::ast::Expr) -> Result<RawIrRef<'ir>> {
         let body = root.downgrade(&mut self)?;
         let thunks = self
@@ -496,18 +557,18 @@ enum Scope<'ctx> {
     Param { sym: StringId, abs_layer: usize },
 }
 
-struct ScopeGuard<'a, 'ctx, 'id, 'ir> {
-    ctx: &'a mut DowngradeCtx<'ctx, 'id, 'ir>,
+struct ScopeGuard<'a, 'ctx, 'id, 'ir, R: VmRuntimeCtx> {
+    ctx: &'a mut DowngradeCtx<'ctx, 'id, 'ir, R>,
 }
 
-impl Drop for ScopeGuard<'_, '_, '_, '_> {
+impl<'a, 'ctx, 'id, 'ir, R: VmRuntimeCtx> Drop for ScopeGuard<'a, 'ctx, 'id, 'ir, R> {
     fn drop(&mut self) {
         self.ctx.scopes.pop();
     }
 }
 
-impl<'id, 'ir, 'ctx> ScopeGuard<'_, 'ctx, 'id, 'ir> {
-    fn as_ctx(&mut self) -> &mut DowngradeCtx<'ctx, 'id, 'ir> {
+impl<'a, 'ctx, 'id, 'ir, R: VmRuntimeCtx> ScopeGuard<'a, 'ctx, 'id, 'ir, R> {
+    fn as_ctx(&mut self) -> &mut DowngradeCtx<'ctx, 'id, 'ir, R> {
         self.ctx
     }
 }
@@ -518,9 +579,6 @@ struct OwnedIr {
 }
 
 impl OwnedIr {
-    /// # Safety
-    ///
-    /// `ir` must be allocated from `bump`.
     unsafe fn new(ir: RawIrRef<'_>, bump: Bump) -> Self {
         Self {
             _bump: bump,
@@ -533,52 +591,14 @@ impl OwnedIr {
     }
 }
 
-impl BytecodeContext for Evaluator {
-    fn intern_string(&mut self, s: &str) -> StringId {
-        StringId(self.strings.get_or_intern(s))
-    }
-
-    fn register_span(&mut self, range: rnix::TextRange) -> u32 {
-        let id = self.spans.len();
-        let source_id = self
-            .sources
-            .len()
-            .checked_sub(1)
-            .expect("current_source not set");
-        self.spans.push((source_id, range));
-        id as u32
-    }
-
-    fn get_code(&self) -> &[u8] {
-        &self.bytecode
-    }
-
-    fn get_code_mut(&mut self) -> &mut Vec<u8> {
-        &mut self.bytecode
-    }
-
-    fn add_constant(&mut self, val: fix_codegen::Const) -> u32 {
-        use fix_codegen::Const::*;
-        let val = match val {
-            Smi(x) => StaticValue::new_inline(x),
-            Float(x) => StaticValue::new_float(x),
-            Bool(x) => StaticValue::new_inline(x),
-            String(x) => StaticValue::new_inline(x),
-            PrimOp { id, arity } => StaticValue::new_primop(id, arity),
-            Null => StaticValue::default(),
-        };
-        self.constants.insert(val)
-    }
-}
-
 impl DisassemblerContext for Evaluator {
     fn get_code(&self) -> &[u8] {
-        &self.bytecode
+        &self.code.bytecode
     }
 
     #[allow(clippy::unwrap_used)]
     fn resolve_string(&self, id: u32) -> &str {
         let id = string_interner::symbol::SymbolU32::try_from_usize(id as usize).unwrap();
-        self.strings.resolve(id).unwrap()
+        self.runtime.strings.resolve(id).unwrap()
     }
 }