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

   1 package main
   2
   3 import (
   4         "cmp"
   5         "context"
   6         "encoding/binary"
   7         "encoding/json"
   8         "errors"
   9         "fmt"
  10         "math/bits"
  11         "os"
  12         "path/filepath"
  13         "slices"
  14         "time"
  15         "unicode/utf8"
  16 )
  17
  18 // nodeBytes is the size limit of a node text in UTF-8 bytes. A model can go a
  19 // few bytes over it, because the view measures the real sizes.
  20 const nodeBytes = 512
  21
  22 // Node is one line of the summary tree. Node (l, i) covers the 2^l messages
  23 // from i*2^l on. A node of level 0 holds one message. A node of a higher level
  24 // merges its two children, (l-1, 2i) and (l-1, 2i+1). Size is the UTF-8 byte
  25 // length of Text.
  26 type Node struct {
  27         Level int    `json:"l"`
  28         Index int    `json:"i"`
  29         Text  string `json:"text"`
  30         Size  int    `json:"size"`
  31 }
  32
  33 // nodeKey names the node (l, i).
  34 type nodeKey struct {
  35         level int
  36         index int
  37 }
  38
  39 // treeRecordBytes is the size of a record of index/tree.idx: the position of
  40 // the node line, the node index, and the node level.
  41 const treeRecordBytes = 32
  42
  43 // Tree holds the built nodes of one chat. It writes one JSON line for each node
  44 // to tree/YYYY-MM-DD.jsonl, named for the local day of the build. Each node is
  45 // built once and never again. A node of a higher level is built only after both
  46 // of its children. After a failed write, the tree refuses all writes until it
  47 // is opened again. The caller holds the lock of the chat and serializes all
  48 // calls.
  49 //
  50 // The tree keeps no node in memory. The index index/tree.idx lists the node
  51 // lines in file order, and index/tree-L.idx holds the position of node (L, i)
  52 // in record i, with a zero record for a node that is not built. At open the
  53 // tree indexes the lines after the last indexed one, so a missing index is
  54 // built again from the day files.
  55 type Tree struct {
  56         clock   func() time.Time
  57         journal journal
  58         dir     string
  59         order   *records
  60         levels  map[int]*records
  61         failure error
  62 }
  63
  64 // openTree opens the tree of the chat in root and indexes its new lines. A chat
  65 // of total messages cannot have a node that covers a later message.
  66 func openTree(ctx context.Context, root string, total int, clock func() time.Time) (*Tree, error) {
  67         for _, dir := range []string{"tree", "index"} {
  68                 err := os.MkdirAll(filepath.Join(root, dir), 0o700)
  69                 if err != nil {
  70                         return nil, fmt.Errorf("create %s directory: %w", dir, err)
  71                 }
  72         }
  73
  74         t := &Tree{
  75                 clock: clock,
  76                 journal: journal{
  77                         dir: filepath.Join(root, "tree"),
  78                 },
  79                 dir:    filepath.Join(root, "index"),
  80                 levels: make(map[int]*records),
  81         }
  82
  83         err := t.open(ctx, total)
  84         if err != nil {
  85                 return nil, errors.Join(fmt.Errorf("index tree: %w", err), t.Close())
  86         }
  87
  88         return t, nil
  89 }
  90
  91 func (t *Tree) open(ctx context.Context, total int) error {
  92         var err error
  93         t.order, err = openRecords(filepath.Join(t.dir, "tree.idx"), treeRecordBytes)
  94         if err != nil {
  95                 return err
  96         }
  97
  98         start := position{}
  99         if t.order.count == 0 {
 100                 err = t.removeLevels()
 101                 if err != nil {
 102                         return err
 103                 }
 104         } else {
 105                 record, _, err := t.order.Read(t.order.count - 1)
 106                 if err != nil {
 107                         return err
 108                 }
 109
 110                 // A crash can stop an add between the two index writes, so the
 111                 // last node gets its level record again.
 112                 node, last := decodeTreeRecord(record)
 113                 err = t.writeLevel(node, last)
 114                 if err != nil {
 115                         return err
 116                 }
 117
 118                 start = position{
 119                         day:    last.day,
 120                         offset: last.offset + uint64(last.length),
 121                 }
 122         }
 123
 124         return t.journal.Scan(ctx, start, func(line []byte, place position) error {
 125                 var node Node
 126                 err := json.Unmarshal(line, &node)
 127                 if err != nil {
 128                         return err
 129                 }
 130
 131                 err = t.check(node)
 132                 if err != nil {
 133                         return err
 134                 }
 135
 136                 if nodeEnd(node.Level, node.Index) > total {
 137                         return fmt.Errorf("node (%d, %d) covers messages after the log", node.Level, node.Index)
 138                 }
 139
 140                 return t.addIndex(nodeKey{node.Level, node.Index}, place)
 141         })
 142 }
 143
 144 // removeLevels removes the level indexes, which are not valid without the order
 145 // index.
 146 func (t *Tree) removeLevels() error {
 147         paths, err := filepath.Glob(filepath.Join(t.dir, "tree-*.idx"))
 148         if err != nil {
 149                 return fmt.Errorf("list level indexes: %w", err)
 150         }
 151
 152         for _, path := range paths {
 153                 err = os.Remove(path)
 154                 if err != nil {
 155                         return fmt.Errorf("remove level index: %w", err)
 156                 }
 157         }
 158
 159         return nil
 160 }
 161
 162 // writeLevel writes the position of node to its level index.
 163 func (t *Tree) writeLevel(node nodeKey, place position) error {
 164         level, err := t.level(node.level)
 165         if err != nil {
 166                 return err
 167         }
 168
 169         record := make([]byte, positionBytes)
 170         encodePosition(record, place)
 171
 172         return level.Write(node.index, record)
 173 }
 174
 175 // level returns the index of level, and opens it at first use.
 176 func (t *Tree) level(level int) (*records, error) {
 177         index, open := t.levels[level]
 178         if open {
 179                 return index, nil
 180         }
 181
 182         index, err := openRecords(filepath.Join(t.dir, fmt.Sprintf("tree-%d.idx", level)), positionBytes)
 183         if err != nil {
 184                 return nil, err
 185         }
 186
 187         t.levels[level] = index
 188
 189         return index, nil
 190 }
 191
 192 // check makes sure that node is valid and can join the tree now.
 193 func (t *Tree) check(node Node) error {
 194         if node.Level < 0 || node.Level >= bits.UintSize-2 || node.Index < 0 {
 195                 return fmt.Errorf("node (%d, %d) is out of range", node.Level, node.Index)
 196         }
 197
 198         if node.Text == "" || !utf8.ValidString(node.Text) || node.Size != len(node.Text) {
 199                 return fmt.Errorf("node (%d, %d) has no valid text of its size", node.Level, node.Index)
 200         }
 201
 202         built, err := t.Built(node.Level, node.Index)
 203         if err != nil {
 204                 return err
 205         }
 206
 207         if built {
 208                 return fmt.Errorf("node (%d, %d) is already built", node.Level, node.Index)
 209         }
 210
 211         if node.Level == 0 {
 212                 return nil
 213         }
 214
 215         for _, index := range []int{2 * node.Index, 2*node.Index + 1} {
 216                 built, err = t.Built(node.Level-1, index)
 217                 if err != nil {
 218                         return err
 219                 }
 220
 221                 if !built {
 222                         return fmt.Errorf("node (%d, %d) has a child that is not built", node.Level, node.Index)
 223                 }
 224         }
 225
 226         return nil
 227 }
 228
 229 // addIndex adds the node line at place to both indexes.
 230 func (t *Tree) addIndex(node nodeKey, place position) error {
 231         record := make([]byte, treeRecordBytes)
 232         encodePosition(record, place)
 233         binary.LittleEndian.PutUint64(record[positionBytes:], uint64(node.index))
 234         binary.LittleEndian.PutUint32(record[positionBytes+8:], uint32(node.level))
 235
 236         err := t.order.Append(record)
 237         if err != nil {
 238                 return err
 239         }
 240
 241         return t.writeLevel(node, place)
 242 }
 243
 244 // Close closes the files of the tree.
 245 func (t *Tree) Close() error {
 246         err := t.journal.Close()
 247         if t.order != nil {
 248                 err = errors.Join(err, t.order.Close())
 249                 t.order = nil
 250         }
 251
 252         for level, index := range t.levels {
 253                 err = errors.Join(err, index.Close())
 254                 delete(t.levels, level)
 255         }
 256
 257         return err
 258 }
 259
 260 // Add writes a new node to the tree.
 261 func (t *Tree) Add(node Node) error {
 262         if t.failure != nil {
 263                 return fmt.Errorf("tree needs a reopen after a failed write: %w", t.failure)
 264         }
 265
 266         err := t.check(node)
 267         if err != nil {
 268                 return err
 269         }
 270
 271         line, err := json.Marshal(node)
 272         if err != nil {
 273                 return fmt.Errorf("encode node: %w", err)
 274         }
 275
 276         place, err := t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
 277         if err == nil {
 278                 err = t.addIndex(nodeKey{node.Level, node.Index}, place)
 279         }
 280
 281         if err != nil {
 282                 t.failure = err
 283                 return fmt.Errorf("write node (%d, %d): %w", node.Level, node.Index, err)
 284         }
 285
 286         return nil
 287 }
 288
 289 // Built tells if the node (l, i) is built.
 290 func (t *Tree) Built(level, index int) (bool, error) {
 291         _, built, err := t.place(level, index)
 292         return built, err
 293 }
 294
 295 // place returns the position of the node (l, i) if it is built.
 296 func (t *Tree) place(level, index int) (position, bool, error) {
 297         if level < 0 || level >= bits.UintSize-2 || index < 0 {
 298                 return position{}, false, nil
 299         }
 300
 301         records, err := t.level(level)
 302         if err != nil {
 303                 return position{}, false, err
 304         }
 305
 306         record, found, err := records.Read(index)
 307         if err != nil || !found {
 308                 return position{}, false, err
 309         }
 310
 311         place := decodePosition(record)
 312
 313         return place, place.length > 0, nil
 314 }
 315
 316 // Node reads the node (l, i) if it is built.
 317 func (t *Tree) Node(level, index int) (Node, bool, error) {
 318         place, built, err := t.place(level, index)
 319         if err != nil || !built {
 320                 return Node{}, false, err
 321         }
 322
 323         line, err := t.journal.Read(place)
 324         if err != nil {
 325                 return Node{}, false, fmt.Errorf("read node (%d, %d): %w", level, index, err)
 326         }
 327
 328         var node Node
 329         err = json.Unmarshal(line, &node)
 330         if err != nil {
 331                 return Node{}, false, fmt.Errorf("decode node (%d, %d): %w", level, index, err)
 332         }
 333
 334         if node.Level != level || node.Index != index {
 335                 return Node{}, false, fmt.Errorf("index of node (%d, %d) points to node (%d, %d)", level, index, node.Level, node.Index)
 336         }
 337
 338         return node, true, nil
 339 }
 340
 341 // Count returns the number of built nodes.
 342 func (t *Tree) Count() int {
 343         return t.order.count
 344 }
 345
 346 // Keys returns the keys of the built nodes, ordered by level and then by index.
 347 func (t *Tree) Keys() ([]nodeKey, error) {
 348         keys := make([]nodeKey, 0, t.order.count)
 349         for n := range t.order.count {
 350                 record, _, err := t.order.Read(n)
 351                 if err != nil {
 352                         return nil, err
 353                 }
 354
 355                 key, _ := decodeTreeRecord(record)
 356                 keys = append(keys, key)
 357         }
 358
 359         slices.SortFunc(keys, func(a, b nodeKey) int {
 360                 return cmp.Or(cmp.Compare(a.level, b.level), cmp.Compare(a.index, b.index))
 361         })
 362
 363         return keys, nil
 364 }
 365
 366 // leafSource returns the text that the leaf of message holds or summarizes.
 367 func leafSource(message Message) string {
 368         return message.Kind + ": " + message.Text
 369 }
 370
 371 // freeLeaf returns the leaf of message when its source fits in a node. Such a
 372 // leaf holds the source word for word and needs no model call.
 373 func freeLeaf(message Message) (Node, bool) {
 374         text := leafSource(message)
 375         if len(text) > nodeBytes {
 376                 return Node{}, false
 377         }
 378
 379         return Node{
 380                 Level: 0,
 381                 Index: message.ID,
 382                 Text:  text,
 383                 Size:  len(text),
 384         }, true
 385 }
 386
 387 // freeMerge returns the parent of two sibling nodes when both texts, joined by
 388 // a newline, fit in a node. Such a parent needs no model call.
 389 func freeMerge(left, right Node) (Node, bool) {
 390         if left.Size+1+right.Size > nodeBytes {
 391                 return Node{}, false
 392         }
 393
 394         text := left.Text + "\n" + right.Text
 395
 396         return Node{
 397                 Level: left.Level + 1,
 398                 Index: left.Index / 2,
 399                 Text:  text,
 400                 Size:  len(text),
 401         }, true
 402 }
 403
 404 // decodeTreeRecord reads the key and the position of a record of
 405 // index/tree.idx.
 406 func decodeTreeRecord(record []byte) (nodeKey, position) {
 407         key := nodeKey{
 408                 level: int(binary.LittleEndian.Uint32(record[positionBytes+8:])),
 409                 index: int(binary.LittleEndian.Uint64(record[positionBytes:])),
 410         }
 411
 412         return key, decodePosition(record)
 413 }
 414
 415 // nodeEnd returns the ID after the last message of node (l, i).
 416 func nodeEnd(level, index int) int {
 417         return (index + 1) << level
 418 }