feat: add integration test, fix edge cases (#15)

kevmo314 · Jan 2, 2024 · 02a56c9 · 02a56c9
1 parent da42537
commit 02a56c9
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 41 deletions.
diff --git a/pkg/btree/bptree.go b/pkg/btree/bptree.go
@@ -239,10 +239,10 @@ func (t *BPTree) Insert(key []byte, value MemoryPointer) error {
 				poffset := noffset + nsize
 				// create a new root
 				p := &BPTreeNode{Pointers: []MemoryPointer{rootOffset}}
-				p.Keys = [][]byte{m.Keys[0]}
+				p.Keys = [][]byte{midKey}
 				p.Pointers = []MemoryPointer{
-					MemoryPointer{Offset: uint64(noffset), Length: uint32(nsize)},
-					MemoryPointer{Offset: uint64(moffset), Length: uint32(msize)},
+					{Offset: uint64(noffset), Length: uint32(nsize)},
+					{Offset: uint64(moffset), Length: uint32(msize)},
 				}
 
 				psize, err := p.WriteTo(t.tree)
@@ -519,10 +519,15 @@ func (t *BPTree) String() string {
 				return err.Error()
 			}
 		}
-		for _, p := range node.Pointers {
-			if _, err := buf.Write([]byte(fmt.Sprintf("%04d ", p))); err != nil {
+		for i := 0; i < len(node.Pointers); i++ {
+			if _, err := buf.Write([]byte(fmt.Sprintf("%04d ", node.Pointers[i]))); err != nil {
 				return err.Error()
 			}
+			if i < len(node.Keys) {
+				if _, err := buf.Write([]byte(fmt.Sprintf("%02d ", node.Keys[i]))); err != nil {
+					return err.Error()
+				}
+			}
 		}
 		if _, err := buf.Write([]byte("\n")); err != nil {
 			return err.Error()

diff --git a/pkg/btree/bptree_test.go b/pkg/btree/bptree_test.go
@@ -1,6 +1,7 @@
 package btree
 
 import (
+	"encoding/binary"
 	"fmt"
 	"testing"
 )
@@ -200,30 +201,34 @@ func TestBPTree(t *testing.T) {
 		fmt.Printf(tree.String())
 	})
 
-	// t.Run("insertion test", func(t *testing.T) {
-	// 	b := newSeekableBuffer()
-	// 	tree, err := NewBPTree(b, 4)
-	// 	if err != nil {
-	// 		t.Fatal(err)
-	// 	}
-	// 	for i := 0; i < 1000; i++ {
-	// 		if err := tree.Insert([]byte{byte(i)}, uint64(i)); err != nil {
-	// 			t.Fatal(err)
-	// 		}
-	// 	}
-	// 	for i := 0; i < 1000; i++ {
-	// 		v, found, err := tree.Find([]byte{byte(i)})
-	// 		if err != nil {
-	// 			t.Fatal(err)
-	// 		}
-	// 		if !found {
-	// 			t.Fatalf("expected to find key %d", i)
-	// 		}
-	// 		if v != uint64(i) {
-	// 			t.Fatalf("expected value %d, got %d", i, v)
-	// 		}
-	// 	}
-	// })
+	t.Run("insertion test", func(t *testing.T) {
+		b := newSeekableBuffer()
+		tree, err := NewBPTree(b, 512)
+		if err != nil {
+			t.Fatal(err)
+		}
+		for i := 0; i < 10240; i++ {
+			buf := make([]byte, 4)
+			binary.BigEndian.PutUint32(buf, uint32(i))
+			if err := tree.Insert(buf, MemoryPointer{Offset: uint64(i)}); err != nil {
+				t.Fatal(err)
+			}
+		}
+		for i := 0; i < 10240; i++ {
+			buf := make([]byte, 4)
+			binary.BigEndian.PutUint32(buf, uint32(i))
+			v, found, err := tree.Find(buf)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !found {
+				t.Fatalf("expected to find key %d", i)
+			}
+			if v.Offset != uint64(i) {
+				t.Fatalf("expected value %d, got %d", i, v)
+			}
+		}
+	})
 
 	// t.Run("bulk insert", func(t *testing.T) {
 	// 	b := newSeekableBuffer()

diff --git a/pkg/btree/node.go b/pkg/btree/node.go
@@ -25,15 +25,18 @@ func (n *BPTreeNode) leaf() bool {
 }
 
 func (n *BPTreeNode) WriteTo(w io.Writer) (int64, error) {
-	size := byte(len(n.Keys))
+	size := int32(len(n.Keys))
 	// set the first bit to 1 if it's a leaf
 	if n.leaf() {
-		size |= 1 << 7
-	}
-	if err := binary.Write(w, binary.BigEndian, size); err != nil {
-		return 0, err
+		if err := binary.Write(w, binary.BigEndian, -size); err != nil {
+			return 0, err
+		}
+	} else {
+		if err := binary.Write(w, binary.BigEndian, size); err != nil {
+			return 0, err
+		}
 	}
-	ct := 1
+	ct := 4
 	for _, k := range n.Keys {
 		if err := binary.Write(w, binary.BigEndian, uint32(len(k))); err != nil {
 			return 0, err
@@ -57,18 +60,19 @@ func (n *BPTreeNode) WriteTo(w io.Writer) (int64, error) {
 }
 
 func (n *BPTreeNode) ReadFrom(r io.Reader) (int64, error) {
-	var size byte
+	var size int32
 	if err := binary.Read(r, binary.BigEndian, &size); err != nil {
 		return 0, err
 	}
-	leaf := size&(1<<7) != 0
+	leaf := size < 0
 	if leaf {
-		n.Pointers = make([]MemoryPointer, size&0x7f)
+		n.Pointers = make([]MemoryPointer, -size)
+		n.Keys = make([][]byte, -size)
 	} else {
-		n.Pointers = make([]MemoryPointer, (size&0x7f)+1)
+		n.Pointers = make([]MemoryPointer, size+1)
+		n.Keys = make([][]byte, size)
 	}
-	n.Keys = make([][]byte, size&0x7f)
-	m := 1
+	m := 4
 	for i := range n.Keys {
 		var l uint32
 		if err := binary.Read(r, binary.BigEndian, &l); err != nil {