--- /dev/null
+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
+}
--- /dev/null
+package main
+
+import (
+ "os"
+ "path/filepath"
+ "strings"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestFreeLeafHoldsAShortMessageWordForWord(t *testing.T) {
+ message := Message{
+ ID: 7,
+ Kind: "user",
+ Text: strings.Repeat("a", nodeBytes-len("user: ")),
+ }
+
+ node, free := freeLeaf(message)
+ require.True(t, free)
+ require.Equal(t, Node{
+ Level: 0,
+ Index: 7,
+ Text: "user: " + message.Text,
+ Size: nodeBytes,
+ }, node)
+
+ message.Text += "a"
+ _, free = freeLeaf(message)
+ require.False(t, free)
+}
+
+func TestFreeMergeJoinsShortSiblings(t *testing.T) {
+ left := testNode(1, 2, "left")
+ right := testNode(1, 3, "right")
+
+ node, free := freeMerge(left, right)
+ require.True(t, free)
+ require.Equal(t, testNode(2, 1, "left\nright"), node)
+
+ right = testNode(1, 3, strings.Repeat("r", nodeBytes-len("left")))
+ _, free = freeMerge(left, right)
+ require.False(t, free)
+}
+
+func TestTreeBuildsEachNodeOnce(t *testing.T) {
+ tree, err := openTree(t.Context(), t.TempDir(), 2, newTestClock().Now)
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, tree.Close())
+ }()
+
+ require.NoError(t, tree.Add(testNode(0, 0, "user: a")))
+ require.ErrorContains(t, tree.Add(testNode(0, 0, "user: b")), "already built")
+ require.ErrorContains(t, tree.Add(testNode(1, 0, "a b")), "child that is not built")
+
+ require.NoError(t, tree.Add(testNode(0, 1, "unii: b")))
+ require.NoError(t, tree.Add(testNode(1, 0, "a b")))
+
+ node, built := tree.Node(0, 0)
+ require.True(t, built)
+ require.Equal(t, "user: a", node.Text)
+}
+
+func TestTreeReopenLoadsTheSameNodes(t *testing.T) {
+ root := t.TempDir()
+ clock := newTestClock()
+ tree, err := openTree(t.Context(), root, 2, clock.Now)
+ require.NoError(t, err)
+
+ nodes := []Node{
+ testNode(0, 1, "unii: b"),
+ testNode(0, 0, "user: a"),
+ testNode(1, 0, "user: a\nunii: b"),
+ }
+ for _, node := range nodes {
+ require.NoError(t, tree.Add(node))
+ }
+
+ require.NoError(t, tree.Close())
+
+ data, err := os.ReadFile(filepath.Join(root, "tree", "2026-01-01.jsonl"))
+ require.NoError(t, err)
+ require.Equal(t, `{"l":0,"i":1,"text":"unii: b","size":7}`+"\n", strings.SplitAfter(string(data), "\n")[0])
+
+ tree, err = openTree(t.Context(), root, 2, clock.Now)
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, tree.Close())
+ }()
+
+ require.Len(t, tree.nodes, len(nodes))
+ for _, node := range nodes {
+ loaded, built := tree.Node(node.Level, node.Index)
+ require.True(t, built)
+ require.Equal(t, node, loaded)
+ }
+
+ _, err = openTree(t.Context(), root, 1, clock.Now)
+ require.ErrorContains(t, err, "covers messages after the log")
+}
+
+func testNode(level, index int, text string) Node {
+ return Node{
+ Level: level,
+ Index: index,
+ Text: text,
+ Size: len(text),
+ }
+}
projects / onechat / commitdiff