Merge branch 'FI-733' into 'master'

feat(FI-733) [icrc index] add list_subaccounts

 

See merge request dfinity-lab/public/ic!12548

Author: MarioDfinity

rs/rosetta-api/icrc1/index-ng/index-ng.did

@@ -94,9 +94,15 @@ type GetTransactionsResult = variant {
   Err : GetTransactionsErr;
 };
 
+type ListSubaccountsArgs = record {
+    owner: principal;
+    start: opt SubAccount;
+};
+
 service : (index_arg: IndexArg) -> {
     get_account_transactions : (GetAccountTransactionsArgs) -> (GetTransactionsResult);
     get_blocks : (GetBlocksRequest) -> (GetBlocksResponse) query;
     icrc1_balance_of : (Account) -> (Tokens) query;
     ledger_id : () -> (principal) query;
+    list_subaccounts : (ListSubaccountsArgs) -> (vec SubAccount) query;
 }

rs/rosetta-api/icrc1/index-ng/src/lib.rs

@@ -1,9 +1,12 @@
 use candid::{CandidType, Deserialize, Nat, Principal};
-use icrc_ledger_types::icrc1::account::Account;
+use icrc_ledger_types::icrc1::account::{Account, Subaccount};
 use icrc_ledger_types::icrc1::transfer::BlockIndex;
 use icrc_ledger_types::icrc3::blocks::GenericBlock;
 use icrc_ledger_types::icrc3::transactions::Transaction;
 
+/// The maximum number of blocks to return in a single [get_blocks] request.
+pub const DEFAULT_MAX_BLOCKS_PER_RESPONSE: u64 = 2000;
+
 #[derive(CandidType, Debug, Deserialize)]
 pub enum IndexArg {
     InitArg(InitArg),
@@ -56,3 +59,13 @@ pub struct GetAccountTransactionsError {
 
 pub type GetAccountTransactionsResult =
     Result<GetAccountTransactionsResponse, GetAccountTransactionsError>;
+
+#[derive(CandidType, Debug, Deserialize, PartialEq, Eq)]
+pub struct ListSubaccountsArgs {
+    pub owner: Principal,
+    // The last subaccount seen by the client for the given principal.
+    // This subaccount is excluded in the result.
+    // If None then the results will start from the first
+    // in natural order.
+    pub start: Option<Subaccount>,
+}

rs/rosetta-api/icrc1/index-ng/src/main.rs

@@ -7,7 +7,7 @@ use ic_icrc1::blocks::{encoded_block_to_generic_block, generic_block_to_encoded_
 use ic_icrc1::{Block, Operation};
 use ic_icrc1_index_ng::{
     GetAccountTransactionsArgs, GetAccountTransactionsResponse, GetAccountTransactionsResult,
-    IndexArg, TransactionWithId,
+    IndexArg, ListSubaccountsArgs, TransactionWithId, DEFAULT_MAX_BLOCKS_PER_RESPONSE,
 };
 use ic_ledger_core::block::{BlockIndex as BlockIndex64, BlockType, EncodedBlock};
 use ic_stable_structures::memory_manager::{MemoryId, VirtualMemory};
@@ -16,7 +16,7 @@ use ic_stable_structures::{
     memory_manager::MemoryManager, BoundedStorable, DefaultMemoryImpl, StableBTreeMap, StableCell,
     StableLog, Storable,
 };
-use icrc_ledger_types::icrc1::account::Account;
+use icrc_ledger_types::icrc1::account::{Account, Subaccount};
 use icrc_ledger_types::icrc3::archive::{ArchivedRange, QueryBlockArchiveFn};
 use icrc_ledger_types::icrc3::blocks::{
     BlockRange, GenericBlock, GetBlocksRequest, GetBlocksResponse,
@@ -30,11 +30,9 @@ use std::cell::RefCell;
 use std::cmp::Reverse;
 use std::convert::TryFrom;
 use std::hash::Hash;
+use std::ops::Bound::{Excluded, Included};
 use std::time::Duration;
 
-/// The maximum number of blocks to return in a single [get_blocks] request.
-const DEFAULT_MAX_BLOCKS_PER_RESPONSE: u64 = 2000;
-
 const STATE_MEMORY_ID: MemoryId = MemoryId::new(0);
 const BLOCK_LOG_INDEX_MEMORY_ID: MemoryId = MemoryId::new(1);
 const BLOCK_LOG_DATA_MEMORY_ID: MemoryId = MemoryId::new(2);
@@ -602,6 +600,33 @@ fn icrc1_balance_of(account: Account) -> Nat {
     get_balance(account).into()
 }
 
+#[query]
+#[candid_method(query)]
+fn list_subaccounts(args: ListSubaccountsArgs) -> Vec<Subaccount> {
+    let start_key = balance_key(Account {
+        owner: args.owner,
+        subaccount: args.start,
+    });
+    let end_key = balance_key(Account {
+        owner: args.owner,
+        subaccount: Some([u8::MAX; 32]),
+    });
+    let range = (
+        if args.start.is_none() {
+            Included(start_key)
+        } else {
+            Excluded(start_key)
+        },
+        Included(end_key),
+    );
+    with_account_data(|data| {
+        data.range(range)
+            .take(DEFAULT_MAX_BLOCKS_PER_RESPONSE as usize)
+            .map(|((_, (_, subaccount)), _)| subaccount)
+            .collect()
+    })
+}
+
 fn main() {}
 
 #[cfg(test)]

rs/rosetta-api/icrc1/index-ng/tests/tests.rs

@@ -2,14 +2,15 @@ use candid::{Decode, Encode, Nat};
 use ic_base_types::{CanisterId, PrincipalId};
 use ic_icrc1_index_ng::{
     GetAccountTransactionsArgs, GetAccountTransactionsResponse, GetAccountTransactionsResult,
-    GetBlocksResponse, IndexArg, InitArg as IndexInitArg, TransactionWithId,
+    GetBlocksResponse, IndexArg, InitArg as IndexInitArg, ListSubaccountsArgs, TransactionWithId,
+    DEFAULT_MAX_BLOCKS_PER_RESPONSE,
 };
 use ic_icrc1_ledger::{InitArgs as LedgerInitArgs, LedgerArgument};
 use ic_icrc1_test_utils::{valid_transactions_strategy, CallerTransferArg};
 use ic_ledger_canister_core::archive::ArchiveOptions;
 use ic_state_machine_tests::StateMachine;
 use icrc_ledger_types::icrc::generic_metadata_value::MetadataValue as Value;
-use icrc_ledger_types::icrc1::account::Account;
+use icrc_ledger_types::icrc1::account::{Account, Subaccount};
 use icrc_ledger_types::icrc1::transfer::{BlockIndex, TransferArg, TransferError};
 use icrc_ledger_types::icrc3::blocks::{BlockRange, GenericBlock, GetBlocksRequest};
 use icrc_ledger_types::icrc3::transactions::{Mint, Transaction, Transfer};
@@ -229,6 +230,30 @@ fn get_account_transactions(
         .expect("Failed to perform GetAccountTransactionsArgs")
 }
 
+fn list_subaccounts(
+    env: &StateMachine,
+    index: CanisterId,
+    principal: PrincipalId,
+    start: Option<Subaccount>,
+) -> Vec<Subaccount> {
+    Decode!(
+        &env.execute_ingress_as(
+            principal,
+            index,
+            "list_subaccounts",
+            Encode!(&ListSubaccountsArgs {
+                owner: principal.into(),
+                start,
+            })
+            .unwrap()
+        )
+        .expect("failed to list_subaccounts")
+        .bytes(),
+        Vec<Subaccount>
+    )
+    .expect("failed to decode list_subaccounts response")
+}
+
 // Assert that the index canister contains the same blocks as the ledger
 fn assert_ledger_index_parity(env: &StateMachine, ledger_id: CanisterId, index_id: CanisterId) {
     let ledger_blocks = icrc1_get_blocks(env, ledger_id);
@@ -617,3 +642,95 @@ fn test_icrc1_balance_of() {
         )
         .unwrap();
 }
+
+#[test]
+fn test_list_subaccounts() {
+    // For this test, we add minting operations for some principals:
+    // - The principal 1 has one account with the last possible
+    // subaccount.
+    // - The principal 2 has a number of subaccounts equals to
+    // two times the DEFAULT_MAX_BLOCKS_PER_RESPONSE. Therefore fetching
+    // its subaccounts will trigger pagination.
+    // - The principal 3 has one account with the first possible
+    // subaccount.
+    // - The principal 4 has one account with the default subaccount,
+    // which should map to [0;32] in the index.
+
+    let account_1 = Account {
+        owner: PrincipalId::new_user_test_id(1).into(),
+        subaccount: Some([u8::MAX; 32]),
+    };
+    let accounts_2: Vec<_> = (0..(DEFAULT_MAX_BLOCKS_PER_RESPONSE * 2))
+        .map(|i| account(2, i as u128))
+        .collect();
+    let account_3 = account(3, 0);
+    let account_4 = Account {
+        owner: PrincipalId::new_user_test_id(4).into(),
+        subaccount: None,
+    };
+
+    let mut initial_balances: Vec<_> = vec![
+        (account_1, 10_000),
+        (account_3, 10_000),
+        (account_4, 40_000),
+    ];
+    initial_balances.extend(accounts_2.iter().map(|account| (*account, 10_000)));
+
+    let env = &StateMachine::new();
+    let ledger_id = install_ledger(env, initial_balances, default_archive_options());
+    let index_id = install_index(env, ledger_id);
+
+    trigger_heartbeat(env);
+
+    // list account_1.owner subaccounts when no starting subaccount is specified
+    assert_eq!(
+        vec![*account_1.effective_subaccount()],
+        list_subaccounts(env, index_id, PrincipalId(account_1.owner), None)
+    );
+
+    // list account_3.owner subaccounts when no starting subaccount is specified
+    assert_eq!(
+        vec![*account_3.effective_subaccount()],
+        list_subaccounts(env, index_id, PrincipalId(account_3.owner), None)
+    );
+
+    // list account_3.owner subaccounts when an existing starting subaccount is specified but no subaccount is in that range
+    assert!(list_subaccounts(
+        env,
+        index_id,
+        PrincipalId(account_3.owner),
+        Some(*account(3, 1).effective_subaccount())
+    )
+    .is_empty());
+
+    // list acccount_4.owner subaccounts should return the default subaccount
+    // mapped to [0;32]
+    assert_eq!(
+        vec![[0; 32]],
+        list_subaccounts(env, index_id, PrincipalId(account_4.owner), None)
+    );
+
+    // account_2.owner should have two batches of subaccounts
+    let principal_2 = accounts_2.get(0).unwrap().owner;
+    let batch_1 = list_subaccounts(env, index_id, PrincipalId(principal_2), None);
+    let expected_batch_1: Vec<_> = accounts_2
+        .iter()
+        .take(DEFAULT_MAX_BLOCKS_PER_RESPONSE as usize)
+        .map(|account| *account.effective_subaccount())
+        .collect();
+    assert_eq!(expected_batch_1, batch_1);
+
+    let batch_2 = list_subaccounts(
+        env,
+        index_id,
+        PrincipalId(principal_2),
+        Some(*batch_1.last().unwrap()),
+    );
+    let expected_batch_2: Vec<_> = accounts_2
+        .iter()
+        .skip(DEFAULT_MAX_BLOCKS_PER_RESPONSE as usize)
+        .take(DEFAULT_MAX_BLOCKS_PER_RESPONSE as usize)
+        .map(|account| *account.effective_subaccount())
+        .collect();
+    assert_eq!(expected_batch_2, batch_2);
+}