base tree creation

client-network/src/data.rs

@@ -1,7 +1,7 @@
+use rand::{Rng, rng};
+use sha2::{Digest, Sha256};
 use std::collections::HashMap;
 use std::hash::{DefaultHasher, Hash, Hasher};
-use rand::{rng, Rng};
-use sha2::{Digest, Sha256};
 
 // --- Constants ---
 const MAX_CHUNK_DATA_SIZE: usize = 1024;
@@ -53,9 +53,7 @@ fn generate_random_filename() -> [u8; FILENAME_HASH_SIZE] {
 pub type NodeHash = [u8; FILENAME_HASH_SIZE];
 
 pub fn node_hash_to_hex_string(hash: &NodeHash) -> String {
-    hash.iter()
-        .map(|b| format!("{:02x}", b))
-        .collect()
+    hash.iter().map(|b| format!("{:02x}", b)).collect()
 }
 
 #[repr(u8)]
@@ -83,7 +81,9 @@ impl MerkleTree {
     }
 }
 
-fn generate_random_file_node(storage: &mut HashMap<NodeHash, MerkleNode>) -> Result<NodeHash, String> {
+fn generate_random_file_node(
+    storage: &mut HashMap<NodeHash, MerkleNode>,
+) -> Result<NodeHash, String> {
     let mut rng = rng();
     let is_big = rng.random_bool(0.2); // 20% chance of being a big file
 
@@ -116,7 +116,7 @@ fn generate_random_file_node(storage: &mut HashMap<NodeHash, MerkleNode>) -> Res
 fn generate_random_directory_node(
     depth: u32,
     max_depth: u32,
-    storage: &mut HashMap<NodeHash, MerkleNode>
+    storage: &mut HashMap<NodeHash, MerkleNode>,
 ) -> Result<NodeHash, String> {
     let mut rng = rng();
     let current_depth = depth + 1;
@@ -157,7 +157,6 @@ fn generate_random_directory_node(
         let hash = hash(&node.serialize());
         storage.insert(hash, node);
         Ok(hash)
-
     } else {
         // Generate a BigDirectory Node (internal directory structure)
         let num_children = rng.random_range(MIN_BIG_CHILDREN..=MAX_BIG_CHILDREN.min(4)); // Limit children count
@@ -215,7 +214,10 @@ pub struct DirectoryEntry {
 }
 
 pub fn filename_to_string(filename: [u8; FILENAME_HASH_SIZE]) -> String {
-    let end_index = filename.iter().position(|&b| b == 0).unwrap_or(FILENAME_HASH_SIZE);
+    let end_index = filename
+        .iter()
+        .position(|&b| b == 0)
+        .unwrap_or(FILENAME_HASH_SIZE);
     String::from_utf8_lossy(&filename[..end_index]).to_string()
 }
 
@@ -312,7 +314,9 @@ impl MerkleNode {
         bytes
     }
 
-    pub fn generate_random_tree(max_depth: u32) -> Result<(NodeHash, HashMap<NodeHash, MerkleNode>), String> {
+    pub fn generate_random_tree(
+        max_depth: u32,
+    ) -> Result<(NodeHash, HashMap<NodeHash, MerkleNode>), String> {
         let mut storage = HashMap::new();
 
         // Start tree generation from the root directory at depth 0
@@ -320,4 +324,56 @@ impl MerkleNode {
 
         Ok((root_hash, storage))
     }
+
+    pub fn generate_base_tree() -> (NodeHash, HashMap<NodeHash, MerkleNode>) {
+        let mut res = HashMap::new();
+
+        let node1 = MerkleNode::Chunk(ChunkNode::new_random());
+        let hash1 = hash(&node1.serialize());
+
+        let node2 = MerkleNode::Chunk(ChunkNode::new_random());
+        let hash2 = hash(&node2.serialize());
+
+        res.insert(hash1, node1);
+        res.insert(hash2, node2);
+
+        let node3 = MerkleNode::Chunk(ChunkNode::new_random());
+        let hash3 = hash(&node3.serialize());
+
+        res.insert(hash3, node3);
+
+        let dir1 = MerkleNode::Directory(DirectoryNode {
+            entries: [DirectoryEntry {
+                filename: generate_random_filename(),
+                content_hash: hash3,
+            }]
+            .to_vec(),
+        });
+        let hash_dir1 = hash(&dir1.serialize());
+
+        res.insert(hash_dir1, dir1);
+
+        let root = MerkleNode::Directory(DirectoryNode {
+            entries: [
+                DirectoryEntry {
+                    filename: generate_random_filename(),
+                    content_hash: hash1,
+                },
+                DirectoryEntry {
+                    filename: generate_random_filename(),
+                    content_hash: hash2,
+                },
+                DirectoryEntry {
+                    filename: generate_random_filename(),
+                    content_hash: hash_dir1,
+                },
+            ]
+            .to_vec(),
+        });
+
+        let root_hash = hash(&root.serialize());
+        res.insert(root_hash, root);
+
+        (root_hash, res)
+    }
 }