+package main
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "math/bits"
+ "os"
+ "path/filepath"
+ "time"
+ "unicode/utf8"
+)
+
+// nodeBytes is the size limit of a node text in UTF-8 bytes. A model can go a
+// few bytes over it, because the view measures the real sizes.
+const nodeBytes = 512
+
+// Node is one line of the summary tree. Node (l, i) covers the 2^l messages
+// from i*2^l on. A node of level 0 holds one message. A node of a higher level
+// merges its two children, (l-1, 2i) and (l-1, 2i+1). Size is the UTF-8 byte
+// length of Text.
+type Node struct {
+ Level int `json:"l"`
+ Index int `json:"i"`
+ Text string `json:"text"`
+ Size int `json:"size"`
+}
+
+// nodeKey names the node (l, i).
+type nodeKey struct {
+ level int
+ index int
+}
+
+// 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.
+type Tree struct {
+ clock func() time.Time
+ journal journal
+ nodes map[nodeKey]Node
+ 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.
+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)
+ }
+
+ t := &Tree{
+ clock: clock,
+ journal: journal{
+ dir: filepath.Join(root, "tree"),
+ },
+ nodes: make(map[nodeKey]Node),
+ }
+
+ err = readLines(ctx, t.journal.dir, func(node Node) error {
+ err := t.check(node)
+ if err != nil {
+ return err
+ }
+
+ if nodeEnd(node.Level, node.Index) > total {
+ 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
+ })
+ if err != nil {
+ return nil, fmt.Errorf("load tree: %w", err)
+ }
+
+ return t, nil
+}
+
+// check makes sure that node is valid and can join the tree now.
+func (t *Tree) check(node Node) error {
+ if node.Level < 0 || node.Level >= bits.UintSize-2 || node.Index < 0 {
+ return fmt.Errorf("node (%d, %d) is out of range", node.Level, node.Index)
+ }
+
+ if node.Text == "" || !utf8.ValidString(node.Text) || node.Size != len(node.Text) {
+ return fmt.Errorf("node (%d, %d) has no valid text of its size", node.Level, node.Index)
+ }
+
+ if t.Built(node.Level, node.Index) {
+ 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)
+ }
+
+ return nil
+}
+
+// Close closes the open day file.
+func (t *Tree) Close() error {
+ return t.journal.Close()
+}
+
+// Add writes a new node to the tree.
+func (t *Tree) Add(node Node) error {
+ if t.failure != nil {
+ return fmt.Errorf("tree needs a reopen after a failed write: %w", t.failure)
+ }
+
+ err := t.check(node)
+ if err != nil {
+ return err
+ }
+
+ line, err := json.Marshal(node)
+ if err != nil {
+ return fmt.Errorf("encode node: %w", err)
+ }
+
+ err = t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
+ 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
+}
+
+// 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
+}
+
+// leafSource returns the text that the leaf of message holds or summarizes.
+func leafSource(message Message) string {
+ return message.Kind + ": " + message.Text
+}
+
+// freeLeaf returns the leaf of message when its source fits in a node. Such a
+// leaf holds the source word for word and needs no model call.
+func freeLeaf(message Message) (Node, bool) {
+ text := leafSource(message)
+ if len(text) > nodeBytes {
+ return Node{}, false
+ }
+
+ return Node{
+ Level: 0,
+ Index: message.ID,
+ Text: text,
+ Size: len(text),
+ }, true
+}
+
+// freeMerge returns the parent of two sibling nodes when both texts, joined by
+// a newline, fit in a node. Such a parent needs no model call.
+func freeMerge(left, right Node) (Node, bool) {
+ if left.Size+1+right.Size > nodeBytes {
+ return Node{}, false
+ }
+
+ text := left.Text + "\n" + right.Text
+
+ return Node{
+ Level: left.Level + 1,
+ Index: left.Index / 2,
+ Text: text,
+ Size: len(text),
+ }, true
+}
+
+// nodeEnd returns the ID after the last message of node (l, i).
+func nodeEnd(level, index int) int {
+ return (index + 1) << level
+}
projects / onechat / commitdiff