X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/6e20c4df4a0eac7f993d7783e82eb4c1ca6d8c70..150e81034dc65943803924ef29962e2481ffe771:/tree.go?ds=sidebyside

diff --git a/tree.go b/tree.go
index ee3ff68..9fb033f 100644
--- a/tree.go
+++ b/tree.go
@@ -1,12 +1,16 @@
 package main
 
 import (
+	"cmp"
 	"context"
+	"encoding/binary"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/bits"
 	"os"
 	"path/filepath"
+	"slices"
 	"time"
 	"unicode/utf8"
 )
@@ -32,25 +36,39 @@ type nodeKey struct {
 	index int
 }
 
+// treeRecordBytes is the size of a record of index/tree.idx: the position of
+// the node line, the node index, and the node level.
+const treeRecordBytes = 32
+
 // Tree holds the built nodes of one chat. It writes one JSON line for each node
-// to tree/YYYY-MM-DD.jsonl, named for the local day of the build, and loads all
-// lines again at start. Each node is built once and never again. A node of a
-// higher level is built only after both of its children. After a failed write,
-// the tree refuses all writes until it is opened again. The caller holds the
-// lock of the chat and serializes all calls.
+// to tree/YYYY-MM-DD.jsonl, named for the local day of the build. Each node is
+// built once and never again. A node of a higher level is built only after both
+// of its children. After a failed write, the tree refuses all writes until it
+// is opened again. The caller holds the lock of the chat and serializes all
+// calls.
+//
+// The tree keeps no node in memory. The index index/tree.idx lists the node
+// lines in file order, and index/tree-L.idx holds the position of node (L, i)
+// in record i, with a zero record for a node that is not built. At open the
+// tree indexes the lines after the last indexed one, so a missing index is
+// built again from the day files.
 type Tree struct {
 	clock   func() time.Time
 	journal journal
-	nodes   map[nodeKey]Node
+	dir     string
+	order   *records
+	levels  map[int]*records
 	failure error
 }
 
-// openTree loads the nodes of the chat in root. A chat of total messages cannot
-// have a node that covers a later message.
+// openTree opens the tree of the chat in root and indexes its new lines. A chat
+// of total messages cannot have a node that covers a later message.
 func openTree(ctx context.Context, root string, total int, clock func() time.Time) (*Tree, error) {
-	err := os.MkdirAll(filepath.Join(root, "tree"), 0o700)
-	if err != nil {
-		return nil, fmt.Errorf("create node directory: %w", err)
+	for _, dir := range []string{"tree", "index"} {
+		err := os.MkdirAll(filepath.Join(root, dir), 0o700)
+		if err != nil {
+			return nil, fmt.Errorf("create %s directory: %w", dir, err)
+		}
 	}
 
 	t := &Tree{
@@ -58,11 +76,59 @@ func openTree(ctx context.Context, root string, total int, clock func() time.Tim
 		journal: journal{
 			dir: filepath.Join(root, "tree"),
 		},
-		nodes: make(map[nodeKey]Node),
+		dir:    filepath.Join(root, "index"),
+		levels: make(map[int]*records),
+	}
+
+	err := t.open(ctx, total)
+	if err != nil {
+		return nil, errors.Join(fmt.Errorf("index tree: %w", err), t.Close())
+	}
+
+	return t, nil
+}
+
+func (t *Tree) open(ctx context.Context, total int) error {
+	var err error
+	t.order, err = openRecords(filepath.Join(t.dir, "tree.idx"), treeRecordBytes)
+	if err != nil {
+		return err
+	}
+
+	start := position{}
+	if t.order.count == 0 {
+		err = t.removeLevels()
+		if err != nil {
+			return err
+		}
+	} else {
+		record, _, err := t.order.Read(t.order.count - 1)
+		if err != nil {
+			return err
+		}
+
+		// A crash can stop an add between the two index writes, so the
+		// last node gets its level record again.
+		node, last := decodeTreeRecord(record)
+		err = t.writeLevel(node, last)
+		if err != nil {
+			return err
+		}
+
+		start = position{
+			day:    last.day,
+			offset: last.offset + uint64(last.length),
+		}
 	}
 
-	err = readLines(ctx, t.journal.dir, func(node Node) error {
-		err := t.check(node)
+	return t.journal.Scan(ctx, start, func(line []byte, place position) error {
+		var node Node
+		err := json.Unmarshal(line, &node)
+		if err != nil {
+			return err
+		}
+
+		err = t.check(node)
 		if err != nil {
 			return err
 		}
@@ -71,15 +137,56 @@ func openTree(ctx context.Context, root string, total int, clock func() time.Tim
 			return fmt.Errorf("node (%d, %d) covers messages after the log", node.Level, node.Index)
 		}
 
-		t.nodes[nodeKey{node.Level, node.Index}] = node
-
-		return nil
+		return t.addIndex(nodeKey{node.Level, node.Index}, place)
 	})
+}
+
+// removeLevels removes the level indexes, which are not valid without the order
+// index.
+func (t *Tree) removeLevels() error {
+	paths, err := filepath.Glob(filepath.Join(t.dir, "tree-*.idx"))
 	if err != nil {
-		return nil, fmt.Errorf("load tree: %w", err)
+		return fmt.Errorf("list level indexes: %w", err)
 	}
 
-	return t, nil
+	for _, path := range paths {
+		err = os.Remove(path)
+		if err != nil {
+			return fmt.Errorf("remove level index: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// writeLevel writes the position of node to its level index.
+func (t *Tree) writeLevel(node nodeKey, place position) error {
+	level, err := t.level(node.level)
+	if err != nil {
+		return err
+	}
+
+	record := make([]byte, positionBytes)
+	encodePosition(record, place)
+
+	return level.Write(node.index, record)
+}
+
+// level returns the index of level, and opens it at first use.
+func (t *Tree) level(level int) (*records, error) {
+	index, open := t.levels[level]
+	if open {
+		return index, nil
+	}
+
+	index, err := openRecords(filepath.Join(t.dir, fmt.Sprintf("tree-%d.idx", level)), positionBytes)
+	if err != nil {
+		return nil, err
+	}
+
+	t.levels[level] = index
+
+	return index, nil
 }
 
 // check makes sure that node is valid and can join the tree now.
@@ -92,20 +199,62 @@ func (t *Tree) check(node Node) error {
 		return fmt.Errorf("node (%d, %d) has no valid text of its size", node.Level, node.Index)
 	}
 
-	if t.Built(node.Level, node.Index) {
+	built, err := t.Built(node.Level, node.Index)
+	if err != nil {
+		return err
+	}
+
+	if built {
 		return fmt.Errorf("node (%d, %d) is already built", node.Level, node.Index)
 	}
 
-	if node.Level > 0 && (!t.Built(node.Level-1, 2*node.Index) || !t.Built(node.Level-1, 2*node.Index+1)) {
-		return fmt.Errorf("node (%d, %d) has a child that is not built", node.Level, node.Index)
+	if node.Level == 0 {
+		return nil
+	}
+
+	for _, index := range []int{2 * node.Index, 2*node.Index + 1} {
+		built, err = t.Built(node.Level-1, index)
+		if err != nil {
+			return err
+		}
+
+		if !built {
+			return fmt.Errorf("node (%d, %d) has a child that is not built", node.Level, node.Index)
+		}
 	}
 
 	return nil
 }
 
-// Close closes the open day file.
+// addIndex adds the node line at place to both indexes.
+func (t *Tree) addIndex(node nodeKey, place position) error {
+	record := make([]byte, treeRecordBytes)
+	encodePosition(record, place)
+	binary.LittleEndian.PutUint64(record[positionBytes:], uint64(node.index))
+	binary.LittleEndian.PutUint32(record[positionBytes+8:], uint32(node.level))
+
+	err := t.order.Append(record)
+	if err != nil {
+		return err
+	}
+
+	return t.writeLevel(node, place)
+}
+
+// Close closes the files of the tree.
 func (t *Tree) Close() error {
-	return t.journal.Close()
+	err := t.journal.Close()
+	if t.order != nil {
+		err = errors.Join(err, t.order.Close())
+		t.order = nil
+	}
+
+	for level, index := range t.levels {
+		err = errors.Join(err, index.Close())
+		delete(t.levels, level)
+	}
+
+	return err
 }
 
 // Add writes a new node to the tree.
@@ -124,27 +273,94 @@ func (t *Tree) Add(node Node) error {
 		return fmt.Errorf("encode node: %w", err)
 	}
 
-	err = t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
+	place, err := t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
+	if err == nil {
+		err = t.addIndex(nodeKey{node.Level, node.Index}, place)
+	}
+
 	if err != nil {
 		t.failure = err
 		return fmt.Errorf("write node (%d, %d): %w", node.Level, node.Index, err)
 	}
 
-	t.nodes[nodeKey{node.Level, node.Index}] = node
-
 	return nil
 }
 
 // Built tells if the node (l, i) is built.
-func (t *Tree) Built(level, index int) bool {
-	_, built := t.nodes[nodeKey{level, index}]
-	return built
+func (t *Tree) Built(level, index int) (bool, error) {
+	_, built, err := t.place(level, index)
+	return built, err
 }
 
-// Node returns the node (l, i) if it is built.
-func (t *Tree) Node(level, index int) (Node, bool) {
-	node, built := t.nodes[nodeKey{level, index}]
-	return node, built
+// place returns the position of the node (l, i) if it is built.
+func (t *Tree) place(level, index int) (position, bool, error) {
+	if level < 0 || level >= bits.UintSize-2 || index < 0 {
+		return position{}, false, nil
+	}
+
+	records, err := t.level(level)
+	if err != nil {
+		return position{}, false, err
+	}
+
+	record, found, err := records.Read(index)
+	if err != nil || !found {
+		return position{}, false, err
+	}
+
+	place := decodePosition(record)
+
+	return place, place.length > 0, nil
+}
+
+// Node reads the node (l, i) if it is built.
+func (t *Tree) Node(level, index int) (Node, bool, error) {
+	place, built, err := t.place(level, index)
+	if err != nil || !built {
+		return Node{}, false, err
+	}
+
+	line, err := t.journal.Read(place)
+	if err != nil {
+		return Node{}, false, fmt.Errorf("read node (%d, %d): %w", level, index, err)
+	}
+
+	var node Node
+	err = json.Unmarshal(line, &node)
+	if err != nil {
+		return Node{}, false, fmt.Errorf("decode node (%d, %d): %w", level, index, err)
+	}
+
+	if node.Level != level || node.Index != index {
+		return Node{}, false, fmt.Errorf("index of node (%d, %d) points to node (%d, %d)", level, index, node.Level, node.Index)
+	}
+
+	return node, true, nil
+}
+
+// Count returns the number of built nodes.
+func (t *Tree) Count() int {
+	return t.order.count
+}
+
+// Keys returns the keys of the built nodes, ordered by level and then by index.
+func (t *Tree) Keys() ([]nodeKey, error) {
+	keys := make([]nodeKey, 0, t.order.count)
+	for n := range t.order.count {
+		record, _, err := t.order.Read(n)
+		if err != nil {
+			return nil, err
+		}
+
+		key, _ := decodeTreeRecord(record)
+		keys = append(keys, key)
+	}
+
+	slices.SortFunc(keys, func(a, b nodeKey) int {
+		return cmp.Or(cmp.Compare(a.level, b.level), cmp.Compare(a.index, b.index))
+	})
+
+	return keys, nil
 }
 
 // leafSource returns the text that the leaf of message holds or summarizes.
@@ -185,6 +401,17 @@ func freeMerge(left, right Node) (Node, bool) {
 	}, true
 }
 
+// decodeTreeRecord reads the key and the position of a record of
+// index/tree.idx.
+func decodeTreeRecord(record []byte) (nodeKey, position) {
+	key := nodeKey{
+		level: int(binary.LittleEndian.Uint32(record[positionBytes+8:])),
+		index: int(binary.LittleEndian.Uint64(record[positionBytes:])),
+	}
+
+	return key, decodePosition(record)
+}
+
 // nodeEnd returns the ID after the last message of node (l, i).
 func nodeEnd(level, index int) int {
 	return (index + 1) << level