Add local context to memory operation failures

[onechat] / tree.go

diff --git a/tree.go b/tree.go
index ee3ff6807dbe8c660b8aba4cec90e6eb095c9867..db93b1d70357bb98e893e6ee6484b05570b24885 100644 (file)
--- a/tree.go
+++ b/tree.go
@@ -1,12 +1,16 @@
 package main
 
 import (
 package main
 
 import (

+       "cmp"

        "context"
        "context"

+       "encoding/binary"

        "encoding/json"
        "encoding/json"

+       "errors"

        "fmt"
        "math/bits"
        "os"
        "path/filepath"
        "fmt"
        "math/bits"
        "os"
        "path/filepath"

+       "slices"

        "time"
        "unicode/utf8"
 )
        "time"
        "unicode/utf8"
 )


@@ -15,6 +19,31 @@ import (
 // few bytes over it, because the view measures the real sizes.
 const nodeBytes = 512
 
 // few bytes over it, because the view measures the real sizes.
 const nodeBytes = 512
 

+// treeRecordBytes is the size of a record of index/tree.idx: the position of
+// the node line, the node index, and the node level.
+const treeRecordBytes = 32
+
+// 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. 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.
+//
+// The tree keeps no node in memory. The index index/tree.idx lists the node
+// lines in file order, and index/tree-L.idx holds the position of node (L, i)
+// in record i, with a zero record for a node that is not built. At open the
+// tree indexes the lines after the last indexed one, so a missing index is
+// built again from the day files.
+type Tree struct {
+       clock   func() time.Time
+       journal Journal
+       dir     string
+       order   *Records
+       levels  map[int]*Records
+       failure error
+}
+

 // 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
 // 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


@@ -26,43 +55,82 @@ type Node struct {
        Size  int    `json:"size"`
 }
 
        Size  int    `json:"size"`
 }
 

-// nodeKey names the node (l, i).
-type nodeKey struct {
+// NodeKey names the node (l, i).
+type NodeKey struct {

        level int
        index int
 }
 
        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.
+// openTree opens the tree of the chat in root and indexes its new lines. 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) {
 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)
+       for _, dir := range []string{"tree", "index"} {
+               err := os.MkdirAll(filepath.Join(root, dir), 0o700)
+               if err != nil {
+                       return nil, fmt.Errorf("create %s directory: %w", dir, err)
+               }

        }
 
        t := &Tree{
                clock: clock,
        }
 
        t := &Tree{
                clock: clock,

-               journal: journal{
+               journal: Journal{

                        dir: filepath.Join(root, "tree"),
                },
                        dir: filepath.Join(root, "tree"),
                },

-               nodes: make(map[nodeKey]Node),
+               dir:    filepath.Join(root, "index"),
+               levels: make(map[int]*Records),
+       }
+
+       err := t.open(ctx, total)
+       if err != nil {
+               return nil, errors.Join(fmt.Errorf("index tree: %w", err), t.Close())

        }
 
        }
 

-       err = readLines(ctx, t.journal.dir, func(node Node) error {
-               err := t.check(node)
+       return t, nil
+}
+
+func (t *Tree) open(ctx context.Context, total int) error {
+       var err error
+
+       t.order, err = openRecords(filepath.Join(t.dir, "tree.idx"), treeRecordBytes)
+       if err != nil {
+               return fmt.Errorf("open node order index: %w", err)
+       }
+
+       start := Position{}
+       if t.order.count == 0 {
+               err = t.removeLevels()
+               if err != nil {
+                       return fmt.Errorf("remove empty level indexes: %w", err)
+               }
+       } else {
+               record, _, err := t.order.Read(t.order.count - 1)
+               if err != nil {
+                       return fmt.Errorf("read last node order: %w", err)
+               }
+
+               // A crash can stop an add between the two index writes, so the
+               // last node gets its level record again.
+               node, last := treeDecodeRecord(record)
+               err = t.writeLevel(node, last)
+               if err != nil {
+                       return fmt.Errorf("restore last node level: %w", err)
+               }
+
+               start = Position{
+                       day:    last.day,
+                       offset: last.offset + uint64(last.length),
+               }
+       }
+
+       return t.journal.Scan(ctx, start, func(line []byte, place Position) error {
+               var node Node
+
+               err := json.Unmarshal(line, &node)
+               if err != nil {
+                       return fmt.Errorf("decode tree node: %w", err)
+               }
+
+               err = t.check(node)

                if err != nil {
                        return err
                }
                if err != nil {
                        return err
                }


@@ -71,15 +139,56 @@ func openTree(ctx context.Context, root string, total int, clock func() time.Tim
                        return fmt.Errorf("node (%d, %d) covers messages after the log", node.Level, node.Index)
                }
 
                        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
+               return t.addIndex(NodeKey{node.Level, node.Index}, place)

        })
        })

+}
+
+// Close closes the files of the tree.
+func (t *Tree) Close() error {
+       err := t.journal.Close()
+       if t.order != nil {
+               err = errors.Join(err, t.order.Close())
+               t.order = nil
+       }
+
+       for level, index := range t.levels {
+               err = errors.Join(err, index.Close())
+
+               delete(t.levels, level)
+       }
+
+       return err
+}
+
+// removeLevels removes the level indexes, which are not valid without the order
+// index.
+func (t *Tree) removeLevels() error {
+       paths, err := filepath.Glob(filepath.Join(t.dir, "tree-*.idx"))

        if err != nil {
        if err != nil {

-               return nil, fmt.Errorf("load tree: %w", err)
+               return fmt.Errorf("list level indexes: %w", err)

        }
 
        }
 

-       return t, nil
+       for _, path := range paths {
+               err = os.Remove(path)
+               if err != nil {
+                       return fmt.Errorf("remove level index: %w", err)
+               }
+       }
+
+       return nil
+}
+
+// writeLevel writes the position of node to its level index.
+func (t *Tree) writeLevel(node NodeKey, place Position) error {
+       level, err := t.level(node.level)
+       if err != nil {
+               return fmt.Errorf("open node level index: %w", err)
+       }
+
+       record := make([]byte, positionBytes)
+       encodePosition(record, place)
+
+       return level.Write(node.index, record)

 }
 
 // check makes sure that node is valid and can join the tree now.
 }
 
 // check makes sure that node is valid and can join the tree now.


@@ -92,20 +201,94 @@ func (t *Tree) check(node Node) error {
                return fmt.Errorf("node (%d, %d) has no valid text of its size", node.Level, node.Index)
        }
 
                return fmt.Errorf("node (%d, %d) has no valid text of its size", node.Level, node.Index)
        }
 

-       if t.Built(node.Level, node.Index) {
+       built, err := t.Built(node.Level, node.Index)
+       if err != nil {
+               return fmt.Errorf("check target node: %w", err)
+       }
+
+       if built {

                return fmt.Errorf("node (%d, %d) is already 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)
+       if node.Level == 0 {
+               return nil
+       }
+
+       for _, index := range []int{2 * node.Index, 2*node.Index + 1} {
+               built, err = t.Built(node.Level-1, index)
+               if err != nil {
+                       return fmt.Errorf("check child node: %w", err)
+               }
+
+               if !built {
+                       return fmt.Errorf("node (%d, %d) has a child that is not built", node.Level, node.Index)
+               }

        }
 
        return nil
 }
 
        }
 
        return nil
 }
 

-// Close closes the open day file.
-func (t *Tree) Close() error {
-       return t.journal.Close()
+// addIndex adds the node line at place to both indexes.
+func (t *Tree) addIndex(node NodeKey, place Position) error {
+       record := make([]byte, treeRecordBytes)
+       encodePosition(record, place)
+       binary.LittleEndian.PutUint64(record[positionBytes:], uint64(node.index))
+       binary.LittleEndian.PutUint32(record[positionBytes+8:], uint32(node.level))
+
+       err := t.order.Append(record)
+       if err != nil {
+               return fmt.Errorf("append node order: %w", err)
+       }
+
+       return t.writeLevel(node, place)
+}
+
+// level returns the index of level, and opens it at first use.
+func (t *Tree) level(level int) (*Records, error) {
+       index, open := t.levels[level]
+       if open {
+               return index, nil
+       }
+
+       index, err := openRecords(filepath.Join(t.dir, fmt.Sprintf("tree-%d.idx", level)), positionBytes)
+       if err != nil {
+               return nil, err
+       }
+
+       t.levels[level] = index
+
+       return index, nil
+}
+
+// Built tells if the node (l, i) is built.
+func (t *Tree) Built(level, index int) (bool, error) {
+       _, built, err := t.place(level, index)
+       return built, err
+}
+
+// place returns the position of the node (l, i) if it is built.
+func (t *Tree) place(level, index int) (Position, bool, error) {
+       if level < 0 || level >= bits.UintSize-2 || index < 0 {
+               return Position{}, false, nil
+       }
+
+       records, err := t.level(level)
+       if err != nil {
+               return Position{}, false, fmt.Errorf("open node level index: %w", err)
+       }
+
+       record, found, err := records.Read(index)
+       if err != nil || !found {
+               if err != nil {
+                       return Position{}, false, fmt.Errorf("read node level record: %w", err)
+               }
+
+               return Position{}, false, nil
+       }
+
+       place := decodePosition(record)
+
+       return place, place.length > 0, nil

 }
 
 // Add writes a new node to the tree.
 }
 
 // Add writes a new node to the tree.


@@ -116,7 +299,7 @@ func (t *Tree) Add(node Node) error {
 
        err := t.check(node)
        if err != nil {
 
        err := t.check(node)
        if err != nil {

-               return err
+               return fmt.Errorf("validate node: %w", err)

        }
 
        line, err := json.Marshal(node)
        }
 
        line, err := json.Marshal(node)


@@ -124,39 +307,95 @@ func (t *Tree) Add(node Node) error {
                return fmt.Errorf("encode node: %w", err)
        }
 
                return fmt.Errorf("encode node: %w", err)
        }
 

-       err = t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
+       place, err := t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
+       if err == nil {
+               err = t.addIndex(NodeKey{node.Level, node.Index}, place)
+       }
+

        if err != nil {
                t.failure = err
                return fmt.Errorf("write node (%d, %d): %w", node.Level, node.Index, err)
        }
 
        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
 }
 
        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 reads the node (l, i) if it is built.
+func (t *Tree) Node(level, index int) (Node, bool, error) {
+       place, built, err := t.place(level, index)
+       if err != nil || !built {
+               if err != nil {
+                       return Node{}, false, fmt.Errorf("find node position: %w", err)
+               }
+
+               return Node{}, false, nil
+       }
+
+       line, err := t.journal.Read(place)
+       if err != nil {
+               return Node{}, false, fmt.Errorf("read node (%d, %d): %w", level, index, err)
+       }
+
+       var node Node
+
+       err = json.Unmarshal(line, &node)
+       if err != nil {
+               return Node{}, false, fmt.Errorf("decode node (%d, %d): %w", level, index, err)
+       }
+
+       if node.Level != level || node.Index != index {
+               return Node{}, false, fmt.Errorf("index of node (%d, %d) points to node (%d, %d)", level, index, node.Level, node.Index)
+       }
+
+       return node, true, nil

 }
 
 }
 

-// 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
+// Keys returns the keys of the built nodes, ordered by level and then by index.
+func (t *Tree) Keys() ([]NodeKey, error) {
+       keys := make([]NodeKey, 0, t.order.count)
+       for n := range t.order.count {
+               record, _, err := t.order.Read(n)
+               if err != nil {
+                       return nil, err
+               }
+
+               key, _ := treeDecodeRecord(record)
+               keys = append(keys, key)
+       }
+
+       slices.SortFunc(keys, func(a, b NodeKey) int {
+               return cmp.Or(cmp.Compare(a.level, b.level), cmp.Compare(a.index, b.index))
+       })
+
+       return keys, nil

 }
 
 }
 

-// leafSource returns the text that the leaf of message holds or summarizes.
-func leafSource(message Message) string {
-       return message.Kind + ": " + message.Text
+// Count returns the number of built nodes.
+func (t *Tree) Count() int {
+       return t.order.count

 }
 
 }
 

-// 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.
+// treeDecodeRecord reads the key and the position of a record of
+// index/tree.idx.
+func treeDecodeRecord(record []byte) (NodeKey, Position) {
+       key := NodeKey{
+               level: int(binary.LittleEndian.Uint32(record[positionBytes+8:])),
+               index: int(binary.LittleEndian.Uint64(record[positionBytes:])),
+       }
+
+       return key, decodePosition(record)
+}
+
+// nodeEnd returns the ID after the last message of node (l, i).
+func nodeEnd(level, index int) int {
+       return (index + 1) << level
+}
+
+// freeLeaf returns a text-only leaf when its source fits in a node. Images need
+// a model description even when their text fits.

 func freeLeaf(message Message) (Node, bool) {
        text := leafSource(message)
 func freeLeaf(message Message) (Node, bool) {
        text := leafSource(message)

-       if len(text) > nodeBytes {
+       if len(text) > nodeBytes || len(imageFiles(message.Text)) > 0 {

                return Node{}, false
        }
 
                return Node{}, false
        }
 


@@ -168,6 +407,11 @@ func freeLeaf(message Message) (Node, bool) {
        }, true
 }
 
        }, true
 }
 

+// leafSource returns the text that the leaf of message holds or summarizes.
+func leafSource(message Message) string {
+       return message.Kind + ": " + message.Text
+}
+

 // 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) {
 // 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) {


@@ -184,8 +428,3 @@ func freeMerge(left, right Node) (Node, bool) {
                Size:  len(text),
        }, true
 }
                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
-}