git.giacomolli.uk - onechat/blob - tree.go

   1 package main
   2
   3 import (
   4         "cmp"
   5         "context"
   6         "encoding/json"
   7         "fmt"
   8         "math/bits"
   9         "os"
  10         "path/filepath"
  11         "slices"
  12         "time"
  13         "unicode/utf8"
  14 )
  15
  16 // nodeBytes is the size limit of a node text in UTF-8 bytes. A model can go a
  17 // few bytes over it, because the view measures the real sizes.
  18 const nodeBytes = 512
  19
  20 // Node is one line of the summary tree. Node (l, i) covers the 2^l messages
  21 // from i*2^l on. A node of level 0 holds one message. A node of a higher level
  22 // merges its two children, (l-1, 2i) and (l-1, 2i+1). Size is the UTF-8 byte
  23 // length of Text.
  24 type Node struct {
  25         Level int    `json:"l"`
  26         Index int    `json:"i"`
  27         Text  string `json:"text"`
  28         Size  int    `json:"size"`
  29 }
  30
  31 // nodeKey names the node (l, i).
  32 type nodeKey struct {
  33         level int
  34         index int
  35 }
  36
  37 // Tree holds the built nodes of one chat. It writes one JSON line for each node
  38 // to tree/YYYY-MM-DD.jsonl, named for the local day of the build, and loads all
  39 // lines again at start. Each node is built once and never again. A node of a
  40 // higher level is built only after both of its children. After a failed write,
  41 // the tree refuses all writes until it is opened again. The caller holds the
  42 // lock of the chat and serializes all calls.
  43 type Tree struct {
  44         clock   func() time.Time
  45         journal journal
  46         nodes   map[nodeKey]Node
  47         failure error
  48 }
  49
  50 // openTree loads the nodes of the chat in root. A chat of total messages cannot
  51 // have a node that covers a later message.
  52 func openTree(ctx context.Context, root string, total int, clock func() time.Time) (*Tree, error) {
  53         err := os.MkdirAll(filepath.Join(root, "tree"), 0o700)
  54         if err != nil {
  55                 return nil, fmt.Errorf("create node directory: %w", err)
  56         }
  57
  58         t := &Tree{
  59                 clock: clock,
  60                 journal: journal{
  61                         dir: filepath.Join(root, "tree"),
  62                 },
  63                 nodes: make(map[nodeKey]Node),
  64         }
  65
  66         err = readLines(ctx, t.journal.dir, func(node Node) error {
  67                 err := t.check(node)
  68                 if err != nil {
  69                         return err
  70                 }
  71
  72                 if nodeEnd(node.Level, node.Index) > total {
  73                         return fmt.Errorf("node (%d, %d) covers messages after the log", node.Level, node.Index)
  74                 }
  75
  76                 t.nodes[nodeKey{node.Level, node.Index}] = node
  77
  78                 return nil
  79         })
  80         if err != nil {
  81                 return nil, fmt.Errorf("load tree: %w", err)
  82         }
  83
  84         return t, nil
  85 }
  86
  87 // check makes sure that node is valid and can join the tree now.
  88 func (t *Tree) check(node Node) error {
  89         if node.Level < 0 || node.Level >= bits.UintSize-2 || node.Index < 0 {
  90                 return fmt.Errorf("node (%d, %d) is out of range", node.Level, node.Index)
  91         }
  92
  93         if node.Text == "" || !utf8.ValidString(node.Text) || node.Size != len(node.Text) {
  94                 return fmt.Errorf("node (%d, %d) has no valid text of its size", node.Level, node.Index)
  95         }
  96
  97         if t.Built(node.Level, node.Index) {
  98                 return fmt.Errorf("node (%d, %d) is already built", node.Level, node.Index)
  99         }
 100
 101         if node.Level > 0 && (!t.Built(node.Level-1, 2*node.Index) || !t.Built(node.Level-1, 2*node.Index+1)) {
 102                 return fmt.Errorf("node (%d, %d) has a child that is not built", node.Level, node.Index)
 103         }
 104
 105         return nil
 106 }
 107
 108 // Close closes the open day file.
 109 func (t *Tree) Close() error {
 110         return t.journal.Close()
 111 }
 112
 113 // Add writes a new node to the tree.
 114 func (t *Tree) Add(node Node) error {
 115         if t.failure != nil {
 116                 return fmt.Errorf("tree needs a reopen after a failed write: %w", t.failure)
 117         }
 118
 119         err := t.check(node)
 120         if err != nil {
 121                 return err
 122         }
 123
 124         line, err := json.Marshal(node)
 125         if err != nil {
 126                 return fmt.Errorf("encode node: %w", err)
 127         }
 128
 129         err = t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
 130         if err != nil {
 131                 t.failure = err
 132                 return fmt.Errorf("write node (%d, %d): %w", node.Level, node.Index, err)
 133         }
 134
 135         t.nodes[nodeKey{node.Level, node.Index}] = node
 136
 137         return nil
 138 }
 139
 140 // Built tells if the node (l, i) is built.
 141 func (t *Tree) Built(level, index int) bool {
 142         _, built := t.nodes[nodeKey{level, index}]
 143         return built
 144 }
 145
 146 // Node returns the node (l, i) if it is built.
 147 func (t *Tree) Node(level, index int) (Node, bool) {
 148         node, built := t.nodes[nodeKey{level, index}]
 149         return node, built
 150 }
 151
 152 // Nodes returns the built nodes, ordered by level and then by index.
 153 func (t *Tree) Nodes() []Node {
 154         nodes := make([]Node, 0, len(t.nodes))
 155         for _, node := range t.nodes {
 156                 nodes = append(nodes, node)
 157         }
 158
 159         slices.SortFunc(nodes, func(a, b Node) int {
 160                 return cmp.Or(cmp.Compare(a.Level, b.Level), cmp.Compare(a.Index, b.Index))
 161         })
 162
 163         return nodes
 164 }
 165
 166 // leafSource returns the text that the leaf of message holds or summarizes.
 167 func leafSource(message Message) string {
 168         return message.Kind + ": " + message.Text
 169 }
 170
 171 // freeLeaf returns the leaf of message when its source fits in a node. Such a
 172 // leaf holds the source word for word and needs no model call.
 173 func freeLeaf(message Message) (Node, bool) {
 174         text := leafSource(message)
 175         if len(text) > nodeBytes {
 176                 return Node{}, false
 177         }
 178
 179         return Node{
 180                 Level: 0,
 181                 Index: message.ID,
 182                 Text:  text,
 183                 Size:  len(text),
 184         }, true
 185 }
 186
 187 // freeMerge returns the parent of two sibling nodes when both texts, joined by
 188 // a newline, fit in a node. Such a parent needs no model call.
 189 func freeMerge(left, right Node) (Node, bool) {
 190         if left.Size+1+right.Size > nodeBytes {
 191                 return Node{}, false
 192         }
 193
 194         text := left.Text + "\n" + right.Text
 195
 196         return Node{
 197                 Level: left.Level + 1,
 198                 Index: left.Index / 2,
 199                 Text:  text,
 200                 Size:  len(text),
 201         }, true
 202 }
 203
 204 // nodeEnd returns the ID after the last message of node (l, i).
 205 func nodeEnd(level, index int) int {
 206         return (index + 1) << level
 207 }