TsDB2/DB/core/index/btree/secondary_test.go

package btree

import (
	"bytes"
	"encoding/binary"
	"math/rand"
	"sort"
	"testing"

	"TsDB2/DB/core/page"
)

// ============================================================================
// 复合分隔键 (key, pkID) 回归测试。
//
// 二级树的全局不变量: 所有叶条目按复合键 (key, pkID) 有序; 内部条目的前
// EntryLen-4 字节 = 其子树的最小复合键。一旦破坏:
//   - sbtChildIdxPK 的点路由会落到错误子树 → (key,pkID) 精确删除漏删;
//   - sbtLeafLowerBoundPK 的叶内二分会漏条目 → SBTSearchAll 返回不全。
// 本测试用「乱序插入 → 乱序删除 → 重插入 → 删空」全流程 + 全树不变量校验钉死它。
// ============================================================================

// checkSubtree 递归校验以 pid 为根的子树:
//   - 叶内/内部条目按复合键严格递增;
//   - 内部条目复合键 == 其子树的最小复合键;
//   - 不存在 "空的非根子树" (删除路径的下溢处理应消掉它们);
//   - 叶子链 (PrePage/NextPage) 与 DFS 顺序一致。
// 返回该子树的最小复合键; 空树返回 nil。
func checkSubtree(t *testing.T, h page.HandleDb, pid uint32) []byte {
	t.Helper()
	n := sbtLoad(h, pid)
	if n == nil {
		t.Fatalf("checkSubtree: 读不出节点 pid=%d", pid)
	}
	cnt := int(n.info.Count)
	if n.isLeaf() {
		for i := 1; i < cnt; i++ {
			a, b := sbtEntry(n, i-1), sbtEntry(n, i)
			if sbtCmpKeyPK(sbtLeafKey(a), sbtPKID(a), sbtLeafKey(b), sbtPKID(b)) >= 0 {
				t.Fatalf("叶 %d 内复合键乱序 @%d", pid, i)
			}
		}
		if cnt == 0 {
			return nil
		}
		return sbtEntry(n, 0) // 叶条目本身就是 [key|pkID]
	}
	if cnt == 0 {
		return nil // 仅允许根为空内部节点 (上层归一化会换掉它)
	}
	var first []byte
	for i := 0; i < cnt; i++ {
		e := sbtEntry(n, i)
		if i > 0 {
			p := sbtEntry(n, i-1)
			if sbtCmpKeyPK(sbtIntKey(p), sbtIntPKID(p), sbtIntKey(e), sbtIntPKID(e)) >= 0 {
				t.Fatalf("内部节点 %d 复合键乱序 @%d", pid, i)
			}
		}
		childPID := sbtChildPID(e)
		child := sbtLoad(h, childPID)
		if child == nil {
			t.Fatalf("内部节点 %d 条目 %d 指向读不出的子页 %d", pid, i, childPID)
		}
		if child.info.Count == 0 {
			t.Fatalf("存在空的非根子树: 父 %d 条目 %d → 子 %d (下溢处理漏了它)", pid, i, childPID)
		}
		cm := checkSubtree(t, h, childPID)
		if !bytes.Equal(cm, e[:len(e)-4]) {
			t.Fatalf("父条目复合键 != 子树最小复合键: 父 %d 条目 %d → 子 %d\n父=%v 子=%v",
				pid, i, childPID, e[:len(e)-4], cm)
		}
		if i == 0 {
			first = append([]byte(nil), cm...)
		}
	}
	return first
}

// dfsLeafOrder DFS 收集叶页号 (复合键序)。
func dfsLeafOrder(h page.HandleDb, pid uint32) []uint32 {
	n := sbtLoad(h, pid)
	if n == nil {
		return nil
	}
	if n.isLeaf() {
		return []uint32{pid}
	}
	var out []uint32
	for i := 0; i < int(n.info.Count); i++ {
		out = append(out, dfsLeafOrder(h, sbtChildPID(sbtEntry(n, i)))...)
	}
	return out
}

func checkLeafChain(t *testing.T, h page.HandleDb, root uint32) {
	t.Helper()
	leaves := dfsLeafOrder(h, root)
	if len(leaves) == 0 {
		return
	}
	for i, pid := range leaves {
		n := sbtLoad(h, pid)
		if n == nil || !n.isLeaf() {
			t.Fatalf("DFS 叶序里出现非叶/读不出的页 %d", pid)
		}
		wantPre, wantNext := uint32(0), uint32(0)
		if i > 0 {
			wantPre = leaves[i-1]
		}
		if i+1 < len(leaves) {
			wantNext = leaves[i+1]
		}
		if n.info.Hdr.PrePage != wantPre || n.info.Hdr.NextPage != wantNext {
			t.Fatalf("叶子链断裂: 叶[%d]=%d want(pre=%d,next=%d) got(pre=%d,next=%d)",
				i, pid, wantPre, wantNext, n.info.Hdr.PrePage, n.info.Hdr.NextPage)
		}
	}
}

// TestSecondaryCompositeOrderStress 乱序插入/删除/重插入全程校验复合键不变量。
func TestSecondaryCompositeOrderStress(t *testing.T) {
	const (
		keyLen = 8
		keys   = 50  // 每个 key 400 条重复 → 单 key 区段 6400B, 必然跨多个叶
		perKey = 400
	)
	total := keys * perKey

	m := newMemHandle()
	h := m.handle()

	type pair struct {
		key uint64
		pk  uint64
	}
	all := make([]pair, 0, total)
	for k := uint64(1); k <= keys; k++ {
		for p := uint64(1); p <= perKey; p++ {
			all = append(all, pair{key: k, pk: p*k + p%97}) // pkID 乱序于插入序列
		}
	}
	rng := rand.New(rand.NewSource(20260914))
	rng.Shuffle(len(all), func(i, j int) { all[i], all[j] = all[j], all[i] })

	mkEntry := func(p pair) []byte {
		e := make([]byte, keyLen+8)
		binary.BigEndian.PutUint64(e[:keyLen], p.key)
		binary.BigEndian.PutUint64(e[keyLen:], p.pk)
		return e
	}
	mkKey := func(k uint64) []byte {
		b := make([]byte, keyLen)
		binary.BigEndian.PutUint64(b, k)
		return b
	}

	// 活跃集合 (期望状态): key → 有序 pkID 列表
	live := make(map[uint64][]uint64)
	var root uint32

	// ---- 阶段 1: 乱序全量插入 ----
	for _, p := range all {
		var err error
		root, err = SBTInsert(root, mkEntry(p), h)
		if err != nil {
			t.Fatalf("插入 (%d,%d) 失败: %v", p.key, p.pk, err)
		}
		live[p.key] = append(live[p.key], p.pk)
	}
	checkSubtree(t, h, root)
	checkLeafChain(t, h, root)

	// live[k] 保持有序, 供 sort.Search 精确移除
	for k := range live {
		sort.Slice(live[k], func(a, b int) bool { return live[k][a] < live[k][b] })
	}

	// ---- 阶段 2: 乱序删除一半 ----
	rng.Shuffle(len(all), func(i, j int) { all[i], all[j] = all[j], all[i] })
	deleted := make(map[pair]bool)
	for i, p := range all[:total/2] {
		newRoot, err := SBTDeleteWithPKID(root, mkKey(p.key), p.pk, h)
		if err != nil {
			t.Fatalf("删除 (%d,%d) 失败: %v", p.key, p.pk, err)
		}
		root = newRoot
		pks := live[p.key]
		idx := sort.Search(len(pks), func(x int) bool { return pks[x] >= p.pk })
		live[p.key] = append(pks[:idx], pks[idx+1:]...)
		deleted[p] = true

		if i%500 == 499 {
			checkSubtree(t, h, root)
			checkLeafChain(t, h, root)
			// 抽查被删 key 与随机 key 的 SBTSearchAll
			for _, k := range []uint64{p.key, uint64(rng.Intn(keys) + 1), 1, uint64(keys)} {
				want := append([]uint64(nil), live[k]...)
				sort.Slice(want, func(a, b int) bool { return want[a] < want[b] })
				got, ok := SBTSearchAll(root, mkKey(k), h)
				if len(want) == 0 {
					if ok {
						t.Fatalf("SBTSearchAll(key=%d) 应为空, 返回 %v", k, got)
					}
					continue
				}
				if !ok || !equalU64(got, want) {
					t.Fatalf("SBTSearchAll(key=%d) 不符: want %d 条, got %v (ok=%v)", k, len(want), got, ok)
				}
			}
		}
	}
	checkSubtree(t, h, root)

	// ---- 阶段 3: 重插入部分已删条目 (混合分裂/合并后的结构) ----
	reins := 0
	for _, p := range all[:total/2] {
		if reins >= total/8 {
			break
		}
		if !deleted[p] {
			continue
		}
		var err error
		root, err = SBTInsert(root, mkEntry(p), h)
		if err != nil {
			t.Fatalf("重插入 (%d,%d) 失败: %v", p.key, p.pk, err)
		}
		live[p.key] = append(live[p.key], p.pk)
		sort.Slice(live[p.key], func(a, b int) bool { return live[p.key][a] < live[p.key][b] })
		delete(deleted, p)
		reins++
	}
	checkSubtree(t, h, root)
	checkLeafChain(t, h, root)

	// ---- 阶段 4: 全量删除至空 ----
	remaining := make([]pair, 0, total)
	for k := uint64(1); k <= keys; k++ {
		for _, pk := range live[k] {
			remaining = append(remaining, pair{key: k, pk: pk})
		}
	}
	rng.Shuffle(len(remaining), func(i, j int) { remaining[i], remaining[j] = remaining[j], remaining[i] })
	for _, p := range remaining {
		newRoot, err := SBTDeleteWithPKID(root, mkKey(p.key), p.pk, h)
		if err != nil {
			t.Fatalf("清空阶段删除 (%d,%d) 失败: %v", p.key, p.pk, err)
		}
		root = newRoot
	}
	// 删空后: 根必须是 (空) 叶, 全树只剩 1 页
	n := sbtLoad(h, root)
	if n == nil || !n.isLeaf() || n.info.Count != 0 {
		t.Fatalf("删空后根不是空叶: %+v", n.info)
	}
	// 页收支平衡: 曾分配的每一页, 要么是最终根, 要么在下溢/合并中已归还
	freedSoFar := sortedSet(m.freed)
	if len(freedSoFar)+1 != len(m.allocated()) {
		t.Fatalf("页收支不平衡: 分配 %d 页, 已归还 %d, 最终根 1 页",
			len(m.allocated()), len(freedSoFar))
	}
	// 释放单页空根应恰好只归还根本身 (其余 34 页已在途中归还, 不得二次归还)
	m.freed = nil
	btReleaseSubtree(h, root)
	if freed := sortedSet(m.freed); !sameSet(freed, []uint32{root}) {
		t.Fatalf("释放单页空根应只归还根 %d, 实际归还 %v", root, freed)
	}
}

func equalU64(a, b []uint64) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// TestSecondaryBuildFromSortedComposite 批量建树要求条目按复合键 (key, pkID) 有序;
// 同 key 条目 pkID 无序时, 复合路由/二分都会静默漏条目 —— 这里验证有序输入的正确性,
// 并抽查精确删除落在叶中部 (非区段端点) 的场景。
func TestSecondaryBuildFromSortedComposite(t *testing.T) {
	const keyLen = 8
	m := newMemHandle()
	h := m.handle()

	// 3 个 key, 每个 800 条, 按 (key, pkID) 排序
	var entries [][]byte
	for k := uint64(1); k <= 3; k++ {
		for p := uint64(1); p <= 800; p++ {
			e := make([]byte, keyLen+8)
			binary.BigEndian.PutUint64(e[:keyLen], k)
			binary.BigEndian.PutUint64(e[keyLen:], p)
			entries = append(entries, e)
		}
	}
	root := SBTBuildFromSorted(entries, h)
	if root == 0 {
		t.Fatal("SBTBuildFromSorted 返回 0")
	}
	checkSubtree(t, h, root)

	// SBTSearchAll 应按 pkID 升序返回完整区段
	for k := uint64(1); k <= 3; k++ {
		b := make([]byte, keyLen)
		binary.BigEndian.PutUint64(b, k)
		got, ok := SBTSearchAll(root, b, h)
		if !ok || len(got) != 800 {
			t.Fatalf("key=%d: want 800 条, got %d (ok=%v)", k, len(got), ok)
		}
		for i := 1; i < len(got); i++ {
			if got[i-1] >= got[i] {
				t.Fatalf("key=%d: 返回未按 pkID 有序", k)
			}
		}
	}

	// 删掉中间一段 pkID, 区段两端不受影响
	for p := uint64(300); p <= 500; p++ {
		b := make([]byte, keyLen)
		binary.BigEndian.PutUint64(b, 2)
		var err error
		root, err = SBTDeleteWithPKID(root, b, p, h)
		if err != nil {
			t.Fatalf("删除 (2,%d) 失败: %v", p, err)
		}
	}
	checkSubtree(t, h, root)
	b2 := make([]byte, keyLen)
	binary.BigEndian.PutUint64(b2, 2)
	got, ok := SBTSearchAll(root, b2, h)
	if !ok || len(got) != 599 {
		t.Fatalf("删除中段 (300..500 共 201 条) 后 key=2 应剩 599 条, got %d", len(got))
	}
}

// TestSecondaryDeleteRoutingExact 点删除不再向左试探兄弟子树:
// 构造 "同 key 区段横跨多个叶" 后, 删除区段最左侧的条目必须一次命中。
// (旧实现从 key 路由落点向左扫, 删左端条目要试完整段; 若复合路由错位, 这里会漏删报
// ErrNotFound 或删错条目。)
func TestSecondaryDeleteRoutingExact(t *testing.T) {
	const keyLen = 8
	m := newMemHandle()
	h := m.handle()

	// key=7 的 3000 条 + 两侧各 1 条异键哨兵
	var entries [][]byte
	add := func(k, p uint64) {
		e := make([]byte, keyLen+8)
		binary.BigEndian.PutUint64(e[:keyLen], k)
		binary.BigEndian.PutUint64(e[keyLen:], p)
		entries = append(entries, e)
	}
	add(6, 1)
	for p := uint64(1); p <= 3000; p++ {
		add(7, p)
	}
	add(8, 1)
	sort.Slice(entries, func(i, j int) bool {
		c := bytes.Compare(entries[i][:keyLen], entries[j][:keyLen])
		if c != 0 {
			return c < 0
		}
		return binary.BigEndian.Uint64(entries[i][keyLen:]) < binary.BigEndian.Uint64(entries[j][keyLen:])
	})
	root := SBTBuildFromSorted(entries, h)

	// 从区段最左 (pkID=1) 删到最右 (pkID=3000): 任何一步 ErrNotFound 都是路由错位
	for p := uint64(1); p <= 3000; p++ {
		b := make([]byte, keyLen)
		binary.BigEndian.PutUint64(b, 7)
		var err error
		root, err = SBTDeleteWithPKID(root, b, p, h)
		if err != nil {
			t.Fatalf("顺序删除 (7,%d) 失败: %v", p, err)
		}
	}
	// 区段删空, 哨兵仍在
	for k, want := range map[uint64]int{6: 1, 7: 0, 8: 1} {
		b := make([]byte, keyLen)
		binary.BigEndian.PutUint64(b, k)
		got, ok := SBTSearchAll(root, b, h)
		if want == 0 && ok {
			t.Fatalf("key=%d 应已删空, got %v", k, got)
		}
		if want > 0 && (!ok || len(got) != want) {
			t.Fatalf("key=%d want %d 条, got %v (ok=%v)", k, want, got, ok)
		}
	}
	checkSubtree(t, h, root)
}