feat: tree construction

This commit is contained in:
Pravdyvy 2025-11-11 17:25:08 +02:00
parent fcb3993a47
commit 6cbc5028cf
3 changed files with 40 additions and 1 deletions

View File

@ -127,6 +127,23 @@ impl<Node: KeyNode> KeyTree<Node> {
self.addr_map.insert(addr, chain_index.clone());
self.key_map.insert(chain_index, node);
}
pub fn generate_tree_for_depth(&mut self, depth: u32) {
let mut id_stack = vec![ChainIndex::root()];
while !id_stack.is_empty() {
let curr_id = id_stack.pop().unwrap();
self.generate_new_node(curr_id.clone());
let mut next_id = curr_id.n_th_child(0);
while (next_id.chain().iter().sum::<u32>()) < depth - 1 {
id_stack.push(next_id.clone());
next_id = next_id.next_in_line();
}
}
}
}
#[cfg(test)]

View File

@ -259,6 +259,13 @@ pub enum OverCommand {
#[arg(short, long)]
password: String,
},
/// !!!WARNING!!! will rewrite current storage
RestoreKeys {
#[arg(short, long)]
password: String,
#[arg(short, long)]
depth: u32,
}
}
///To execute commands, env var NSSA_WALLET_HOME_DIR must be set into directory with config
@ -511,3 +518,15 @@ pub async fn execute_setup(password: String) -> Result<()> {
Ok(())
}
pub async fn execute_keys_restoration(password: String, depth: u32) -> Result<()> {
let config = fetch_config().await?;
let mut wallet_core = WalletCore::start_from_config_new_storage(config.clone(), password.clone()).await?;
wallet_core.storage.user_data.public_key_tree.generate_tree_for_depth(depth);
wallet_core.storage.user_data.private_key_tree.generate_tree_for_depth(depth);
wallet_core.store_persistent_data().await?;
Ok(())
}

View File

@ -1,7 +1,7 @@
use anyhow::Result;
use clap::{CommandFactory, Parser};
use tokio::runtime::Builder;
use wallet::{Args, OverCommand, execute_continious_run, execute_setup, execute_subcommand};
use wallet::{Args, OverCommand, execute_continious_run, execute_keys_restoration, execute_setup, execute_subcommand};
pub const NUM_THREADS: usize = 2;
@ -25,6 +25,9 @@ fn main() -> Result<()> {
OverCommand::Setup { password } => {
execute_setup(password).await.unwrap();
}
OverCommand::RestoreKeys { password, depth } => {
execute_keys_restoration(password, depth).await.unwrap();
}
}
} else if args.continious_run {
execute_continious_run().await.unwrap();