1
package sqlite3
2
 
3
import (
4
	"context"
5
	"database/sql"
6
 
7
	"go.miragespace.co/specter/spec/chord"
8
)
9
 
10
const (
11
	SimpleFlag = 1 << iota
12
	PrefixFlag
13
	LeaseFlag
14
)
15
 
16
type KeyTracker struct {
17
	Key   []byte
18
	Hash  uint64
19
	Flags uint8 // Bitmask: 1=Simple, 2=Prefix, 4=Lease
20
}
21
 
22
// optimized by GPT 4o-mini via flags
23
func (s *SqliteKV) updateKeyTracker(ctx context.Context, tx *sql.Tx, key []byte, addFlags, removeFlags uint8) error {
24
	var (
25
		hash  int64
26
		flags uint8
27
	)
28
 
29
	// Fetch existing tracker
30
	lookupErr := tx.StmtContext(ctx, s.stmts.trackerLookup).QueryRow(key).Scan(&hash, &flags)
31
	if lookupErr != nil && lookupErr != sql.ErrNoRows {
32
		return lookupErr
33
	}
34
 
35
	found := lookupErr == nil
36
 
37
	// consistency check: if this instance has a different hash function, refuse to modify
38
	if found && scanInt64AsUint64(hash) != s.hashFn(key) {
39
		return chord.ErrKVHashFnChanged
40
	}
41
 
42
	// If tracker does not exist but we're adding flags, create a new entry
43
	if !found {
44
		if addFlags != 0 {
45
			_, err := tx.StmtContext(ctx, s.stmts.trackerInsert).Exec(key, bindUint64AsInt64(s.hashFn(key)), addFlags)
46
			return err
47
		}
48
		return nil
49
	}
50
 
51
	// Check for remaining values before removing flags
52
	if removeFlags&PrefixFlag != 0 {
53
		var count int64
54
		if err := tx.StmtContext(ctx, s.stmts.prefixCount).QueryRow(key).Scan(&count); err != nil {
55
			return err
56
		}
57
		if count > 0 {
58
			removeFlags &^= PrefixFlag // Cancel removing PrefixFlag
59
		}
60
	}
61
 
62
	// Update bitmask flags
63
	newFlags := (flags | addFlags) &^ removeFlags
64
 
65
	// If no values remain, delete tracker
66
	if newFlags == 0 {
67
		_, err := tx.StmtContext(ctx, s.stmts.trackerDelete).Exec(key)
68
		return err
69
	}
70
 
71
	// Otherwise, update tracker
72
	_, err := tx.StmtContext(ctx, s.stmts.trackerUpdate).Exec(newFlags, key)
73
	return err
74
}