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