TsDB2/DB/index/btree/secondary.go

package btree

import (
	"TsDB2/DB/page"
	"encoding/binary"
	"errors"
	"unsafe"
)

// ======================== 浜岀骇 B+ 鏍戠储寮曟搷浣?========================
// 鍙惰妭鐐?   [key(keyLen) + PKID(8B)]     entryLen = keyLen + 8
// 鍐呰妭鐐?   [key(keyLen) + childPID(4B)] entryLen = keyLen + 4
// 鍒嗚: 缁熶竴 50/50 (涓嶅仛椤哄簭浼樺寲, 浜岀骇绱㈠紩閿彲鑳介潪鑷)
// 鍚堝苟: InnoDB 椋庢牸 鈥?鍒犻櫎鍚庢鏌ュ綋鍓嶉〉+鍏勫紵椤佃兘鍚︽斁杩涗竴椤?

// ----- 鑺傜偣 I/O -----

func sbtNew(handle page.HandleDb, isLeaf bool, entryLen int) *btNode {
	pid := handle.NewPage()
	buf := make([]byte, page.IndexPageSize)
	info := (*page.IndexHead)(unsafe.Pointer(&buf[0]))
	info.Hdr.Type = page.IndexDataIndex
	info.Hdr.PageID = pid
	info.Hdr.UseSize = uint32(page.IndexHeadSize)
	info.IndexType = page.IndexTypeBTree
	info.EntryLen = uint16(entryLen)
	info.Count = 0
	if isLeaf {
		info.IsLeaf = 1
	}
	return &btNode{buf: buf, info: info, entries: buf[page.IndexHeadSize:]}
}

func sbtLoad(handle page.HandleDb, pid uint32) *btNode {
	if pid == 0 {
		return nil
	}
	buf := make([]byte, page.IndexPageSize)
	if err := handle.ReadPage(pid, buf); err != nil {
		return nil
	}
	info := (*page.IndexHead)(unsafe.Pointer(&buf[0]))
	return &btNode{buf: buf, info: info, entries: buf[page.IndexHeadSize:]}
}

func sbtSave(handle page.HandleDb, n *btNode) {
	_ = handle.SavePage(n.info.Hdr.PageID, n.buf)
}

// ----- 鏉＄洰鎿嶄綔 -----

func sbtEntry(n *btNode, i int) []byte {
	el := int(n.info.EntryLen)
	return n.entries[i*el : (i+1)*el]
}

func sbtLeafKey(e []byte) []byte          { return e[:len(e)-8] }
func sbtPKID(e []byte) uint64             { return binary.BigEndian.Uint64(e[len(e)-8:]) }
func sbtIntKey(e []byte) []byte           { return e[:len(e)-4] }
func sbtChildPID(e []byte) uint32         { return binary.BigEndian.Uint32(e[len(e)-4:]) }
func sbtSetChildPID(e []byte, pid uint32) { binary.BigEndian.PutUint32(e[len(e)-4:], pid) }

func sbtCmp(a, b []byte) int {
	ml := len(a)
	if len(b) < ml {
		ml = len(b)
	}
	for i := 0; i < ml; i++ {
		if a[i] < b[i] {
			return -1
		}
		if a[i] > b[i] {
			return 1
		}
	}
	return 0
}

// ----- 鍐呰妭鐐硅矾鐢?-----

func sbtRoute(n *btNode, key []byte) uint32 {
	cnt := int(n.info.Count)
	result := 0
	for i := 0; i < cnt; i++ {
		if sbtCmp(sbtIntKey(sbtEntry(n, i)), key) <= 0 {
			result = i
		}
	}
	return sbtChildPID(sbtEntry(n, result))
}

func sbtChildIdx(n *btNode, key []byte) int {
	cnt := int(n.info.Count)
	result := 0
	for i := 0; i < cnt; i++ {
		if sbtCmp(sbtEntry(n, i)[:int(n.info.EntryLen)-4], key) <= 0 {
			result = i
		}
	}
	return result
}

// ----- 鍙惰妭鐐规搷浣?-----

func sbtInsEntry(n *btNode, entry []byte, pos int) {
	el := int(n.info.EntryLen)
	cnt := int(n.info.Count)
	for i := cnt; i > pos; i-- {
		copy(sbtEntry(n, i), sbtEntry(n, i-1))
	}
	copy(sbtEntry(n, pos), entry)
	n.info.Count++
	n.info.Hdr.UseSize += uint32(el)
}

func sbtRmEntry(n *btNode, pos int) {
	el := int(n.info.EntryLen)
	cnt := int(n.info.Count)
	for i := pos; i < cnt-1; i++ {
		copy(sbtEntry(n, i), sbtEntry(n, i+1))
	}
	n.info.Count--
	n.info.Hdr.UseSize -= uint32(el)
}

func sbtLeafFind(n *btNode, key []byte) (int, bool) {
	el := int(n.info.EntryLen)
	cnt := int(n.info.Count)
	for i := 0; i < cnt; i++ {
		if sbtCmp(sbtEntry(n, i)[:el-8], key) == 0 {
			return i, true
		}
	}
	return 0, false
}

func sbtLeafPos(n *btNode, key []byte) int {
	el := int(n.info.EntryLen)
	cnt := int(n.info.Count)
	for i := 0; i < cnt; i++ {
		if sbtCmp(sbtEntry(n, i)[:el-8], key) >= 0 {
			return i
		}
	}
	return cnt
}

// ======================== 鎼滅储 ========================

func SBTSearch(rootPID uint32, key []byte, handle page.HandleDb) (uint64, bool) {
	if rootPID == 0 {
		return 0, false
	}
	cur := rootPID
	for {
		n := sbtLoad(handle, cur)
		if n == nil {
			return 0, false
		}
		if n.isLeaf() {
			pos, ok := sbtLeafFind(n, key)
			if !ok {
				return 0, false
			}
			return sbtPKID(sbtEntry(n, pos)), true
		}
		cur = sbtRoute(n, key)
	}
}

// ======================== 鑼冨洿鎵弿 ========================

func SBTRangeScan(rootPID uint32, handle page.HandleDb,
	startKey, endKey []byte, startInc, endInc bool, reverse bool,
	callback func(row []byte) bool) error {
	if rootPID == 0 {
		return nil
	}
	if !reverse {
		cur := rootPID
		for {
			n := sbtLoad(handle, cur)
			if n == nil {
				return nil
			}
			if n.isLeaf() {
				break
			}
			cur = sbtRoute(n, startKey)
		}
		for cur != 0 {
			n := sbtLoad(handle, cur)
			if n == nil {
				break
			}
			el := int(n.info.EntryLen)
			for i := 0; i < int(n.info.Count); i++ {
				entry := sbtEntry(n, i)
				key := entry[:el-8]
				if endKey != nil {
					cmp := sbtCmp(key, endKey)
					if cmp > 0 || (cmp == 0 && !endInc) {
						return nil
					}
				}
				if startKey != nil {
					cmp := sbtCmp(key, startKey)
					if cmp < 0 || (cmp == 0 && !startInc) {
						continue
					}
				}
				pkID := entry[el-8:]
				if !callback(pkID) {
					return nil
				}
			}
			cur = n.info.Hdr.NextPage
		}
	} else {
		// 反向: 路由到 endKey 或最右叶, 沿 PrePage 遍历
		cur := rootPID
		if endKey != nil {
			for {
				n := sbtLoad(handle, cur)
				if n == nil {
					return nil
				}
				if n.isLeaf() {
					break
				}
				cur = sbtRoute(n, endKey)
			}
		} else {
			for {
				n := sbtLoad(handle, cur)
				if n == nil {
					return nil
				}
				if n.isLeaf() {
					break
				}
				cur = sbtChildPID(sbtEntry(n, int(n.info.Count)-1))
			}
		}
		for cur != 0 {
			n := sbtLoad(handle, cur)
			if n == nil {
				break
			}
			el := int(n.info.EntryLen)
			for i := int(n.info.Count) - 1; i >= 0; i-- {
				entry := sbtEntry(n, i)
				key := entry[:el-8]
				if startKey != nil {
					cmp := sbtCmp(key, startKey)
					if cmp < 0 || (cmp == 0 && !startInc) {
						return nil
					}
				}
				if endKey != nil {
					cmp := sbtCmp(key, endKey)
					if cmp > 0 || (cmp == 0 && !endInc) {
						continue
					}
				}
				pkID := entry[el-8:]
				if !callback(pkID) {
					return nil
				}
			}
			cur = n.info.Hdr.PrePage
		}
	}
	return nil
}

// ======================== 鎻掑叆 ========================

func SBTInsert(rootPID uint32, entry []byte, handle page.HandleDb) (uint32, error) {
	el := len(entry)
	if rootPID == 0 {
		leaf := sbtNew(handle, true, el)
		sbtInsEntry(leaf, entry, 0)
		sbtSave(handle, leaf)
		return leaf.pageID(), nil
	}
	var path []btPathEntry
	cur := rootPID
	for {
		n := sbtLoad(handle, cur)
		if n == nil {
			return rootPID, errors.New("sbt load failed")
		}
		if n.isLeaf() {
			key := entry[:el-8]
			ip := sbtLeafPos(n, key)
			// 二级索引允许重复键, 插入到相同 key 的条目之后
			for ip < int(n.info.Count) && sbtCmp(sbtEntry(n, ip)[:el-8], key) == 0 {
				ip++
			}
			if int(n.info.Hdr.UseSize)+el <= btSplitThreshold() {
				sbtInsEntry(n, entry, ip)
				sbtSave(handle, n)
				return rootPID, nil
			}
			cnt := int(n.info.Count)
			total := cnt + 1
			all := make([]byte, total*el)
			if ip > 0 {
				copy(all, n.entries[:ip*el])
			}
			copy(all[ip*el:], entry)
			if cnt > ip {
				copy(all[(ip+1)*el:], n.entries[ip*el:cnt*el])
			}
			mid := total / 2
			rl := sbtNew(handle, true, el)
			rc := total - mid
			copy(rl.entries, all[mid*el:total*el])
			rl.info.Count = uint32(rc)
			rl.info.Hdr.UseSize = uint32(page.IndexHeadSize) + uint32(rc*el)
			rl.info.Hdr.NextPage = n.info.Hdr.NextPage
			rl.info.Hdr.PrePage = n.info.Hdr.PageID
			n.info.Hdr.NextPage = rl.info.Hdr.PageID
			copy(n.entries, all[:mid*el])
			n.info.Count = uint32(mid)
			n.info.Hdr.UseSize = uint32(page.IndexHeadSize) + uint32(mid*el)
			sbtSave(handle, n)
			sbtSave(handle, rl)
			if rl.info.Hdr.NextPage != 0 {
				nl := sbtLoad(handle, rl.info.Hdr.NextPage)
				if nl != nil {
					nl.info.Hdr.PrePage = rl.info.Hdr.PageID
					sbtSave(handle, nl)
				}
			}
			// 使用右子树最左叶的第一个 key 作为 promoKey
			pk := sbtLeftmostKey(handle, rl.info.Hdr.PageID, el)
			nr, err := sbtPromote(handle, rootPID, path, n.info.Hdr.PageID, pk, rl.info.Hdr.PageID, el)
			return nr, err
		}
		ci := sbtChildIdx(n, entry[:el-8])
		cpid := sbtChildPID(sbtEntry(n, ci))
		path = append(path, btPathEntry{nodePID: cur, childIdx: ci})
		cur = cpid
	}
}

func sbtLeftmostKey(handle page.HandleDb, nodePID uint32, el int) []byte {
	cur := nodePID
	for {
		n := sbtLoad(handle, cur)
		if n == nil {
			return nil
		}
		if n.isLeaf() {
			e := sbtEntry(n, 0)
			keyLen := el - 8
			result := make([]byte, el-4)
			copy(result, e[:keyLen])
			return result
		}
		cur = sbtChildPID(sbtEntry(n, 0))
	}
}

func sbtPromote(handle page.HandleDb, rootPID uint32, path []btPathEntry, leftPID uint32, promoKey []byte, rightPID uint32, el int) (uint32, error) {
	if len(path) == 0 {
		nr := sbtNew(handle, false, el)
		leftKey := sbtLeftmostKey(handle, leftPID, el)
		le := make([]byte, el)
		copy(le[:el-4], leftKey)
		sbtSetChildPID(le, leftPID)
		re := make([]byte, el)
		copy(re[:el-4], promoKey)
		sbtSetChildPID(re, rightPID)
		sbtInsEntry(nr, le, 0)
		sbtInsEntry(nr, re, 1)
		sbtSave(handle, nr)
		return nr.pageID(), nil
	}
	pi := len(path) - 1
	parent := sbtLoad(handle, path[pi].nodePID)
	if parent == nil {
		return rootPID, errors.New("sbt parent load failed")
	}
	ci := path[pi].childIdx
	cnt := int(parent.info.Count)
	pos := cnt
	for i := ci + 1; i < cnt; i++ {
		if sbtCmp(sbtIntKey(sbtEntry(parent, i)), promoKey) >= 0 {
			pos = i
			break
		}
	}
	for i := cnt; i > pos; i-- {
		copy(sbtEntry(parent, i), sbtEntry(parent, i-1))
	}
	ne := make([]byte, el)
	copy(ne[:el-4], promoKey)
	sbtSetChildPID(ne, rightPID)
	copy(sbtEntry(parent, pos), ne)
	parent.info.Count++
	parent.info.Hdr.UseSize += uint32(el)
	if int(parent.info.Hdr.UseSize)+el <= btSplitThreshold() {
		sbtSave(handle, parent)
		return rootPID, nil
	}
	cnt = int(parent.info.Count)
	mid := cnt / 2
	pk2 := make([]byte, el-4)
	copy(pk2, sbtIntKey(sbtEntry(parent, mid)))
	rn := sbtNew(handle, false, el)
	rc := cnt - mid - 1
	if rc > 0 {
		copy(rn.entries, parent.entries[(mid+1)*el:cnt*el])
	}
	rn.info.Count = uint32(rc)
	rn.info.Hdr.UseSize = uint32(page.IndexHeadSize) + uint32(rc*el)
	parent.info.Count = uint32(mid)
	parent.info.Hdr.UseSize = uint32(page.IndexHeadSize) + uint32(mid*el)
	sbtSave(handle, parent)
	sbtSave(handle, rn)
	return sbtPromote(handle, rootPID, path[:pi], parent.info.Hdr.PageID, pk2, rn.info.Hdr.PageID, el)
}

// ======================== 鍒犻櫎 ========================

func SBTDelete(rootPID uint32, key []byte, handle page.HandleDb) (uint32, error) {
	if rootPID == 0 {
		return 0, page.ErrNotFound
	}
	return sbtDeleteRec(rootPID, key, nil, handle)
}

func sbtDeleteRec(nodePID uint32, key []byte, path []btPathEntry, handle page.HandleDb) (uint32, error) {
	node := sbtLoad(handle, nodePID)
	if node == nil {
		return nodePID, page.ErrNotFound
	}
	el := int(node.info.EntryLen)
	if node.isLeaf() {
		pos, ok := sbtLeafFind(node, key)
		if !ok {
			return nodePID, page.ErrNotFound
		}
		sbtRmEntry(node, pos)
		if node.info.Count == 0 && len(path) == 0 {
			sbtSave(handle, node)
			return nodePID, nil
		}
		sbtSave(handle, node)
		if len(path) > 0 && node.info.Count > 0 {
			sbtTryMerge(handle, path, node, el)
		}
		return nodePID, nil
	}
	ci := sbtChildIdx(node, key)
	childPID := sbtChildPID(sbtEntry(node, ci))
	newPath := make([]btPathEntry, len(path)+1)
	copy(newPath, path)
	newPath[len(path)] = btPathEntry{nodePID: nodePID, childIdx: ci}
	newChildPID, err := sbtDeleteRec(childPID, key, newPath, handle)
	if err != nil {
		return nodePID, err
	}
	if newChildPID != childPID {
		sbtSetChildPID(sbtEntry(node, ci), newChildPID)
	}
	if node.info.Count > 0 {
		ci2 := sbtChildIdx(node, key)
		if ci2 >= 0 && ci2 < int(node.info.Count) {
			cpID := sbtChildPID(sbtEntry(node, ci2))
			child := sbtLoad(handle, cpID)
			if child != nil && child.info.Count == 0 {
				sbtHandleUnderflow(node, ci2, el, handle)
			}
		}
	}
	if node.info.Count == 0 {
		return nodePID, nil
	}
	sbtSave(handle, node)
	return nodePID, nil
}

func sbtTryMerge(handle page.HandleDb, path []btPathEntry, node *btNode, el int) {
	if len(path) == 0 {
		return
	}
	pp := path[len(path)-1].nodePID
	parent := sbtLoad(handle, pp)
	if parent == nil {
		return
	}
	ci := -1
	for i := 0; i < int(parent.info.Count); i++ {
		if sbtChildPID(sbtEntry(parent, i)) == node.info.Hdr.PageID {
			ci = i
			break
		}
	}
	if ci < 0 {
		return
	}
	if ci > 0 {
		leftPID := sbtChildPID(sbtEntry(parent, ci-1))
		left := sbtLoad(handle, leftPID)
		if left != nil && left.isLeaf() {
			combined := uint32(page.IndexHeadSize) + (left.info.Count+node.info.Count)*uint32(el)
			threshold := uint32(btSplitThreshold())
			if combined <= threshold {
				rc := int(node.info.Count)
				copy(left.entries[int(left.info.Count)*el:], node.entries[:rc*el])
				left.info.Count += node.info.Count
				left.info.Hdr.UseSize += uint32(rc * el)
				left.info.Hdr.NextPage = node.info.Hdr.NextPage
				sbtSave(handle, left)
				if left.info.Hdr.NextPage != 0 {
					nl := sbtLoad(handle, left.info.Hdr.NextPage)
					if nl != nil {
						nl.info.Hdr.PrePage = leftPID
						sbtSave(handle, nl)
					}
				}
				sbtRmEntry(parent, ci)
				sbtSave(handle, parent)
				return
			}
		}
	}
	if ci < int(parent.info.Count)-1 {
		rightPID := sbtChildPID(sbtEntry(parent, ci+1))
		right := sbtLoad(handle, rightPID)
		if right != nil && right.isLeaf() {
			combined := uint32(page.IndexHeadSize) + (node.info.Count+right.info.Count)*uint32(el)
			threshold := uint32(btSplitThreshold())
			if combined <= threshold {
				rc := int(right.info.Count)
				copy(node.entries[int(node.info.Count)*el:], right.entries[:rc*el])
				node.info.Count += right.info.Count
				node.info.Hdr.UseSize += uint32(rc * el)
				node.info.Hdr.NextPage = right.info.Hdr.NextPage
				sbtSave(handle, node)
				if node.info.Hdr.NextPage != 0 {
					nl := sbtLoad(handle, node.info.Hdr.NextPage)
					if nl != nil {
						nl.info.Hdr.PrePage = node.info.Hdr.PageID
						sbtSave(handle, nl)
					}
				}
				sbtRmEntry(parent, ci+1)
				sbtSave(handle, parent)
				return
			}
		}
	}
}

func sbtHandleUnderflow(parent *btNode, childIdx int, el int, handle page.HandleDb) {
	childPID := sbtChildPID(sbtEntry(parent, childIdx))
	child := sbtLoad(handle, childPID)
	if child == nil {
		return
	}
	if childIdx > 0 {
		leftPID := sbtChildPID(sbtEntry(parent, childIdx-1))
		leftSib := sbtLoad(handle, leftPID)
		if leftSib != nil && int(leftSib.info.Count) > 1 {
			borrowed := make([]byte, el)
			copy(borrowed, sbtEntry(leftSib, int(leftSib.info.Count)-1))
			sbtRmEntry(leftSib, int(leftSib.info.Count)-1)
			if child.isLeaf() {
				sbtInsEntry(child, borrowed, 0)
				copy(sbtEntry(parent, childIdx)[:el-4], sbtLeafKey(sbtEntry(child, 0)))
			} else {
				sbtInsEntry(child, borrowed, 0)
				copy(sbtEntry(parent, childIdx)[:el-4], sbtIntKey(borrowed))
			}
			sbtSave(handle, leftSib)
			sbtSave(handle, child)
			sbtSave(handle, parent)
			return
		}
	}
	if childIdx < int(parent.info.Count)-1 {
		rightPID := sbtChildPID(sbtEntry(parent, childIdx+1))
		rightSib := sbtLoad(handle, rightPID)
		if rightSib != nil && int(rightSib.info.Count) > 1 {
			borrowed := make([]byte, el)
			copy(borrowed, sbtEntry(rightSib, 0))
			sbtRmEntry(rightSib, 0)
			if child.isLeaf() {
				sbtInsEntry(child, borrowed, int(child.info.Count))
				copy(sbtEntry(parent, childIdx)[:el-4], sbtLeafKey(sbtEntry(rightSib, 0)))
			} else {
				sbtInsEntry(child, borrowed, int(child.info.Count))
				copy(sbtEntry(parent, childIdx)[:el-4], sbtIntKey(borrowed))
			}
			sbtSave(handle, rightSib)
			sbtSave(handle, child)
			sbtSave(handle, parent)
			return
		}
	}
	if childIdx > 0 {
		sbtMergeChildren(handle, parent, childIdx-1, el)
	} else if childIdx < int(parent.info.Count)-1 {
		sbtMergeChildren(handle, parent, childIdx, el)
	}
}

func sbtMergeChildren(handle page.HandleDb, parent *btNode, sepIdx int, el int) {
	leftPID := sbtChildPID(sbtEntry(parent, sepIdx))
	rightPID := sbtChildPID(sbtEntry(parent, sepIdx+1))
	left := sbtLoad(handle, leftPID)
	right := sbtLoad(handle, rightPID)
	if left == nil || right == nil {
		return
	}
	if left.isLeaf() {
		rc := int(right.info.Count)
		copy(left.entries[int(left.info.Count)*el:], right.entries[:rc*el])
		left.info.Count += right.info.Count
		left.info.Hdr.UseSize += uint32(rc * el)
		left.info.Hdr.NextPage = right.info.Hdr.NextPage
	} else {
		sepEntry := make([]byte, el)
		copy(sepEntry, sbtEntry(parent, sepIdx))
		rc := int(right.info.Count)
		off := int(left.info.Count) * el
		copy(left.entries[off:], sepEntry)
		off += el
		copy(left.entries[off:], right.entries[:rc*el])
		left.info.Count += 1 + right.info.Count
		left.info.Hdr.UseSize += uint32((1 + rc) * el)
	}
	sbtSave(handle, left)
	if left.isLeaf() && left.info.Hdr.NextPage != 0 {
		nl := sbtLoad(handle, left.info.Hdr.NextPage)
		if nl != nil {
			nl.info.Hdr.PrePage = leftPID
			sbtSave(handle, nl)
		}
	}
	sbtRmEntry(parent, sepIdx+1)
	sbtSetChildPID(sbtEntry(parent, sepIdx), leftPID)
	sbtSave(handle, left)
}