Skip to content

Commit

Permalink
feat: add tree node encode/decode (#8)
Browse files Browse the repository at this point in the history
  • Loading branch information
kevmo314 authored Dec 31, 2023
1 parent 6adbe40 commit 8d5e078
Show file tree
Hide file tree
Showing 4 changed files with 177 additions and 0 deletions.
83 changes: 83 additions & 0 deletions pkg/btree/node.go
Original file line number Diff line number Diff line change
@@ -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
}
67 changes: 67 additions & 0 deletions pkg/btree/node_test.go
Original file line number Diff line number Diff line change
@@ -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")
}
})
}
12 changes: 12 additions & 0 deletions pkg/encoding/encoding.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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 {
Expand Down
15 changes: 15 additions & 0 deletions pkg/encoding/encoding_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{}
Expand Down

0 comments on commit 8d5e078

Please sign in to comment.