TsDB2/DB/core/index/primary_handler.go

package index

import (
	"TsDB2/DB/core/index/btree"
	"TsDB2/DB/core/index/datapage"
	"TsDB2/DB/core/page"
	"encoding/binary"
)

// PrimaryHandler — 主键 B+ 树索引 (IndexTypePrimary)
// 存储: storeRow = [pkID(8B LE) + rowData]
type PrimaryHandler struct{}

func init() {
	Compose.RegisterHandler(page.IndexTypePrimary, &PrimaryHandler{})
}

func (h *PrimaryHandler) Insert(handle page.HandleDb, entryPage *uint32, indexKey []byte, pkID uint64, rowData []byte, keyLen int) (uint32, error) {
	pkBuf := make([]byte, 8)
	binary.BigEndian.PutUint64(pkBuf, pkID)
	storeRowData := append(pkBuf, rowData...)
	// 主键 B+ 树始终以 pkID(8B BE) 为 key, 与 Search/Delete/Update 保持一致
	return h.insertPrebuilt(handle, entryPage, pkBuf, storeRowData)
}

// InsertPrebuilt 插入已构建好的 storeRowData (避免重复分配)
// storeRowData = [pkID(8B BE) + rowData]
func (h *PrimaryHandler) InsertPrebuilt(handle page.HandleDb, entryPage *uint32, indexKey []byte, storeRowData []byte) (uint32, error) {
	return h.insertPrebuilt(handle, entryPage, indexKey, storeRowData)
}

// BuildStoreRow 用复用缓冲区构建 storeRowData = [pkID(8B BE) + rowData]
// 返回的切片与 buf 共享底层数组, 调用方需立即使用
func (h *PrimaryHandler) BuildStoreRow(buf []byte, pkID uint64, rowData []byte) []byte {
	binary.BigEndian.PutUint64(buf[:8], pkID)
	copy(buf[8:], rowData)
	return buf[:8+len(rowData)]
}

func (h *PrimaryHandler) insertPrebuilt(handle page.HandleDb, entryPage *uint32, indexKey []byte, storeRowData []byte) (uint32, error) {
	rootPID := *entryPage
	rowSize := uint32(len(storeRowData)) + 4 // 4B len prefix + row data

	// 尝试复用 DataPage: 只在路由目标叶子内查找
	if rootPID != 0 {
		if dpID, found := btree.FindLastLeafDP(handle, rootPID, rowSize, indexKey); found {
			dp, err := datapage.Load(handle, dpID)
			if err == nil {
				dp.AppendRow(storeRowData)
				dp.Save()
				newRoot, _, dup, err := btree.BTreePKInsert(rootPID, indexKey, dpID, handle)
				if err != nil {
					return 0, err
				}
				if !dup {
					*entryPage = newRoot
				}
				return newRoot, nil
			}
		}
	}

	// 无可用空间或树为空: 创建新 DataPage
	dp := datapage.New(handle)
	dp.AppendRow(storeRowData)
	dp.Save()
	newRoot, _, _, err := btree.BTreePKInsert(rootPID, indexKey, dp.PageID(), handle)
	if err != nil {
		return 0, err
	}
	*entryPage = newRoot
	return newRoot, nil
}

func (h *PrimaryHandler) Search(handle page.HandleDb, entryPage uint32, indexKey []byte, keyLen int) ([]byte, bool) {
	rootPID := entryPage
	if rootPID == 0 {
		return nil, false
	}
	dpID, found := btree.BTreePKSearch(rootPID, indexKey, handle)
	if !found {
		return nil, false
	}
	dp, err := datapage.Load(handle, dpID)
	if err != nil {
		return nil, false
	}
	return dp.FindRowByIndexKey(indexKey, 9) // skip pkID(8B) + schemaVer(1B)
}

func (h *PrimaryHandler) SearchByPK(handle page.HandleDb, entryPage uint32, pkID uint64) ([]byte, bool) {
	rootPID := entryPage
	if rootPID == 0 {
		return nil, false
	}
	pkKeyBE := make([]byte, 8)
	binary.BigEndian.PutUint64(pkKeyBE, pkID)
	dpID, found := btree.BTreePKSearch(rootPID, pkKeyBE, handle)
	if found {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			if rowData, ok := dp.FindRowByPK(pkKeyBE); ok {
				return rowData, true
			}
		}
	}
	// 路由失败或未找到, 扫描所有叶子
	return btree.ScanAllLeavesForPK(rootPID, handle, pkKeyBE)
}

func (h *PrimaryHandler) Delete(handle page.HandleDb, entryPage *uint32, indexKey []byte, keyLen int) error {
	rootPID := *entryPage
	if rootPID == 0 {
		return page.ErrNotFound
	}
	dpID, found := btree.BTreePKSearch(rootPID, indexKey, handle)
	if found {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			deleted := false
			if dp.DeleteByIndexKey(indexKey, 8) {
				deleted = true
			} else if dp.DeleteByPK(indexKey) {
				deleted = true
			}
			if deleted {
				if dp.AllDeleted() {
					newRoot, _, err := btree.BTreePKDelete(rootPID, indexKey, handle)
					if err != nil {
						return err
					}
					if newRoot != rootPID {
						*entryPage = newRoot
					}
					if handle.FreePage != nil {
						handle.FreePage(dpID)
					}
				} else {
					dp.Save()
				}
				return nil
			}
		}
	}
	// 索引键路由失败, 扫描所有叶子
	return btree.ScanAndDelete(rootPID, handle, indexKey, indexKey)
}

func (h *PrimaryHandler) DeleteByPK(handle page.HandleDb, entryPage *uint32, pkID uint64) error {
	rootPID := *entryPage
	if rootPID == 0 {
		return page.ErrNotFound
	}
	pkKeyBE := make([]byte, 8)
	binary.BigEndian.PutUint64(pkKeyBE, pkID)
	dpID, found := btree.BTreePKSearch(rootPID, pkKeyBE, handle)
	if found {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			if dp.DeleteByPK(pkKeyBE) {
				if dp.AllDeleted() {
					newRoot, _, err := btree.BTreePKDelete(rootPID, pkKeyBE, handle)
					if err != nil {
						return err
					}
					if newRoot != rootPID {
						*entryPage = newRoot
					}
					if handle.FreePage != nil {
						handle.FreePage(dpID)
					}
				} else {
					dp.Save()
				}
				return nil
			}
		}
	}
	// 路由失败, 扫描所有叶子
	return btree.ScanAndDeleteByPK(rootPID, handle, pkKeyBE)
}

func (h *PrimaryHandler) Update(handle page.HandleDb, entryPage *uint32, oldKey []byte, rowData []byte, pkID uint64, keyLen int) error {
	rootPID := *entryPage
	if rootPID == 0 {
		return page.ErrNotFound
	}
	pkBuf := make([]byte, 8)
	binary.BigEndian.PutUint64(pkBuf, pkID)
	storeRowData := append(pkBuf, rowData...)
	pkKeyBE := make([]byte, 8)
	binary.BigEndian.PutUint64(pkKeyBE, pkID)
	// 先用 B+ 树路由定位 DataPage
	dpID, found := btree.BTreePKSearch(rootPID, pkKeyBE, handle)
	if found {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			if err := dp.UpdateByPK(pkKeyBE, storeRowData); err == nil {
				return nil
			}
		}
	}
	// 路由失败, 回退扫描
	return btree.ScanAndUpdate(rootPID, handle, pkKeyBE, storeRowData)
}

func (h *PrimaryHandler) GetDataPageIDs(handle page.HandleDb, entryPage uint32) []uint32 {
	return btree.GetDataPageIDs(handle, entryPage)
}

func (h *PrimaryHandler) ReadDataPageRows(handle page.HandleDb, dpID uint32) [][]byte {
	dp, err := datapage.Load(handle, dpID)
	if err != nil {
		return nil
	}
	return dp.GetAllRows()
}

func (h *PrimaryHandler) RangeScan(handle page.HandleDb, entryPage uint32,
	startKey, endKey []byte, startInclusive, endInclusive bool, reverse bool,
	callback func(row []byte) bool) error {
	return btree.BTreePKRangeScan(entryPage, handle, startKey, endKey, startInclusive, endInclusive, reverse, callback)
}