From 3170c65d901582d0909f556781d7c271b091ac99 Mon Sep 17 00:00:00 2001
From: Toru Maesaka <toru@chronohq.com>
Date: Fri, 4 Oct 2024 21:14:17 -0700
Subject: [PATCH] radixdb: support the parent restructure case

Resolves #1
---
 radixdb.go      | 37 +++++++++++++++++++++++++++++++++----
 radixdb_test.go | 33 ++++++++++++++++++++++++++++-----
 2 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/radixdb.go b/radixdb.go
index 1e7ee38..2efb87b 100644
--- a/radixdb.go
+++ b/radixdb.go
@@ -87,6 +87,30 @@ func (rdb *RadixDB) Insert(key []byte, value any) error {
 			return ErrDuplicateKey
 		}
 
+		// newNode's key matches the longest common prefix and is shorter than
+		// the current node's key. Therefore, newNode logically becomes the
+		// parent of the current node, which requires restructuring the tree.
+		//
+		// For example, suppose newNode.key is "app" and current.key is "apple".
+		// The common prefix is "app", and thus "app" becomes the parent of "le".
+		if len(prefix) == len(newNode.key) && len(prefix) < len(current.key) {
+			current.key = current.key[len(newNode.key):]
+			newNode.children = append(newNode.children, current)
+
+			if parent == nil {
+				rdb.root = newNode
+			} else {
+				for i, child := range parent.children {
+					if child == current {
+						parent.children[i] = newNode
+					}
+				}
+			}
+
+			rdb.numNodes += 1
+			return nil
+		}
+
 		// Partial match: Insert newNode by splitting the curent node.
 		// Meeting this condition means that the key has been exhausted.
 		if len(prefix) > 0 && len(prefix) < len(current.key) {
@@ -101,13 +125,18 @@ func (rdb *RadixDB) Insert(key []byte, value any) error {
 		nextNode := current.findCompatibleChild(key)
 
 		if nextNode == nil {
-			if len(current.children) > 0 {
-				current.children = append(current.children, newNode)
-				rdb.numNodes += 1
+			newNode.key = key
+
+			if parent == nil {
+				rdb.root = &node{
+					key:      prefix,
+					children: []*node{current, newNode},
+				}
 			} else {
-				rdb.splitNode(parent, current, newNode, prefix)
+				parent.children = append(parent.children, newNode)
 			}
 
+			rdb.numNodes += 1
 			return nil
 		}
 
diff --git a/radixdb_test.go b/radixdb_test.go
index 81c1cc7..2b962b8 100644
--- a/radixdb_test.go
+++ b/radixdb_test.go
@@ -2,6 +2,8 @@ package radixdb
 
 import (
 	"bytes"
+	"crypto/rand"
+	mrand "math/rand/v2"
 	"testing"
 )
 
@@ -78,7 +80,7 @@ func TestSplitNode(t *testing.T) {
 	rdb.splitNode(nil, rdb.root, newNode, commonPrefix)
 
 	if rdb.Len() != 1 && len(rdb.root.children) != 1 {
-		t.Errorf("tree size: got:%d, want:1", rdb.Len())
+		t.Errorf("Len(): got:%d, want:1", rdb.Len())
 	}
 
 	if !bytes.Equal(rdb.root.key, commonPrefix) {
@@ -100,7 +102,7 @@ func TestSplitNode(t *testing.T) {
 	rdb.splitNode(rdb.root, newNode, strawberryNode, commonPrefix)
 
 	if rdb.Len() != 2 && len(rdb.root.children) != 2 {
-		t.Errorf("tree size: got:%d, want:2", rdb.Len())
+		t.Errorf("Len(): got:%d, want:2", rdb.Len())
 	}
 
 	// newNode should now be further split to "st[ore]".
@@ -136,7 +138,7 @@ func TestInsert(t *testing.T) {
 	}
 
 	if len := rdb.Len(); len != 1 {
-		t.Errorf("expected Len(): got:%d, want:1", len)
+		t.Errorf("Len(): got:%d, want:1", len)
 	}
 
 	// Test non-common key insertion. The node should be a direct child of root.
@@ -157,7 +159,7 @@ func TestInsert(t *testing.T) {
 	}
 
 	if len := rdb.Len(); len != 2 {
-		t.Errorf("expected Len(): got:%d, want:2", len)
+		t.Errorf("Len(): got:%d, want:2", len)
 	}
 
 	// Test common prefix insertion.
@@ -174,6 +176,27 @@ func TestInsert(t *testing.T) {
 	}
 
 	if len := rdb.Len(); len != 5 {
-		t.Errorf("expected Len(): got:%d, want:5", len)
+		t.Errorf("Len(): got:%d, want:5", len)
+	}
+
+	// Mild fuzzing: Insert random keys for memory errors.
+	numRandomInserts := 2000
+	numRecordsBefore := rdb.Len()
+	numRecordsExpected := uint64(numRandomInserts + int(numRecordsBefore))
+
+	for i := 0; i < numRandomInserts; i++ {
+		// Random key length between 1 and 32 bytes.
+		keyLength := mrand.IntN(32-1) + 1
+		randomKey := make([]byte, keyLength)
+
+		if _, err := rand.Read(randomKey); err != nil {
+			t.Fatal(err)
+		}
+
+		rdb.Insert(randomKey, randomKey)
+	}
+
+	if len := rdb.Len(); len != numRecordsExpected {
+		t.Errorf("Len(): got:%d, want:%d", len, numRecordsExpected)
 	}
 }