consciousness/src/hippocampus/store/index.rs
Kent Overstreet 4696bb8b7d store: index ops take WriteTransaction, mutations batch properly
Index functions now take &WriteTransaction instead of &Database,
allowing callers to batch multiple index operations in a single
transaction. Store mutations (upsert, delete, rename, etc.) now
begin_write/commit their own transactions, ensuring atomicity.

- replay_relations uses single txn for all relation indexing
- Store::db() exposes Database for callers needing txn control
- Convenience wrappers open their own txn for simple cases

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
2026-04-13 21:44:20 -04:00

210 lines
7.5 KiB
Rust

// redb index tables
//
// capnp logs are source of truth; redb provides indexed access.
//
// Node tables:
// NODES: key → offset (current version)
// KEY_TO_UUID: key → uuid
// UUID_OFFSETS: uuid → offsets (multimap, all versions)
// NODES_BY_PROVENANCE: provenance → keys (multimap)
// NODES_BY_TYPE: [type_byte][timestamp_be] → key (for range queries by type+date)
//
// Relation tables:
// RELS: node_uuid → (other_uuid, strength, rel_type, is_outgoing) packed (multimap)
// Each relation stored twice — once per endpoint with direction bit.
//
// To get key from uuid: UUID_OFFSETS → read_node_at_offset() → node.key
use anyhow::{Context, Result};
use redb::{Database, MultimapTableDefinition, ReadableDatabase, ReadableTable, TableDefinition, WriteTransaction};
use std::path::Path;
// Node tables
pub const NODES: TableDefinition<&str, u64> = TableDefinition::new("nodes");
pub const KEY_TO_UUID: TableDefinition<&str, &[u8]> = TableDefinition::new("key_to_uuid");
pub const UUID_OFFSETS: MultimapTableDefinition<&[u8], u64> = MultimapTableDefinition::new("uuid_offsets");
pub const NODES_BY_PROVENANCE: MultimapTableDefinition<&str, &str> = MultimapTableDefinition::new("nodes_by_provenance");
// Composite key: [node_type: u8][timestamp: i64 BE] for range queries
pub const NODES_BY_TYPE: TableDefinition<&[u8], &str> = TableDefinition::new("nodes_by_type");
// Relations table - each relation stored twice (once per endpoint)
// Value: (other_uuid: [u8;16], strength: f32, rel_type: u8, is_outgoing: bool)
// Packed as 22 bytes: [other_uuid:16][strength:4][rel_type:1][is_outgoing:1]
pub const RELS: MultimapTableDefinition<&[u8], &[u8]> = MultimapTableDefinition::new("rels");
/// Open or create the redb database, ensuring all tables exist.
pub fn open_db(path: &Path) -> Result<Database> {
let db = Database::create(path)
.with_context(|| format!("create redb {}", path.display()))?;
// Ensure tables exist by opening a write transaction
let txn = db.begin_write()?;
{
// Node tables
let _ = txn.open_table(NODES)?;
let _ = txn.open_table(KEY_TO_UUID)?;
let _ = txn.open_multimap_table(UUID_OFFSETS)?;
let _ = txn.open_multimap_table(NODES_BY_PROVENANCE)?;
let _ = txn.open_table(NODES_BY_TYPE)?;
// Relations
let _ = txn.open_multimap_table(RELS)?;
}
txn.commit()?;
Ok(db)
}
/// Record a node's location in the index.
pub fn index_node(txn: &WriteTransaction, key: &str, offset: u64, uuid: &[u8; 16]) -> Result<()> {
let mut nodes_table = txn.open_table(NODES)?;
let mut key_uuid_table = txn.open_table(KEY_TO_UUID)?;
let mut uuid_offsets = txn.open_multimap_table(UUID_OFFSETS)?;
nodes_table.insert(key, offset)?;
key_uuid_table.insert(key, uuid.as_slice())?;
uuid_offsets.insert(uuid.as_slice(), offset)?;
Ok(())
}
/// Get offset for a node by key.
pub fn get_offset(db: &Database, key: &str) -> Result<Option<u64>> {
let txn = db.begin_read()?;
let table = txn.open_table(NODES)?;
Ok(table.get(key)?.map(|v| v.value()))
}
/// Check if a key exists in the index.
pub fn contains_key(db: &Database, key: &str) -> Result<bool> {
let txn = db.begin_read()?;
let table = txn.open_table(NODES)?;
Ok(table.get(key)?.is_some())
}
/// Get a node's UUID from its key.
pub fn get_uuid_for_key(db: &Database, key: &str) -> Result<Option<[u8; 16]>> {
let txn = db.begin_read()?;
let table = txn.open_table(KEY_TO_UUID)?;
match table.get(key)? {
Some(uuid) => {
let slice = uuid.value();
let mut arr = [0u8; 16];
arr.copy_from_slice(slice);
Ok(Some(arr))
}
None => Ok(None),
}
}
/// Get all offsets for a UUID (all versions). Returns newest first.
pub fn get_offsets_for_uuid(db: &Database, uuid: &[u8; 16]) -> Result<Vec<u64>> {
let txn = db.begin_read()?;
let table = txn.open_multimap_table(UUID_OFFSETS)?;
let mut offsets = Vec::new();
for entry in table.get(uuid.as_slice())? {
offsets.push(entry?.value());
}
// Sort descending so newest (highest offset) is first
offsets.sort_by(|a, b| b.cmp(a));
Ok(offsets)
}
/// Remove a node from the index (key mappings only; UUID history preserved).
pub fn remove_node(txn: &WriteTransaction, key: &str) -> Result<()> {
let mut nodes_table = txn.open_table(NODES)?;
let mut key_uuid_table = txn.open_table(KEY_TO_UUID)?;
// Note: UUID_OFFSETS is not cleared - preserves version history
nodes_table.remove(key)?;
key_uuid_table.remove(key)?;
Ok(())
}
/// Collect all keys from the index.
pub fn all_keys(db: &Database) -> Result<Vec<String>> {
let txn = db.begin_read()?;
let table = txn.open_table(NODES)?;
let mut keys = Vec::new();
for entry in table.iter()? {
let (key, _) = entry?;
keys.push(key.value().to_string());
}
Ok(keys)
}
// ── Relation index operations ──────────────────────────────────────
//
// RELS value format: [other_uuid:16][strength:4][rel_type:1][is_outgoing:1] = 22 bytes
/// Pack relation data into bytes for RELS table.
fn pack_rel(other_uuid: &[u8; 16], strength: f32, rel_type: u8, is_outgoing: bool) -> [u8; 22] {
let mut buf = [0u8; 22];
buf[0..16].copy_from_slice(other_uuid);
buf[16..20].copy_from_slice(&strength.to_be_bytes());
buf[20] = rel_type;
buf[21] = if is_outgoing { 1 } else { 0 };
buf
}
/// Unpack relation data from RELS table.
pub fn unpack_rel(data: &[u8]) -> ([u8; 16], f32, u8, bool) {
let mut other_uuid = [0u8; 16];
other_uuid.copy_from_slice(&data[0..16]);
let strength = f32::from_be_bytes([data[16], data[17], data[18], data[19]]);
let rel_type = data[20];
let is_outgoing = data[21] != 0;
(other_uuid, strength, rel_type, is_outgoing)
}
/// Index a relation: store twice (once per endpoint).
pub fn index_relation(
txn: &WriteTransaction,
source_uuid: &[u8; 16],
target_uuid: &[u8; 16],
strength: f32,
rel_type: u8,
) -> Result<()> {
let mut rels = txn.open_multimap_table(RELS)?;
// Store outgoing: source → (target, strength, type, true)
let outgoing = pack_rel(target_uuid, strength, rel_type, true);
rels.insert(source_uuid.as_slice(), outgoing.as_slice())?;
// Store incoming: target → (source, strength, type, false)
let incoming = pack_rel(source_uuid, strength, rel_type, false);
rels.insert(target_uuid.as_slice(), incoming.as_slice())?;
Ok(())
}
/// Remove a relation from the index.
pub fn remove_relation(
txn: &WriteTransaction,
source_uuid: &[u8; 16],
target_uuid: &[u8; 16],
strength: f32,
rel_type: u8,
) -> Result<()> {
let mut rels = txn.open_multimap_table(RELS)?;
let outgoing = pack_rel(target_uuid, strength, rel_type, true);
rels.remove(source_uuid.as_slice(), outgoing.as_slice())?;
let incoming = pack_rel(source_uuid, strength, rel_type, false);
rels.remove(target_uuid.as_slice(), incoming.as_slice())?;
Ok(())
}
/// Get all edges for a node. Returns (other_uuid, strength, rel_type, is_outgoing).
pub fn edges_for_node(db: &Database, node_uuid: &[u8; 16]) -> Result<Vec<([u8; 16], f32, u8, bool)>> {
let txn = db.begin_read()?;
let rels = txn.open_multimap_table(RELS)?;
let mut edges = Vec::new();
for entry in rels.get(node_uuid.as_slice())? {
let guard = entry?;
let slice = guard.value();
let mut data = [0u8; 22];
data.copy_from_slice(slice);
edges.push(unpack_rel(&data));
}
Ok(edges)
}