maurice/snix
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor(tvix/eval): non-recursive thunk forcing

Browse source 
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 <tazjin@tvl.su>
Co-authored-by: Griffin Smith <grfn@gws.fyi>
Change-Id: Ifc1747e712663684b2fff53095de62b8459a47f3
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7551
Reviewed-by: grfn <grfn@gws.fyi>
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>
This commit is contained in:
Adam Joseph 2022-12-09 17:27:32 +03:00 • committed by tazjin
parent 4cda236c0c
commit 67d508f2ec
2 changed files with 244 additions and 68 deletions
Download patch file Download diff file Expand all files Collapse all files

144
tvix/eval/src/value/thunk.rs
View file

@ -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<Box<dyn Fn(&mut VM) -> Result<Value, ErrorKind>>>,
) -> Self {
@ -103,7 +99,7 @@ impl Thunk {
None,
move |v: Vec<Value>, 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<Trampoline, ErrorKind> {
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<Trampoline, ErrorKind> {
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<Value> {
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"),
})