From 67d508f2ece710714ce8abf6f7deba1fd2440487 Mon Sep 17 00:00:00 2001 From: Adam Joseph Date: Fri, 9 Dec 2022 17:27:32 +0300 Subject: refactor(tvix/eval): non-recursive thunk forcing Introduces continuation-passing-based trampolining of thunk forcing to avoid recursing when forcing deeply nested expressions. This is required for evaluating large expressions. This change was extracted out of cl/7362. Co-authored-by: Vincent Ambo Co-authored-by: Griffin Smith Change-Id: Ifc1747e712663684b2fff53095de62b8459a47f3 Reviewed-on: https://cl.tvl.fyi/c/depot/+/7551 Reviewed-by: grfn Tested-by: BuildkiteCI Reviewed-by: tazjin --- tvix/eval/src/value/thunk.rs | 144 +++++++++++++++++++++++++++++-------- tvix/eval/src/vm.rs | 166 +++++++++++++++++++++++++++++++++---------- 2 files changed, 243 insertions(+), 67 deletions(-) (limited to 'tvix') diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs index 680cc9b1fb..c7cdfa2441 100644 --- a/tvix/eval/src/value/thunk.rs +++ b/tvix/eval/src/value/thunk.rs @@ -30,7 +30,7 @@ use crate::{ spans::LightSpan, upvalues::Upvalues, value::{Builtin, Closure}, - vm::VM, + vm::{Trampoline, TrampolineAction, VM}, Value, }; @@ -84,10 +84,6 @@ impl Thunk { }))) } - /// Create a new thunk from suspended Rust code. - /// - /// The suspended code will be executed and expected to return a - /// value whenever the thunk is forced like any other thunk. pub fn new_suspended_native( native: Rc Result>>, ) -> Self { @@ -103,7 +99,7 @@ impl Thunk { None, move |v: Vec, vm: &mut VM| { // sanity check that only the dummy argument was popped - assert_eq!(v.len(), 1); + assert!(v.len() == 1); assert!(matches!(v[0], Value::Null)); native(vm) }, @@ -127,41 +123,103 @@ impl Thunk { }))) } + /// Force a thunk from a context that can't handle trampoline + /// continuations, eg outside the VM's normal execution loop. Calling + /// `force_trampoline()` instead should be preferred whenever possible. + pub fn force(&self, vm: &mut VM) -> Result<(), ErrorKind> { + if self.is_forced() { + return Ok(()); + } + vm.push(Value::Thunk(self.clone())); + let mut trampoline = Self::force_trampoline(vm)?; + loop { + match trampoline.action { + None => (), + Some(TrampolineAction::EnterFrame { + lambda, + upvalues, + arg_count, + light_span: _, + }) => vm.enter_frame(lambda, upvalues, arg_count)?, + } + match trampoline.continuation { + None => break (), + Some(cont) => { + trampoline = cont(vm)?; + continue; + } + } + } + vm.pop(); + Ok(()) + } + /// Evaluate the content of a thunk, potentially repeatedly, until a /// non-thunk value is returned. /// /// This will change the existing thunk (and thus all references to it, /// providing memoization) through interior mutability. In case of nested /// thunks, the intermediate thunk representations are replaced. - pub fn force(&self, vm: &mut VM) -> Result<(), ErrorKind> { + /// + /// The thunk to be forced should be at the top of the VM stack, + /// and will be left there (but possibly partially forced) when + /// this function returns. + pub fn force_trampoline(vm: &mut VM) -> Result { + match vm.pop() { + Value::Thunk(thunk) => thunk.force_trampoline_self(vm), + v => { + vm.push(v); + Ok(Trampoline::default()) + } + } + } + + fn force_trampoline_self(&self, vm: &mut VM) -> Result { loop { - let mut thunk_mut = self.0.borrow_mut(); + if !self.is_suspended() { + let thunk = self.0.borrow(); + match *thunk { + ThunkRepr::Evaluated(Value::Thunk(ref inner_thunk)) => { + let inner_repr = inner_thunk.0.borrow().clone(); + drop(thunk); + self.0.replace(inner_repr); + } - match *thunk_mut { - ThunkRepr::Evaluated(Value::Thunk(ref inner_thunk)) => { - let inner_repr = inner_thunk.0.borrow().clone(); - *thunk_mut = inner_repr; + ThunkRepr::Evaluated(ref v) => { + vm.push(v.clone()); + return Ok(Trampoline::default()); + } + ThunkRepr::Blackhole => return Err(ErrorKind::InfiniteRecursion), + _ => panic!("impossible"), } - - ThunkRepr::Evaluated(_) => return Ok(()), - ThunkRepr::Blackhole => return Err(ErrorKind::InfiniteRecursion), - - ThunkRepr::Suspended { .. } => { - if let ThunkRepr::Suspended { + } else { + match self.0.replace(ThunkRepr::Blackhole) { + ThunkRepr::Suspended { lambda, upvalues, light_span, - } = std::mem::replace(&mut *thunk_mut, ThunkRepr::Blackhole) - { - drop(thunk_mut); - vm.enter_frame(lambda, upvalues, 0).map_err(|e| { - ErrorKind::ThunkForce(Box::new(Error { - span: light_span.span(), - ..e - })) - })?; - (*self.0.borrow_mut()) = ThunkRepr::Evaluated(vm.pop()) + } => { + let self_clone = self.clone(); + return Ok(Trampoline { + action: Some(TrampolineAction::EnterFrame { + lambda, + upvalues: upvalues.clone(), + arg_count: 0, + light_span: light_span.clone(), + }), + continuation: Some(Box::new(move |vm| { + let should_be_blackhole = + self_clone.0.replace(ThunkRepr::Evaluated(vm.pop())); + assert!(matches!(should_be_blackhole, ThunkRepr::Blackhole)); + vm.push(Value::Thunk(self_clone)); + return Self::force_trampoline(vm).map_err(|kind| Error { + kind, + span: light_span.span(), + }); + })), + }); } + _ => panic!("impossible"), } } } @@ -175,6 +233,20 @@ impl Thunk { matches!(*self.0.borrow(), ThunkRepr::Evaluated(_)) } + pub fn is_suspended(&self) -> bool { + matches!(*self.0.borrow(), ThunkRepr::Suspended { .. }) + } + + /// Returns true if forcing this thunk will not change it. + pub fn is_forced(&self) -> bool { + match *self.0.borrow() { + ThunkRepr::Blackhole => panic!("is_forced() called on a blackholed thunk"), + ThunkRepr::Evaluated(Value::Thunk(_)) => false, + ThunkRepr::Evaluated(_) => true, + _ => false, + } + } + /// Returns a reference to the inner evaluated value of a thunk. /// It is an error to call this on a thunk that has not been /// forced, or is not otherwise known to be fully evaluated. @@ -183,7 +255,21 @@ impl Thunk { // API too much. pub fn value(&self) -> Ref { Ref::map(self.0.borrow(), |thunk| match thunk { - ThunkRepr::Evaluated(value) => value, + ThunkRepr::Evaluated(value) => { + /* + #[cfg(debug_assertions)] + if matches!( + value, + Value::Closure(Closure { + is_finalised: false, + .. + }) + ) { + panic!("Thunk::value called on an unfinalised closure"); + } + */ + return value; + } ThunkRepr::Blackhole => panic!("Thunk::value called on a black-holed thunk"), ThunkRepr::Suspended { .. } => panic!("Thunk::value called on a suspended thunk"), }) diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs index b074bd4224..6c0d1157ec 100644 --- a/tvix/eval/src/vm.rs +++ b/tvix/eval/src/vm.rs @@ -18,6 +18,52 @@ use crate::{ warnings::{EvalWarning, WarningKind}, }; +/// Representation of a VM continuation; +/// see: https://en.wikipedia.org/wiki/Continuation-passing_style#CPS_in_Haskell +type Continuation = Box EvalResult>; + +/// A description of how to continue evaluation of a thunk when returned to by the VM +/// +/// This struct is used when forcing thunks to avoid stack-based recursion, which for deeply nested +/// evaluation can easily overflow the stack. +#[must_use = "this `Trampoline` may be a continuation request, which should be handled"] +#[derive(Default)] +pub struct Trampoline { + /// The action to perform upon return to the trampoline + pub action: Option, + + /// The continuation to execute after the action has completed + pub continuation: Option, +} + +impl Trampoline { + /// Add the execution of a new [`Continuation`] to the existing continuation + /// of this `Trampoline`, returning the resulting `Trampoline`. + pub fn append_to_continuation(self, f: Continuation) -> Self { + Trampoline { + action: self.action, + continuation: match self.continuation { + None => Some(f), + Some(f0) => Some(Box::new(move |vm| { + let trampoline = f0(vm)?; + Ok(trampoline.append_to_continuation(f)) + })), + }, + } + } +} + +/// Description of an action to perform upon return to a [`Trampoline`] by the VM +pub enum TrampolineAction { + /// Enter a new stack frame + EnterFrame { + lambda: Rc, + upvalues: Rc, + light_span: LightSpan, + arg_count: usize, + }, +} + struct CallFrame { /// The lambda currently being executed. lambda: Rc, @@ -32,6 +78,8 @@ struct CallFrame { /// Stack offset, i.e. the frames "view" into the VM's full stack. stack_offset: usize, + + continuation: Option, } impl CallFrame { @@ -324,7 +372,6 @@ impl<'o> VM<'o> { Ok(res) } - #[inline(always)] fn tail_call_value(&mut self, callable: Value) -> EvalResult<()> { match callable { Value::Builtin(builtin) => self.call_builtin(builtin), @@ -362,8 +409,8 @@ impl<'o> VM<'o> { } } - /// Execute the given lambda in this VM's context, returning its - /// value after its stack frame completes. + /// Execute the given lambda in this VM's context, leaving the + /// computed value on its stack after the frame completes. pub fn enter_frame( &mut self, lambda: Rc, @@ -378,10 +425,33 @@ impl<'o> VM<'o> { upvalues, ip: CodeIdx(0), stack_offset: self.stack.len() - arg_count, + continuation: None, }; + let starting_frames_depth = self.frames.len(); self.frames.push(frame); - let result = self.run(); + + let result = loop { + let op = self.inc_ip(); + + self.observer + .observe_execute_op(self.frame().ip, &op, &self.stack); + + let res = self.run_op(op); + + let mut retrampoline: Option = None; + + // we need to pop the frame before checking `res` for an + // error in order to implement `tryEval` correctly. + if self.frame().ip.0 == self.chunk().code.len() { + let frame = self.frames.pop(); + retrampoline = frame.and_then(|frame| frame.continuation); + } + self.trampoline_loop(res?, retrampoline)?; + if self.frames.len() == starting_frames_depth { + break Ok(()); + } + }; self.observer .observe_exit_frame(self.frames.len() + 1, &self.stack); @@ -389,35 +459,53 @@ impl<'o> VM<'o> { result } - /// Run the VM's current call frame to completion. - /// - /// On successful return, the top of the stack is the value that - /// the frame evaluated to. The frame itself is popped off. It is - /// up to the caller to consume the value. - fn run(&mut self) -> EvalResult<()> { + fn trampoline_loop( + &mut self, + mut trampoline: Trampoline, + mut retrampoline: Option, + ) -> EvalResult<()> { loop { - // Break the loop if this call frame has already run to - // completion, pop it off, and return the value to the - // caller. - if self.frame().ip.0 == self.chunk().code.len() { - self.frames.pop(); - return Ok(()); + if let Some(TrampolineAction::EnterFrame { + lambda, + upvalues, + arg_count, + light_span: _, + }) = trampoline.action + { + let frame = CallFrame { + lambda, + upvalues, + ip: CodeIdx(0), + stack_offset: self.stack.len() - arg_count, + continuation: match retrampoline { + None => trampoline.continuation, + Some(retrampoline) => match trampoline.continuation { + None => None, + Some(cont) => Some(Box::new(|vm| { + Ok(cont(vm)?.append_to_continuation(retrampoline)) + })), + }, + }, + }; + self.frames.push(frame); + break; } - let op = self.inc_ip(); - - self.observer - .observe_execute_op(self.frame().ip, &op, &self.stack); - - let res = self.run_op(op); - - if self.frame().ip.0 == self.chunk().code.len() { - self.frames.pop(); - return res; - } else { - res?; + match trampoline.continuation { + None => { + if let Some(cont) = retrampoline.take() { + trampoline = cont(self)?; + } else { + break; + } + } + Some(cont) => { + trampoline = cont(self)?; + continue; + } } } + Ok(()) } pub(crate) fn nix_eq( @@ -428,7 +516,8 @@ impl<'o> VM<'o> { ) -> EvalResult { self.push(v1); self.push(v2); - self.nix_op_eq(allow_top_level_pointer_equality_on_functions_and_thunks)?; + let res = self.nix_op_eq(allow_top_level_pointer_equality_on_functions_and_thunks); + self.trampoline_loop(res?, None)?; match self.pop() { Value::Bool(b) => Ok(b), v => panic!("run_op(OpEqual) left a non-boolean on the stack: {v:#?}"), @@ -438,7 +527,7 @@ impl<'o> VM<'o> { pub(crate) fn nix_op_eq( &mut self, allow_top_level_pointer_equality_on_functions_and_thunks: bool, - ) -> EvalResult<()> { + ) -> EvalResult { // This bit gets set to `true` (if it isn't already) as soon // as we start comparing the contents of two // {lists,attrsets} -- but *not* the contents of two thunks. @@ -566,10 +655,10 @@ impl<'o> VM<'o> { }; self.pop_then_drop(numpairs * 2); self.push(Value::Bool(res)); - Ok(()) + Ok(Trampoline::default()) } - fn run_op(&mut self, op: OpCode) -> EvalResult<()> { + pub(crate) fn run_op(&mut self, op: OpCode) -> EvalResult { match op { OpCode::OpConstant(idx) => { let c = self.chunk()[idx].clone(); @@ -918,14 +1007,15 @@ impl<'o> VM<'o> { } OpCode::OpForce => { - let mut value = self.pop(); + let value = self.pop(); if let Value::Thunk(thunk) = value { - fallible!(self, thunk.force(self)); - value = thunk.value().clone(); + self.push(Value::Thunk(thunk)); + let trampoline = fallible!(self, Thunk::force_trampoline(self)); + return Ok(trampoline); + } else { + self.push(value); } - - self.push(value); } OpCode::OpFinalise(StackIdx(idx)) => { @@ -953,7 +1043,7 @@ impl<'o> VM<'o> { } } - Ok(()) + Ok(Trampoline::default()) } fn run_attrset(&mut self, count: usize) -> EvalResult<()> { -- cgit 1.4.1