return account id on new private account creation

key_protocol/src/key_management/key_tree/mod.rs

@@ -255,6 +255,28 @@ impl KeyTree<ChildKeysPublic> {
 }
 
 impl KeyTree<ChildKeysPrivate> {
+    /// Generate a new private key node, registering the identifier=0 account_id immediately.
+    pub fn generate_new_private_node(
+        &mut self,
+        parent_cci: &ChainIndex,
+    ) -> Option<(nssa::AccountId, ChainIndex)> {
+        let cci = self.generate_new_node(parent_cci)?;
+        let node = self.key_map.get(&cci)?;
+        let account_id = nssa::AccountId::from((&node.value.0.nullifier_public_key, 0_u128));
+        self.account_id_map.insert(account_id, cci.clone());
+        Some((account_id, cci))
+    }
+
+    /// Generate a new private key node using layered placement, registering the identifier=0
+    /// account_id immediately.
+    pub fn generate_new_private_node_layered(&mut self) -> Option<(nssa::AccountId, ChainIndex)> {
+        let cci = self.generate_new_node_layered()?;
+        let node = self.key_map.get(&cci)?;
+        let account_id = nssa::AccountId::from((&node.value.0.nullifier_public_key, 0_u128));
+        self.account_id_map.insert(account_id, cci.clone());
+        Some((account_id, cci))
+    }
+
     /// Cleanup of non-initialized accounts in a private tree.
     ///
     /// If account has no synced entries, removes it, stops at first initialized account.

key_protocol/src/key_protocol_core/mod.rs

@@ -123,19 +123,19 @@ impl NSSAUserData {
 
     /// Generated new private key for privacy preserving transactions.
     ///
-    /// Returns the `ChainIndex` of the new node.
+    /// Returns the `AccountId` (for identifier=0) and `ChainIndex` of the new node.
     pub fn generate_new_privacy_preserving_transaction_key_chain(
         &mut self,
         parent_cci: Option<ChainIndex>,
-    ) -> ChainIndex {
+    ) -> (nssa::AccountId, ChainIndex) {
         match parent_cci {
             Some(parent_cci) => self
                 .private_key_tree
-                .generate_new_node(&parent_cci)
+                .generate_new_private_node(&parent_cci)
                 .expect("Parent must be present in a tree"),
             None => self
                 .private_key_tree
-                .generate_new_node_layered()
+                .generate_new_private_node_layered()
                 .expect("Search for new node slot failed"),
         }
     }
@@ -211,13 +211,16 @@ mod tests {
     fn new_account() {
         let mut user_data = NSSAUserData::default();
 
-        let chain_index = user_data
+        let (account_id, chain_index) = user_data
             .generate_new_privacy_preserving_transaction_key_chain(Some(ChainIndex::root()));
 
         let is_key_chain_generated = user_data.private_key_tree.key_map.contains_key(&chain_index);
-
         assert!(is_key_chain_generated);
 
+        let is_account_id_registered =
+            user_data.private_key_tree.account_id_map.contains_key(&account_id);
+        assert!(is_account_id_registered);
+
         let key_chain = &user_data.private_key_tree.key_map[&chain_index].value.0;
         println!("{key_chain:#?}");
     }

wallet-ffi/src/account.rs

@@ -98,16 +98,7 @@ pub unsafe extern "C" fn wallet_ffi_create_account_private(
         }
     };
 
-    let chain_index = wallet.create_new_account_private(None);
+    let (account_id, _chain_index) = wallet.create_new_account_private(None);
-
-    let node = wallet
-        .storage()
-        .user_data
-        .private_key_tree
-        .key_map
-        .get(&chain_index)
-        .expect("Node was just inserted");
-    let account_id = AccountId::from((&node.value.0.nullifier_public_key, 0_u128));
 
     unsafe {
         (*out_account_id).data = *account_id.value();

wallet/src/cli/account.rs

@@ -147,7 +147,7 @@ impl WalletSubcommand for NewSubcommand {
                     anyhow::bail!("Label '{label}' is already in use by another account");
                 }
 
-                let chain_index = wallet_core.create_new_account_private(cci);
+                let (account_id, chain_index) = wallet_core.create_new_account_private(cci);
 
                 let node = wallet_core
                     .storage
@@ -157,7 +157,6 @@ impl WalletSubcommand for NewSubcommand {
                     .get(&chain_index)
                     .expect("Node was just inserted");
                 let key = &node.value.0;
-                let account_id = AccountId::from((&key.nullifier_public_key, 0_u128));
 
                 if let Some(label) = label {
                     wallet_core

wallet/src/lib.rs

@@ -259,7 +259,7 @@ impl WalletCore {
     pub fn create_new_account_private(
         &mut self,
         chain_index: Option<ChainIndex>,
-    ) -> ChainIndex {
+    ) -> (AccountId, ChainIndex) {
         self.storage
             .user_data
             .generate_new_privacy_preserving_transaction_key_chain(chain_index)