five/hbrdd/ntx/build.go

// Copyright (c) 2026 Charles KWON OhJun (charleskwonohjun@gmail.com)
// All rights reserved.

// NTX index creation (INDEX ON) and key insertion.
package ntx

import (
	"encoding/binary"
	"fmt"
	"os"
	"sort"
)

// KeyRecord pairs a key value with its record number for sorting.
type KeyRecord struct {
	Key   []byte
	RecNo uint32
}

// CreateIndex builds a new NTX index file from sorted key-record pairs.
func CreateIndex(path string, keyExpr string, keyLen int, unique bool, descend bool, keys []KeyRecord) (*Index, error) {
	if len(path) < 4 || path[len(path)-4:] != ".ntx" {
		path += ".ntx"
	}

	f, err := os.Create(path)
	if err != nil {
		return nil, fmt.Errorf("create NTX %s: %w", path, err)
	}

	itemSize := 8 + keyLen
	maxItem := calculateMaxItems(itemSize)
	halfPage := maxItem / 2

	// Phase 1: Build leaf pages and assign file offsets immediately
	var allPages []*buildPage
	nextOffset := int64(HeaderSize)

	leafPages := buildLeafPages(keys, keyLen, itemSize, maxItem, &nextOffset)
	allPages = append(allPages, leafPages...)

	if len(leafPages) == 0 {
		pg := makeEmptyPage(keyLen, itemSize, maxItem, nextOffset)
		nextOffset += BlockSize
		allPages = append(allPages, pg)
		leafPages = append(leafPages, pg)
	}

	// Phase 2: Build internal pages bottom-up
	// Each level's pages have offsets already assigned, so children are resolvable.
	currentLevel := leafPages
	for len(currentLevel) > 1 {
		parentLevel := buildInternalLevel(currentLevel, keyLen, itemSize, maxItem, &nextOffset)
		allPages = append(allPages, parentLevel...)
		currentLevel = parentLevel
	}

	rootOffset := uint32(currentLevel[0].fileOffset)

	// Write header
	hdr := Header{
		Type:     0x0401,
		Version:  1,
		Root:     rootOffset,
		NextPage: uint32(nextOffset),
		ItemSize: uint16(itemSize),
		KeySize:  uint16(keyLen),
		KeyDec:   0,
		MaxItem:  uint16(maxItem),
		HalfPage: uint16(halfPage),
	}
	copy(hdr.KeyExpr[:], keyExpr)
	if unique {
		hdr.Unique = 1
	}
	if descend {
		hdr.Descend = 1
	}

	if err := WriteHeader(f, &hdr); err != nil {
		f.Close()
		return nil, err
	}

	// Write all pages
	for _, pg := range allPages {
		if _, err := f.WriteAt(pg.data[:], pg.fileOffset); err != nil {
			f.Close()
			return nil, fmt.Errorf("write NTX page at %d: %w", pg.fileOffset, err)
		}
	}

	f.Close()
	return OpenIndex(path)
}

// --- Internal build structures ---

type buildPage struct {
	data       [BlockSize]byte
	fileOffset int64
	keyCount   int
	firstKey   []byte
	firstRecNo uint32
}

func calculateMaxItems(itemSize int) int {
	// Page layout: [keyCount:2] [offsets:(max+1)*2] [entries:(max+1)*itemSize]
	// Total must fit in BlockSize (1024).
	// 2 + (max+1)*2 + (max+1)*itemSize <= 1024
	// (max+1) * (itemSize + 2) <= 1022
	// max+1 <= 1022 / (itemSize + 2)
	// max <= 1022/(itemSize+2) - 1
	max := 1022/(itemSize+2) - 1
	if max < 2 {
		max = 2
	}
	if max > 250 {
		max = 250
	}
	return max
}

func makeEmptyPage(keyLen, itemSize, maxItem int, offset int64) *buildPage {
	pg := &buildPage{
		fileOffset: offset,
		firstKey:   make([]byte, keyLen),
	}
	binary.LittleEndian.PutUint16(pg.data[0:2], 0)
	return pg
}

func buildLeafPages(keys []KeyRecord, keyLen, itemSize, maxItem int, nextOffset *int64) []*buildPage {
	if len(keys) == 0 {
		return nil
	}

	var pages []*buildPage
	for i := 0; i < len(keys); i += maxItem {
		end := i + maxItem
		if end > len(keys) {
			end = len(keys)
		}
		chunk := keys[i:end]
		pg := encodeLeafPage(chunk, keyLen, itemSize, maxItem, *nextOffset)
		*nextOffset += BlockSize
		pages = append(pages, pg)
	}
	return pages
}

func encodeLeafPage(keys []KeyRecord, keyLen, itemSize, maxItem int, offset int64) *buildPage {
	pg := &buildPage{
		fileOffset: offset,
		keyCount:   len(keys),
		firstKey:   make([]byte, keyLen),
		firstRecNo: keys[0].RecNo,
	}
	copy(pg.firstKey, keys[0].Key)

	binary.LittleEndian.PutUint16(pg.data[0:2], uint16(len(keys)))
	dataStart := 2 + (maxItem+1)*2

	for i, kr := range keys {
		entryOffset := dataStart + i*itemSize
		binary.LittleEndian.PutUint16(pg.data[2+i*2:4+i*2], uint16(entryOffset))
		binary.LittleEndian.PutUint32(pg.data[entryOffset:entryOffset+4], 0) // leaf
		binary.LittleEndian.PutUint32(pg.data[entryOffset+4:entryOffset+8], kr.RecNo)
		key := make([]byte, keyLen)
		for j := range key {
			key[j] = ' '
		}
		copy(key, kr.Key)
		copy(pg.data[entryOffset+8:entryOffset+8+keyLen], key)
	}
	// Trailing offset
	binary.LittleEndian.PutUint16(
		pg.data[2+len(keys)*2:4+len(keys)*2],
		uint16(dataStart+len(keys)*itemSize),
	)
	return pg
}

func buildInternalLevel(children []*buildPage, keyLen, itemSize, maxItem int, nextOffset *int64) []*buildPage {
	var pages []*buildPage
	// Each internal page holds up to maxItem keys and maxItem+1 child pointers.
	// So each parent page covers maxItem+1 children.
	fanout := maxItem + 1

	for i := 0; i < len(children); i += fanout {
		end := i + fanout
		if end > len(children) {
			end = len(children)
		}
		chunk := children[i:end]
		pg := encodeInternalPage(chunk, keyLen, itemSize, maxItem, *nextOffset)
		*nextOffset += BlockSize
		pages = append(pages, pg)
	}
	return pages
}

func encodeInternalPage(children []*buildPage, keyLen, itemSize, maxItem int, offset int64) *buildPage {
	nKeys := len(children) - 1
	if nKeys < 0 {
		nKeys = 0
	}

	pg := &buildPage{
		fileOffset: offset,
		keyCount:   nKeys,
	}
	if len(children) > 0 {
		pg.firstKey = make([]byte, keyLen)
		copy(pg.firstKey, children[0].firstKey)
		pg.firstRecNo = children[0].firstRecNo
	}

	binary.LittleEndian.PutUint16(pg.data[0:2], uint16(nKeys))
	dataStart := 2 + (maxItem+1)*2

	for i := 0; i <= nKeys; i++ {
		entryOffset := dataStart + i*itemSize
		if entryOffset+itemSize > BlockSize {
			// Page overflow — cap keys
			binary.LittleEndian.PutUint16(pg.data[0:2], uint16(i))
			break
		}
		binary.LittleEndian.PutUint16(pg.data[2+i*2:4+i*2], uint16(entryOffset))

		// Child pointer
		binary.LittleEndian.PutUint32(pg.data[entryOffset:entryOffset+4],
			uint32(children[i].fileOffset))

		if i < nKeys {
			// Separator = first key of child[i+1], promoted (not duplicated)
			binary.LittleEndian.PutUint32(pg.data[entryOffset+4:entryOffset+8],
				children[i+1].firstRecNo)
			key := make([]byte, keyLen)
			for j := range key {
				key[j] = ' '
			}
			copy(key, children[i+1].firstKey)
			copy(pg.data[entryOffset+8:entryOffset+8+keyLen], key)
		} else {
			binary.LittleEndian.PutUint32(pg.data[entryOffset+4:entryOffset+8], 0)
		}
	}

	return pg
}

// --- Single key operations ---

func (idx *Index) InsertKey(key []byte, recNo uint32) error {
	keys := idx.collectAllKeys()
	newKey := make([]byte, idx.keyLen)
	copy(newKey, key)
	keys = append(keys, KeyRecord{Key: newKey, RecNo: recNo})

	sort.Slice(keys, func(i, j int) bool {
		cmp := idx.compareKeys(keys[i].Key, keys[j].Key)
		if cmp == 0 {
			return keys[i].RecNo < keys[j].RecNo
		}
		return cmp < 0
	})

	path := idx.file.Name()
	idx.Close()

	newIdx, err := CreateIndex(path, idx.header.GetKeyExpr(), idx.keyLen,
		idx.uniqueKey, !idx.ascendKey, keys)
	if err != nil {
		return err
	}
	*idx = *newIdx
	return nil
}

func (idx *Index) DeleteKey(recNo uint32) error {
	keys := idx.collectAllKeys()
	filtered := keys[:0]
	for _, kr := range keys {
		if kr.RecNo != recNo {
			filtered = append(filtered, kr)
		}
	}

	path := idx.file.Name()
	idx.Close()

	newIdx, err := CreateIndex(path, idx.header.GetKeyExpr(), idx.keyLen,
		idx.uniqueKey, !idx.ascendKey, filtered)
	if err != nil {
		return err
	}
	*idx = *newIdx
	return nil
}

func (idx *Index) collectAllKeys() []KeyRecord {
	var keys []KeyRecord
	if !idx.GoTop() {
		return keys
	}
	keys = append(keys, KeyRecord{
		Key:   append([]byte{}, idx.curKey[:idx.keyLen]...),
		RecNo: idx.curRecNo,
	})
	for idx.SkipNext() {
		keys = append(keys, KeyRecord{
			Key:   append([]byte{}, idx.curKey[:idx.keyLen]...),
			RecNo: idx.curRecNo,
		})
	}
	return keys
}