From b92ee28bef17eb1deb981edc38b07159cf5391b6 Mon Sep 17 00:00:00 2001
From: yukirij <yukirij@yukiri.jp>
Date: Mon, 21 Apr 2025 11:42:37 -0700
Subject: [PATCH] Add id update and non-update insert to trie.

---
 src/triefile/mod.rs | 103 +++++++++++++++++++++++++-------------------
 1 file changed, 59 insertions(+), 44 deletions(-)

diff --git a/src/triefile/mod.rs b/src/triefile/mod.rs
index 5dad748..6b2a545 100644
--- a/src/triefile/mod.rs
+++ b/src/triefile/mod.rs
@@ -115,7 +115,7 @@ impl<const Z:usize> TrieFile<Z> {
         }
     }
 
-    pub async fn set(&mut self, key:&[u8], data:&[u8]) -> Result<usize,std::io::Error>
+    async fn set_data(&mut self, key:&[u8], data:&[u8], overwrite:bool) -> Result<usize,std::io::Error>
     {
         let mut node = Node::<Z>::new();
 
@@ -165,15 +165,15 @@ impl<const Z:usize> TrieFile<Z> {
                         if key_index == key.len() {
                             //println!(" - found node");
 
-                            node.has_data = true;
-                            node.data = [0; Z];
-                            
-                            for i in 0..Z.min(data.len()) {
-                                node.data[i] = data[i];
-                            }
+                            if overwrite {
+                                node.has_data = true;
+                                Self::copy_to(&mut node.data, data);
 
-                            self.write_node(node_index, &node).await?;
-                            result = node_index as usize;
+                                self.write_node(node_index, &node).await?;
+                                result = node_index as usize;
+                            } else {
+                                return Err(std::io::Error::other("key already exists"));
+                            }
                         }
 
                         // Continue to child node.
@@ -216,27 +216,11 @@ impl<const Z:usize> TrieFile<Z> {
                         node.next = NO_REL;
 
                         // Set length and text.
-                        let mut bindex = 0;
                         node.length = suffix.len() as u8;
-                        while bindex < suffix.len() {
-                            node.bytes[bindex] = suffix[bindex];
-                            bindex += 1;
-                        }
-                        while bindex < node.bytes.len() {
-                            node.bytes[bindex] = 0;
-                            bindex += 1;
-                        }
+                        Self::copy_to(&mut node.bytes, &suffix);
 
-                        bindex = 0;
                         parent_node.length = prefix.len() as u8;
-                        while bindex < prefix.len() {
-                            parent_node.bytes[bindex] = prefix[bindex];
-                            bindex += 1;
-                        }
-                        while bindex < node.bytes.len() {
-                            parent_node.bytes[bindex] = 0;
-                            bindex += 1;
-                        }
+                        Self::copy_to(&mut parent_node.bytes, &prefix);
 
                         // Write data to trailing nodes.
                         if key_index < key.len() {
@@ -247,9 +231,7 @@ impl<const Z:usize> TrieFile<Z> {
                         // Write data to current node.
                         else {
                             parent_node.has_data = true;
-                            for i in 0..Z.min(data.len()) {
-                                parent_node.data[i] = data[i];
-                            }
+                            Self::copy_to(&mut parent_node.data, data);
                             result = new_index as usize;
                         }
 
@@ -298,6 +280,29 @@ impl<const Z:usize> TrieFile<Z> {
         Ok(result)
     }
 
+    pub async fn set(&mut self, key:&[u8], data:&[u8]) -> Result<usize,std::io::Error>
+    {
+        self.set_data(key, data, true).await
+    }
+
+    pub async fn insert(&mut self, key:&[u8], data:&[u8]) -> Result<usize,std::io::Error>
+    {
+        self.set_data(key, data, false).await
+    }
+
+    pub async fn update(&mut self, id:usize, data:&[u8]) -> Result<(),std::io::Error>
+    {
+        let mut node = Node::<Z>::new();
+        self.read_node(id as u32, &mut node).await?;
+
+        node.has_data = true;
+        Self::copy_to(&mut node.data, data);
+
+        self.write_node(id as u32, &node).await?;
+
+        Ok(())
+    }
+
     pub async fn find(&self, key:&[u8]) -> Result<Option<usize>, std::io::Error>
     {
         let mut node = Node::<Z>::new();
@@ -421,19 +426,18 @@ impl<const Z:usize> TrieFile<Z> {
         Ok(output)
     }
 
-    pub async fn unset(&mut self, key:&[u8]) -> Result<(), std::io::Error>
+    pub async fn unset(&mut self, id:usize) -> Result<(), std::io::Error>
     {
         let mut node = Node::<Z>::new();
-        match self.find(key).await {
-            Ok(Some(id)) => {
-                self.read_node(id as u32, &mut node).await?;
-                node.has_data = false;
-                node.data = [0; Z];
-                self.write_node(id as u32, &node).await?;
-                Ok(())
-            }
-            Ok(None) => Ok(()),
-            Err(e) => Err(e),
+        if id < self.block_count().await? as usize {
+            self.read_node(id as u32, &mut node).await?;
+            node.has_data = false;
+            node.data = [0; Z];
+            self.write_node(id as u32, &node).await?;
+
+            Ok(())
+        } else {
+            Err(std::io::Error::other("id out of bounds"))
         }
     }
 
@@ -465,9 +469,7 @@ impl<const Z:usize> TrieFile<Z> {
             // Otherwise, write data to node.
             else {
                 node.has_data = true;
-                for i in 0..Z.min(data.len()) {
-                    node.data[i] = data[i];
-                }
+                Self::copy_to(&mut node.data, data);
             }
 
             self.write_node(node_index, &node).await?;
@@ -511,6 +513,19 @@ impl<const Z:usize> TrieFile<Z> {
         Ok(((file.seek(SeekFrom::End(0)).await? as usize) / Self::block_size()) as u32)
     }
 
+    fn copy_to(dst:&mut [u8], src:&[u8])
+    {
+        let mut index = 0;
+        while index < dst.len() && index < src.len() {
+            dst[index] = src[index];
+            index += 1;
+        }
+        while index < dst.len() {
+            dst[index] = 0;
+            index += 1;
+        }
+    }
+
     const fn block_size() -> usize
     {
         Node::<Z>::len()