Skip to content

Commit

Permalink
feat: use uint64 for internal pointers
Browse files Browse the repository at this point in the history 
Closes Remove `Offset` for pointers #91
  • Loading branch information
@kevmo314
kevmo314 committed Feb 8, 2024
1 parent 4ef9b10 commit cd637b6
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 44 deletions.
48 changes: 22 additions & 26 deletions pkg/btree/bptree.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ func (t *BPTree) Find(key []byte) (MemoryPointer, bool, error) {
n := path[0].node
i, found := n.bsearch(key)
if found {
return n.Pointers[i], true, nil
return n.Pointer(i), true, nil
}
return MemoryPointer{}, false, nil
}
Expand Down Expand Up @@ -85,25 +85,22 @@ func (t *BPTree) traverse(key []byte, node *BPTreeNode, ptr MemoryPointer) ([]Tr
}
for i, k := range node.Keys {
if bytes.Compare(key, k.Value) < 0 {
child, err := t.readNode(node.Pointers[i])
child, err := t.readNode(node.Pointer(i))
if err != nil {
return nil, err
}
path, err := t.traverse(key, child, node.Pointers[i])
path, err := t.traverse(key, child, node.Pointer(i))
if err != nil {
return nil, err
}
return append(path, TraversalRecord{node: node, index: i, ptr: ptr}), nil
}
}
if node.Pointers[len(node.Pointers)-1].Offset == ptr.Offset {
panic("infinite loop 2")
}
child, err := t.readNode(node.Pointers[len(node.Pointers)-1])
child, err := t.readNode(node.Pointer(-1))
if err != nil {
return nil, err
}
path, err := t.traverse(key, child, node.Pointers[len(node.Pointers)-1])
path, err := t.traverse(key, child, node.Pointer(-1))
if err != nil {
return nil, err
}
Expand All @@ -119,7 +116,7 @@ func (t *BPTree) Insert(key ReferencedValue, value MemoryPointer) error {
// special case, create the root as the first node
node := &BPTreeNode{Data: t.Data}
node.Keys = []ReferencedValue{key}
node.Pointers = []MemoryPointer{value}
node.leafPointers = []MemoryPointer{value}
buf, err := node.MarshalBinary()
if err != nil {
return err
Expand All @@ -141,12 +138,12 @@ func (t *BPTree) Insert(key ReferencedValue, value MemoryPointer) error {
j, _ := n.bsearch(key.Value)
if j == len(n.Keys) {
n.Keys = append(n.Keys, key)
n.Pointers = append(n.Pointers, value)
n.leafPointers = append(n.leafPointers, value)
} else {
n.Keys = append(n.Keys[:j+1], n.Keys[j:]...)
n.Keys[j] = key
n.Pointers = append(n.Pointers[:j+1], n.Pointers[j:]...)
n.Pointers[j] = value
n.leafPointers = append(n.leafPointers[:j+1], n.leafPointers[j:]...)
n.leafPointers[j] = value
}

// traverse up the tree and split if necessary
Expand All @@ -162,11 +159,11 @@ func (t *BPTree) Insert(key ReferencedValue, value MemoryPointer) error {
// n is the left node, m the right node
m := &BPTreeNode{Data: t.Data}
if n.leaf() {
m.Pointers = n.Pointers[mid:]
m.leafPointers = n.leafPointers[mid:]
m.Keys = n.Keys[mid:]
} else {
// for non-leaf nodes, the mid key is inserted into the parent
m.Pointers = n.Pointers[mid+1:]
m.internalPointers = n.internalPointers[mid+1:]
m.Keys = n.Keys[mid+1:]
}
mbuf, err := m.MarshalBinary()
Expand All @@ -179,10 +176,10 @@ func (t *BPTree) Insert(key ReferencedValue, value MemoryPointer) error {
}

if n.leaf() {
n.Pointers = n.Pointers[:mid]
n.leafPointers = n.leafPointers[:mid]
n.Keys = n.Keys[:mid]
} else {
n.Pointers = n.Pointers[:mid+1]
n.internalPointers = n.internalPointers[:mid+1]
n.Keys = n.Keys[:mid]
}

Expand All @@ -208,17 +205,16 @@ func (t *BPTree) Insert(key ReferencedValue, value MemoryPointer) error {
p.node.Keys = append(p.node.Keys[:p.index+1], p.node.Keys[p.index:]...)
p.node.Keys[p.index] = midKey
}
p.node.Pointers = append(p.node.Pointers[:p.index+1], p.node.Pointers[p.index:]...)
p.node.Pointers[p.index] = MemoryPointer{Offset: uint64(noffset), Length: uint32(len(nbuf))}
p.node.Pointers[p.index+1] = MemoryPointer{Offset: uint64(moffset), Length: uint32(len(mbuf))}
p.node.internalPointers = append(p.node.internalPointers[:p.index+1], p.node.internalPointers[p.index:]...)
p.node.internalPointers[p.index] = uint64(noffset)
p.node.internalPointers[p.index+1] = uint64(moffset)
// the parent will be written to disk in the next iteration
} else {
// the root split, so create a new root
p := &BPTreeNode{Data: t.Data}
p.Keys = []ReferencedValue{midKey}
p.Pointers = []MemoryPointer{
{Offset: uint64(noffset), Length: uint32(len(nbuf))},
{Offset: uint64(moffset), Length: uint32(len(mbuf))},
p.internalPointers = []uint64{
uint64(noffset), uint64(moffset),
}

pbuf, err := p.MarshalBinary()
Expand Down Expand Up @@ -346,21 +342,21 @@ func (t *BPTree) recursiveString(n *BPTreeNode, indent int) string {
// print the node itself
var buf bytes.Buffer
if !n.leaf() {
for i := range n.Pointers {
child, err := t.readNode(n.Pointers[i])
for i := range n.internalPointers {
child, err := t.readNode(n.Pointer(i))
if err != nil {
return fmt.Sprintf("error: failed to read child node: %v", err)
}
buf.WriteString(t.recursiveString(child, indent+1))
if i < len(n.Pointers)-1 {
if i < len(n.internalPointers)-1 {
for i := 0; i < indent; i++ {
buf.WriteString(" ")
}
buf.WriteString(fmt.Sprintf("key %v\n", n.Keys[i]))
}
}
} else {
for i := range n.Pointers {
for i := range n.leafPointers {
for i := 0; i < indent; i++ {
buf.WriteString(" ")
}
Expand Down
51 changes: 40 additions & 11 deletions pkg/btree/node.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,31 @@ type BPTreeNode struct {
Data []byte
// contains the offset of the child node or the offset of the record for leaf
// if the node is a leaf, the last pointer is the offset of the next leaf
Pointers []MemoryPointer
Keys []ReferencedValue
leafPointers []MemoryPointer
internalPointers []uint64
Keys []ReferencedValue
}

func (n *BPTreeNode) leaf() bool {
// leafs contain the same number of pointers as keys
return len(n.Pointers) == len(n.Keys)
return len(n.leafPointers) > 0
}

func (n *BPTreeNode) Pointers() []MemoryPointer {
if n.leaf() {
return n.leafPointers
}
pointers := make([]MemoryPointer, len(n.internalPointers))
for i, p := range n.internalPointers {
pointers[i].Offset = p
}
return pointers
}

func (n *BPTreeNode) Pointer(i int) MemoryPointer {
if n.leaf() {
return n.leafPointers[(len(n.leafPointers)+i)%len(n.leafPointers)]
}
return MemoryPointer{Offset: n.internalPointers[(len(n.internalPointers)+i)%len(n.internalPointers)]}
}

func (n *BPTreeNode) Size() int64 {
Expand All @@ -44,9 +62,12 @@ func (n *BPTreeNode) Size() int64 {
size += 4 + len(k.Value)
}
}
for range n.Pointers {
for range n.leafPointers {
size += 12
}
for range n.internalPointers {
size += 8
}
return int64(size)
}

Expand Down Expand Up @@ -78,11 +99,15 @@ func (n *BPTreeNode) MarshalBinary() ([]byte, error) {
ct += m + 4
}
}
for _, p := range n.Pointers {
for _, p := range n.leafPointers {
binary.BigEndian.PutUint64(buf[ct:ct+8], p.Offset)
binary.BigEndian.PutUint32(buf[ct+8:ct+12], p.Length)
ct += 12
}
for _, p := range n.internalPointers {
binary.BigEndian.PutUint64(buf[ct:ct+8], p)
ct += 8
}
if ct != int(n.Size()) {
panic("size mismatch")
}
Expand All @@ -102,10 +127,10 @@ func (n *BPTreeNode) UnmarshalBinary(buf []byte) error {
size := int32(binary.BigEndian.Uint32(buf[:4]))
leaf := size < 0
if leaf {
n.Pointers = make([]MemoryPointer, -size)
n.leafPointers = make([]MemoryPointer, -size)
n.Keys = make([]ReferencedValue, -size)
} else {
n.Pointers = make([]MemoryPointer, size+1)
n.internalPointers = make([]uint64, size+1)
n.Keys = make([]ReferencedValue, size)
}
if size == 0 {
Expand All @@ -127,11 +152,15 @@ func (n *BPTreeNode) UnmarshalBinary(buf []byte) error {
m += 4 + int(l)
}
}
for i := range n.Pointers {
n.Pointers[i].Offset = binary.BigEndian.Uint64(buf[m : m+8])
n.Pointers[i].Length = binary.BigEndian.Uint32(buf[m+8 : m+12])
for i := range n.leafPointers {
n.leafPointers[i].Offset = binary.BigEndian.Uint64(buf[m : m+8])
n.leafPointers[i].Length = binary.BigEndian.Uint32(buf[m+8 : m+12])
m += 12
}
for i := range n.internalPointers {
n.internalPointers[i] = binary.BigEndian.Uint64(buf[m : m+8])
m += 8
}
return nil
}

Expand Down
9 changes: 2 additions & 7 deletions pkg/btree/node_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import (
func TestBPTreeNode_ReadWriteLeaf(t *testing.T) {
// Create a test BPTreeNode
node1 := &BPTreeNode{
Pointers: []MemoryPointer{
leafPointers: []MemoryPointer{
{Offset: 0, Length: 1},
{Offset: 1, Length: 2},
{Offset: 2, Length: 3},
Expand Down Expand Up @@ -43,12 +43,7 @@ func TestBPTreeNode_ReadWriteLeaf(t *testing.T) {
func TestBPTreeNode_ReadWriteIntermediate(t *testing.T) {
// Create a test BPTreeNode
node1 := &BPTreeNode{
Pointers: []MemoryPointer{
{Offset: 0, Length: 1},
{Offset: 1, Length: 2},
{Offset: 2, Length: 3},
{Offset: 3, Length: 4},
},
internalPointers: []uint64{0, 1, 2, 3},
Keys: []ReferencedValue{
{Value: []byte{0}},
{Value: []byte{1, 2}},
Expand Down

0 comments on commit cd637b6

Please sign in to comment.