feat: add batched state Merkle tree support

[ananas-block]

Overview

This pr adds support for V2 (batched) state and address Merkle tree types. Existing V1 state and address Merkle trees are concurrent Merkle trees. Changes include the indexing logic, and new api endpoints.

Motivation for the V2 update are costs associated with V1 Merkle trees:

1. V1 (concurrent) Merkle trees are expensive for users in terms of CU, and for the forester in terms of number of transactions.

2. V2 (batched) Merkle trees reduce the number of forester transactions by 1250x by updating Merkle tree in batches with zero knowledge proofs. Users transactions are cheaper because compressed account hashes and nullifiers are both inserted into queues (in V1 only nullifiers were inserted into queues).

Additionally, the V2 update includes CU performance optimizations in the system program, ie removal of noop event cpi, use of pinocchio instead of anchor, and zero copies instead of borsh for instruction data serialization.

No breaking changes for existing programs.

Life cycle Concurrent Merkle tree:

1. user tx create -> append account hash to Merkle tree
2. user tx spend -> insert account hash into nullifier queue
3. forester -> insert nullifier into Merkle tree

Lifecycle batched Merkle tree:

1. user tx create -> insert account hash into output queue
2. forester -> appends account hashes in batch zero knowledge proof to Merkle tree
3. user tx spend -> insert nullifier into input queue (nullifier queue)
4. forester -> inserts nullifiers in batch zero knowledge proof into Merkle tree

Note, a created account can be spent before step 2 is completed. The order of steps can be any of the following combinations:

• 1 2 3 4
• 1 3 4 2
• 1 3 2 4

Changes:

Api

1. get_compressed_account_v2
   returns AccountV2
2. get_compressed_account_proof_v2
   returns GetCompressedAccountProofResponseValueV2 which includes:
   2.1. prove_by_index is set to true if account is in_output_queue
   2.2. tree_context (tree pubkey, queue pubkey, tree type, cpi context pubkey (placeholder))
3. get_multiple_compressed_account_proofs_v2
   returns GetCompressedAccountProofResponseValueV2
4. get_compressed_token_accounts_by_owner_v2
   returns PaginatedAccountListV2 of AccountV2
5. get_compressed_token_accounts_by_delegate_v2
   returns TokenAccountListResponseV2, contains AccountV2
6. get_queue_elements
   • used by foresters to empty queues and update batched Merkle trees
   • returns queue elements including Merkle proofs for the index of the queue element
     (See notion doc for more details.)
7. get_compressed_accounts_by_owner_v2
   return TokenAccountListV2 which contains AccountV2
8. get_multiple_compressed_accounts_v2
   return AccountListV2 which contains AccountV2
9. get_transaction_with_compression_info_v2
   returns CompressionInfoV2 which includes opened and closed accounts v2
10. get_validity_proof_v2
    • returns None as root index for elements in batched output queues
11. get_batch_address_update_info
    • used by forester for batched address tree updates
    • returns multiple queue elements, and their address Merkle proofs

Account Model

src/dao/generated/accounts.rs
Changed Properties:

1. seq: i64 -> seq: Option<i64>
   In concurrent trees seq is always some since account hashes are directly appended to the Merkle tree by the user transaction.
   Accounts that are inserted into batched trees are first inserted into the trees output queue, at that time the account has no sequence number.

New Fields:

1. queue
   Queue pubkey, for batched trees the output queue pubkey, for concurrent trees the nullifier queue pubkey.
   (For batched trees the nullifier queue is part of the Merkle tree account.)
   We added the output queue property so that we can return the queue efficiently in get_compressed_account_v2.
   We need the output queue to spend accounts in batched trees.
2. in_output_queue
   If true, account is in output queue and not in the tree yet.
   Once the account is inserted into the tree by a forester event.
   Note, the account can be in the nullifier queue or nullified in the Merkle tree and still be in the output queue.
3. nullifier_queue_index
   Index of value in a batched nullifier queue.
   Used by get_queue_elements.
4. nullified_in_tree
true -> account is nullified in tree -> nullifier_queue_index is None
   is a separate field for // could we just say nullifier_queue_index.is_none && spent ?
5. tree_type
   Either concurrent or batched.
   Clients need to know which tree type an account is in to deserialize accounts correctly.
6. nullifier
   Only used to nullify accounts in batched trees.
   Accounts in concurrent trees are nullified by updating the leaf the account hash is stored in with zero
   (the nullifier is constant for all account hashes).
   For accounts in batched trees the nullifier is H(account_hash, leaf_index, tx_hash). This allows to do a zk proof about
   the transaction an account was spent in.
7. tx_hash
   Unique per public transaction event.
   tx_hash = H(hash_chain_input_account_hashes, hash_chain_output_account_hashes, slot)
   Tx hash is only created when a transaction has input compressed accounts.

Parser

1. src/ingester/parser/indexer_events.rs
   1.1. MerkleTreeEvent
   Added events emitted by forester transactions updating batched Merkle trees.
   1.1.1. BatchAppend(BatchEvent)
   1.1.2. BatchNullify(BatchEvent)
   1.1.3. BatchAddressAppend(BatchEvent)
   1.2. BatchPublicTransactionEvent
   Parsed from user transactions system program and account compression program instructions.
   Wraps PublicTransactionEvent.
2. src/ingester/parser/merkle_tree_events_parser.rs
   Extracts logic to parse Merkle tree events from parser/mod.rs
   2.1. parse_merkle_tree_event
   Deserializes event from an instruction.
   2.2. parse_legacy_nullifier_event
   Parses legacy (concurrent) nullifier event.
   2.3. parse_indexed_merkle_tree_update
   Parses an indexed Merkle tree event.
3. src/ingester/parser/mod.rs
   3.1. add if branch to parse_public_transaction_event_v2
   3.2. refactor existing event parsing logic into parse_legacy_public_transaction_event and parse_merkle_tree_event
   3.3. move helper functions unmodified into separate files
4. src/ingester/parser/state_update.rs
   4.1. pub out_accounts: Vec<Account> -> Vec<AccountWithContext>
   4.2. pub batch_events: IndexedBatchEvents,
   Contains both batch append and batch nullify events.
   4.3. pub input_context: Vec<BatchNullifyContext>
   Contains context to insert accounts into the batched input queue.
   4.4. merge_updates include added fields
5. src/ingester/parser/tree_info.rs
   Adds global hash set to map tree pubkey -> queue pubkey and queue pubkey -> tree pubkey
   We need the mapping to determine the tree pubkey for accounts created in a batched Merkle tree since these account hashes are inserted into the output queue -> the event doesn't contain the tree pubkey.
   As the number of trees grows this should likely become a separate table instead of a hash map that needs to be filled manually.
6. src/ingester/parser/tx_event_parser.rs
   Contains parsing logic moved from mod.rs to parse PublicTransactionEvent from legacy transactions.
7. src/ingester/parser/tx_event_parser_v2.rs
   Contains new parsing logic to parse events from a light transaction (one or multiple system program instructions and one account compression program instruction) without the event emitted in an explicit noop cpi.
   7.1. event_from_light_transaction does the parsing and returns Vec<BatchPublicTransactionEvent>

Persist

1. src/ingester/persist/mod.rs
   1.1. remove hardcoded tree height
   1.2. add logic to insert into input queues spend_input_accounts_batched
   1.3. add logic to persist batched events persist_batch_events
   1.4. append_output_accounts adapted to fill all new fields
   1.5. move spend_input_accounts unmodified to src/ingester/persist/spend.rs
2. src/ingester/persist/leaf_node.rs unchanged logic moved from src/ingester/persist/persisted_state_tree.rs
3. src/ingester/persist/leaf_node_proof.rs unchanged logic moved from src/ingester/persist/persisted_state_tree.rs
   3.1. proof validation is commented because it makes new tests 20x slower
4. src/ingester/persist/merkle_proof_with_context.rs
   Contains MerkleProofWithContext and validate method
5. src/ingester/persist/persisted_batch_event.rs
   Persists both batch append and batch nullify events.
6. src/ingester/persist/persisted_indexed_merkle_tree.rs
   6.1. adapt filling seq to Option<i64>
7. src/ingester/persist/persisted_state_tree.rs
   7.1. get_proof_nodes include empty leaves in proof needed for get_queue_elements
8. src/ingester/persist/spend.rs
   8.1. unchanged spend_input_accounts moved from mod.rs
   8.2. spend_input_accounts_batched
   Update the nullifier queue index and nullifier of the input accounts in batched trees.

Migrations

Added migration from current to new account model.

Tests

1. tests/integration_tests/batched_state_tree_tests.rs
   Contains 3 new tests, test data is generated with export-photon-test-transactions.sh in light-protocol.
   1.1. test_batched_tree_transactions
   Tested API calls with data in a batched state tree:
   - get compressed account by owner
   - get compressed account proofs
   - correct root update after batch append and batch nullify events
   - get_validity_proof_v2
   - get_queue_elements (for input and output queues)
   Test data:
   - 50 active compressed accounts with 1_000_000 each owned by Pubkey::new_unique()
   - 50 nullified compressed accounts
   - All accounts are inserted into the batched Merkle tree
   - 10 append events and 5 nullify events (zero knowledge proof size 10)
   - Queues are empty once all transactions are indexed
   1.2. test_batched_tree_token_transactions
   Test correct indexing of token accounts in a batched state Merkle tree.
   Data:
   - 4 recipients with 1 token account each
   - 1 sender with 3 token accounts
   Asserts:
   - Sender has 3 token accounts with 12341 balance each.
   - Recipients have 1 token account each with 9255, 9255, 9255, 9258 balance.
   - Sender's token balances are correct.
   - Recipients' token balances are correct.
   1.3. test_four_cpi_events
   Test indexes a transaction which creates 4 compressed accounts in 4 CPIs that create a transaction event each in one outer instruction.

2. tests/integration_tests/e2e_tests.rs
   (asserted with assert_json_snapshot)
   2.1. get_compressed_token_accounts_by_owner_v2
   2.2. get_validity_proof_v2
   2.3. get_transaction_helper_v2

3. tests/integration_tests/mock_tests.rs
   (V1 tests already existed.)
   3.1. test_multiple_accounts
   Test V1 accounts with V1 and V2 endpoints:
   - get_compressed_accounts_by_owner
   - get_compressed_accounts_by_owner_v2
   - get_multiple_compressed_accounts
   - get_multiple_compressed_accounts_v2
   - get_compressed_account
   - get_compressed_account_v2

   3.2. test_persist_token_data
   Test V1 token accounts with V1 and V2 endpoints:
   - get_compressed_token_accounts_by_owner
   - get_compressed_token_accounts_by_owner_v2
   - get_compressed_token_balances_by_owner
   - get_compressed_token_balances_by_owner_v2
   - get_compressed_token_account_balance
   - get_compressed_token_accounts_by_delegate
   - get_compressed_token_accounts_by_delegate_v2

Notes

1. QUEUE_TREE_MAPPING should probably be a temporary solution and be replaced in a different pr, for example by a tree metadata table.