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