From 8d5e0789fdce36ea3527b21ef69da2908f9ce76e Mon Sep 17 00:00:00 2001
From: Kevin Wang <kevmo314@gmail.com>
Date: Sat, 30 Dec 2023 22:34:36 -0500
Subject: [PATCH] feat: add tree node encode/decode (#8)

---
 pkg/btree/node.go             | 83 +++++++++++++++++++++++++++++++++++
 pkg/btree/node_test.go        | 67 ++++++++++++++++++++++++++++
 pkg/encoding/encoding.go      | 12 +++++
 pkg/encoding/encoding_test.go | 15 +++++++
 4 files changed, 177 insertions(+)
 create mode 100644 pkg/btree/node.go
 create mode 100644 pkg/btree/node_test.go

diff --git a/pkg/btree/node.go b/pkg/btree/node.go
new file mode 100644
index 00000000..6c52ac2b
--- /dev/null
+++ b/pkg/btree/node.go
@@ -0,0 +1,83 @@
+package btree
+
+import (
+	"io"
+
+	"github.com/kevmo314/appendable/pkg/encoding"
+)
+
+type Node struct {
+	Size     uint8
+	Keys     [8]DataPointer
+	Children [9]uint64
+	Leaf     bool
+}
+
+type DataPointer struct {
+	RecordOffset        uint64
+	FieldOffset, Length uint32
+}
+
+func (n *Node) encode(w io.Writer) error {
+	size := n.Size
+	if n.Leaf {
+		// mark the first bit
+		size |= 1 << 7
+	}
+	if err := encoding.WriteUint8(w, size); err != nil {
+		return err
+	}
+	for i := 0; i < int(n.Size); i++ {
+		if err := encoding.WriteUint64(w, n.Keys[i].RecordOffset); err != nil {
+			return err
+		}
+		if err := encoding.WriteUint32(w, n.Keys[i].FieldOffset); err != nil {
+			return err
+		}
+		if err := encoding.WriteUint32(w, n.Keys[i].Length); err != nil {
+			return err
+		}
+	}
+	for i := 0; i < int(n.Size)+1; i++ {
+		if err := encoding.WriteUint64(w, n.Children[i]); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (n *Node) decode(r io.Reader) error {
+	size, err := encoding.ReadByte(r)
+	if err != nil {
+		return err
+	}
+	n.Size = size & 0x7f
+	n.Leaf = size&(1<<7) != 0
+	for i := 0; i < int(n.Size); i++ {
+		recordOffset, err := encoding.ReadUint64(r)
+		if err != nil {
+			return err
+		}
+		fieldOffset, err := encoding.ReadUint32(r)
+		if err != nil {
+			return err
+		}
+		length, err := encoding.ReadUint32(r)
+		if err != nil {
+			return err
+		}
+		n.Keys[i] = DataPointer{
+			RecordOffset: recordOffset,
+			FieldOffset:  fieldOffset,
+			Length:       length,
+		}
+	}
+	for i := 0; i < int(n.Size)+1; i++ {
+		child, err := encoding.ReadUint64(r)
+		if err != nil {
+			return err
+		}
+		n.Children[i] = child
+	}
+	return nil
+}
diff --git a/pkg/btree/node_test.go b/pkg/btree/node_test.go
new file mode 100644
index 00000000..547077e5
--- /dev/null
+++ b/pkg/btree/node_test.go
@@ -0,0 +1,67 @@
+package btree
+
+import (
+	"bytes"
+	"reflect"
+	"testing"
+)
+
+func TestNode(t *testing.T) {
+	t.Run("encode", func(t *testing.T) {
+		n := &Node{
+			Size: 2,
+			Keys: [8]DataPointer{
+				{
+					RecordOffset: 0,
+					FieldOffset:  0,
+					Length:       5,
+				},
+				{
+					RecordOffset: 0,
+					FieldOffset:  5,
+					Length:       5,
+				},
+			},
+			Children: [9]uint64{0, 1, 2},
+			Leaf:     true,
+		}
+		buf := &bytes.Buffer{}
+		if err := n.encode(buf); err != nil {
+			t.Fatal(err)
+		}
+		if buf.Len() != 1+16*2+8*3 {
+			t.Fatalf("expected buffer length to be 1+16*2+8*3, got %d", buf.Len())
+		}
+	})
+
+	t.Run("decode", func(t *testing.T) {
+		n := &Node{
+			Size: 2,
+			Keys: [8]DataPointer{
+				{
+					RecordOffset: 0,
+					FieldOffset:  0,
+					Length:       5,
+				},
+				{
+					RecordOffset: 0,
+					FieldOffset:  5,
+					Length:       5,
+				},
+			},
+			Children: [9]uint64{0, 1, 2},
+			Leaf:     true,
+		}
+		buf := &bytes.Buffer{}
+		if err := n.encode(buf); err != nil {
+			t.Fatal(err)
+		}
+		m := &Node{}
+		if err := m.decode(buf); err != nil {
+			t.Fatal(err)
+		}
+		if !reflect.DeepEqual(n, m) {
+			t.Fatalf("expected decoded node to be equal to original node")
+		}
+	})
+}
diff --git a/pkg/encoding/encoding.go b/pkg/encoding/encoding.go
index b5eee172..d36f08d0 100644
--- a/pkg/encoding/encoding.go
+++ b/pkg/encoding/encoding.go
@@ -12,6 +12,10 @@ func WriteByte(w io.Writer, b byte) error {
 	return err
 }
 
+func WriteUint8(w io.Writer, u uint8) error {
+	return binary.Write(w, binary.BigEndian, u)
+}
+
 func WriteUint16(w io.Writer, u uint16) error {
 	return binary.Write(w, binary.BigEndian, u)
 }
@@ -51,6 +55,14 @@ func ReadByte(r io.Reader) (byte, error) {
 	return b[0], nil
 }
 
+func ReadUint8(r io.Reader) (uint8, error) {
+	var u uint8
+	if err := binary.Read(r, binary.BigEndian, &u); err != nil {
+		return 0, err
+	}
+	return u, nil
+}
+
 func ReadUint16(r io.Reader) (uint16, error) {
 	var u uint16
 	if err := binary.Read(r, binary.BigEndian, &u); err != nil {
diff --git a/pkg/encoding/encoding_test.go b/pkg/encoding/encoding_test.go
index 28eb7719..f5d7af43 100644
--- a/pkg/encoding/encoding_test.go
+++ b/pkg/encoding/encoding_test.go
@@ -21,6 +21,21 @@ func TestEncoding(t *testing.T) {
 		}
 	})
 
+	t.Run("uint8 encoding", func(t *testing.T) {
+		u := uint8(1)
+		buf := &bytes.Buffer{}
+		if err := WriteUint8(buf, u); err != nil {
+			t.Fatal(err)
+		}
+		u2, err := ReadUint8(buf)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if u != u2 {
+			t.Errorf("expected %v, got %v", u, u2)
+		}
+	})
+
 	t.Run("uint32 encoding", func(t *testing.T) {
 		u := uint32(1)
 		buf := &bytes.Buffer{}