TsDB2/DB/core/index/compose_handler.go

package index

import (
	"TsDB2/DB/core/index/composepage"
	"TsDB2/DB/core/page"
	"TsDB2/DB/core/util"
	"bytes"
	"encoding/binary"
	"unsafe"
)

func init() {
	Compose.RegisterHandler(page.IndexTypeCompose, &ComposeChainHandler{})
}

// ======================== IEntry: 复合索引条目 ========================

// ComposeEntry 通用复合索引结构体
// Keys: 存储 N 个索引键 (uint64)
// PKID: 主键ID（8字节），通过 B+ 树解析数据位置
type ComposeEntry struct {
	Keys []uint64 // 任意多个索引键（1个、2个、3..都支持）
	PKID uint64   // 主键ID (8字节)
}

func (e *ComposeEntry) DataEntry() IDateEntry {
	return &composepage.ComposeEntryPage{}
}

// EntryLength 自动计算长度：8*键数量 + 8 (PK_ID)
func (e *ComposeEntry) EntryLength() uint16 {
	return uint16(len(e.Keys)*8 + 8)
}

// Compare 通用复合索引比较逻辑：自动按顺序比较所有键（不含末尾 PK_ID）
func (e *ComposeEntry) Compare(midEntry []byte, key []byte) int {
	entryLen := int(e.EntryLength())
	keyCount := (entryLen - 8) / 8
	for i := 0; i < keyCount; i++ {
		offset := i * 8
		f1 := binary.BigEndian.Uint64(midEntry[offset : offset+8])
		f2 := binary.BigEndian.Uint64(key[offset : offset+8])
		cmp := util.CompareUint64(f1, f2)
		if cmp != 0 {
			return cmp
		}
	}
	return 0
}

// Serialize 自动序列化 N 个键 + PK_ID
func (e *ComposeEntry) Serialize(ie IEntry) []byte {
	ent, ok := ie.(*ComposeEntry)
	if !ok {
		return nil
	}
	keyCount := len(ent.Keys)
	buf := make([]byte, keyCount*8+8)
	for i, k := range ent.Keys {
		binary.BigEndian.PutUint64(buf[i*8:(i+1)*8], k)
	}
	binary.BigEndian.PutUint64(buf[keyCount*8:], ent.PKID)
	return buf
}

// Deserialize 自动反序列化 N 个键 + PK_ID
func (e *ComposeEntry) Deserialize(buf []byte) IEntry {
	keyCount := (len(buf) - 8) / 8
	keys := make([]uint64, 0, keyCount)
	for i := 0; i < keyCount; i++ {
		keys = append(keys, binary.BigEndian.Uint64(buf[i*8:(i+1)*8]))
	}
	return &ComposeEntry{
		Keys: keys,
		PKID: binary.BigEndian.Uint64(buf[keyCount*8:]),
	}
}

// Insert 插入索引
func (e *ComposeEntry) Insert(entries *[]byte, data []byte) error {
	*entries = append(*entries, data...)
	return nil
}

// Search 二分查找（通用 N 字段复合索引）
func (e *ComposeEntry) Search(entries []byte, key []byte) (uint32, bool) {
	entryLen := int(e.EntryLength())
	pos, found := util.BinarySearch(entries, entryLen, key, e.Compare)
	if !found {
		return 0, false
	}
	return pos, true
}

// ======================== IndexHandler: 复合二级索引处理器 ========================

// ComposeChainHandler — 复合二级索引 (page.IndexTypeCompose)
// 扁平索引页链, 条目 [composeKey(N*8B) + PKID(8B)]
type ComposeChainHandler struct{}

func (h *ComposeChainHandler) Insert(handle page.HandleDb, entryPage *uint32, indexKey []byte, pkID uint64, rowData []byte, keyLen int) (uint32, error) {
	dataEntry := h.getDataEntry()
	curPageID := *entryPage
	lastIdxPageID := uint32(0)
	for curPageID != 0 {
		buf := make([]byte, page.IndexPageSize)
		if err := handle.ReadPage(curPageID, buf); err != nil {
			return 0, err
		}
		hdr := (*page.IndexHead)(unsafe.Pointer(&buf[0]))
		pageEntryLen := int(hdr.EntryLen)
		count := int(hdr.Count)
		idxPageHasRoom := hdr.Hdr.UseSize+uint32(pageEntryLen) <= page.IndexPageSize
		if !idxPageHasRoom {
			lastIdxPageID = curPageID
			curPageID = hdr.Hdr.NextPage
			continue
		}
		entries := buf[page.IndexHeadSize : page.IndexHeadSize+(count+1)*pageEntryLen]
		err := dataEntry.Insert(handle, entries, pageEntryLen, count, indexKey, pkID, rowData)
		if err == nil {
			hdr.Count++
			hdr.Hdr.UseSize += uint32(pageEntryLen)
			if err := handle.SavePage(curPageID, buf); err != nil {
				return 0, err
			}
			return *entryPage, nil
		}
		lastIdxPageID = curPageID
		curPageID = hdr.Hdr.NextPage
	}
	entryLen := keyLen + 8
	newEntry := make([]byte, entryLen)
	copy(newEntry[:keyLen], indexKey)
	binary.BigEndian.PutUint64(newEntry[keyLen:], pkID)
	return *entryPage, writeToNewIndexPage(handle, entryPage, newEntry, lastIdxPageID)
}

func (h *ComposeChainHandler) Search(handle page.HandleDb, entryPage uint32, indexKey []byte, keyLen int, resolve RowResolver) ([]byte, bool) {
	dataEntry := h.getDataEntry()
	curPageID := entryPage
	for curPageID != 0 {
		buf := make([]byte, page.IndexPageSize)
		if err := handle.ReadPage(curPageID, buf); err != nil {
			break
		}
		hdr := (*page.IndexHead)(unsafe.Pointer(&buf[0]))
		entries := buf[page.IndexHeadSize : page.IndexHeadSize+int(hdr.Count)*int(hdr.EntryLen)]
		entryLen := int(hdr.EntryLen)
		pkBuf, found := dataEntry.Search(handle, entries, entryLen, indexKey)
		if found {
			// 二级索引不存行数据: 内部回表取整行
			if resolve != nil && len(pkBuf) == 8 {
				if row, ok := resolve(binary.BigEndian.Uint64(pkBuf)); ok {
					return row, true
				}
				return nil, false
			}
			return pkBuf, true
		}
		curPageID = hdr.Hdr.NextPage
	}
	return nil, false
}

// Delete 按 (key, pkID) 精确删除复合索引条目。
// 链页里的条目是 [composeKey(N*8B) + pkID(8B)] 的扁平数组, 直接原地摘除即可。
func (h *ComposeChainHandler) Delete(handle page.HandleDb, entryPage *uint32, indexKey []byte, pkID uint64, keyLen int) error {
	curPageID := *entryPage
	for curPageID != 0 {
		buf := make([]byte, page.IndexPageSize)
		if err := handle.ReadPage(curPageID, buf); err != nil {
			return err
		}
		hdr := (*page.IndexHead)(unsafe.Pointer(&buf[0]))
		entryLen := int(hdr.EntryLen)
		count := int(hdr.Count)
		if entryLen > 8 && count > 0 {
			base := page.IndexHeadSize
			keyPart := entryLen - 8
			for i := 0; i < count; i++ {
				off := base + i*entryLen
				if off+entryLen > len(buf) {
					break
				}
				if composeCmp(buf[off:off+keyPart], indexKey) != 0 {
					continue
				}
				if pkID != 0 && binary.BigEndian.Uint64(buf[off+keyPart:off+entryLen]) != pkID {
					continue
				}
				copy(buf[off:base+count*entryLen], buf[off+entryLen:base+count*entryLen])
				copy(buf[base+(count-1)*entryLen:base+count*entryLen], make([]byte, entryLen))
				hdr.Count--
				hdr.Hdr.UseSize -= uint32(entryLen)
				if err := handle.SavePage(curPageID, buf); err != nil {
					return err
				}
				return nil
			}
		}
		curPageID = hdr.Hdr.NextPage
	}
	return page.ErrNotFound
}

func (h *ComposeChainHandler) Update(handle page.HandleDb, entryPage *uint32, oldKey, newKey, rowData []byte, pkID uint64, keyLen int) error {
	// 二级索引条目只存 [key + pkID]: 键未变则无事可做, 键变了则删旧插新
	if bytes.Equal(oldKey, newKey) {
		return nil
	}
	if oldKey != nil {
		if err := h.Delete(handle, entryPage, oldKey, pkID, len(oldKey)); err != nil && err != page.ErrNotFound {
			return err
		}
	}
	_, err := h.Insert(handle, entryPage, newKey, pkID, rowData, len(newKey))
	return err
}

// IsRowStore 复合二级索引不承载行数据
func (h *ComposeChainHandler) IsRowStore() bool { return false }

// ResolveRow 复合二级索引无法按 pkID 取行, 由调用方的回表函数负责
func (h *ComposeChainHandler) ResolveRow(handle page.HandleDb, entryPage uint32, pkID uint64, cache *[]DPRange) ([]byte, bool) {
	return nil, false
}

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

func (h *ComposeChainHandler) ReadDataPageRows(handle page.HandleDb, dpID uint32) [][]byte {
	return nil
}

// RangeScan 复合二级索引范围扫描
// 条目格式: [key(keyLen) + pkID(8B BE)]
// 回调收到 8B BigEndian pkID
func (h *ComposeChainHandler) RangeScan(handle page.HandleDb, entryPage uint32,
	startKey, endKey []byte, startInclusive, endInclusive bool, reverse bool,
	callback func(row []byte) bool) error {
	if entryPage == 0 {
		return nil
	}
	type pageData struct {
		entries []byte
		el      int
		count   int
	}
	var pages []pageData
	curPageID := entryPage
	for curPageID != 0 {
		buf := make([]byte, page.IndexPageSize)
		if err := handle.ReadPage(curPageID, buf); err != nil {
			break
		}
		hdr := (*page.IndexHead)(unsafe.Pointer(&buf[0]))
		el := int(hdr.EntryLen)
		count := int(hdr.Count)
		entries := make([]byte, count*el)
		copy(entries, buf[page.IndexHeadSize:page.IndexHeadSize+count*el])
		pages = append(pages, pageData{entries: entries, el: el, count: count})
		curPageID = hdr.Hdr.NextPage
	}
	if !reverse {
		for p := 0; p < len(pages); p++ {
			pg := &pages[p]
			for i := 0; i < pg.count; i++ {
				entry := pg.entries[i*pg.el : (i+1)*pg.el]
				key := entry[:pg.el-8]
				if endKey != nil {
					cmp := composeCmp(key, endKey)
					if cmp > 0 || (cmp == 0 && !endInclusive) {
						return nil
					}
				}
				if startKey != nil {
					cmp := composeCmp(key, startKey)
					if cmp < 0 || (cmp == 0 && !startInclusive) {
						continue
					}
				}
				pkID := entry[pg.el-8:]
				if !callback(pkID) {
					return nil
				}
			}
		}
	} else {
		for p := len(pages) - 1; p >= 0; p-- {
			pg := &pages[p]
			for i := pg.count - 1; i >= 0; i-- {
				entry := pg.entries[i*pg.el : (i+1)*pg.el]
				key := entry[:pg.el-8]
				if startKey != nil {
					cmp := composeCmp(key, startKey)
					if cmp < 0 || (cmp == 0 && !startInclusive) {
						return nil
					}
				}
				if endKey != nil {
					cmp := composeCmp(key, endKey)
					if cmp > 0 || (cmp == 0 && !endInclusive) {
						continue
					}
				}
				pkID := entry[pg.el-8:]
				if !callback(pkID) {
					return nil
				}
			}
		}
	}
	return nil
}

// composeCmp 字节切片比较
func composeCmp(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
		}
	}
	if len(a) < len(b) {
		return -1
	}
	if len(a) > len(b) {
		return 1
	}
	return 0
}

func (h *ComposeChainHandler) deleteEntry(handle page.HandleDb, entryPage *uint32, indexKey []byte) error {
	dataEntry := h.getDataEntry()
	curPageID := *entryPage
	for curPageID != 0 {
		buf := make([]byte, page.IndexPageSize)
		if err := handle.ReadPage(curPageID, buf); err != nil {
			return err
		}
		hdr := (*page.IndexHead)(unsafe.Pointer(&buf[0]))
		entries := buf[page.IndexHeadSize : page.IndexHeadSize+int(hdr.Count)*int(hdr.EntryLen)]
		entryLen := int(hdr.EntryLen)
		err := dataEntry.Delete(handle, entries, entryLen, indexKey)
		if err == nil {
			hdr.Hdr.UseSize -= uint32(entryLen)
			if err := handle.SavePage(curPageID, buf); err != nil {
				return err
			}
			return nil
		}
		curPageID = hdr.Hdr.NextPage
	}
	return page.ErrNotFound
}

func (h *ComposeChainHandler) getDataEntry() IDateEntry {
	handler, _ := Compose.GetEntryHandler(page.IndexTypeCompose)
	return handler.DataEntry()
}

// CollectPageIDs 实现 index.PageCollector: 顺着 NextPage 走完整条扁平页链。
// seen 兜底防环 (链被写坏时 NextPage 可能绕回自身, 否则这里会死循环)。
func (h *ComposeChainHandler) CollectPageIDs(handle page.HandleDb, entryPage uint32) []uint32 {
	var out []uint32
	seen := make(map[uint32]bool)
	for cur := entryPage; cur != 0 && !seen[cur]; {
		seen[cur] = true
		out = append(out, cur)
		buf := make([]byte, page.IndexPageSize)
		if err := handle.ReadPage(cur, buf); err != nil {
			// 读不到就到此为止: 继续按 buf 里的残留指针走会收集到无关的页号
			break
		}
		cur = (*page.IndexHead)(unsafe.Pointer(&buf[0])).Hdr.NextPage
	}
	return out
}