Multi-Level BTree

mTree.go

package main

import (
    "bytes"
    "encoding/gob"
    "fmt"
    "github.com/google/btree"
    "log"
    "sync"
)

const Degree = 32 // Degree of the B-tree

// Item represents a single item in the B-tree.
type Item[T any] struct {
    Key   uint32
    Value T
}

// Less method for Item to satisfy btree.Item interface.
func (a Item[T]) Less(b btree.Item) bool {
    return a.Key < b.(Item[T]).Key
}

// Node represents a node in the multi-level binary tree of B-trees.
type Node[T btree.Item] struct {
    Tree  *btree.BTreeG[T]
    Left  *Node[T]
    Right *Node[T]
    dirty bool // Indicates if the node has been updated
    mu    sync.RWMutex
}

// MultiLevelBTree represents a multi-level binary tree of B-trees.
type MultiLevelBTree[T btree.Item] struct {
    root *Node[T]
    mu   sync.RWMutex
    splitThreshold int
    mergeThreshold int
}

// NewMultiLevelBTree initializes a new MultiLevelBTree.
func NewMultiLevelBTree[T btree.Item](splitThreshold, mergeThreshold int) *MultiLevelBTree[T] {
    return &MultiLevelBTree[T]{
        root: &Node[T]{Tree: btree.NewG[T](Degree, nil)},
        splitThreshold: splitThreshold,
        mergeThreshold: mergeThreshold,
    }
}

// insertItem inserts an item into the appropriate B-tree in the binary tree.
func (mlbt *MultiLevelBTree[T]) insertItem(item T) {
    mlbt.mu.Lock()
    defer mlbt.mu.Unlock()
    mlbt.root = insertIntoNode(mlbt.root, item, mlbt.splitThreshold)
}

func insertIntoNode[T btree.Item](node *Node[T], item T, splitThreshold int) *Node[T] {
    if node == nil {
        newNode := &Node[T]{Tree: btree.NewG[T](Degree, nil)}
        newNode.Tree.ReplaceOrInsert(item)
        newNode.dirty = true
        return newNode
    }

    node.mu.Lock()
    defer node.mu.Unlock()

    // Insert into the current B-tree
    node.Tree.ReplaceOrInsert(item)
    node.dirty = true

    // If the current tree grows beyond a threshold, split it
    if node.Tree.Len() > splitThreshold {
        middle := node.Tree.DeleteMax().(T)
        newNode := &Node[T]{Tree: btree.NewG[T](Degree, nil)}
        newNode.Tree.ReplaceOrInsert(middle)
        newNode.Left = node
        newNode.Right = insertIntoNode(nil, middle, splitThreshold)
        newNode.dirty = true
        return newNode
    }

    if item.Less(node.Tree.Min().(T)) {
        node.Left = insertIntoNode(node.Left, item, splitThreshold)
    } else {
        node.Right = insertIntoNode(node.Right, item, splitThreshold)
    }

    return node
}

// searchItem searches for an item across all B-trees in the binary tree.
func (mlbt *MultiLevelBTree[T]) searchItem(key T) *T {
    mlbt.mu.RLock()
    defer mlbt.mu.RUnlock()
    return searchInNode(mlbt.root, key)
}

func searchInNode[T btree.Item](node *Node[T], key T) *T {
    if node == nil {
        return nil
    }

    node.mu.RLock()
    defer node.mu.RUnlock()

    item := node.Tree.Get(key)
    if item != nil {
        result := item.(T)
        return &result
    }

    if key.Less(node.Tree.Min().(T)) {
        return searchInNode(node.Left, key)
    } else {
        return searchInNode(node.Right, key)
    }
}

// deleteItem deletes an item from the appropriate B-tree in the binary tree.
func (mlbt *MultiLevelBTree[T]) deleteItem(key T) {
    mlbt.mu.Lock()
    defer mlbt.mu.Unlock()
    mlbt.root = deleteFromNode(mlbt.root, key, mlbt.mergeThreshold)
}

func deleteFromNode[T btree.Item](node *Node[T], key T, mergeThreshold int) *Node[T] {
    if node == nil {
        return nil
    }

    node.mu.Lock()
    defer node.mu.Unlock()

    node.Tree.Delete(key)
    node.dirty = true

    // If the current tree becomes too small, handle balancing or merging
    if node.Tree.Len() < mergeThreshold {
        // Merge with left child if possible
        if node.Left != nil {
            node.Left.Tree.Ascend(func(i T) bool {
                node.Tree.ReplaceOrInsert(i)
                return true
            })
            node.Left = nil
            node.dirty = true
        }
        // Merge with right child if possible
        if node.Right != nil {
            node.Right.Tree.Ascend(func(i T) bool {
                node.Tree.ReplaceOrInsert(i)
                return true
            })
            node.Right = nil
            node.dirty = true
        }
    }

    if key.Less(node.Tree.Min().(T)) {
        node.Left = deleteFromNode(node.Left, key, mergeThreshold)
    } else {
        node.Right = deleteFromNode(node.Right, key, mergeThreshold)
    }

    return node
}

// serializeNode serializes a Node to a byte slice in a thread-safe manner.
func serializeNode[T btree.Item](node *Node[T]) ([]byte, error) {
    node.mu.RLock()
    defer node.mu.RUnlock()

    var buf bytes.Buffer
    encoder := gob.NewEncoder(&buf)
    err := encoder.Encode(node)
    if err != nil {
        return nil, err
    }
    return buf.Bytes(), nil
}

// deserializeNode deserializes a byte slice into a Node in a thread-safe manner.
func deserializeNode[T btree.Item](data []byte) (*Node[T], error) {
    var node Node[T]
    buf := bytes.NewBuffer(data)
    decoder := gob.NewDecoder(buf)
    err := decoder.Decode(&node)
    if err != nil {
        return nil, err
    }
    return &node, nil
}

func main() {
    mlbt := NewMultiLevelBTree[Item[string]](100, 20) // Configurable thresholds

    // Insert items
    for i := 0; i < 500; i++ {
        mlbt.insertItem(Item[string]{Key: uint32(i), Value: fmt.Sprintf("Value %d", i)})
    }

    // Serialize the root node
    serializedData, err := serializeNode(mlbt.root)
    if err != nil {
        log.Fatalf("Failed to serialize node: %v", err)
    }

    fmt.Println("Serialized data length:", len(serializedData))

    // Deserialize the root node
    deserializedNode, err := deserializeNode[Item[string]](serializedData)
    if err != nil {
        log.Fatalf("Failed to deserialize node: %v", err)
    }

    mlbt.root = deserializedNode

    // Search for an item
    key := Item[string]{Key: 250}
    item := mlbt.searchItem(key)
    if item != nil {
        fmt.Printf("Found item: Key = %d, Value = %s\n", item.Key, item.Value)
    } else {
        fmt.Printf("Item with Key = %d not found\n", key.Key)
    }

    // Delete an item
    mlbt.deleteItem(key)
    item = mlbt.searchItem(key)
    if item != nil {
        fmt.Printf("Found item: Key = %d, Value = %s\n", item.Key, item.Value)
    } else {
        fmt.Printf("Item with Key = %d not found\n", key.Key)
    }
}