Skip to content

fix: lambdas capture their defining environment #883

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mwaddip wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

fix: lambdas capture their defining environment #883

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions ergotree-interpreter/src/eval.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use ergotree_ir::sigma_protocol::sigma_boolean::SigmaProp;
use snumeric::numeric_method_evalfn;

use ergotree_ir::mir::expr::Expr;
use ergotree_ir::mir::value::Lambda;
use ergotree_ir::mir::value::Value;
use ergotree_ir::sigma_protocol::sigma_boolean::SigmaBoolean;

Expand Down Expand Up @@ -205,6 +206,43 @@ pub(crate) trait Evaluable {
) -> Result<Value<'ctx>, EvalError>;
}

/// Per-lambda invoker mirroring the JVM's closure semantics: `FuncValue.eval`
/// returns a closure over the *defining* environment, and each application
/// evaluates the body in that captured env extended with the argument
/// bindings (`env1 = env + (argId -> value)`) — the caller's environment at
/// application time plays no role.
///
/// The captured base is materialized once per lambda value; each `invoke`
/// overwrites the argument slot(s), which matches Scala's per-call extension
/// because an argument binding always shadows a same-id captured binding.
pub(crate) struct LambdaInvoker<'l, 'ctx> {
lambda: &'l Lambda<'ctx>,
env: Env<'ctx>,
}

impl<'l, 'ctx> LambdaInvoker<'l, 'ctx> {
pub(crate) fn new(lambda: &'l Lambda<'ctx>) -> Self {
let mut env = Env::empty();
for (idx, v) in &lambda.captured {
env.insert(*idx, v.clone());
}
Self { lambda, env }
}

/// Bind `args` positionally to the lambda's parameters and evaluate the
/// body in the captured environment.
pub(crate) fn invoke(
&mut self,
ctx: &Context<'ctx>,
args: Vec<Value<'ctx>>,
) -> Result<Value<'ctx>, EvalError> {
for (arg, v) in self.lambda.args.iter().zip(args) {
self.env.insert(arg.idx, v);
}
self.lambda.body.eval(&mut self.env, ctx)
}
}

type EvalFn = for<'ctx> fn(
mc: &SMethod,
env: &mut Env<'ctx>,
Expand Down
170 changes: 146 additions & 24 deletions ergotree-interpreter/src/eval/apply.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,12 @@
use alloc::vec::Vec;
use ergotree_ir::mir::apply::Apply;
use ergotree_ir::mir::val_def::ValId;
use ergotree_ir::mir::value::Value;
use hashbrown::HashMap;

use crate::eval::env::Env;
use crate::eval::Context;
use crate::eval::EvalError;
use crate::eval::Evaluable;
use crate::eval::LambdaInvoker;

impl Evaluable for Apply {
fn eval<'ctx>(
Expand All @@ -23,28 +22,13 @@ impl Evaluable for Apply {
.collect::<Result<_, EvalError>>()?;
match func_v {
Value::Lambda(fv) => {
let arg_ids: Vec<ValId> = fv.args.iter().map(|a| a.idx).collect();
let mut existing_variables = HashMap::new();
let mut new_variables = vec![];
arg_ids.iter().zip(args_v).for_each(|(idx, arg_v)| {
if let Some(old_val) = env.get(*idx) {
existing_variables.insert(idx, old_val.clone());
} else {
new_variables.push(*idx);
}
env.insert(*idx, arg_v);
});
let res = fv.body.eval(env, ctx);
new_variables.into_iter().for_each(|idx| {
env.remove(&idx);
});
existing_variables
.into_iter()
.for_each(|(idx, orig_value)| {
env.insert(*idx, orig_value);
});

res
// The body evaluates in the lambda's CAPTURED environment
// (extended with the argument bindings) — not in the caller's
// env. Mirrors the JVM, where `FuncValue.eval` returns a
// closure over the defining env; the previous bind-into-the-
// caller's-env dance lost the outer binding of a curried
// lambda (`add(3)(1)`) as soon as the outer `Apply` returned.
LambdaInvoker::new(&fv).invoke(ctx, args_v)
}
_ => Err(EvalError::UnexpectedValue(format!(
"expected func_v to be Value::FuncValue got: {0:?}",
Expand Down Expand Up @@ -72,6 +56,144 @@ mod tests {

use super::*;

// SANTA HOF regression (`HOF_currying_Apply_of_Apply` vector): a lambda
// returned from another lambda must capture its defining environment —
// JVM `FuncValue.eval` returns a closure over the env it was created in.
// Pre-fix, `Apply` bound arguments into the CALLER's env and removed them
// after the body ran, so the inner lambda of
// `add = (a: Int) => (b: Int) => a + b` lost `a` the moment the outer
// Apply returned, and `add(3)(1)` errored instead of evaluating to 4.
#[test]
fn curried_lambda_captures_environment() {
use crate::eval::test_util::try_eval_out_with_version;
use ergotree_ir::chain::context::Context;
use ergotree_ir::ergo_tree::ErgoTree;
use ergotree_ir::serialization::SigmaSerializable;
use sigma_test_util::force_any_val;

fn hx(s: &str) -> alloc::vec::Vec<u8> {
(0..s.len())
.step_by(2)
.map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
.collect()
}
// { val add = {(a: Int) => {(b: Int) => a + b}}; add(3)(1) } — the
// vector tree (v3, segregated, sized). Clear the size bit and drop
// the size byte: the eval-tier corpus carries an Int-typed root,
// which the sized parse path rejects.
let bytes = {
let b = hx("1b200204060402d801d601d9010204d90103049a72027203dada7201017300017301");
let mut out = alloc::vec::Vec::with_capacity(b.len() - 1);
out.push(b[0] & !0x08);
out.extend_from_slice(&b[2..]);
out
};
let tree = ErgoTree::sigma_parse_bytes(&bytes).unwrap();
let expr = tree.proposition().unwrap();
let ctx = force_any_val::<Context>();
assert_eq!(
try_eval_out_with_version::<i32>(&expr, &ctx, 3, 3).unwrap(),
4,
"add(3)(1) must evaluate to 4 — the returned lambda must keep its captured `a`"
);
}

// The capture must hold uniformly across invocation sites: a lambda that
// escapes its creation site and is then fed to a HOF must still see its
// captured bindings. `{ val mk = (a: Int) => (b: Int) => a + b;
// val f = mk(3); Coll(1, 2).map(f) }` → [4, 5]. Pre-fix, `Coll.map`
// bound the argument into the caller's env, where `a` no longer existed.
#[test]
fn escaped_lambda_into_hof_keeps_captured_env() {
use ergotree_ir::mir::apply::Apply;
use ergotree_ir::mir::bin_op::ArithOp;
use ergotree_ir::mir::coll_map::Map;
use ergotree_ir::mir::collection::Collection;

let inner_lambda: Expr = FuncValue::new(
vec![FuncArg {
idx: 3.into(),
tpe: SType::SInt,
}],
Expr::BinOp(
BinOp {
kind: ArithOp::Plus.into(),
left: Box::new(
ValUse {
val_id: 2.into(),
tpe: SType::SInt,
}
.into(),
),
right: Box::new(
ValUse {
val_id: 3.into(),
tpe: SType::SInt,
}
.into(),
),
}
.into(),
),
)
.into();
let mk: Expr = FuncValue::new(
vec![FuncArg {
idx: 2.into(),
tpe: SType::SInt,
}],
inner_lambda,
)
.into();
let f: Expr = Apply::new(
ValUse {
val_id: 1.into(),
tpe: mk.tpe(),
}
.into(),
vec![Expr::Const(3i32.into())],
)
.unwrap()
.into();
let block: Expr = BlockValue {
items: vec![
ValDef {
id: 1.into(),
rhs: Box::new(mk),
}
.into(),
ValDef {
id: 4.into(),
rhs: Box::new(f.clone()),
}
.into(),
],
result: Box::new(
Map::new(
Collection::new(
SType::SInt,
vec![Expr::Const(1i32.into()), Expr::Const(2i32.into())],
)
.unwrap()
.into(),
ValUse {
val_id: 4.into(),
tpe: f.tpe(),
}
.into(),
)
.unwrap()
.into(),
),
}
.into();
assert_eq!(
eval_out_wo_ctx::<alloc::vec::Vec<i32>>(&block),
alloc::vec![4, 5],
"the escaped lambda must keep `a = 3` when invoked by Coll.map"
);
}

#[test]
fn eval_user_defined_func_call() {
let arg = Expr::Const(1i32.into());
Expand Down
28 changes: 13 additions & 15 deletions ergotree-interpreter/src/eval/coll_exists.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use crate::eval::env::Env;
use crate::eval::Context;
use crate::eval::EvalError;
use crate::eval::Evaluable;
use crate::eval::LambdaInvoker;

impl Evaluable for Exists {
fn eval<'ctx>(
Expand All @@ -18,28 +19,25 @@ impl Evaluable for Exists {
let input_v = self.input.eval(env, ctx)?;
let condition_v = self.condition.eval(env, ctx)?;
let input_v_clone = input_v.clone();
let mut condition_call = |arg: Value<'ctx>| match &condition_v {
let mut invoker = match &condition_v {
Value::Lambda(func_value) => {
let func_arg = func_value.args.first().ok_or_else(|| {
func_value.args.first().ok_or_else(|| {
EvalError::NotFound(
"Exists: evaluated condition has empty arguments list".to_string(),
)
})?;
let orig_val = env.get(func_arg.idx).cloned();
env.insert(func_arg.idx, arg);
let res = func_value.body.eval(env, ctx);
if let Some(orig_val) = orig_val {
env.insert(func_arg.idx, orig_val);
} else {
env.remove(&func_arg.idx);
}
res
// The body evaluates in the lambda's CAPTURED environment
// (JVM closure semantics) — not in the caller's env.
LambdaInvoker::new(func_value)
}
_ => {
return Err(EvalError::UnexpectedValue(format!(
"expected Exists::condition to be Value::FuncValue got: {0:?}",
input_v_clone
)))
}
_ => Err(EvalError::UnexpectedValue(format!(
"expected Exists::condition to be Value::FuncValue got: {0:?}",
input_v_clone
))),
};
let mut condition_call = |arg: Value<'ctx>| invoker.invoke(ctx, vec![arg]);
let normalized_input_vals: Vec<Value> = match input_v {
Value::Coll(coll) => {
if coll.elem_tpe() != &*self.elem_tpe {
Expand Down
28 changes: 13 additions & 15 deletions ergotree-interpreter/src/eval/coll_filter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use crate::eval::env::Env;
use crate::eval::Context;
use crate::eval::EvalError;
use crate::eval::Evaluable;
use crate::eval::LambdaInvoker;

impl Evaluable for Filter {
fn eval<'ctx>(
Expand All @@ -21,28 +22,25 @@ impl Evaluable for Filter {
let input_v = self.input.eval(env, ctx)?;
let condition_v = self.condition.eval(env, ctx)?;
let input_v_clone = input_v.clone();
let mut condition_call = |arg: Value<'ctx>| match &condition_v {
let mut invoker = match &condition_v {
Value::Lambda(func_value) => {
let func_arg = func_value.args.first().ok_or_else(|| {
func_value.args.first().ok_or_else(|| {
EvalError::NotFound(
"Filter: evaluated condition has empty arguments list".to_string(),
)
})?;
let orig_val = env.get(func_arg.idx).cloned();
env.insert(func_arg.idx, arg);
let res = func_value.body.eval(env, ctx);
if let Some(orig_val) = orig_val {
env.insert(func_arg.idx, orig_val);
} else {
env.remove(&func_arg.idx);
}
res
// The body evaluates in the lambda's CAPTURED environment
// (JVM closure semantics) — not in the caller's env.
LambdaInvoker::new(func_value)
}
_ => {
return Err(EvalError::UnexpectedValue(format!(
"expected Filter::condition to be Value::FuncValue got: {0:?}",
input_v_clone
)))
}
_ => Err(EvalError::UnexpectedValue(format!(
"expected Filter::condition to be Value::FuncValue got: {0:?}",
input_v_clone
))),
};
let mut condition_call = |arg: Value<'ctx>| invoker.invoke(ctx, vec![arg]);
let normalized_input_vals: Vec<Value> = match input_v {
Value::Coll(coll) => {
if coll.elem_tpe() != &*self.elem_tpe {
Expand Down
28 changes: 13 additions & 15 deletions ergotree-interpreter/src/eval/coll_fold.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use crate::eval::env::Env;
use crate::eval::Context;
use crate::eval::EvalError;
use crate::eval::Evaluable;
use crate::eval::LambdaInvoker;

impl Evaluable for Fold {
fn eval<'ctx>(
Expand All @@ -19,27 +20,24 @@ impl Evaluable for Fold {
let zero_v = self.zero.eval(env, ctx)?;
let fold_op_v = self.fold_op.eval(env, ctx)?;
let input_v_clone = input_v.clone();
let mut fold_op_call = |arg: Value<'ctx>| match &fold_op_v {
let mut invoker = match &fold_op_v {
Value::Lambda(func_value) => {
let func_arg = func_value
func_value
.args
.first()
.ok_or_else(|| EvalError::NotFound("empty argument for fold op".to_string()))?;
let orig_val = env.get(func_arg.idx).cloned();
env.insert(func_arg.idx, arg);
let res = func_value.body.eval(env, ctx);
if let Some(orig_val) = orig_val {
env.insert(func_arg.idx, orig_val);
} else {
env.remove(&func_arg.idx);
}
res
// The body evaluates in the lambda's CAPTURED environment
// (JVM closure semantics) — not in the caller's env.
LambdaInvoker::new(func_value)
}
_ => {
return Err(EvalError::UnexpectedValue(format!(
"expected fold_op to be Value::FuncValue got: {0:?}",
input_v_clone
)))
}
_ => Err(EvalError::UnexpectedValue(format!(
"expected fold_op to be Value::FuncValue got: {0:?}",
input_v_clone
))),
};
let mut fold_op_call = |arg: Value<'ctx>| invoker.invoke(ctx, vec![arg]);
match input_v {
Value::Coll(coll) => match coll {
CollKind::NativeColl(NativeColl::CollByte(coll_byte)) => {
Expand Down
Loading
Loading