IndexedMerkleMap does not update the root inside ZkProgram method

[dfstio]

@mitschabaude

According to https://github.com/o1-labs/o1js/blob/main/src/lib/provable/merkle-tree-indexed.ts#L41, we can pass MerkleMap directly to the ZkProgram method. Doing so works correctly when used with rawMethods (without proof calculation) but does not update the root and length of the Indexed MerkleMap when used with proofs.

The code to reproduce the issue:

import { describe, expect, it } from "@jest/globals";
import { Experimental, Field, ZkProgram, Cache } from "o1js";

const { IndexedMerkleMap } = Experimental;
const height = 3;
class MerkleMap extends IndexedMerkleMap(height) {}

const MapProgram = ZkProgram({
  name: "MapProgram",
  publicInput: Field,
  publicOutput: Field,
  methods: {
    insert: {
      privateInputs: [MerkleMap, Field, Field],

      async method(oldRoot: Field, map: MerkleMap, key: Field, value: Field) {
        map.root.assertEquals(oldRoot);
        map.insert(key, value);
        return map.root;
      },
    },
  },
});

describe("Map", () => {
  it(`should build the Indexed Merkle Map without proofs`, async () => {
    const map = new MerkleMap();
    const root = map.root;
    const step1 = await MapProgram.rawMethods.insert(
      root,
      map,
      Field(1),
      Field(2)
    );
    expect(step1.toBigInt()).toBe(map.root.toBigInt());
  });
  it(`should build the Indexed Merkle Map with proofs`, async () => {
    const map = new MerkleMap();
    const root = map.root;
    await MapProgram.compile({ cache: Cache.FileSystem("./cache") });
    console.log("map before:", {
      root: map.root.toJSON(),
      length: map.length.toJSON(),
      nodes: map.data.get().nodes,
      sortedLeaves: map.data.get().sortedLeaves,
    });
    const step1 = await MapProgram.insert(root, map, Field(1), Field(2));
    console.log("map after:", {
      root: map.root.toJSON(),
      length: map.length.toJSON(),
      nodes: map.data.get().nodes,
      sortedLeaves: map.data.get().sortedLeaves,
    });
    expect(step1.publicOutput.toBigInt()).toBe(map.root.toBigInt());
  });
});

The log:

[4:20:57 PM] map before: {
  root: '11424064169442499656248270967957466020056181284936027899258689782550997266764',
  length: '1',
  nodes: [
    [
      15632594125205012933270775512041100774922902129630782398682867600973342943653n
    ],
    [
      24167050500389110549889417520872236040881848908173047560091884295208630545258n
    ],
    [
      11424064169442499656248270967957466020056181284936027899258689782550997266764n
    ]
  ],
  sortedLeaves: [ { key: 0n, value: 0n, nextKey: 0n, index: 0 } ]
}
[4:21:09 PM] map after: {
  root: '11424064169442499656248270967957466020056181284936027899258689782550997266764',
  length: '1',
  nodes: [
    [
      23233281230813636529734057503003001761255023005895160991036382347279357583241n,
      26492836593797301377607267624407155506937202095849486381650364063349767768086n
    ],
    [
      1337184076127276975935409848604683921263008880902249293868243195127772462995n
    ],
    [
      27914741461995568510587207701502469538471560106234225357427537690948886846724n
    ]
  ],
  sortedLeaves: [
    { key: 0n, value: 0n, nextKey: 1n, index: 0 },
    { key: 1n, value: 2n, nextKey: 0n, index: 1 }
  ]
}
 FAIL  tests/map.issue.test.ts
  Map
    ✓ should build the Indexed Merkle Map without proofs (11 ms)
    ✕ should build the Indexed Merkle Map with proofs (13861 ms)

  ● Map › should build the Indexed Merkle Map with proofs

    expect(received).toBe(expected) // Object.is equality

    Expected: 11424064169442499656248270967957466020056181284936027899258689782550997266764n
    Received: 27914741461995568510587207701502469538471560106234225357427537690948886846724n

      52 |       sortedLeaves: map.data.get().sortedLeaves,
      53 |     });
    > 54 |     expect(step1.publicOutput.toBigInt()).toBe(map.root.toBigInt());
         |                                           ^
      55 |   });
      56 | });
      57 |

      at Object.<anonymous> (tests/map.issue.test.ts:54:43)

Test Suites: 1 failed, 1 total
Tests:       1 failed, 1 passed, 2 total
Snapshots:   0 total
Time:        14.492 s

[mitschabaude]

Yes I'm aware of this, had the same problem in the offchain state implementation (this also shows a workaround)

o1js/src/lib/mina/actions/offchain-state-rollup.ts

Lines 287 to 292 in 72a2779

	let proof = await offchainStateRollup.firstBatch(inputState, slice, tree);
	// update tree root/length again, they aren't mutated :(
	// TODO: this shows why the full tree should be the public output
	tree.root = proof.publicOutput.root;
	tree.length = proof.publicOutput.length;

I think it's ok that private proof inputs are not designed to be mutated - that's what return values are for.
So the solution for this issue will come once we enable zkprogram public outputs with auxiliary data. Then we can just return the entire merkle tree as output from the zkprogram, and it will have the updated content.

[dfstio]

Thank you; this workaround works. I also got correct results by cloning the map, passing the cloned map to the ZkProgram, and repeating the same insert operation (that can be done with rawMethods to make sure that the code is the same) on the main map. But your workaround is much more efficient.