use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use plugmem_arena::{
Arena, ArenaCfg, BlobHeap, BlobHeapCfg, BlobId, ChunkPool, ChunkPoolCfg, Interner, ListHandle,
ShardMode, Slot, TermId, key,
};
use crate::config::Config;
use crate::error::Error;
use crate::id::{EdgeId, EntityId, FactId, NONE_U32};
use crate::index::IdListIndex;
use crate::index::bm25::Bm25Index;
use crate::index::hnsw::HnswGraph;
use crate::index::vecpool::VecPool;
use crate::journal::{JournalScan, Op, scan};
use crate::model::{
EdgeHistorySlot, EdgeSlot, EntityByName, EntityRecord, FactAux, FactRecord, TemporalSlot,
VALID_TO_OPEN, close_edge_history_payload, edge_history_key, edge_key, fact_flags,
};
use crate::storage::Storage;
use crate::tokenizer::Tokenizer;
use maintain::TOKENIZER_INDEX_VERSION;
const SIMILAR_CANDIDATE_CAP: usize = 32;
const MAX_TAGS: usize = 32;
mod maintain;
mod migrations;
mod persist;
mod recall;
mod reembed;
mod shards;
mod tags;
pub use maintain::{MaintainReport, MaintenanceMode, MaintenanceOptions};
pub use recall::{RecallQuery, RecallResult, RecallScratch, RecalledEdge, RecalledFact, source};
pub use reembed::{ReembedError, ReembedReport};
pub use shards::ShardLayout;
pub use tags::{DEFAULT_TAG_PAGE_LIMIT, MAX_TAG_PAGE_LIMIT, TagPage, TagQuery, TagSummary};
pub const MAX_VECTOR_SPACE_ID_BYTES: usize = 256;
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct RememberInput<'a> {
pub now: u64,
pub text: &'a str,
pub entity: Option<&'a str>,
pub tags: &'a [&'a str],
pub links: &'a [(&'a str, &'a str)],
pub vector: Option<&'a [f32]>,
pub valid_from: Option<u64>,
pub metadata: Option<&'a [(&'a str, &'a str)]>,
}
impl<'a> RememberInput<'a> {
pub fn text(now: u64, text: &'a str) -> Self {
Self {
now,
text,
entity: None,
tags: &[],
links: &[],
vector: None,
valid_from: None,
metadata: None,
}
}
}
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct LinkInput<'a> {
pub now: u64,
pub src: &'a str,
pub rel: &'a str,
pub dst: &'a str,
pub provenance: Option<FactId>,
}
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct UnlinkInput<'a> {
pub now: u64,
pub src: &'a str,
pub rel: &'a str,
pub dst: &'a str,
}
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct RememberOutcome {
pub id: FactId,
pub entity: Option<EntityId>,
pub similar: Vec<Similar>,
}
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "serde", serde(tag = "status", rename_all = "snake_case"))]
pub enum GuardedRememberOutcome {
Stored {
outcome: RememberOutcome,
checked: bool,
},
Blocked {
similar: Vec<Similar>,
},
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct RemoveTagReport {
pub affected: u32,
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Similar {
pub id: FactId,
pub score: f32,
pub reason: SimilarReason,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum SimilarReason {
LexicalOverlap,
VectorCosine,
}
#[derive(Debug, Default)]
struct SimilarityScratch {
new_terms: Vec<u32>,
unknown: String,
unknown_ranges: Vec<(usize, usize)>,
candidate_terms: Vec<u32>,
vector: Vec<u8>,
}
#[derive(Clone, Copy)]
enum SimilarVector<'s> {
None,
Stored(u32),
Encoded(&'s [u8]),
}
#[derive(Clone, Copy)]
struct SimilarityQuery<'s> {
entity: EntityId,
exclude: Option<FactId>,
terms: &'s [u32],
term_count: usize,
vector: SimilarVector<'s>,
}
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct FactView<'a> {
pub record: FactRecord,
pub text: &'a str,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum FactFault {
Text,
Vector,
Metadata,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[non_exhaustive]
pub struct Stats {
pub facts: usize,
pub entities: usize,
pub terms: usize,
pub edges: usize,
pub edge_versions: usize,
pub vectors: usize,
pub tombstones: usize,
pub hnsw_indexed: u32,
pub next_fact: u32,
pub next_entity: u32,
pub next_edge: u32,
pub db_uuid: u128,
pub pool_bytes: usize,
pub shards: ShardLayout,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct OpenReport {
pub replayed: usize,
pub skipped: usize,
pub truncated_tail: bool,
}
pub struct Memory<'a> {
cfg: Config,
facts: Arena<'a, FactRecord>,
fact_aux: Arena<'a, FactAux>,
entities: Arena<'a, EntityRecord>,
by_name: Arena<'a, EntityByName>,
edges_out: Arena<'a, EdgeSlot>,
edges_in: Arena<'a, EdgeSlot>,
edges_hist_out: Arena<'a, EdgeHistorySlot>,
edges_hist_in: Arena<'a, EdgeHistorySlot>,
temporal: Arena<'a, TemporalSlot>,
texts: BlobHeap<'a>,
metas: BlobHeap<'a>,
terms: Interner<'a>,
tag_lists: ChunkPool<'a>,
bm25: Bm25Index<'a>,
tags_idx: IdListIndex<'a>,
tag_catalog: tags::TagCatalog,
entity_facts: IdListIndex<'a>,
vecs: VecPool<'a>,
hnsw: HnswGraph<'a>,
vector_space: Option<String>,
next_fact: u32,
next_entity: u32,
next_edge: u32,
tombstones: usize,
bm25_tokenizer_version: u32,
tokenizer: Tokenizer,
tf_scratch: Vec<(u32, u8)>,
tf_probe: Vec<u32>,
name_scratch: String,
similarity_scratch: SimilarityScratch,
}
impl<'a> Memory<'a> {
pub fn new(cfg: Config) -> Result<Self, Error> {
cfg.validate()?;
let uni =
|shards: usize| ArenaCfg::new(shards, ShardMode::Uniform).with_max_bytes(cfg.max_bytes);
let ord =
|shards: usize| ArenaCfg::new(shards, ShardMode::Ordered).with_max_bytes(cfg.max_bytes);
let blob = BlobHeapCfg::new()
.with_max_bytes(cfg.max_bytes)
.with_max_blob(cfg.max_blob);
Ok(Self {
facts: Arena::new(uni(cfg.shards_facts))?,
fact_aux: Arena::new(uni(cfg.shards_facts))?,
entities: Arena::new(uni(cfg.shards_entities))?,
by_name: Arena::new(ord(cfg.shards_entities))?,
edges_out: Arena::new(ord(cfg.shards_edges))?,
edges_in: Arena::new(ord(cfg.shards_edges))?,
edges_hist_out: Arena::new(ord(cfg.shards_edges))?,
edges_hist_in: Arena::new(ord(cfg.shards_edges))?,
temporal: Arena::new(ord(cfg.shards_temporal))?,
texts: BlobHeap::new(blob),
metas: BlobHeap::new(blob),
terms: Interner::new(blob),
tag_lists: ChunkPool::new(ChunkPoolCfg::new().with_max_bytes(cfg.max_bytes)),
bm25: Bm25Index::new(cfg.shards_postings, cfg.max_bytes)?,
tags_idx: IdListIndex::new(cfg.shards_postings, cfg.max_bytes)?,
tag_catalog: tags::TagCatalog::new(),
entity_facts: IdListIndex::new(cfg.shards_entities, cfg.max_bytes)?,
vecs: VecPool::new(cfg.dim, cfg.max_bytes),
hnsw: HnswGraph::new(cfg.hnsw_m, cfg.hnsw_m0, cfg.max_bytes)?,
vector_space: None,
next_fact: 0,
next_entity: 0,
next_edge: 0,
tombstones: 0,
bm25_tokenizer_version: maintain::TOKENIZER_INDEX_VERSION,
tokenizer: Tokenizer::new(),
tf_scratch: Vec::new(),
tf_probe: Vec::new(),
name_scratch: String::new(),
similarity_scratch: SimilarityScratch::default(),
cfg,
})
}
pub fn open<S: Storage>(store: &mut S, cfg: Config) -> Result<(Self, OpenReport), Error> {
let snapshot = store
.read_snapshot()
.map_err(|e| Error::Storage(format!("{e:?}")))?;
let journal = store
.read_journal()
.map_err(|e| Error::Storage(format!("{e:?}")))?;
Self::from_bytes(snapshot.as_deref(), &journal, cfg)
}
pub fn from_bytes(
snapshot: Option<&[u8]>,
journal: &[u8],
cfg: Config,
) -> Result<(Self, OpenReport), Error> {
let mut mem = match snapshot {
Some(bytes) => Self::load_snapshot(bytes, cfg)?,
None => Self::new(cfg)?,
};
let report = mem.replay(journal)?;
Ok((mem, report))
}
pub fn from_bytes_borrowed(
snapshot: &'a [u8],
journal: &[u8],
cfg: Config,
) -> Result<Self, Error> {
let mem = Self::load_snapshot_borrowed(snapshot, cfg)?;
let JournalScan { entries, .. } = scan(journal)?;
if !entries.is_empty() {
return Err(Error::Invalid(
"read-only open requires a checkpointed (empty) journal",
));
}
Ok(mem)
}
pub fn from_bytes_overlay(
snapshot: &'a [u8],
journal: &[u8],
cfg: Config,
) -> Result<(Self, OpenReport), Error> {
let mut mem = Self::load_snapshot_borrowed(snapshot, cfg)?;
let report = mem.replay(journal)?;
Ok((mem, report))
}
fn replay(&mut self, journal: &[u8]) -> Result<OpenReport, Error> {
let JournalScan {
entries,
truncated_tail,
} = scan(journal)?;
let mut report = OpenReport {
truncated_tail,
..OpenReport::default()
};
for entry in entries {
let op = Op::decode(entry.op, entry.payload)?;
match op {
Op::Remember {
now,
valid_from,
entity,
text,
ref tags,
ref links,
ref vector,
ref metadata,
revises,
assigned,
} => {
if assigned.0 < self.next_fact {
report.skipped += 1;
continue;
}
if assigned.0 != self.next_fact {
return Err(Error::Corrupt("journal fact ids are not contiguous"));
}
if !vector.is_empty() && vector.len() != self.cfg.dim {
return Err(Error::Corrupt(
"journal vector dimension disagrees with dim",
));
}
if tags.len() > MAX_TAGS {
return Err(Error::Corrupt("journal record carries too many tags"));
}
if let Some(target) = revises.some() {
self.check_revisable(target)
.map_err(|_| Error::Corrupt("journal revises an unrevisable fact"))?;
}
self.apply_remember(
&RememberInput {
now,
text,
entity,
tags: &tags.to_vec(),
links: &links.to_vec(),
vector: (!vector.is_empty()).then_some(vector.as_slice()),
valid_from: Some(valid_from),
metadata: (!metadata.is_empty()).then_some(metadata.as_slice()),
},
revises,
None,
)?;
if let Some(target) = revises.some() {
self.close_target(target, valid_from);
}
report.replayed += 1;
}
Op::Forget { fact, .. } => {
match self.apply_forget(fact) {
Ok(_) => report.replayed += 1,
Err(Error::NotFound(_)) => {
return Err(Error::Corrupt("journal forgets an unknown fact"));
}
Err(e) => return Err(e),
}
}
Op::Link {
now,
src,
rel,
dst,
provenance,
} => {
self.apply_link(now, src, rel, dst, provenance)?;
report.replayed += 1;
}
Op::Unlink { now, src, rel, dst } => {
self.apply_unlink(now, src, rel, dst)?;
report.replayed += 1;
}
Op::RemoveTag { now, tag } => {
self.apply_remove_tag(now, tag)?;
report.replayed += 1;
}
Op::SetVectorSpace { space } => {
self.apply_set_vector_space(space)?;
report.replayed += 1;
}
Op::Maintain {
mode,
max_hnsw_inserts,
..
} => {
let options =
maintain::MaintenanceOptions::from_journal(mode, max_hnsw_inserts)?;
self.replay_maintain_with_options(options)?;
report.replayed += 1;
}
}
}
Ok(report)
}
pub fn vector_space(&self) -> Option<&str> {
self.vector_space.as_deref()
}
pub fn claim_vector_space<S: Storage>(
&mut self,
store: &mut S,
space: &str,
) -> Result<bool, Error> {
let needs_claim = self.check_vector_space_claim(space)?;
if !needs_claim {
return Ok(false);
}
let mut entry = Vec::new();
Op::SetVectorSpace { space }.encode(&mut entry);
store
.append_journal(&entry)
.map_err(|e| Error::Storage(format!("{e:?}")))?;
self.vector_space = Some(space.into());
Ok(true)
}
fn check_vector_space_claim(&self, space: &str) -> Result<bool, Error> {
Self::validate_vector_space(space)?;
if let Some(stored) = &self.vector_space {
if stored == space {
return Ok(false);
}
if !self.vecs.is_empty() {
return Err(Error::VectorSpaceMismatch {
stored: stored.clone(),
requested: space.into(),
});
}
}
if !self.vecs.is_empty() {
return Err(Error::UntrackedVectorSpace);
}
Ok(true)
}
fn apply_set_vector_space(&mut self, space: &str) -> Result<(), Error> {
Self::validate_vector_space(space)?;
if let Some(stored) = &self.vector_space {
if stored == space {
return Ok(());
}
if !self.vecs.is_empty() {
return Err(Error::Corrupt(
"journal changes an established vector space",
));
}
}
if !self.vecs.is_empty() {
return Err(Error::Corrupt(
"journal assigns a vector space after vector records",
));
}
self.vector_space = Some(space.into());
Ok(())
}
pub(super) fn validate_vector_space(space: &str) -> Result<(), Error> {
if space.is_empty() {
return Err(Error::Invalid("vector space must not be empty"));
}
if space.len() > MAX_VECTOR_SPACE_ID_BYTES {
return Err(Error::TooLarge {
what: "vector space",
len: space.len(),
max: MAX_VECTOR_SPACE_ID_BYTES,
});
}
if space.bytes().any(|b| b < b' ' || b == 0x7f) {
return Err(Error::Invalid(
"vector space must not contain control bytes",
));
}
Ok(())
}
pub fn remember<S: Storage>(
&mut self,
store: &mut S,
input: RememberInput<'_>,
) -> Result<RememberOutcome, Error> {
self.validate_input(&input)?;
let mut outcome = self.apply_remember(&input, FactId::NONE, None)?;
self.find_similar(&mut outcome);
self.journal_remember(store, &input, FactId::NONE, outcome.id)?;
Ok(outcome)
}
pub fn remember_guarded<S: Storage>(
&mut self,
store: &mut S,
input: RememberInput<'_>,
) -> Result<GuardedRememberOutcome, Error> {
self.remember_guarded_with_vector_space(store, input, None)
}
#[doc(hidden)]
pub fn remember_guarded_with_vector_space<S: Storage>(
&mut self,
store: &mut S,
input: RememberInput<'_>,
vector_space: Option<&str>,
) -> Result<GuardedRememberOutcome, Error> {
self.validate_input(&input)?;
if let Some(space) = vector_space {
self.check_vector_space_claim(space)?;
}
let checked = input.entity.is_some();
let similar = self.find_similar_input(&input)?;
if !similar.is_empty() {
return Ok(GuardedRememberOutcome::Blocked { similar });
}
if let Some(space) = vector_space {
self.claim_vector_space(store, space)?;
}
let outcome = self.apply_remember(&input, FactId::NONE, None)?;
self.journal_remember(store, &input, FactId::NONE, outcome.id)?;
Ok(GuardedRememberOutcome::Stored { outcome, checked })
}
pub fn remember_batch<S: Storage>(
&mut self,
store: &mut S,
inputs: &[RememberInput<'_>],
skip_similar: bool,
) -> Result<Vec<RememberOutcome>, Error> {
let mut out = Vec::with_capacity(inputs.len());
for input in inputs {
self.validate_input(input)?;
let mut outcome = self.apply_remember(input, FactId::NONE, None)?;
if !skip_similar {
self.find_similar(&mut outcome);
}
self.journal_remember(store, input, FactId::NONE, outcome.id)?;
out.push(outcome);
}
Ok(out)
}
fn find_similar_input(&mut self, input: &RememberInput<'_>) -> Result<Vec<Similar>, Error> {
let Some(name) = input.entity else {
return Ok(Vec::new());
};
let Some(entity) = self.lookup_entity_name(name) else {
return Ok(Vec::new());
};
let mut scratch = core::mem::take(&mut self.similarity_scratch);
scratch.new_terms.clear();
scratch.unknown.clear();
scratch.unknown_ranges.clear();
scratch.candidate_terms.clear();
scratch.vector.clear();
let terms = &self.terms;
let new_terms = &mut scratch.new_terms;
let unknown = &mut scratch.unknown;
let unknown_ranges = &mut scratch.unknown_ranges;
self.tokenizer.tokenize(input.text, &mut |token| {
if let Some(term) = terms.lookup(token) {
if !new_terms.contains(&term.0) {
new_terms.push(term.0);
}
return;
}
if unknown_ranges
.iter()
.any(|&(start, end)| &unknown[start..end] == token)
{
return;
}
let start = unknown.len();
unknown.push_str(token);
unknown_ranges.push((start, unknown.len()));
});
let new_term_count = scratch.new_terms.len() + scratch.unknown_ranges.len();
let mut similar = Vec::new();
let result = (|| {
let encoded = match input.vector {
Some(vector) => {
self.vecs
.encode_slot_into(FactId::NONE, vector, &mut scratch.vector)?;
SimilarVector::Encoded(&scratch.vector)
}
None => SimilarVector::None,
};
self.scan_similar(
SimilarityQuery {
entity,
exclude: None,
terms: &scratch.new_terms,
term_count: new_term_count,
vector: encoded,
},
&mut scratch.candidate_terms,
&mut similar,
);
Ok::<(), Error>(())
})();
self.similarity_scratch = scratch;
result.map(|()| similar)
}
fn find_similar(&mut self, outcome: &mut RememberOutcome) {
let Some(entity) = outcome.entity else { return };
let new_vec = self
.fact(outcome.id)
.filter(|record| record.has_vector())
.map(|record| record.vector);
let mut scratch = core::mem::take(&mut self.similarity_scratch);
scratch.new_terms.clear();
scratch
.new_terms
.extend(self.tf_scratch.iter().map(|&(term, _)| term));
let new_term_count = scratch.new_terms.len();
self.scan_similar(
SimilarityQuery {
entity,
exclude: Some(outcome.id),
terms: &scratch.new_terms,
term_count: new_term_count,
vector: new_vec.map_or(SimilarVector::None, SimilarVector::Stored),
},
&mut scratch.candidate_terms,
&mut outcome.similar,
);
self.similarity_scratch = scratch;
}
fn scan_similar(
&mut self,
query: SimilarityQuery<'_>,
candidate_terms: &mut Vec<u32>,
out: &mut Vec<Similar>,
) {
out.clear();
if query.term_count == 0 && matches!(query.vector, SimilarVector::None) {
return;
}
let mut ring = [FactId::NONE; SIMILAR_CANDIDATE_CAP];
let mut n = 0usize;
for (fact, _) in self.entity_facts.entries(query.entity.0) {
if query.exclude == Some(fact) {
continue;
}
ring[n % SIMILAR_CANDIDATE_CAP] = fact;
n += 1;
}
let summaries_trustworthy = self.bm25_tokenizer_version == TOKENIZER_INDEX_VERSION;
let lexical_only = matches!(query.vector, SimilarVector::None);
for &fact in ring.iter().take(n.min(SIMILAR_CANDIDATE_CAP)) {
let may_overlap = !query.terms.is_empty()
&& self.overlap_possible(
fact,
query.terms,
query.term_count,
summaries_trustworthy,
);
if lexical_only && !may_overlap {
continue;
}
let Some(record) = self.fact(fact) else {
continue;
};
if record.is_tombstone() || record.is_closed() {
continue;
}
let mut lexical = None;
if may_overlap && let Ok(text) = core::str::from_utf8(self.texts.get(record.text)) {
candidate_terms.clear();
let terms = &self.terms;
let cand = &mut *candidate_terms;
self.tokenizer.tokenize(text, &mut |token| {
if let Some(term) = terms.lookup(token)
&& !cand.contains(&term.0)
{
cand.push(term.0);
}
});
if !candidate_terms.is_empty() {
let both = candidate_terms
.iter()
.filter(|term| query.terms.contains(term))
.count();
let union = candidate_terms.len() + query.term_count - both;
let jaccard = both as f32 / union as f32;
if jaccard > self.cfg.similar_jaccard {
lexical = Some(jaccard);
}
}
}
let mut vector = None;
if record.has_vector() {
let cos = match query.vector {
SimilarVector::None => 0.0,
SimilarVector::Stored(slot) => self.vecs.cosine_slots(slot, record.vector),
SimilarVector::Encoded(encoded) => {
self.vecs.cosine_encoded_slot(encoded, record.vector)
}
};
if cos > self.cfg.similar_cos {
vector = Some(cos);
}
}
let best = match (lexical, vector) {
(Some(l), Some(v)) if v > l => Some((v, SimilarReason::VectorCosine)),
(Some(l), _) => Some((l, SimilarReason::LexicalOverlap)),
(None, Some(v)) => Some((v, SimilarReason::VectorCosine)),
(None, None) => None,
};
if let Some((score, reason)) = best {
out.push(Similar {
id: fact,
score,
reason,
});
}
}
out.sort_unstable_by(|a, b| b.score.total_cmp(&a.score).then(a.id.cmp(&b.id)));
out.truncate(8);
}
fn overlap_possible(
&self,
candidate: FactId,
new_terms: &[u32],
new_term_count: usize,
trust_summary: bool,
) -> bool {
if !trust_summary {
return true;
}
let Some(doc) = self.bm25.doc(candidate) else {
return true;
};
if !doc.has_signature() {
return true;
}
let bound = doc.overlap_bound(new_terms);
debug_assert!(
bound <= new_terms.len(),
"the overlap bound counts query terms, so it cannot exceed them"
);
let union = (new_term_count - bound) + usize::from(doc.distinct);
bound as f32 / union as f32 > self.cfg.similar_jaccard
}
pub fn revise<S: Storage>(
&mut self,
store: &mut S,
target: FactId,
input: RememberInput<'_>,
) -> Result<RememberOutcome, Error> {
self.validate_input(&input)?;
self.check_revisable(target)?;
let outcome = self.apply_remember(&input, target, None)?;
let valid_from = input.valid_from.unwrap_or(input.now);
self.close_target(target, valid_from);
self.journal_remember(store, &input, target, outcome.id)?;
Ok(outcome)
}
pub fn forget<S: Storage>(
&mut self,
store: &mut S,
now: u64,
id: FactId,
) -> Result<bool, Error> {
let fresh = self.apply_forget(id)?;
let mut entry = Vec::new();
Op::Forget { now, fact: id }.encode(&mut entry);
store
.append_journal(&entry)
.map_err(|e| Error::Storage(format!("{e:?}")))?;
Ok(fresh)
}
pub fn remove_tag<S: Storage>(
&mut self,
store: &mut S,
now: u64,
tag: &str,
) -> Result<RemoveTagReport, Error> {
if tag.is_empty() {
return Err(Error::Invalid("empty tag"));
}
let report = self.apply_remove_tag(now, tag)?;
if report.affected != 0 {
let mut entry = Vec::new();
Op::RemoveTag { now, tag }.encode(&mut entry);
store
.append_journal(&entry)
.map_err(|e| Error::Storage(format!("{e:?}")))?;
}
Ok(report)
}
pub fn link<S: Storage>(&mut self, store: &mut S, input: LinkInput<'_>) -> Result<(), Error> {
self.apply_link(
input.now,
input.src,
input.rel,
input.dst,
FactId::from_opt(input.provenance),
)?;
let mut entry = Vec::new();
Op::Link {
now: input.now,
src: input.src,
rel: input.rel,
dst: input.dst,
provenance: FactId::from_opt(input.provenance),
}
.encode(&mut entry);
store
.append_journal(&entry)
.map_err(|e| Error::Storage(format!("{e:?}")))?;
Ok(())
}
pub fn unlink<S: Storage>(
&mut self,
store: &mut S,
input: UnlinkInput<'_>,
) -> Result<bool, Error> {
let fresh = self.apply_unlink(input.now, input.src, input.rel, input.dst)?;
let mut entry = Vec::new();
Op::Unlink {
now: input.now,
src: input.src,
rel: input.rel,
dst: input.dst,
}
.encode(&mut entry);
store
.append_journal(&entry)
.map_err(|e| Error::Storage(format!("{e:?}")))?;
Ok(fresh)
}
pub fn get(&self, id: FactId) -> Option<FactView<'_>> {
let record = self.fact(id)?;
if record.is_tombstone() {
return None;
}
let text = core::str::from_utf8(self.texts.get(record.text)).ok()?;
Some(FactView { record, text })
}
pub fn tags_of(&self, id: FactId, out: &mut Vec<TermId>) {
let Some(record) = self.fact(id) else { return };
if record.is_tombstone() {
return;
}
let Some(aux) = self.fact_aux.get(&id.0.to_be_bytes()) else {
return;
};
for chunk in self.tag_lists.iter(&aux.tags) {
for raw in chunk.chunks_exact(4) {
out.push(TermId(u32::from_be_bytes(raw.try_into().unwrap())));
}
}
}
pub fn list_tags(&self, query: TagQuery<'_>) -> Result<TagPage, Error> {
self.tag_catalog.page(&self.terms, self.cfg.db_uuid, query)
}
pub fn metadata_of<'s>(&'s self, id: FactId, out: &mut Vec<(&'s str, &'s str)>) -> bool {
out.clear();
let Some(record) = self.fact(id) else {
return false;
};
if record.is_tombstone() {
return false;
}
let Some(aux) = self.fact_aux.get(&id.0.to_be_bytes()) else {
return false;
};
if aux.meta.0 == NONE_U32 || aux.meta.0 >= self.metas.len() as u32 {
return false;
}
crate::metadata::decode(self.metas.get(aux.meta), out).is_ok() && !out.is_empty()
}
pub fn entity(&mut self, name: &str) -> Option<EntityId> {
let mut norm = core::mem::take(&mut self.name_scratch);
normalize_name(&mut self.tokenizer, name, &mut norm);
let found = if norm.is_empty() {
None
} else {
self.lookup_entity_by_norm(&norm)
};
self.name_scratch = norm;
found
}
pub fn term(&self, id: TermId) -> &str {
self.terms.resolve(id)
}
pub fn entity_name(&self, id: EntityId) -> Option<&str> {
let record = self.entities.get(&id.0.to_be_bytes())?;
core::str::from_utf8(self.texts.get(record.name)).ok()
}
pub fn edges_each(&self, mut visit: impl FnMut(&str, &str, &str, FactId) -> bool) {
for slot in self.edges_out.iter() {
let (Some(src), Some(dst)) = (self.entity_name(slot.a), self.entity_name(slot.b))
else {
continue;
};
if !visit(src, self.term(slot.rel), dst, slot.fact) {
return;
}
}
}
pub fn facts_len(&self) -> usize {
self.facts.len()
}
pub fn entities_len(&self) -> usize {
self.entities.len()
}
pub fn cfg(&self) -> &Config {
&self.cfg
}
pub fn stats(&self) -> Stats {
Stats {
facts: self.facts.len(),
entities: self.entities.len(),
terms: self.terms.len(),
edges: self.edges_out.len(),
edge_versions: self.edges_hist_out.len(),
vectors: self.vecs.len(),
tombstones: self.tombstones,
hnsw_indexed: self.hnsw.indexed(),
next_fact: self.next_fact,
next_entity: self.next_entity,
next_edge: self.next_edge,
db_uuid: self.cfg.db_uuid,
pool_bytes: self.facts.pool_bytes()
+ self.fact_aux.pool_bytes()
+ self.entities.pool_bytes()
+ self.by_name.pool_bytes()
+ self.edges_out.pool_bytes()
+ self.edges_in.pool_bytes()
+ self.edges_hist_out.pool_bytes()
+ self.edges_hist_in.pool_bytes()
+ self.temporal.pool_bytes()
+ self.texts.pool_bytes()
+ self.terms.pool_bytes()
+ self.tag_lists.pool_bytes()
+ self.bm25.pool_bytes()
+ self.tags_idx.pool_bytes()
+ self.tag_catalog.pool_bytes()
+ self.entity_facts.pool_bytes()
+ self.vecs.pool_bytes()
+ self.hnsw.pool_bytes(),
shards: shards::ShardLayout::of_config(&self.cfg),
}
}
fn fact(&self, id: FactId) -> Option<FactRecord> {
self.facts.get(&id.0.to_be_bytes())
}
fn validate_input(&self, input: &RememberInput<'_>) -> Result<(), Error> {
if input.text.len() > self.cfg.max_text {
return Err(Error::TooLarge {
what: "text",
len: input.text.len(),
max: self.cfg.max_text,
});
}
if input.tags.len() > MAX_TAGS {
return Err(Error::TooLarge {
what: "tags",
len: input.tags.len(),
max: MAX_TAGS,
});
}
if input.links.len() > 16 {
return Err(Error::TooLarge {
what: "links",
len: input.links.len(),
max: 16,
});
}
if input.tags.iter().any(|t| t.is_empty()) {
return Err(Error::Invalid("empty tag"));
}
if !input.links.is_empty() && input.entity.is_none() {
return Err(Error::Invalid("links require a subject entity"));
}
if let Some(v) = input.vector {
if self.cfg.dim == 0 {
return Err(Error::Invalid("vector given but dim is 0"));
}
if v.len() != self.cfg.dim {
return Err(Error::DimMismatch {
got: v.len(),
want: self.cfg.dim,
});
}
}
Ok(())
}
fn check_revisable(&self, target: FactId) -> Result<(), Error> {
let record = self.fact(target).ok_or(Error::NotFound(target))?;
if record.is_tombstone() {
return Err(Error::NotFound(target));
}
if record.is_closed() {
return Err(Error::AlreadyClosed(target));
}
Ok(())
}
fn close_target(&mut self, target: FactId, valid_to: u64) {
let record = self.fact(target).expect("checked revisable");
let payload = self
.facts
.payload_mut(&target.0.to_be_bytes())
.expect("record fetched above");
let flags = record.flags | fact_flags::CLOSED;
payload[4..6].copy_from_slice(&flags.to_be_bytes());
payload[36..44].copy_from_slice(&valid_to.to_be_bytes());
self.change_catalog_for_fact(target, -1);
}
fn apply_remember(
&mut self,
input: &RememberInput<'_>,
revises: FactId,
copy_vector: Option<u32>,
) -> Result<RememberOutcome, Error> {
let id = FactId(self.next_fact);
let entity = match input.entity {
Some(name) => Some(self.resolve_or_create_entity(name, input.now)?),
None => None,
};
let text_id = self.texts.push(input.text.as_bytes())?;
let mut tfs = core::mem::take(&mut self.tf_scratch);
let mut probe = core::mem::take(&mut self.tf_probe);
tfs.clear();
probe.fill(0);
let terms = &mut self.terms;
let mut intern_err = None;
self.tokenizer.tokenize(input.text, &mut |token| {
if intern_err.is_some() {
return;
}
match terms.intern(token) {
Ok(term) => tally_term(&mut tfs, &mut probe, term.0),
Err(e) => intern_err = Some(e),
}
});
self.tf_probe = probe;
if let Some(e) = intern_err {
self.tf_scratch = tfs;
return Err(Error::Arena(e));
}
self.bm25.index_doc(id, &tfs)?;
self.tf_scratch = tfs;
let meta = match input.metadata {
Some(pairs) if !pairs.is_empty() => {
self.metas.push(&crate::metadata::encode(pairs)?)?
}
_ => BlobId(NONE_U32),
};
let mut aux = FactAux {
id,
tags: ListHandle::EMPTY,
meta,
};
let mut seen_tags: [u32; MAX_TAGS] = [NONE_U32; MAX_TAGS];
let mut seen_cnt = 0usize;
for tag in input.tags {
let term = self.terms.intern(tag)?;
if seen_tags[..seen_cnt].contains(&term.0) {
continue;
}
seen_tags[seen_cnt] = term.0;
seen_cnt += 1;
self.tag_lists.push(&mut aux.tags, &term.0.to_be_bytes())?;
self.tags_idx.push(term.0, id, 0)?;
}
self.fact_aux.insert(&aux)?;
if let Some(src) = entity {
for &(rel, dst_name) in input.links {
let dst = self.resolve_or_create_entity(dst_name, input.now)?;
let rel = self.terms.intern(rel)?;
self.open_edge(input.now, src, rel, dst, id)?;
}
self.entity_facts.push(src.0, id, 0)?;
}
let (vector, flags) = match (input.vector, copy_vector) {
(Some(v), None) => (self.vecs.push(id, v)?, fact_flags::HAS_VECTOR),
(None, Some(source)) => (
self.vecs.clone_slot_for_fact(id, source)?,
fact_flags::HAS_VECTOR,
),
(None, None) => (NONE_U32, 0),
(Some(_), Some(_)) => unreachable!("retag does not provide a raw vector"),
};
let recorded_at = input.now;
let valid_from = input.valid_from.unwrap_or(input.now);
self.facts.insert(&FactRecord {
id,
entity: EntityId::from_opt(entity),
flags,
kind: 0,
text: text_id,
vector,
revises,
recorded_at,
valid_from,
valid_to: VALID_TO_OPEN,
})?;
self.temporal.insert(&TemporalSlot {
recorded_at,
fact: id,
})?;
for &term in &seen_tags[..seen_cnt] {
self.tag_catalog.change(&self.terms, TermId(term), 1);
}
self.next_fact += 1;
Ok(RememberOutcome {
id,
entity,
similar: Vec::new(),
})
}
fn apply_forget(&mut self, id: FactId) -> Result<bool, Error> {
let record = self.fact(id).ok_or(Error::NotFound(id))?;
if record.is_tombstone() {
return Ok(false);
}
let payload = self
.facts
.payload_mut(&id.0.to_be_bytes())
.expect("record fetched above");
let flags = record.flags | fact_flags::TOMBSTONE;
payload[4..6].copy_from_slice(&flags.to_be_bytes());
self.tombstones += 1;
if !record.is_closed() {
self.change_catalog_for_fact(id, -1);
}
Ok(true)
}
fn apply_remove_tag(&mut self, now: u64, tag: &str) -> Result<RemoveTagReport, Error> {
let Some(term) = self.terms.lookup(tag) else {
return Ok(RemoveTagReport::default());
};
let targets: Vec<FactId> = self
.tags_idx
.entries(term.0)
.map(|(id, _)| id)
.filter(|&id| {
self.fact(id)
.is_some_and(|record| !record.is_tombstone() && !record.is_closed())
})
.collect();
let mut affected = 0u32;
for target in targets {
self.retag_without(now, target, tag)?;
affected = affected.saturating_add(1);
}
Ok(RemoveTagReport { affected })
}
fn retag_without(&mut self, now: u64, target: FactId, removed: &str) -> Result<(), Error> {
let record = self.fact(target).ok_or(Error::NotFound(target))?;
let view = self.get(target).ok_or(Error::NotFound(target))?;
let text = view.text.to_string();
let entity = record
.entity
.some()
.and_then(|id| self.entity_name(id))
.map(ToString::to_string);
let mut tag_terms = Vec::new();
self.tags_of(target, &mut tag_terms);
let tags: Vec<String> = tag_terms
.into_iter()
.map(|term| self.term(term))
.filter(|name| *name != removed)
.map(ToString::to_string)
.collect();
let tag_refs: Vec<&str> = tags.iter().map(String::as_str).collect();
let mut metadata = Vec::new();
self.metadata_of(target, &mut metadata);
let metadata: Vec<(String, String)> = metadata
.into_iter()
.map(|(key, value)| (key.to_string(), value.to_string()))
.collect();
let metadata_refs: Vec<(&str, &str)> = metadata
.iter()
.map(|(key, value)| (key.as_str(), value.as_str()))
.collect();
let vector = (record.flags & fact_flags::HAS_VECTOR != 0).then_some(record.vector);
let input = RememberInput {
now,
text: &text,
entity: entity.as_deref(),
tags: &tag_refs,
links: &[],
vector: None,
valid_from: Some(now),
metadata: (!metadata_refs.is_empty()).then_some(metadata_refs.as_slice()),
};
self.apply_remember(&input, target, vector)?;
self.close_target(target, now);
Ok(())
}
fn change_catalog_for_fact(&mut self, id: FactId, delta: i32) {
let Some(aux) = self.fact_aux.get(&id.0.to_be_bytes()) else {
return;
};
let mut ids = [NONE_U32; MAX_TAGS];
let mut len = 0usize;
for chunk in self.tag_lists.iter(&aux.tags) {
for raw in chunk.chunks_exact(4) {
if len == ids.len() {
break;
}
ids[len] = u32::from_be_bytes(raw.try_into().unwrap());
len += 1;
}
}
for &term in &ids[..len] {
self.tag_catalog.change(&self.terms, TermId(term), delta);
}
}
fn apply_link(
&mut self,
now: u64,
src: &str,
rel: &str,
dst: &str,
provenance: FactId,
) -> Result<(), Error> {
let src = self.resolve_or_create_entity(src, now)?;
let dst = self.resolve_or_create_entity(dst, now)?;
let rel = self.terms.intern(rel)?;
self.open_edge(now, src, rel, dst, provenance)
}
fn apply_unlink(&mut self, now: u64, src: &str, rel: &str, dst: &str) -> Result<bool, Error> {
let Some(src) = self.lookup_entity_name(src) else {
return Ok(false);
};
let Some(dst) = self.lookup_entity_name(dst) else {
return Ok(false);
};
let Some(rel) = self.terms.lookup(rel) else {
return Ok(false);
};
self.close_current_edge(now, src, rel, dst)
}
fn open_edge(
&mut self,
now: u64,
src: EntityId,
rel: TermId,
dst: EntityId,
fact: FactId,
) -> Result<(), Error> {
if let Some(current) = self.current_edge(src, rel, dst) {
if current.fact == fact {
return Ok(());
}
self.close_current_edge(now, src, rel, dst)?;
}
let edge = EdgeId(self.next_edge);
let history = EdgeHistorySlot {
a: src,
rel,
b: dst,
edge,
fact,
flags: 0,
kind: 0,
recorded_at: now,
valid_from: now,
valid_to: VALID_TO_OPEN,
};
self.insert_history_edge(history)?;
self.insert_current_edge(src, rel, dst, fact, edge, now)?;
self.next_edge += 1;
Ok(())
}
fn insert_current_edge(
&mut self,
src: EntityId,
rel: TermId,
dst: EntityId,
fact: FactId,
edge: EdgeId,
valid_from: u64,
) -> Result<(), Error> {
for (arena, a, b) in [
(&mut self.edges_out, src, dst),
(&mut self.edges_in, dst, src),
] {
let slot = EdgeSlot {
a,
rel,
b,
fact,
edge,
valid_from,
};
if !arena.insert(&slot)? {
let payload = arena
.payload_mut(&edge_key(a, rel, b))
.expect("insert reported a duplicate");
let mut full = [0u8; EdgeSlot::SIZE];
slot.write(&mut full);
payload.copy_from_slice(&full[EdgeSlot::KEY_LEN..]);
}
}
Ok(())
}
fn insert_history_edge(&mut self, edge: EdgeHistorySlot) -> Result<(), Error> {
self.edges_hist_out.insert(&edge)?;
self.edges_hist_in.insert(&EdgeHistorySlot {
a: edge.b,
b: edge.a,
..edge
})?;
Ok(())
}
fn close_current_edge(
&mut self,
now: u64,
src: EntityId,
rel: TermId,
dst: EntityId,
) -> Result<bool, Error> {
let Some(current) = self.current_edge(src, rel, dst) else {
return Ok(false);
};
let close_at = now.max(current.valid_from);
let out_key = edge_history_key(src, current.valid_from, current.edge);
let in_key = edge_history_key(dst, current.valid_from, current.edge);
close_edge_history_payload(
self.edges_hist_out
.payload_mut(&out_key)
.ok_or(Error::Corrupt("missing outgoing edge history"))?,
close_at,
);
close_edge_history_payload(
self.edges_hist_in
.payload_mut(&in_key)
.ok_or(Error::Corrupt("missing incoming edge history"))?,
close_at,
);
let out_removed = self.edges_out.remove(&edge_key(src, rel, dst));
let in_removed = self.edges_in.remove(&edge_key(dst, rel, src));
if out_removed != in_removed {
return Err(Error::Corrupt("edge mirrors disagree"));
}
Ok(out_removed)
}
fn current_edge(&self, src: EntityId, rel: TermId, dst: EntityId) -> Option<EdgeSlot> {
self.edges_out
.get_slot(&edge_key(src, rel, dst))
.map(EdgeSlot::read)
}
fn lookup_entity_by_norm(&self, norm: &str) -> Option<EntityId> {
let term = self.terms.lookup(norm)?;
let mut from = [0u8; 8];
key::write_u32(&mut from, term.0);
let mut to = [0u8; 8];
key::write_u32(&mut to, term.0);
to[4..].copy_from_slice(&u32::MAX.to_be_bytes());
self.by_name.range(&from, &to).next().map(|e| e.id)
}
fn lookup_entity_name(&mut self, name: &str) -> Option<EntityId> {
let mut norm = core::mem::take(&mut self.name_scratch);
normalize_name(&mut self.tokenizer, name, &mut norm);
let result = (!norm.is_empty())
.then(|| self.lookup_entity_by_norm(&norm))
.flatten();
self.name_scratch = norm;
result
}
fn resolve_or_create_entity(&mut self, name: &str, now: u64) -> Result<EntityId, Error> {
let mut norm = core::mem::take(&mut self.name_scratch);
normalize_name(&mut self.tokenizer, name, &mut norm);
if norm.is_empty() {
self.name_scratch = norm;
return Err(Error::Invalid("entity name has no indexable characters"));
}
let result = (|| {
if let Some(found) = self.lookup_entity_by_norm(&norm) {
return Ok(found);
}
let term = self.terms.intern(&norm)?;
let id = EntityId(self.next_entity);
let name_id = self.texts.push(name.as_bytes())?;
self.entities.insert(&EntityRecord {
id,
name: name_id,
name_term: term,
created_at: now,
flags: 0,
})?;
self.by_name.insert(&EntityByName {
name_term: term,
id,
})?;
self.next_entity += 1;
Ok(id)
})();
self.name_scratch = norm;
result
}
fn journal_remember<S: Storage>(
&mut self,
store: &mut S,
input: &RememberInput<'_>,
revises: FactId,
assigned: FactId,
) -> Result<(), Error> {
let mut entry = Vec::new();
Op::Remember {
now: input.now,
valid_from: input.valid_from.unwrap_or(input.now),
entity: input.entity,
text: input.text,
tags: input.tags.to_vec(),
links: input.links.to_vec(),
vector: input.vector.map(<[f32]>::to_vec).unwrap_or_default(),
metadata: input.metadata.map(<[_]>::to_vec).unwrap_or_default(),
revises,
assigned,
}
.encode(&mut entry);
store
.append_journal(&entry)
.map_err(|e| Error::Storage(format!("{e:?}")))
}
}
impl core::fmt::Debug for Memory<'_> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Memory")
.field("facts", &self.facts.len())
.field("entities", &self.entities.len())
.field("terms", &self.terms.len())
.finish()
}
}
fn normalize_name(tokenizer: &mut Tokenizer, name: &str, out: &mut String) {
out.clear();
tokenizer.tokenize(name, &mut |token| {
if !out.is_empty() {
out.push(' ');
}
out.push_str(token);
});
}
#[inline]
fn probe_at(term: u32, mask: usize) -> usize {
(u64::from(term).wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 32) as usize & mask
}
fn tally_term(tfs: &mut Vec<(u32, u8)>, probe: &mut Vec<u32>, term: u32) {
if (tfs.len() + 1) * 2 > probe.len() {
let grown = (probe.len() * 2).max(64);
probe.clear();
probe.resize(grown, 0);
let mask = grown - 1;
for (i, &(seen, _)) in tfs.iter().enumerate() {
let mut at = probe_at(seen, mask);
while probe[at] != 0 {
at = (at + 1) & mask;
}
probe[at] = i as u32 + 1;
}
}
let mask = probe.len() - 1;
let mut at = probe_at(term, mask);
loop {
match probe[at] {
0 => {
probe[at] = tfs.len() as u32 + 1;
tfs.push((term, 1));
return;
}
slot => {
let entry = &mut tfs[slot as usize - 1];
if entry.0 == term {
entry.1 = entry.1.saturating_add(1);
return;
}
}
}
at = (at + 1) & mask;
}
}