Merge pull request #1638 from o1-labs/fix/dynamic-call

Dynamic call unit test

src/lib/mina/test/dynamic-call.unit-test.ts

@@ -0,0 +1,92 @@
+/**
+ * Tests that shows we can call a subcontract dynamically based on the address,
+ * as long as its signature matches the signature our contract was compiled against.
+ *
+ * In other words, the exact implementation/constraints of zkApp methods we call are not hard-coded in the caller contract.
+ */
+import {
+  Bool,
+  UInt64,
+  SmartContract,
+  method,
+  PublicKey,
+  Mina,
+} from '../../../index.js';
+
+type Subcontract = SmartContract & {
+  submethod(a: UInt64, b: UInt64): Promise<Bool>;
+};
+
+// two implementations with same signature of the called method, but different provable logic
+
+class SubcontractA extends SmartContract implements Subcontract {
+  @method.returns(Bool)
+  async submethod(a: UInt64, b: UInt64): Promise<Bool> {
+    return a.greaterThan(b);
+  }
+}
+
+class SubcontractB extends SmartContract implements Subcontract {
+  @method.returns(Bool)
+  async submethod(a: UInt64, b: UInt64): Promise<Bool> {
+    return a.mul(b).equals(UInt64.from(42));
+  }
+}
+
+// caller contract that calls the subcontract
+
+class Caller extends SmartContract {
+  @method
+  async call(a: UInt64, b: UInt64, address: PublicKey) {
+    const subcontract = new Caller.Subcontract(address);
+    await subcontract.submethod(a, b);
+  }
+
+  // subcontract to call. this property is changed below
+  // TODO: having to set this property is a hack, it would be nice to pass the contract as parameter
+  static Subcontract: new (...args: any) => Subcontract = SubcontractA;
+}
+
+// TEST BELOW
+
+// setup
+
+let Local = await Mina.LocalBlockchain({ proofsEnabled: true });
+Mina.setActiveInstance(Local);
+
+let [sender, callerAccount, aAccount, bAccount] = Local.testAccounts;
+
+await SubcontractA.compile();
+await SubcontractB.compile();
+await Caller.compile();
+
+let caller = new Caller(callerAccount);
+let a = new SubcontractA(aAccount);
+let b = new SubcontractB(bAccount);
+
+await Mina.transaction(sender, async () => {
+  await caller.deploy();
+  await a.deploy();
+  await b.deploy();
+})
+  .sign([callerAccount.key, aAccount.key, bAccount.key, sender.key])
+  .send();
+
+// subcontract A call
+
+let x = UInt64.from(10);
+let y = UInt64.from(5);
+
+await Mina.transaction(sender, () => caller.call(x, y, aAccount))
+  .prove()
+  .sign([sender.key])
+  .send();
+
+// subcontract B call
+
+Caller.Subcontract = SubcontractB;
+
+await Mina.transaction(sender, () => caller.call(x, y, bAccount))
+  .prove()
+  .sign([sender.key])
+  .send();

src/lib/mina/zkapp.ts

@@ -1612,13 +1612,23 @@ function diffRecursive(
   let { transaction, index, accountUpdate: input } = inputData;
   diff(transaction, index, prover.toPretty(), input.toPretty());
   // TODO
-  let inputChildren = accountUpdateLayout()!.get(input)!.children.mutable!;
-  let proverChildren = accountUpdateLayout()!.get(prover)!.children.mutable!;
+  let proverChildren = accountUpdateLayout()?.get(prover)?.children.mutable;
+  if (proverChildren === undefined) return;
+
+  // collect input children
+  let inputChildren: AccountUpdate[] = [];
+  let callDepth = input.body.callDepth;
+  for (let i = index; i < transaction.accountUpdates.length; i++) {
+    let update = transaction.accountUpdates[i];
+    if (update.body.callDepth <= callDepth) break;
+    if (update.body.callDepth === callDepth + 1) inputChildren.push(update);
+  }
+
   let nChildren = inputChildren.length;
   for (let i = 0; i < nChildren; i++) {
-    let inputChild = inputChildren[i].mutable;
+    let inputChild = inputChildren[i];
     let child = proverChildren[i].mutable;
-    if (!inputChild || !child) return;
+    if (!child) return;
     diffRecursive(child, { transaction, index, accountUpdate: inputChild });
   }
 }