fix(aiur): replay constrained calls cached as hints

Ix/Aiur/Stages/Codegen.lean

@@ -398,8 +398,9 @@ private def emitCall (out : Nat) (callee : FunIdx) (args : Array ValIdx)
  s!"\{ let __args: [G; IN_{callee}] = {argsStr};" ++
  s!" let __cu = {cuExpr};" ++
  s!" if let Some(result) = record.function_queries[{callee}].get_mut(&__args[..]) \{" ++
- bumpStmt ++
- retExpr ++
+ s!" if !__cu && *result.multiplicity == G::ZERO \{" ++
+ s!" aiur_fn_{callee}(__args, record, io_buffer, false)?" ++
+ s!" } else \{" ++ bumpStmt ++ retExpr ++ " }" ++
  s!" } else \{ aiur_fn_{callee}(__args, record, io_buffer, __cu)? } }"
  let mut stmts : Array RustStmt := #[
  .letStmt false "__r_arr" (some s!"[G; OUT_{callee}]") (.lit blockExpr)
@@ -816,15 +817,13 @@ partial def emitCtrl (funIdx : FunIdx) (mcLabel? : Option String)
  let outArr : RustStmt :=
  .letStmt false "__ret" (some s!"[G; OUT_{funIdx}]")
  (.arrayLit (outs.map valVar))
- let insertCall : RustStmt :=
- .exprStmt (.call
- (.field
- (.index (.field (.var "record") "function_queries")
- (.lit (toString funIdx)))
- "insert")
- #[.ref (.index (.var "inp") (.lit "..")),
- .ref (.index (.var "__ret") (.lit "..")),
- gFromBool (.lit "!unconstrained")])
+ let insertCall : RustStmt := .exprStmt (.lit <|
+ s!"if let Some(result) = record.function_queries[{funIdx}].get_mut(&inp[..]) \{" ++
+ " debug_assert_eq!(result.output, &__ret[..]);" ++
+ " if !unconstrained { *result.multiplicity += G::ONE; }" ++
+ " } else {" ++
+ s!" record.function_queries[{funIdx}].insert(&inp[..], &__ret[..], G::from_bool(!unconstrained));" ++
+ " }")
  -- Wrap in Ok(...) since fn now returns Result<[G; OUT_N], ExecError>.
  return #[outArr, insertCall,
  .returnStmt (.call (.var "Ok") #[.var "__ret"])]

Tests/Aiur/Aiur.lean

@@ -215,6 +215,25 @@ def toplevel := ⟦
  }
  }
 
+ -- Regression: a constrained call must recursively promote an identical
+ -- query previously cached by an unconstrained call. `promotion_nested`
+ -- calls another function so promoting only its own multiplicity leaves
+ -- the nested function channel unbalanced.
+ fn promotion_leaf(x: G) -> G {
+ x + 1
+ }
+
+ fn promotion_nested(x: G) -> G {
+ promotion_leaf(x)
+ }
+
+ pub fn unconstrained_call_promotion(x: G) -> G {
+ let hinted = #promotion_nested(x);
+ let constrained = promotion_nested(x);
+ assert_eq!(hinted, constrained);
+ constrained
+ }
+
  ---------------------------------------------------------------------------
  -- IO
  ---------------------------------------------------------------------------
@@ -794,6 +813,10 @@ def aiurTestCases : List AiurTestCase := [
  -- Unconstrained recursion: mixed constrained/unconstrained calls
  .prove `unconstrained_fibonacci #[6] #[13],
 
+ -- A constrained call reuses an identical unconstrained cached query and
+ -- must recursively promote the nested callee's lookup multiplicity.
+ .prove `unconstrained_call_promotion #[3] #[4],
+
  -- IO
  { functionName := `read_write_io
  inputIOBuffer :=

crates/aiur/src/execute.rs

@@ -327,13 +327,26 @@ impl Function {
  },
  ExecEntry::Op(Op::Call(callee_idx, args, _, op_unconstrained)) => {
  let args: Vec<G> = args.iter().map(|i| map[*i]).collect();
- if let Some(result) =
- record.function_queries[*callee_idx].get_mut(&args)
- {
- if !unconstrained && !op_unconstrained {
- *result.multiplicity += G::ONE;
- }
- map.extend_from_slice(result.output);
+ let callee_unconstrained = unconstrained || *op_unconstrained;
+ let cached_output = record.function_queries[*callee_idx]
+ .get_mut(&args)
+ .and_then(|result| {
+ // A zero-multiplicity entry was computed only as an
+ // unconstrained hint. Promoting just this row would omit all
+ // of the callee's child lookups and unbalance their channels;
+ // replay the body constrained so promotion recurses through
+ // the whole dependency tree.
+ if !callee_unconstrained && result.multiplicity.is_zero() {
+ None
+ } else {
+ if !callee_unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ Some(result.output.to_vec())
+ }
+ });
+ if let Some(output) = cached_output {
+ map.extend(output);
  } else {
  let saved_map = std::mem::replace(&mut map, args);
  callers_states_stack.push(CallerState {
@@ -343,7 +356,7 @@ impl Function {
  continuation_depth: continuation_stack.len(),
  });
  fun_idx = *callee_idx;
- unconstrained = unconstrained || *op_unconstrained;
+ unconstrained = callee_unconstrained;
  push_block_exec_entries!(&toplevel.functions[fun_idx].body);
  }
  },
@@ -699,11 +712,22 @@ impl Function {
  // Register the query.
  let input_size = toplevel.functions[fun_idx].layout.input_size;
  let output = output.iter().map(|i| map[*i]).collect::<Vec<_>>();
- record.function_queries[fun_idx].insert(
- &map[..input_size],
- &output,
- G::from_bool(!unconstrained),
- );
+ if let Some(result) =
+ record.function_queries[fun_idx].get_mut(&map[..input_size])
+ {
+ // The only ordinary way to execute an already cached function
+ // is constrained promotion of an unconstrained hint entry.
+ debug_assert_eq!(result.output, output);
+ if !unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ } else {
+ record.function_queries[fun_idx].insert(
+ &map[..input_size],
+ &output,
+ G::from_bool(!unconstrained),
+ );
+ }
  if let Some(CallerState {
  fun_idx: caller_idx,
  map: caller_map,

crates/aiur/src/synthesis.rs

@@ -562,6 +562,82 @@ mod tests {
  );
  }
 
+ /// Exercise promotion of a cached unconstrained call through nested calls.
+ ///
+ /// `f` first computes `g(x)` as an unconstrained hint, then calls `g(x)`
+ /// constrained. Since `g` calls `h`, promoting the cached `g` query must
+ /// replay its body and promote the cached `h` query as well. Merely bumping
+ /// `g`'s multiplicity leaves the `g -> h` function channel unbalanced.
+ fn unconstrained_call_promotion_toplevel() -> Toplevel {
+ let f = Function {
+ body: Block {
+ ops: vec![
+ Op::Call(1, vec![0], 1, true),
+ Op::Call(1, vec![0], 1, false),
+ ],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 3,
+ lookups: 2,
+ },
+ entry: true,
+ constrained: true,
+ };
+
+ let g = Function {
+ body: Block {
+ ops: vec![Op::Call(2, vec![0], 1, false)],
+ ctrl: Ctrl::Return(0, vec![1]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 2,
+ lookups: 2,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ let h = Function {
+ body: Block {
+ ops: vec![Op::Const(G::ONE), Op::Add(0, 1)],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 1,
+ lookups: 1,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ Toplevel { functions: vec![f, g, h], memory_sizes: vec![] }
+ }
+
+ #[test]
+ fn prove_verify_promotes_nested_unconstrained_call() {
+ let (cp, fp) = test_parameters();
+ let system =
+ AiurSystem::build(unconstrained_call_promotion_toplevel(), cp, fp);
+ let input = [G::from_u64(3)];
+ let mut io_buffer = empty_io_buffer();
+
+ let (claim, proof) = system.prove(0, &input, &mut io_buffer);
+ assert_eq!(
+ claim,
+ vec![function_channel(), G::ZERO, input[0], input[0] + G::ONE]
+ );
+ system
+ .verify(&claim, &proof)
+ .expect("nested constrained promotion must balance function channels");
+ }
+
  #[test]
  fn prove_verify_mul_roundtrip() {
  let (cp, fp) = test_parameters();

Ix/Aiur/Stages/Codegen.lean

@@ -398,8 +398,9 @@ private def emitCall (out : Nat) (callee : FunIdx) (args : Array ValIdx)
  s!"\{ let __args: [G; IN_{callee}] = {argsStr};" ++
  s!" let __cu = {cuExpr};" ++
  s!" if let Some(result) = record.function_queries[{callee}].get_mut(&__args[..]) \{" ++
- bumpStmt ++
- retExpr ++
+ s!" if !__cu && *result.multiplicity == G::ZERO \{" ++
+ s!" aiur_fn_{callee}(__args, record, io_buffer, false)?" ++
+ s!" } else \{" ++ bumpStmt ++ retExpr ++ " }" ++
  s!" } else \{ aiur_fn_{callee}(__args, record, io_buffer, __cu)? } }"
  let mut stmts : Array RustStmt := #[
  .letStmt false "__r_arr" (some s!"[G; OUT_{callee}]") (.lit blockExpr)
@@ -816,15 +817,13 @@ partial def emitCtrl (funIdx : FunIdx) (mcLabel? : Option String)
  let outArr : RustStmt :=
  .letStmt false "__ret" (some s!"[G; OUT_{funIdx}]")
  (.arrayLit (outs.map valVar))
- let insertCall : RustStmt :=
- .exprStmt (.call
- (.field
- (.index (.field (.var "record") "function_queries")
- (.lit (toString funIdx)))
- "insert")
- #[.ref (.index (.var "inp") (.lit "..")),
- .ref (.index (.var "__ret") (.lit "..")),
- gFromBool (.lit "!unconstrained")])
+ let insertCall : RustStmt := .exprStmt (.lit <|
+ s!"if let Some(result) = record.function_queries[{funIdx}].get_mut(&inp[..]) \{" ++
+ " debug_assert_eq!(result.output, &__ret[..]);" ++
+ " if !unconstrained { *result.multiplicity += G::ONE; }" ++
+ " } else {" ++
+ s!" record.function_queries[{funIdx}].insert(&inp[..], &__ret[..], G::from_bool(!unconstrained));" ++
+ " }")
  -- Wrap in Ok(...) since fn now returns Result<[G; OUT_N], ExecError>.
  return #[outArr, insertCall,
  .returnStmt (.call (.var "Ok") #[.var "__ret"])]

Tests/Aiur/Aiur.lean

@@ -215,6 +215,25 @@ def toplevel := ⟦
  }
  }
 
+ -- Regression: a constrained call must recursively promote an identical
+ -- query previously cached by an unconstrained call. `promotion_nested`
+ -- calls another function so promoting only its own multiplicity leaves
+ -- the nested function channel unbalanced.
+ fn promotion_leaf(x: G) -> G {
+ x + 1
+ }
+
+ fn promotion_nested(x: G) -> G {
+ promotion_leaf(x)
+ }
+
+ pub fn unconstrained_call_promotion(x: G) -> G {
+ let hinted = #promotion_nested(x);
+ let constrained = promotion_nested(x);
+ assert_eq!(hinted, constrained);
+ constrained
+ }
+
  ---------------------------------------------------------------------------
  -- IO
  ---------------------------------------------------------------------------
@@ -794,6 +813,10 @@ def aiurTestCases : List AiurTestCase := [
  -- Unconstrained recursion: mixed constrained/unconstrained calls
  .prove `unconstrained_fibonacci #[6] #[13],
 
+ -- A constrained call reuses an identical unconstrained cached query and
+ -- must recursively promote the nested callee's lookup multiplicity.
+ .prove `unconstrained_call_promotion #[3] #[4],
+
  -- IO
  { functionName := `read_write_io
  inputIOBuffer :=

crates/aiur/src/execute.rs

@@ -327,13 +327,26 @@ impl Function {
  },
  ExecEntry::Op(Op::Call(callee_idx, args, _, op_unconstrained)) => {
  let args: Vec<G> = args.iter().map(|i| map[*i]).collect();
- if let Some(result) =
- record.function_queries[*callee_idx].get_mut(&args)
- {
- if !unconstrained && !op_unconstrained {
- *result.multiplicity += G::ONE;
- }
- map.extend_from_slice(result.output);
+ let callee_unconstrained = unconstrained || *op_unconstrained;
+ let cached_output = record.function_queries[*callee_idx]
+ .get_mut(&args)
+ .and_then(|result| {
+ // A zero-multiplicity entry was computed only as an
+ // unconstrained hint. Promoting just this row would omit all
+ // of the callee's child lookups and unbalance their channels;
+ // replay the body constrained so promotion recurses through
+ // the whole dependency tree.
+ if !callee_unconstrained && result.multiplicity.is_zero() {
+ None
+ } else {
+ if !callee_unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ Some(result.output.to_vec())
+ }
+ });
+ if let Some(output) = cached_output {
+ map.extend(output);
  } else {
  let saved_map = std::mem::replace(&mut map, args);
  callers_states_stack.push(CallerState {
@@ -343,7 +356,7 @@ impl Function {
  continuation_depth: continuation_stack.len(),
  });
  fun_idx = *callee_idx;
- unconstrained = unconstrained || *op_unconstrained;
+ unconstrained = callee_unconstrained;
  push_block_exec_entries!(&toplevel.functions[fun_idx].body);
  }
  },
@@ -699,11 +712,22 @@ impl Function {
  // Register the query.
  let input_size = toplevel.functions[fun_idx].layout.input_size;
  let output = output.iter().map(|i| map[*i]).collect::<Vec<_>>();
- record.function_queries[fun_idx].insert(
- &map[..input_size],
- &output,
- G::from_bool(!unconstrained),
- );
+ if let Some(result) =
+ record.function_queries[fun_idx].get_mut(&map[..input_size])
+ {
+ // The only ordinary way to execute an already cached function
+ // is constrained promotion of an unconstrained hint entry.
+ debug_assert_eq!(result.output, output);
+ if !unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ } else {
+ record.function_queries[fun_idx].insert(
+ &map[..input_size],
+ &output,
+ G::from_bool(!unconstrained),
+ );
+ }
  if let Some(CallerState {
  fun_idx: caller_idx,
  map: caller_map,

crates/aiur/src/synthesis.rs

@@ -562,6 +562,82 @@ mod tests {
  );
  }
 
+ /// Exercise promotion of a cached unconstrained call through nested calls.
+ ///
+ /// `f` first computes `g(x)` as an unconstrained hint, then calls `g(x)`
+ /// constrained. Since `g` calls `h`, promoting the cached `g` query must
+ /// replay its body and promote the cached `h` query as well. Merely bumping
+ /// `g`'s multiplicity leaves the `g -> h` function channel unbalanced.
+ fn unconstrained_call_promotion_toplevel() -> Toplevel {
+ let f = Function {
+ body: Block {
+ ops: vec![
+ Op::Call(1, vec![0], 1, true),
+ Op::Call(1, vec![0], 1, false),
+ ],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 3,
+ lookups: 2,
+ },
+ entry: true,
+ constrained: true,
+ };
+
+ let g = Function {
+ body: Block {
+ ops: vec![Op::Call(2, vec![0], 1, false)],
+ ctrl: Ctrl::Return(0, vec![1]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 2,
+ lookups: 2,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ let h = Function {
+ body: Block {
+ ops: vec![Op::Const(G::ONE), Op::Add(0, 1)],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 1,
+ lookups: 1,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ Toplevel { functions: vec![f, g, h], memory_sizes: vec![] }
+ }
+
+ #[test]
+ fn prove_verify_promotes_nested_unconstrained_call() {
+ let (cp, fp) = test_parameters();
+ let system =
+ AiurSystem::build(unconstrained_call_promotion_toplevel(), cp, fp);
+ let input = [G::from_u64(3)];
+ let mut io_buffer = empty_io_buffer();
+
+ let (claim, proof) = system.prove(0, &input, &mut io_buffer);
+ assert_eq!(
+ claim,
+ vec![function_channel(), G::ZERO, input[0], input[0] + G::ONE]
+ );
+ system
+ .verify(&claim, &proof)
+ .expect("nested constrained promotion must balance function channels");
+ }
+
  #[test]
  fn prove_verify_mul_roundtrip() {
  let (cp, fp) = test_parameters();

Ix/Aiur/Stages/Codegen.lean

@@ -398,8 +398,9 @@ private def emitCall (out : Nat) (callee : FunIdx) (args : Array ValIdx)
  s!"\{ let __args: [G; IN_{callee}] = {argsStr};" ++
  s!" let __cu = {cuExpr};" ++
  s!" if let Some(result) = record.function_queries[{callee}].get_mut(&__args[..]) \{" ++
- bumpStmt ++
- retExpr ++
+ s!" if !__cu && *result.multiplicity == G::ZERO \{" ++
+ s!" aiur_fn_{callee}(__args, record, io_buffer, false)?" ++
+ s!" } else \{" ++ bumpStmt ++ retExpr ++ " }" ++
  s!" } else \{ aiur_fn_{callee}(__args, record, io_buffer, __cu)? } }"
  let mut stmts : Array RustStmt := #[
  .letStmt false "__r_arr" (some s!"[G; OUT_{callee}]") (.lit blockExpr)
@@ -816,15 +817,13 @@ partial def emitCtrl (funIdx : FunIdx) (mcLabel? : Option String)
  let outArr : RustStmt :=
  .letStmt false "__ret" (some s!"[G; OUT_{funIdx}]")
  (.arrayLit (outs.map valVar))
- let insertCall : RustStmt :=
- .exprStmt (.call
- (.field
- (.index (.field (.var "record") "function_queries")
- (.lit (toString funIdx)))
- "insert")
- #[.ref (.index (.var "inp") (.lit "..")),
- .ref (.index (.var "__ret") (.lit "..")),
- gFromBool (.lit "!unconstrained")])
+ let insertCall : RustStmt := .exprStmt (.lit <|
+ s!"if let Some(result) = record.function_queries[{funIdx}].get_mut(&inp[..]) \{" ++
+ " debug_assert_eq!(result.output, &__ret[..]);" ++
+ " if !unconstrained { *result.multiplicity += G::ONE; }" ++
+ " } else {" ++
+ s!" record.function_queries[{funIdx}].insert(&inp[..], &__ret[..], G::from_bool(!unconstrained));" ++
+ " }")
  -- Wrap in Ok(...) since fn now returns Result<[G; OUT_N], ExecError>.
  return #[outArr, insertCall,
  .returnStmt (.call (.var "Ok") #[.var "__ret"])]

Tests/Aiur/Aiur.lean

@@ -215,6 +215,25 @@ def toplevel := ⟦
  }
  }
 
+ -- Regression: a constrained call must recursively promote an identical
+ -- query previously cached by an unconstrained call. `promotion_nested`
+ -- calls another function so promoting only its own multiplicity leaves
+ -- the nested function channel unbalanced.
+ fn promotion_leaf(x: G) -> G {
+ x + 1
+ }
+
+ fn promotion_nested(x: G) -> G {
+ promotion_leaf(x)
+ }
+
+ pub fn unconstrained_call_promotion(x: G) -> G {
+ let hinted = #promotion_nested(x);
+ let constrained = promotion_nested(x);
+ assert_eq!(hinted, constrained);
+ constrained
+ }
+
  ---------------------------------------------------------------------------
  -- IO
  ---------------------------------------------------------------------------
@@ -794,6 +813,10 @@ def aiurTestCases : List AiurTestCase := [
  -- Unconstrained recursion: mixed constrained/unconstrained calls
  .prove `unconstrained_fibonacci #[6] #[13],
 
+ -- A constrained call reuses an identical unconstrained cached query and
+ -- must recursively promote the nested callee's lookup multiplicity.
+ .prove `unconstrained_call_promotion #[3] #[4],
+
  -- IO
  { functionName := `read_write_io
  inputIOBuffer :=

crates/aiur/src/execute.rs

@@ -327,13 +327,26 @@ impl Function {
  },
  ExecEntry::Op(Op::Call(callee_idx, args, _, op_unconstrained)) => {
  let args: Vec<G> = args.iter().map(|i| map[*i]).collect();
- if let Some(result) =
- record.function_queries[*callee_idx].get_mut(&args)
- {
- if !unconstrained && !op_unconstrained {
- *result.multiplicity += G::ONE;
- }
- map.extend_from_slice(result.output);
+ let callee_unconstrained = unconstrained || *op_unconstrained;
+ let cached_output = record.function_queries[*callee_idx]
+ .get_mut(&args)
+ .and_then(|result| {
+ // A zero-multiplicity entry was computed only as an
+ // unconstrained hint. Promoting just this row would omit all
+ // of the callee's child lookups and unbalance their channels;
+ // replay the body constrained so promotion recurses through
+ // the whole dependency tree.
+ if !callee_unconstrained && result.multiplicity.is_zero() {
+ None
+ } else {
+ if !callee_unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ Some(result.output.to_vec())
+ }
+ });
+ if let Some(output) = cached_output {
+ map.extend(output);
  } else {
  let saved_map = std::mem::replace(&mut map, args);
  callers_states_stack.push(CallerState {
@@ -343,7 +356,7 @@ impl Function {
  continuation_depth: continuation_stack.len(),
  });
  fun_idx = *callee_idx;
- unconstrained = unconstrained || *op_unconstrained;
+ unconstrained = callee_unconstrained;
  push_block_exec_entries!(&toplevel.functions[fun_idx].body);
  }
  },
@@ -699,11 +712,22 @@ impl Function {
  // Register the query.
  let input_size = toplevel.functions[fun_idx].layout.input_size;
  let output = output.iter().map(|i| map[*i]).collect::<Vec<_>>();
- record.function_queries[fun_idx].insert(
- &map[..input_size],
- &output,
- G::from_bool(!unconstrained),
- );
+ if let Some(result) =
+ record.function_queries[fun_idx].get_mut(&map[..input_size])
+ {
+ // The only ordinary way to execute an already cached function
+ // is constrained promotion of an unconstrained hint entry.
+ debug_assert_eq!(result.output, output);
+ if !unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ } else {
+ record.function_queries[fun_idx].insert(
+ &map[..input_size],
+ &output,
+ G::from_bool(!unconstrained),
+ );
+ }
  if let Some(CallerState {
  fun_idx: caller_idx,
  map: caller_map,

crates/aiur/src/synthesis.rs

@@ -562,6 +562,82 @@ mod tests {
  );
  }
 
+ /// Exercise promotion of a cached unconstrained call through nested calls.
+ ///
+ /// `f` first computes `g(x)` as an unconstrained hint, then calls `g(x)`
+ /// constrained. Since `g` calls `h`, promoting the cached `g` query must
+ /// replay its body and promote the cached `h` query as well. Merely bumping
+ /// `g`'s multiplicity leaves the `g -> h` function channel unbalanced.
+ fn unconstrained_call_promotion_toplevel() -> Toplevel {
+ let f = Function {
+ body: Block {
+ ops: vec![
+ Op::Call(1, vec![0], 1, true),
+ Op::Call(1, vec![0], 1, false),
+ ],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 3,
+ lookups: 2,
+ },
+ entry: true,
+ constrained: true,
+ };
+
+ let g = Function {
+ body: Block {
+ ops: vec![Op::Call(2, vec![0], 1, false)],
+ ctrl: Ctrl::Return(0, vec![1]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 2,
+ lookups: 2,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ let h = Function {
+ body: Block {
+ ops: vec![Op::Const(G::ONE), Op::Add(0, 1)],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 1,
+ lookups: 1,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ Toplevel { functions: vec![f, g, h], memory_sizes: vec![] }
+ }
+
+ #[test]
+ fn prove_verify_promotes_nested_unconstrained_call() {
+ let (cp, fp) = test_parameters();
+ let system =
+ AiurSystem::build(unconstrained_call_promotion_toplevel(), cp, fp);
+ let input = [G::from_u64(3)];
+ let mut io_buffer = empty_io_buffer();
+
+ let (claim, proof) = system.prove(0, &input, &mut io_buffer);
+ assert_eq!(
+ claim,
+ vec![function_channel(), G::ZERO, input[0], input[0] + G::ONE]
+ );
+ system
+ .verify(&claim, &proof)
+ .expect("nested constrained promotion must balance function channels");
+ }
+
  #[test]
  fn prove_verify_mul_roundtrip() {
  let (cp, fp) = test_parameters();

Ix/Aiur/Stages/Codegen.lean

@@ -398,8 +398,9 @@ private def emitCall (out : Nat) (callee : FunIdx) (args : Array ValIdx)
  s!"\{ let __args: [G; IN_{callee}] = {argsStr};" ++
  s!" let __cu = {cuExpr};" ++
  s!" if let Some(result) = record.function_queries[{callee}].get_mut(&__args[..]) \{" ++
- bumpStmt ++
- retExpr ++
+ s!" if !__cu && *result.multiplicity == G::ZERO \{" ++
+ s!" aiur_fn_{callee}(__args, record, io_buffer, false)?" ++
+ s!" } else \{" ++ bumpStmt ++ retExpr ++ " }" ++
  s!" } else \{ aiur_fn_{callee}(__args, record, io_buffer, __cu)? } }"
  let mut stmts : Array RustStmt := #[
  .letStmt false "__r_arr" (some s!"[G; OUT_{callee}]") (.lit blockExpr)
@@ -816,15 +817,13 @@ partial def emitCtrl (funIdx : FunIdx) (mcLabel? : Option String)
  let outArr : RustStmt :=
  .letStmt false "__ret" (some s!"[G; OUT_{funIdx}]")
  (.arrayLit (outs.map valVar))
- let insertCall : RustStmt :=
- .exprStmt (.call
- (.field
- (.index (.field (.var "record") "function_queries")
- (.lit (toString funIdx)))
- "insert")
- #[.ref (.index (.var "inp") (.lit "..")),
- .ref (.index (.var "__ret") (.lit "..")),
- gFromBool (.lit "!unconstrained")])
+ let insertCall : RustStmt := .exprStmt (.lit <|
+ s!"if let Some(result) = record.function_queries[{funIdx}].get_mut(&inp[..]) \{" ++
+ " debug_assert_eq!(result.output, &__ret[..]);" ++
+ " if !unconstrained { *result.multiplicity += G::ONE; }" ++
+ " } else {" ++
+ s!" record.function_queries[{funIdx}].insert(&inp[..], &__ret[..], G::from_bool(!unconstrained));" ++
+ " }")
  -- Wrap in Ok(...) since fn now returns Result<[G; OUT_N], ExecError>.
  return #[outArr, insertCall,
  .returnStmt (.call (.var "Ok") #[.var "__ret"])]

Tests/Aiur/Aiur.lean

@@ -215,6 +215,25 @@ def toplevel := ⟦
  }
  }
 
+ -- Regression: a constrained call must recursively promote an identical
+ -- query previously cached by an unconstrained call. `promotion_nested`
+ -- calls another function so promoting only its own multiplicity leaves
+ -- the nested function channel unbalanced.
+ fn promotion_leaf(x: G) -> G {
+ x + 1
+ }
+
+ fn promotion_nested(x: G) -> G {
+ promotion_leaf(x)
+ }
+
+ pub fn unconstrained_call_promotion(x: G) -> G {
+ let hinted = #promotion_nested(x);
+ let constrained = promotion_nested(x);
+ assert_eq!(hinted, constrained);
+ constrained
+ }
+
  ---------------------------------------------------------------------------
  -- IO
  ---------------------------------------------------------------------------
@@ -794,6 +813,10 @@ def aiurTestCases : List AiurTestCase := [
  -- Unconstrained recursion: mixed constrained/unconstrained calls
  .prove `unconstrained_fibonacci #[6] #[13],
 
+ -- A constrained call reuses an identical unconstrained cached query and
+ -- must recursively promote the nested callee's lookup multiplicity.
+ .prove `unconstrained_call_promotion #[3] #[4],
+
  -- IO
  { functionName := `read_write_io
  inputIOBuffer :=

crates/aiur/src/execute.rs

@@ -327,13 +327,26 @@ impl Function {
  },
  ExecEntry::Op(Op::Call(callee_idx, args, _, op_unconstrained)) => {
  let args: Vec<G> = args.iter().map(|i| map[*i]).collect();
- if let Some(result) =
- record.function_queries[*callee_idx].get_mut(&args)
- {
- if !unconstrained && !op_unconstrained {
- *result.multiplicity += G::ONE;
- }
- map.extend_from_slice(result.output);
+ let callee_unconstrained = unconstrained || *op_unconstrained;
+ let cached_output = record.function_queries[*callee_idx]
+ .get_mut(&args)
+ .and_then(|result| {
+ // A zero-multiplicity entry was computed only as an
+ // unconstrained hint. Promoting just this row would omit all
+ // of the callee's child lookups and unbalance their channels;
+ // replay the body constrained so promotion recurses through
+ // the whole dependency tree.
+ if !callee_unconstrained && result.multiplicity.is_zero() {
+ None
+ } else {
+ if !callee_unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ Some(result.output.to_vec())
+ }
+ });
+ if let Some(output) = cached_output {
+ map.extend(output);
  } else {
  let saved_map = std::mem::replace(&mut map, args);
  callers_states_stack.push(CallerState {
@@ -343,7 +356,7 @@ impl Function {
  continuation_depth: continuation_stack.len(),
  });
  fun_idx = *callee_idx;
- unconstrained = unconstrained || *op_unconstrained;
+ unconstrained = callee_unconstrained;
  push_block_exec_entries!(&toplevel.functions[fun_idx].body);
  }
  },
@@ -699,11 +712,22 @@ impl Function {
  // Register the query.
  let input_size = toplevel.functions[fun_idx].layout.input_size;
  let output = output.iter().map(|i| map[*i]).collect::<Vec<_>>();
- record.function_queries[fun_idx].insert(
- &map[..input_size],
- &output,
- G::from_bool(!unconstrained),
- );
+ if let Some(result) =
+ record.function_queries[fun_idx].get_mut(&map[..input_size])
+ {
+ // The only ordinary way to execute an already cached function
+ // is constrained promotion of an unconstrained hint entry.
+ debug_assert_eq!(result.output, output);
+ if !unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ } else {
+ record.function_queries[fun_idx].insert(
+ &map[..input_size],
+ &output,
+ G::from_bool(!unconstrained),
+ );
+ }
  if let Some(CallerState {
  fun_idx: caller_idx,
  map: caller_map,

crates/aiur/src/synthesis.rs

@@ -562,6 +562,82 @@ mod tests {
  );
  }
 
+ /// Exercise promotion of a cached unconstrained call through nested calls.
+ ///
+ /// `f` first computes `g(x)` as an unconstrained hint, then calls `g(x)`
+ /// constrained. Since `g` calls `h`, promoting the cached `g` query must
+ /// replay its body and promote the cached `h` query as well. Merely bumping
+ /// `g`'s multiplicity leaves the `g -> h` function channel unbalanced.
+ fn unconstrained_call_promotion_toplevel() -> Toplevel {
+ let f = Function {
+ body: Block {
+ ops: vec![
+ Op::Call(1, vec![0], 1, true),
+ Op::Call(1, vec![0], 1, false),
+ ],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 3,
+ lookups: 2,
+ },
+ entry: true,
+ constrained: true,
+ };
+
+ let g = Function {
+ body: Block {
+ ops: vec![Op::Call(2, vec![0], 1, false)],
+ ctrl: Ctrl::Return(0, vec![1]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 2,
+ lookups: 2,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ let h = Function {
+ body: Block {
+ ops: vec![Op::Const(G::ONE), Op::Add(0, 1)],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 1,
+ lookups: 1,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ Toplevel { functions: vec![f, g, h], memory_sizes: vec![] }
+ }
+
+ #[test]
+ fn prove_verify_promotes_nested_unconstrained_call() {
+ let (cp, fp) = test_parameters();
+ let system =
+ AiurSystem::build(unconstrained_call_promotion_toplevel(), cp, fp);
+ let input = [G::from_u64(3)];
+ let mut io_buffer = empty_io_buffer();
+
+ let (claim, proof) = system.prove(0, &input, &mut io_buffer);
+ assert_eq!(
+ claim,
+ vec![function_channel(), G::ZERO, input[0], input[0] + G::ONE]
+ );
+ system
+ .verify(&claim, &proof)
+ .expect("nested constrained promotion must balance function channels");
+ }
+
  #[test]
  fn prove_verify_mul_roundtrip() {
  let (cp, fp) = test_parameters();

Ix/Aiur/Stages/Codegen.lean

@@ -398,8 +398,9 @@ private def emitCall (out : Nat) (callee : FunIdx) (args : Array ValIdx)
  s!"\{ let __args: [G; IN_{callee}] = {argsStr};" ++
  s!" let __cu = {cuExpr};" ++
  s!" if let Some(result) = record.function_queries[{callee}].get_mut(&__args[..]) \{" ++
- bumpStmt ++
- retExpr ++
+ s!" if !__cu && *result.multiplicity == G::ZERO \{" ++
+ s!" aiur_fn_{callee}(__args, record, io_buffer, false)?" ++
+ s!" } else \{" ++ bumpStmt ++ retExpr ++ " }" ++
  s!" } else \{ aiur_fn_{callee}(__args, record, io_buffer, __cu)? } }"
  let mut stmts : Array RustStmt := #[
  .letStmt false "__r_arr" (some s!"[G; OUT_{callee}]") (.lit blockExpr)
@@ -816,15 +817,13 @@ partial def emitCtrl (funIdx : FunIdx) (mcLabel? : Option String)
  let outArr : RustStmt :=
  .letStmt false "__ret" (some s!"[G; OUT_{funIdx}]")
  (.arrayLit (outs.map valVar))
- let insertCall : RustStmt :=
- .exprStmt (.call
- (.field
- (.index (.field (.var "record") "function_queries")
- (.lit (toString funIdx)))
- "insert")
- #[.ref (.index (.var "inp") (.lit "..")),
- .ref (.index (.var "__ret") (.lit "..")),
- gFromBool (.lit "!unconstrained")])
+ let insertCall : RustStmt := .exprStmt (.lit <|
+ s!"if let Some(result) = record.function_queries[{funIdx}].get_mut(&inp[..]) \{" ++
+ " debug_assert_eq!(result.output, &__ret[..]);" ++
+ " if !unconstrained { *result.multiplicity += G::ONE; }" ++
+ " } else {" ++
+ s!" record.function_queries[{funIdx}].insert(&inp[..], &__ret[..], G::from_bool(!unconstrained));" ++
+ " }")
  -- Wrap in Ok(...) since fn now returns Result<[G; OUT_N], ExecError>.
  return #[outArr, insertCall,
  .returnStmt (.call (.var "Ok") #[.var "__ret"])]

Tests/Aiur/Aiur.lean

@@ -215,6 +215,25 @@ def toplevel := ⟦
  }
  }
 
+ -- Regression: a constrained call must recursively promote an identical
+ -- query previously cached by an unconstrained call. `promotion_nested`
+ -- calls another function so promoting only its own multiplicity leaves
+ -- the nested function channel unbalanced.
+ fn promotion_leaf(x: G) -> G {
+ x + 1
+ }
+
+ fn promotion_nested(x: G) -> G {
+ promotion_leaf(x)
+ }
+
+ pub fn unconstrained_call_promotion(x: G) -> G {
+ let hinted = #promotion_nested(x);
+ let constrained = promotion_nested(x);
+ assert_eq!(hinted, constrained);
+ constrained
+ }
+
  ---------------------------------------------------------------------------
  -- IO
  ---------------------------------------------------------------------------
@@ -794,6 +813,10 @@ def aiurTestCases : List AiurTestCase := [
  -- Unconstrained recursion: mixed constrained/unconstrained calls
  .prove `unconstrained_fibonacci #[6] #[13],
 
+ -- A constrained call reuses an identical unconstrained cached query and
+ -- must recursively promote the nested callee's lookup multiplicity.
+ .prove `unconstrained_call_promotion #[3] #[4],
+
  -- IO
  { functionName := `read_write_io
  inputIOBuffer :=

crates/aiur/src/execute.rs

@@ -327,13 +327,26 @@ impl Function {
  },
  ExecEntry::Op(Op::Call(callee_idx, args, _, op_unconstrained)) => {
  let args: Vec<G> = args.iter().map(|i| map[*i]).collect();
- if let Some(result) =
- record.function_queries[*callee_idx].get_mut(&args)
- {
- if !unconstrained && !op_unconstrained {
- *result.multiplicity += G::ONE;
- }
- map.extend_from_slice(result.output);
+ let callee_unconstrained = unconstrained || *op_unconstrained;
+ let cached_output = record.function_queries[*callee_idx]
+ .get_mut(&args)
+ .and_then(|result| {
+ // A zero-multiplicity entry was computed only as an
+ // unconstrained hint. Promoting just this row would omit all
+ // of the callee's child lookups and unbalance their channels;
+ // replay the body constrained so promotion recurses through
+ // the whole dependency tree.
+ if !callee_unconstrained && result.multiplicity.is_zero() {
+ None
+ } else {
+ if !callee_unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ Some(result.output.to_vec())
+ }
+ });
+ if let Some(output) = cached_output {
+ map.extend(output);
  } else {
  let saved_map = std::mem::replace(&mut map, args);
  callers_states_stack.push(CallerState {
@@ -343,7 +356,7 @@ impl Function {
  continuation_depth: continuation_stack.len(),
  });
  fun_idx = *callee_idx;
- unconstrained = unconstrained || *op_unconstrained;
+ unconstrained = callee_unconstrained;
  push_block_exec_entries!(&toplevel.functions[fun_idx].body);
  }
  },
@@ -699,11 +712,22 @@ impl Function {
  // Register the query.
  let input_size = toplevel.functions[fun_idx].layout.input_size;
  let output = output.iter().map(|i| map[*i]).collect::<Vec<_>>();
- record.function_queries[fun_idx].insert(
- &map[..input_size],
- &output,
- G::from_bool(!unconstrained),
- );
+ if let Some(result) =
+ record.function_queries[fun_idx].get_mut(&map[..input_size])
+ {
+ // The only ordinary way to execute an already cached function
+ // is constrained promotion of an unconstrained hint entry.
+ debug_assert_eq!(result.output, output);
+ if !unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ } else {
+ record.function_queries[fun_idx].insert(
+ &map[..input_size],
+ &output,
+ G::from_bool(!unconstrained),
+ );
+ }
  if let Some(CallerState {
  fun_idx: caller_idx,
  map: caller_map,

crates/aiur/src/synthesis.rs

@@ -562,6 +562,82 @@ mod tests {
  );
  }
 
+ /// Exercise promotion of a cached unconstrained call through nested calls.
+ ///
+ /// `f` first computes `g(x)` as an unconstrained hint, then calls `g(x)`
+ /// constrained. Since `g` calls `h`, promoting the cached `g` query must
+ /// replay its body and promote the cached `h` query as well. Merely bumping
+ /// `g`'s multiplicity leaves the `g -> h` function channel unbalanced.
+ fn unconstrained_call_promotion_toplevel() -> Toplevel {
+ let f = Function {
+ body: Block {
+ ops: vec![
+ Op::Call(1, vec![0], 1, true),
+ Op::Call(1, vec![0], 1, false),
+ ],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 3,
+ lookups: 2,
+ },
+ entry: true,
+ constrained: true,
+ };
+
+ let g = Function {
+ body: Block {
+ ops: vec![Op::Call(2, vec![0], 1, false)],
+ ctrl: Ctrl::Return(0, vec![1]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 2,
+ lookups: 2,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ let h = Function {
+ body: Block {
+ ops: vec![Op::Const(G::ONE), Op::Add(0, 1)],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 1,
+ lookups: 1,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ Toplevel { functions: vec![f, g, h], memory_sizes: vec![] }
+ }
+
+ #[test]
+ fn prove_verify_promotes_nested_unconstrained_call() {
+ let (cp, fp) = test_parameters();
+ let system =
+ AiurSystem::build(unconstrained_call_promotion_toplevel(), cp, fp);
+ let input = [G::from_u64(3)];
+ let mut io_buffer = empty_io_buffer();
+
+ let (claim, proof) = system.prove(0, &input, &mut io_buffer);
+ assert_eq!(
+ claim,
+ vec![function_channel(), G::ZERO, input[0], input[0] + G::ONE]
+ );
+ system
+ .verify(&claim, &proof)
+ .expect("nested constrained promotion must balance function channels");
+ }
+
  #[test]
  fn prove_verify_mul_roundtrip() {
  let (cp, fp) = test_parameters();

Ix/Aiur/Stages/Codegen.lean

@@ -398,8 +398,9 @@ private def emitCall (out : Nat) (callee : FunIdx) (args : Array ValIdx)
  s!"\{ let __args: [G; IN_{callee}] = {argsStr};" ++
  s!" let __cu = {cuExpr};" ++
  s!" if let Some(result) = record.function_queries[{callee}].get_mut(&__args[..]) \{" ++
- bumpStmt ++
- retExpr ++
+ s!" if !__cu && *result.multiplicity == G::ZERO \{" ++
+ s!" aiur_fn_{callee}(__args, record, io_buffer, false)?" ++
+ s!" } else \{" ++ bumpStmt ++ retExpr ++ " }" ++
  s!" } else \{ aiur_fn_{callee}(__args, record, io_buffer, __cu)? } }"
  let mut stmts : Array RustStmt := #[
  .letStmt false "__r_arr" (some s!"[G; OUT_{callee}]") (.lit blockExpr)
@@ -816,15 +817,13 @@ partial def emitCtrl (funIdx : FunIdx) (mcLabel? : Option String)
  let outArr : RustStmt :=
  .letStmt false "__ret" (some s!"[G; OUT_{funIdx}]")
  (.arrayLit (outs.map valVar))
- let insertCall : RustStmt :=
- .exprStmt (.call
- (.field
- (.index (.field (.var "record") "function_queries")
- (.lit (toString funIdx)))
- "insert")
- #[.ref (.index (.var "inp") (.lit "..")),
- .ref (.index (.var "__ret") (.lit "..")),
- gFromBool (.lit "!unconstrained")])
+ let insertCall : RustStmt := .exprStmt (.lit <|
+ s!"if let Some(result) = record.function_queries[{funIdx}].get_mut(&inp[..]) \{" ++
+ " debug_assert_eq!(result.output, &__ret[..]);" ++
+ " if !unconstrained { *result.multiplicity += G::ONE; }" ++
+ " } else {" ++
+ s!" record.function_queries[{funIdx}].insert(&inp[..], &__ret[..], G::from_bool(!unconstrained));" ++
+ " }")
  -- Wrap in Ok(...) since fn now returns Result<[G; OUT_N], ExecError>.
  return #[outArr, insertCall,
  .returnStmt (.call (.var "Ok") #[.var "__ret"])]

Tests/Aiur/Aiur.lean

@@ -215,6 +215,25 @@ def toplevel := ⟦
  }
  }
 
+ -- Regression: a constrained call must recursively promote an identical
+ -- query previously cached by an unconstrained call. `promotion_nested`
+ -- calls another function so promoting only its own multiplicity leaves
+ -- the nested function channel unbalanced.
+ fn promotion_leaf(x: G) -> G {
+ x + 1
+ }
+
+ fn promotion_nested(x: G) -> G {
+ promotion_leaf(x)
+ }
+
+ pub fn unconstrained_call_promotion(x: G) -> G {
+ let hinted = #promotion_nested(x);
+ let constrained = promotion_nested(x);
+ assert_eq!(hinted, constrained);
+ constrained
+ }
+
  ---------------------------------------------------------------------------
  -- IO
  ---------------------------------------------------------------------------
@@ -794,6 +813,10 @@ def aiurTestCases : List AiurTestCase := [
  -- Unconstrained recursion: mixed constrained/unconstrained calls
  .prove `unconstrained_fibonacci #[6] #[13],
 
+ -- A constrained call reuses an identical unconstrained cached query and
+ -- must recursively promote the nested callee's lookup multiplicity.
+ .prove `unconstrained_call_promotion #[3] #[4],
+
  -- IO
  { functionName := `read_write_io
  inputIOBuffer :=

crates/aiur/src/execute.rs

@@ -327,13 +327,26 @@ impl Function {
  },
  ExecEntry::Op(Op::Call(callee_idx, args, _, op_unconstrained)) => {
  let args: Vec<G> = args.iter().map(|i| map[*i]).collect();
- if let Some(result) =
- record.function_queries[*callee_idx].get_mut(&args)
- {
- if !unconstrained && !op_unconstrained {
- *result.multiplicity += G::ONE;
- }
- map.extend_from_slice(result.output);
+ let callee_unconstrained = unconstrained || *op_unconstrained;
+ let cached_output = record.function_queries[*callee_idx]
+ .get_mut(&args)
+ .and_then(|result| {
+ // A zero-multiplicity entry was computed only as an
+ // unconstrained hint. Promoting just this row would omit all
+ // of the callee's child lookups and unbalance their channels;
+ // replay the body constrained so promotion recurses through
+ // the whole dependency tree.
+ if !callee_unconstrained && result.multiplicity.is_zero() {
+ None
+ } else {
+ if !callee_unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ Some(result.output.to_vec())
+ }
+ });
+ if let Some(output) = cached_output {
+ map.extend(output);
  } else {
  let saved_map = std::mem::replace(&mut map, args);
  callers_states_stack.push(CallerState {
@@ -343,7 +356,7 @@ impl Function {
  continuation_depth: continuation_stack.len(),
  });
  fun_idx = *callee_idx;
- unconstrained = unconstrained || *op_unconstrained;
+ unconstrained = callee_unconstrained;
  push_block_exec_entries!(&toplevel.functions[fun_idx].body);
  }
  },
@@ -699,11 +712,22 @@ impl Function {
  // Register the query.
  let input_size = toplevel.functions[fun_idx].layout.input_size;
  let output = output.iter().map(|i| map[*i]).collect::<Vec<_>>();
- record.function_queries[fun_idx].insert(
- &map[..input_size],
- &output,
- G::from_bool(!unconstrained),
- );
+ if let Some(result) =
+ record.function_queries[fun_idx].get_mut(&map[..input_size])
+ {
+ // The only ordinary way to execute an already cached function
+ // is constrained promotion of an unconstrained hint entry.
+ debug_assert_eq!(result.output, output);
+ if !unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ } else {
+ record.function_queries[fun_idx].insert(
+ &map[..input_size],
+ &output,
+ G::from_bool(!unconstrained),
+ );
+ }
  if let Some(CallerState {
  fun_idx: caller_idx,
  map: caller_map,

crates/aiur/src/synthesis.rs

@@ -562,6 +562,82 @@ mod tests {
  );
  }
 
+ /// Exercise promotion of a cached unconstrained call through nested calls.
+ ///
+ /// `f` first computes `g(x)` as an unconstrained hint, then calls `g(x)`
+ /// constrained. Since `g` calls `h`, promoting the cached `g` query must
+ /// replay its body and promote the cached `h` query as well. Merely bumping
+ /// `g`'s multiplicity leaves the `g -> h` function channel unbalanced.
+ fn unconstrained_call_promotion_toplevel() -> Toplevel {
+ let f = Function {
+ body: Block {
+ ops: vec![
+ Op::Call(1, vec![0], 1, true),
+ Op::Call(1, vec![0], 1, false),
+ ],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 3,
+ lookups: 2,
+ },
+ entry: true,
+ constrained: true,
+ };
+
+ let g = Function {
+ body: Block {
+ ops: vec![Op::Call(2, vec![0], 1, false)],
+ ctrl: Ctrl::Return(0, vec![1]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 2,
+ lookups: 2,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ let h = Function {
+ body: Block {
+ ops: vec![Op::Const(G::ONE), Op::Add(0, 1)],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 1,
+ lookups: 1,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ Toplevel { functions: vec![f, g, h], memory_sizes: vec![] }
+ }
+
+ #[test]
+ fn prove_verify_promotes_nested_unconstrained_call() {
+ let (cp, fp) = test_parameters();
+ let system =
+ AiurSystem::build(unconstrained_call_promotion_toplevel(), cp, fp);
+ let input = [G::from_u64(3)];
+ let mut io_buffer = empty_io_buffer();
+
+ let (claim, proof) = system.prove(0, &input, &mut io_buffer);
+ assert_eq!(
+ claim,
+ vec![function_channel(), G::ZERO, input[0], input[0] + G::ONE]
+ );
+ system
+ .verify(&claim, &proof)
+ .expect("nested constrained promotion must balance function channels");
+ }
+
  #[test]
  fn prove_verify_mul_roundtrip() {
  let (cp, fp) = test_parameters();

Ix/Aiur/Stages/Codegen.lean

@@ -398,8 +398,9 @@ private def emitCall (out : Nat) (callee : FunIdx) (args : Array ValIdx)
  s!"\{ let __args: [G; IN_{callee}] = {argsStr};" ++
  s!" let __cu = {cuExpr};" ++
  s!" if let Some(result) = record.function_queries[{callee}].get_mut(&__args[..]) \{" ++
- bumpStmt ++
- retExpr ++
+ s!" if !__cu && *result.multiplicity == G::ZERO \{" ++
+ s!" aiur_fn_{callee}(__args, record, io_buffer, false)?" ++
+ s!" } else \{" ++ bumpStmt ++ retExpr ++ " }" ++
  s!" } else \{ aiur_fn_{callee}(__args, record, io_buffer, __cu)? } }"
  let mut stmts : Array RustStmt := #[
  .letStmt false "__r_arr" (some s!"[G; OUT_{callee}]") (.lit blockExpr)
@@ -816,15 +817,13 @@ partial def emitCtrl (funIdx : FunIdx) (mcLabel? : Option String)
  let outArr : RustStmt :=
  .letStmt false "__ret" (some s!"[G; OUT_{funIdx}]")
  (.arrayLit (outs.map valVar))
- let insertCall : RustStmt :=
- .exprStmt (.call
- (.field
- (.index (.field (.var "record") "function_queries")
- (.lit (toString funIdx)))
- "insert")
- #[.ref (.index (.var "inp") (.lit "..")),
- .ref (.index (.var "__ret") (.lit "..")),
- gFromBool (.lit "!unconstrained")])
+ let insertCall : RustStmt := .exprStmt (.lit <|
+ s!"if let Some(result) = record.function_queries[{funIdx}].get_mut(&inp[..]) \{" ++
+ " debug_assert_eq!(result.output, &__ret[..]);" ++
+ " if !unconstrained { *result.multiplicity += G::ONE; }" ++
+ " } else {" ++
+ s!" record.function_queries[{funIdx}].insert(&inp[..], &__ret[..], G::from_bool(!unconstrained));" ++
+ " }")
  -- Wrap in Ok(...) since fn now returns Result<[G; OUT_N], ExecError>.
  return #[outArr, insertCall,
  .returnStmt (.call (.var "Ok") #[.var "__ret"])]

Tests/Aiur/Aiur.lean

@@ -215,6 +215,25 @@ def toplevel := ⟦
  }
  }
 
+ -- Regression: a constrained call must recursively promote an identical
+ -- query previously cached by an unconstrained call. `promotion_nested`
+ -- calls another function so promoting only its own multiplicity leaves
+ -- the nested function channel unbalanced.
+ fn promotion_leaf(x: G) -> G {
+ x + 1
+ }
+
+ fn promotion_nested(x: G) -> G {
+ promotion_leaf(x)
+ }
+
+ pub fn unconstrained_call_promotion(x: G) -> G {
+ let hinted = #promotion_nested(x);
+ let constrained = promotion_nested(x);
+ assert_eq!(hinted, constrained);
+ constrained
+ }
+
  ---------------------------------------------------------------------------
  -- IO
  ---------------------------------------------------------------------------
@@ -794,6 +813,10 @@ def aiurTestCases : List AiurTestCase := [
  -- Unconstrained recursion: mixed constrained/unconstrained calls
  .prove `unconstrained_fibonacci #[6] #[13],
 
+ -- A constrained call reuses an identical unconstrained cached query and
+ -- must recursively promote the nested callee's lookup multiplicity.
+ .prove `unconstrained_call_promotion #[3] #[4],
+
  -- IO
  { functionName := `read_write_io
  inputIOBuffer :=

crates/aiur/src/execute.rs

@@ -327,13 +327,26 @@ impl Function {
  },
  ExecEntry::Op(Op::Call(callee_idx, args, _, op_unconstrained)) => {
  let args: Vec<G> = args.iter().map(|i| map[*i]).collect();
- if let Some(result) =
- record.function_queries[*callee_idx].get_mut(&args)
- {
- if !unconstrained && !op_unconstrained {
- *result.multiplicity += G::ONE;
- }
- map.extend_from_slice(result.output);
+ let callee_unconstrained = unconstrained || *op_unconstrained;
+ let cached_output = record.function_queries[*callee_idx]
+ .get_mut(&args)
+ .and_then(|result| {
+ // A zero-multiplicity entry was computed only as an
+ // unconstrained hint. Promoting just this row would omit all
+ // of the callee's child lookups and unbalance their channels;
+ // replay the body constrained so promotion recurses through
+ // the whole dependency tree.
+ if !callee_unconstrained && result.multiplicity.is_zero() {
+ None
+ } else {
+ if !callee_unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ Some(result.output.to_vec())
+ }
+ });
+ if let Some(output) = cached_output {
+ map.extend(output);
  } else {
  let saved_map = std::mem::replace(&mut map, args);
  callers_states_stack.push(CallerState {
@@ -343,7 +356,7 @@ impl Function {
  continuation_depth: continuation_stack.len(),
  });
  fun_idx = *callee_idx;
- unconstrained = unconstrained || *op_unconstrained;
+ unconstrained = callee_unconstrained;
  push_block_exec_entries!(&toplevel.functions[fun_idx].body);
  }
  },
@@ -699,11 +712,22 @@ impl Function {
  // Register the query.
  let input_size = toplevel.functions[fun_idx].layout.input_size;
  let output = output.iter().map(|i| map[*i]).collect::<Vec<_>>();
- record.function_queries[fun_idx].insert(
- &map[..input_size],
- &output,
- G::from_bool(!unconstrained),
- );
+ if let Some(result) =
+ record.function_queries[fun_idx].get_mut(&map[..input_size])
+ {
+ // The only ordinary way to execute an already cached function
+ // is constrained promotion of an unconstrained hint entry.
+ debug_assert_eq!(result.output, output);
+ if !unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ } else {
+ record.function_queries[fun_idx].insert(
+ &map[..input_size],
+ &output,
+ G::from_bool(!unconstrained),
+ );
+ }
  if let Some(CallerState {
  fun_idx: caller_idx,
  map: caller_map,

crates/aiur/src/synthesis.rs

@@ -562,6 +562,82 @@ mod tests {
  );
  }
 
+ /// Exercise promotion of a cached unconstrained call through nested calls.
+ ///
+ /// `f` first computes `g(x)` as an unconstrained hint, then calls `g(x)`
+ /// constrained. Since `g` calls `h`, promoting the cached `g` query must
+ /// replay its body and promote the cached `h` query as well. Merely bumping
+ /// `g`'s multiplicity leaves the `g -> h` function channel unbalanced.
+ fn unconstrained_call_promotion_toplevel() -> Toplevel {
+ let f = Function {
+ body: Block {
+ ops: vec![
+ Op::Call(1, vec![0], 1, true),
+ Op::Call(1, vec![0], 1, false),
+ ],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 3,
+ lookups: 2,
+ },
+ entry: true,
+ constrained: true,
+ };
+
+ let g = Function {
+ body: Block {
+ ops: vec![Op::Call(2, vec![0], 1, false)],
+ ctrl: Ctrl::Return(0, vec![1]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 2,
+ lookups: 2,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ let h = Function {
+ body: Block {
+ ops: vec![Op::Const(G::ONE), Op::Add(0, 1)],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 1,
+ lookups: 1,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ Toplevel { functions: vec![f, g, h], memory_sizes: vec![] }
+ }
+
+ #[test]
+ fn prove_verify_promotes_nested_unconstrained_call() {
+ let (cp, fp) = test_parameters();
+ let system =
+ AiurSystem::build(unconstrained_call_promotion_toplevel(), cp, fp);
+ let input = [G::from_u64(3)];
+ let mut io_buffer = empty_io_buffer();
+
+ let (claim, proof) = system.prove(0, &input, &mut io_buffer);
+ assert_eq!(
+ claim,
+ vec![function_channel(), G::ZERO, input[0], input[0] + G::ONE]
+ );
+ system
+ .verify(&claim, &proof)
+ .expect("nested constrained promotion must balance function channels");
+ }
+
  #[test]
  fn prove_verify_mul_roundtrip() {
  let (cp, fp) = test_parameters();

Ix/Aiur/Stages/Codegen.lean

@@ -398,8 +398,9 @@ private def emitCall (out : Nat) (callee : FunIdx) (args : Array ValIdx)
  s!"\{ let __args: [G; IN_{callee}] = {argsStr};" ++
  s!" let __cu = {cuExpr};" ++
  s!" if let Some(result) = record.function_queries[{callee}].get_mut(&__args[..]) \{" ++
- bumpStmt ++
- retExpr ++
+ s!" if !__cu && *result.multiplicity == G::ZERO \{" ++
+ s!" aiur_fn_{callee}(__args, record, io_buffer, false)?" ++
+ s!" } else \{" ++ bumpStmt ++ retExpr ++ " }" ++
  s!" } else \{ aiur_fn_{callee}(__args, record, io_buffer, __cu)? } }"
  let mut stmts : Array RustStmt := #[
  .letStmt false "__r_arr" (some s!"[G; OUT_{callee}]") (.lit blockExpr)
@@ -816,15 +817,13 @@ partial def emitCtrl (funIdx : FunIdx) (mcLabel? : Option String)
  let outArr : RustStmt :=
  .letStmt false "__ret" (some s!"[G; OUT_{funIdx}]")
  (.arrayLit (outs.map valVar))
- let insertCall : RustStmt :=
- .exprStmt (.call
- (.field
- (.index (.field (.var "record") "function_queries")
- (.lit (toString funIdx)))
- "insert")
- #[.ref (.index (.var "inp") (.lit "..")),
- .ref (.index (.var "__ret") (.lit "..")),
- gFromBool (.lit "!unconstrained")])
+ let insertCall : RustStmt := .exprStmt (.lit <|
+ s!"if let Some(result) = record.function_queries[{funIdx}].get_mut(&inp[..]) \{" ++
+ " debug_assert_eq!(result.output, &__ret[..]);" ++
+ " if !unconstrained { *result.multiplicity += G::ONE; }" ++
+ " } else {" ++
+ s!" record.function_queries[{funIdx}].insert(&inp[..], &__ret[..], G::from_bool(!unconstrained));" ++
+ " }")
  -- Wrap in Ok(...) since fn now returns Result<[G; OUT_N], ExecError>.
  return #[outArr, insertCall,
  .returnStmt (.call (.var "Ok") #[.var "__ret"])]

Tests/Aiur/Aiur.lean

@@ -215,6 +215,25 @@ def toplevel := ⟦
  }
  }
 
+ -- Regression: a constrained call must recursively promote an identical
+ -- query previously cached by an unconstrained call. `promotion_nested`
+ -- calls another function so promoting only its own multiplicity leaves
+ -- the nested function channel unbalanced.
+ fn promotion_leaf(x: G) -> G {
+ x + 1
+ }
+
+ fn promotion_nested(x: G) -> G {
+ promotion_leaf(x)
+ }
+
+ pub fn unconstrained_call_promotion(x: G) -> G {
+ let hinted = #promotion_nested(x);
+ let constrained = promotion_nested(x);
+ assert_eq!(hinted, constrained);
+ constrained
+ }
+
  ---------------------------------------------------------------------------
  -- IO
  ---------------------------------------------------------------------------
@@ -794,6 +813,10 @@ def aiurTestCases : List AiurTestCase := [
  -- Unconstrained recursion: mixed constrained/unconstrained calls
  .prove `unconstrained_fibonacci #[6] #[13],
 
+ -- A constrained call reuses an identical unconstrained cached query and
+ -- must recursively promote the nested callee's lookup multiplicity.
+ .prove `unconstrained_call_promotion #[3] #[4],
+
  -- IO
  { functionName := `read_write_io
  inputIOBuffer :=

crates/aiur/src/execute.rs

@@ -327,13 +327,26 @@ impl Function {
  },
  ExecEntry::Op(Op::Call(callee_idx, args, _, op_unconstrained)) => {
  let args: Vec<G> = args.iter().map(|i| map[*i]).collect();
- if let Some(result) =
- record.function_queries[*callee_idx].get_mut(&args)
- {
- if !unconstrained && !op_unconstrained {
- *result.multiplicity += G::ONE;
- }
- map.extend_from_slice(result.output);
+ let callee_unconstrained = unconstrained || *op_unconstrained;
+ let cached_output = record.function_queries[*callee_idx]
+ .get_mut(&args)
+ .and_then(|result| {
+ // A zero-multiplicity entry was computed only as an
+ // unconstrained hint. Promoting just this row would omit all
+ // of the callee's child lookups and unbalance their channels;
+ // replay the body constrained so promotion recurses through
+ // the whole dependency tree.
+ if !callee_unconstrained && result.multiplicity.is_zero() {
+ None
+ } else {
+ if !callee_unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ Some(result.output.to_vec())
+ }
+ });
+ if let Some(output) = cached_output {
+ map.extend(output);
  } else {
  let saved_map = std::mem::replace(&mut map, args);
  callers_states_stack.push(CallerState {
@@ -343,7 +356,7 @@ impl Function {
  continuation_depth: continuation_stack.len(),
  });
  fun_idx = *callee_idx;
- unconstrained = unconstrained || *op_unconstrained;
+ unconstrained = callee_unconstrained;
  push_block_exec_entries!(&toplevel.functions[fun_idx].body);
  }
  },
@@ -699,11 +712,22 @@ impl Function {
  // Register the query.
  let input_size = toplevel.functions[fun_idx].layout.input_size;
  let output = output.iter().map(|i| map[*i]).collect::<Vec<_>>();
- record.function_queries[fun_idx].insert(
- &map[..input_size],
- &output,
- G::from_bool(!unconstrained),
- );
+ if let Some(result) =
+ record.function_queries[fun_idx].get_mut(&map[..input_size])
+ {
+ // The only ordinary way to execute an already cached function
+ // is constrained promotion of an unconstrained hint entry.
+ debug_assert_eq!(result.output, output);
+ if !unconstrained {
+ *result.multiplicity += G::ONE;
+ }
+ } else {
+ record.function_queries[fun_idx].insert(
+ &map[..input_size],
+ &output,
+ G::from_bool(!unconstrained),
+ );
+ }
  if let Some(CallerState {
  fun_idx: caller_idx,
  map: caller_map,

crates/aiur/src/synthesis.rs

@@ -562,6 +562,82 @@ mod tests {
  );
  }
 
+ /// Exercise promotion of a cached unconstrained call through nested calls.
+ ///
+ /// `f` first computes `g(x)` as an unconstrained hint, then calls `g(x)`
+ /// constrained. Since `g` calls `h`, promoting the cached `g` query must
+ /// replay its body and promote the cached `h` query as well. Merely bumping
+ /// `g`'s multiplicity leaves the `g -> h` function channel unbalanced.
+ fn unconstrained_call_promotion_toplevel() -> Toplevel {
+ let f = Function {
+ body: Block {
+ ops: vec![
+ Op::Call(1, vec![0], 1, true),
+ Op::Call(1, vec![0], 1, false),
+ ],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 3,
+ lookups: 2,
+ },
+ entry: true,
+ constrained: true,
+ };
+
+ let g = Function {
+ body: Block {
+ ops: vec![Op::Call(2, vec![0], 1, false)],
+ ctrl: Ctrl::Return(0, vec![1]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 2,
+ lookups: 2,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ let h = Function {
+ body: Block {
+ ops: vec![Op::Const(G::ONE), Op::Add(0, 1)],
+ ctrl: Ctrl::Return(0, vec![2]),
+ },
+ layout: FunctionLayout {
+ input_size: 1,
+ selectors: 1,
+ auxiliaries: 1,
+ lookups: 1,
+ },
+ entry: false,
+ constrained: true,
+ };
+
+ Toplevel { functions: vec![f, g, h], memory_sizes: vec![] }
+ }
+
+ #[test]
+ fn prove_verify_promotes_nested_unconstrained_call() {
+ let (cp, fp) = test_parameters();
+ let system =
+ AiurSystem::build(unconstrained_call_promotion_toplevel(), cp, fp);
+ let input = [G::from_u64(3)];
+ let mut io_buffer = empty_io_buffer();
+
+ let (claim, proof) = system.prove(0, &input, &mut io_buffer);
+ assert_eq!(
+ claim,
+ vec![function_channel(), G::ZERO, input[0], input[0] + G::ONE]
+ );
+ system
+ .verify(&claim, &proof)
+ .expect("nested constrained promotion must balance function channels");
+ }
+
  #[test]
  fn prove_verify_mul_roundtrip() {
  let (cp, fp) = test_parameters();