Add the summary tree nodes

[onechat] / tree.go

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
}