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 }