B+Tree

btree.go

package btree

import (
	"fmt"
	"slices"
	"sort"
	"strings"
)

type BTree struct {
	root   *node
	degree int
}

type node struct {
	isLeaf bool
	keys   []string
	values []string
	ptr    []*node
	next   *node
}

type item struct {
	key   string
	value string
	ptr   *node
}

func New() *BTree {
	return &BTree{
		root:   nil,
		degree: 8,
	}
}

func NewWithDegree(degree int) *BTree {
	return &BTree{
		root:   nil,
		degree: degree,
	}
}

func (b *BTree) Cursor(key string) (*cursor, bool) {
	path := b.lookup(key)
	return &cursor{
		node: path.nodes[len(path.nodes)-1],
		i:    path.i[len(path.i)-1],
	}, path.found
}

func (b *BTree) lookup(key string) path {
	if b.root == nil {
		return path{}
	}

	nodes := make([]*node, 0)
	indexes := make([]int, 0)

	current := b.root
	for !current.isLeaf {
		i, _ := sort.Find(len(current.keys), func(i int) int {
			return strings.Compare(key, current.keys[i])
		})

		nodes = append(nodes, current)
		indexes = append(indexes, i)

		if i == len(current.keys) {
			current = current.next
		} else {
			current = current.ptr[i]
		}
	}

	i, found := sort.Find(len(current.keys), func(i int) int {
		return strings.Compare(key, current.keys[i])
	})

	nodes = append(nodes, current)
	indexes = append(indexes, i)

	return path{
		found: found,
		nodes: nodes,
		i:     indexes,
	}
}

type path struct {
	found bool
	nodes []*node
	i     []int // TODO: no need to be sliec
}

func (b *BTree) Put(key, value string) {
	if b.root == nil {
		b.root = &node{
			isLeaf: true,
			keys:   []string{key},
			values: []string{value},
		}
		return
	}

	path := b.lookup(key)
	leaf := path.nodes[len(path.nodes)-1]
	i := path.i[len(path.i)-1]

	if path.found {
		leaf.values[i] = value
		return
	}

	if len(leaf.keys) < b.degree {
		sliceShift(&leaf.keys, i, key)
		sliceShift(&leaf.values, i, value)
		return
	}

	keep := len(leaf.keys) / 2
	var rightLeaf *node
	if i < keep {
		rightLeaf = &node{
			isLeaf: true,
			keys:   slices.Clone(leaf.keys[keep:]),
			values: slices.Clone(leaf.values[keep:]),
			next:   leaf.next,
		}
		leaf.keys = leaf.keys[:keep]
		leaf.values = leaf.values[:keep]
		leaf.next = rightLeaf
		sliceShift(&leaf.keys, i, key)
		sliceShift(&leaf.values, i, value)
	} else {
		rightLeaf = &node{
			isLeaf: true,
			keys:   slices.Clone(leaf.keys[keep:]),
			values: slices.Clone(leaf.values[keep:]),
			next:   leaf.next,
		}
		leaf.keys = leaf.keys[:keep]
		leaf.values = leaf.values[:keep]
		leaf.next = rightLeaf
		sliceShift(&rightLeaf.keys, i-keep, key)
		sliceShift(&rightLeaf.values, i-keep, value)
	}

	nextNodes := path.nodes[:len(path.nodes)-1]
	b.propagateInterior(rightLeaf.keys[0], leaf, rightLeaf, nextNodes)
}

func sliceShift[T any](s *[]T, i int, value T) {
	var empty T
	slice := *s
	slice = append(slice, empty)
	for j := len(slice) - 1; j > i; j-- {
		slice[j] = slice[j-1]
	}
	slice[i] = value
	*s = slice
}

func (b *BTree) propagateInterior(key string, left, right *node, nextNodes []*node) {
	if len(nextNodes) == 0 {
		b.root = &node{
			keys: []string{key},
			ptr:  []*node{left},
			next: right,
		}
		return
	}

	current := nextNodes[len(nextNodes)-1]
	i, _ := sort.Find(len(current.keys), func(i int) int {
		return strings.Compare(key, current.keys[i])
	})

	if len(current.keys) < b.degree {
		if i == len(current.keys) {
			current.next = right
			sliceShift(&current.keys, i, key)
			sliceShift(&current.ptr, i, left)
		} else {
			sliceShift(&current.keys, i, key)
			sliceShift(&current.ptr, i, left)
			current.ptr[i+1] = right
		}
		return
	}

	mid := (len(current.keys) + 1) / 2
	if i < mid {
		rightNode := &node{
			keys: slices.Clone(current.keys[mid:]),
			ptr:  slices.Clone(current.ptr[mid:]),
			next: current.next,
		}
		evictedKey := current.keys[mid-1]
		evictedPtr := current.ptr[mid-1]
		current.keys = current.keys[:mid-1]
		current.ptr = current.ptr[:mid-1]
		current.next = evictedPtr
		if i == len(current.keys) {
			current.next = right
			sliceShift(&current.keys, i, key)
			sliceShift(&current.ptr, i, left)
		} else {
			sliceShift(&current.keys, i, key)
			sliceShift(&current.ptr, i, left)
			current.ptr[i+1] = right
		}
		b.propagateInterior(evictedKey, current, rightNode, nextNodes[:len(nextNodes)-1])
	} else if i == mid {
		rightNode := &node{
			keys: slices.Clone(current.keys[mid:]),
			ptr:  slices.Clone(current.ptr[mid:]),
			next: current.next,
		}
		rightNode.ptr[0] = right
		current.next = left
		current.keys = current.keys[:mid]
		current.ptr = current.ptr[:mid]
		b.propagateInterior(key, current, rightNode, nextNodes[:len(nextNodes)-1])
	} else {
		rightNode := &node{
			keys: slices.Clone(current.keys[mid+1:]),
			ptr:  slices.Clone(current.ptr[mid+1:]),
			next: current.next,
		}
		evictedKey := current.keys[mid]
		evictedPtr := current.ptr[mid]
		current.keys = current.keys[:mid]
		current.ptr = current.ptr[:mid]
		current.next = evictedPtr
		if i-mid-1 == len(rightNode.keys) {
			rightNode.next = right
			sliceShift(&rightNode.keys, i-mid-1, key)
			sliceShift(&rightNode.ptr, i-mid-1, left)
		} else {
			sliceShift(&rightNode.keys, i-mid-1, key)
			sliceShift(&rightNode.ptr, i-mid-1, left)
			rightNode.ptr[i-mid] = right
		}
		b.propagateInterior(evictedKey, current, rightNode, nextNodes[:len(nextNodes)-1])
	}
}

func (b *BTree) Debug() {
	if b.root == nil {
		fmt.Println("(nil)")
		return
	}

	b.debug(b.root)
}

func (b *BTree) debug(node *node) {
	if node == nil {
		return
	}

	fmt.Printf("%p leaf=%v next=%p\n", node, node.isLeaf, node.next)
	if node.isLeaf {
		for i, key := range node.keys {
			fmt.Printf("- %s %s\n", key, node.values[i])
		}
	} else {
		for i, key := range node.keys {
			fmt.Printf("- %s %p\n", key, node.ptr[i])
		}
	}
	fmt.Println("===============================================")

	if !node.isLeaf {
		for _, ptr := range node.ptr {
			b.debug(ptr)
		}
		b.debug(node.next)
	}
}

type cursor struct {
	node *node
	i    int
}

func (c *cursor) Next() (key string, value string, ok bool) {
	if c.node == nil {
		return "", "", false
	}

	key, value, ok = c.node.keys[c.i], c.node.values[c.i], true

	if c.i >= len(c.node.values)-1 {
		c.node = c.node.next
		c.i = 0
	} else {
		c.i++
	}

	return key, value, ok
}