//! Storage backend using RocksDB
//!
//! Handles billions of atoms on disk with efficient indexes.
use crate::atom::{Atom, Hash, Value};
use std::path::Path;
#[cfg(feature = "rocksdb-storage")]
use rocksdb::{DB, Options, ColumnFamilyDescriptor, IteratorMode, Direction};
/// Column family names
const CF_ATOMS: &str = "atoms"; // hash -> atom bytes
const CF_BY_TYPE: &str = "by_type"; // type:hash -> ()
const CF_BY_NAME: &str = "by_name"; // type:name -> hash
const CF_INCOMING: &str = "incoming"; // target:link -> ()
const CF_VALUES: &str = "values"; // hash:key -> value bytes
/// Storage trait for different backends
pub trait AtomStore {
/// Add an atom, returns its hash
fn add(&self, atom: &Atom) -> Hash;
/// Get atom by hash
fn get(&self, hash: &Hash) -> Option<Atom>;
/// Get node by type and name
fn get_node(&self, typ: &str, name: &str) -> Option<(Hash, Atom)>;
/// Check if atom exists
fn contains(&self, hash: &Hash) -> bool;
/// Iterate atoms of a type (streaming)
fn iter_by_type<'a>(&'a self, typ: &str) -> Box<dyn Iterator<Item = (Hash, Atom)> + 'a>;
/// Get incoming links
fn get_incoming(&self, target: &Hash) -> Vec<Hash>;
/// Set a value on an atom
fn set_value(&self, hash: &Hash, key: &str, value: &Value);
/// Get a value from an atom
fn get_value(&self, hash: &Hash, key: &str) -> Option<Value>;
/// Number of atoms
fn len(&self) -> usize;
/// Is empty?
fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[cfg(feature = "rocksdb-storage")]
pub struct RocksStore {
db: DB,
count: std::sync::atomic::AtomicUsize,
}
#[cfg(feature = "rocksdb-storage")]
impl RocksStore {
pub fn open(path: impl AsRef<Path>) -> Result<Self, rocksdb::Error> {
let mut opts = Options::default();
opts.create_if_missing(true);
opts.create_missing_column_families(true);
opts.set_max_open_files(10_000);
opts.set_keep_log_file_num(3);
// Optimize for point lookups (most common operation)
opts.optimize_for_point_lookup(128); // 128MB block cache
// Column families
let cf_opts = Options::default();
let cfs = vec![
ColumnFamilyDescriptor::new(CF_ATOMS, cf_opts.clone()),
ColumnFamilyDescriptor::new(CF_BY_TYPE, cf_opts.clone()),
ColumnFamilyDescriptor::new(CF_BY_NAME, cf_opts.clone()),
ColumnFamilyDescriptor::new(CF_INCOMING, cf_opts.clone()),
ColumnFamilyDescriptor::new(CF_VALUES, cf_opts),
];
let db = DB::open_cf_descriptors(&opts, path, cfs)?;
// Count existing atoms
let count = {
let cf = db.cf_handle(CF_ATOMS).unwrap();
db.iterator_cf(cf, IteratorMode::Start).count()
};
Ok(Self {
db,
count: std::sync::atomic::AtomicUsize::new(count),
})
}
fn hash_to_hex(hash: &Hash) -> String {
hex::encode(hash)
}
fn hex_to_hash(s: &str) -> Option<Hash> {
let bytes = hex::decode(s).ok()?;
if bytes.len() != 16 {
return None;
}
let mut hash = [0u8; 16];
hash.copy_from_slice(&bytes);
Some(hash)
}
}
#[cfg(feature = "rocksdb-storage")]
impl AtomStore for RocksStore {
fn add(&self, atom: &Atom) -> Hash {
let hash = atom.hash();
let hash_hex = Self::hash_to_hex(&hash);
// Check if already exists
let cf_atoms = self.db.cf_handle(CF_ATOMS).unwrap();
if self.db.get_cf(cf_atoms, &hash).unwrap().is_some() {
return hash;
}
// Serialize atom
let bytes = atom.to_bytes();
// Write atom
self.db.put_cf(cf_atoms, &hash, &bytes).unwrap();
// Index by type
let cf_type = self.db.cf_handle(CF_BY_TYPE).unwrap();
let type_key = format!("{}:{}", atom.typ(), hash_hex);
self.db.put_cf(cf_type, type_key.as_bytes(), &[]).unwrap();
// Index by name (for nodes)
if let Some(name) = atom.name() {
let cf_name = self.db.cf_handle(CF_BY_NAME).unwrap();
let name_key = format!("{}:{}", atom.typ(), name);
self.db.put_cf(cf_name, name_key.as_bytes(), &hash).unwrap();
}
// Index incoming (for links)
if let Some(outgoing) = atom.outgoing() {
let cf_incoming = self.db.cf_handle(CF_INCOMING).unwrap();
for target in outgoing {
let inc_key = format!("{}:{}", Self::hash_to_hex(target), hash_hex);
self.db.put_cf(cf_incoming, inc_key.as_bytes(), &[]).unwrap();
}
}
self.count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
hash
}
fn get(&self, hash: &Hash) -> Option<Atom> {
let cf = self.db.cf_handle(CF_ATOMS)?;
let bytes = self.db.get_cf(cf, hash).ok()??;
Atom::from_bytes(&bytes).ok()
}
fn get_node(&self, typ: &str, name: &str) -> Option<(Hash, Atom)> {
let cf = self.db.cf_handle(CF_BY_NAME)?;
let key = format!("{}:{}", typ, name);
let hash_bytes = self.db.get_cf(cf, key.as_bytes()).ok()??;
let mut hash = [0u8; 16];
hash.copy_from_slice(&hash_bytes);
let atom = self.get(&hash)?;
Some((hash, atom))
}
fn contains(&self, hash: &Hash) -> bool {
let cf = self.db.cf_handle(CF_ATOMS).unwrap();
self.db.get_cf(cf, hash).ok().flatten().is_some()
}
fn iter_by_type<'a>(&'a self, typ: &str) -> Box<dyn Iterator<Item = (Hash, Atom)> + 'a> {
let cf = self.db.cf_handle(CF_BY_TYPE).unwrap();
let prefix = format!("{}:", typ);
let prefix_bytes = prefix.as_bytes().to_vec();
let prefix_len = prefix.len();
Box::new(
self.db
.prefix_iterator_cf(cf, &prefix_bytes)
.filter_map(move |result| {
let (key, _) = result.ok()?;
let key_str = std::str::from_utf8(&key).ok()?;
// Stop if we've gone past the prefix
if !key_str.starts_with(&prefix) {
return None;
}
let hash_hex = &key_str[prefix_len..];
let hash = Self::hex_to_hash(hash_hex)?;
let atom = self.get(&hash)?;
Some((hash, atom))
})
)
}
fn get_incoming(&self, target: &Hash) -> Vec<Hash> {
let cf = match self.db.cf_handle(CF_INCOMING) {
Some(cf) => cf,
None => return vec![],
};
let prefix = format!("{}:", Self::hash_to_hex(target));
let prefix_len = prefix.len();
self.db
.prefix_iterator_cf(cf, prefix.as_bytes())
.filter_map(|result| {
let (key, _) = result.ok()?;
let key_str = std::str::from_utf8(&key).ok()?;
if !key_str.starts_with(&prefix) {
return None;
}
let hash_hex = &key_str[prefix_len..];
Self::hex_to_hash(hash_hex)
})
.collect()
}
fn set_value(&self, hash: &Hash, key: &str, value: &Value) {
let cf = self.db.cf_handle(CF_VALUES).unwrap();
let db_key = format!("{}:{}", Self::hash_to_hex(hash), key);
let bytes = bincode::serialize(value).unwrap();
self.db.put_cf(cf, db_key.as_bytes(), &bytes).unwrap();
}
fn get_value(&self, hash: &Hash, key: &str) -> Option<Value> {
let cf = self.db.cf_handle(CF_VALUES)?;
let db_key = format!("{}:{}", Self::hash_to_hex(hash), key);
let bytes = self.db.get_cf(cf, db_key.as_bytes()).ok()??;
bincode::deserialize(&bytes).ok()
}
fn len(&self) -> usize {
self.count.load(std::sync::atomic::Ordering::Relaxed)
}
}
/// In-memory store for testing and small datasets
pub struct MemoryStore {
atoms: parking_lot::RwLock<std::collections::HashMap<Hash, Atom>>,
by_type: parking_lot::RwLock<std::collections::HashMap<String, Vec<Hash>>>,
by_name: parking_lot::RwLock<std::collections::HashMap<String, Hash>>,
incoming: parking_lot::RwLock<std::collections::HashMap<Hash, Vec<Hash>>>,
values: parking_lot::RwLock<std::collections::HashMap<String, Value>>,
}
impl Default for MemoryStore {
fn default() -> Self {
Self::new()
}
}
impl MemoryStore {
pub fn new() -> Self {
Self {
atoms: parking_lot::RwLock::new(std::collections::HashMap::new()),
by_type: parking_lot::RwLock::new(std::collections::HashMap::new()),
by_name: parking_lot::RwLock::new(std::collections::HashMap::new()),
incoming: parking_lot::RwLock::new(std::collections::HashMap::new()),
values: parking_lot::RwLock::new(std::collections::HashMap::new()),
}
}
}
impl AtomStore for MemoryStore {
fn add(&self, atom: &Atom) -> Hash {
let hash = atom.hash();
// Check if exists
if self.atoms.read().contains_key(&hash) {
return hash;
}
// Add atom
self.atoms.write().insert(hash, atom.clone());
// Index by type
self.by_type
.write()
.entry(atom.typ().to_string())
.or_default()
.push(hash);
// Index by name
if let Some(name) = atom.name() {
let key = format!("{}:{}", atom.typ(), name);
self.by_name.write().insert(key, hash);
}
// Index incoming
if let Some(outgoing) = atom.outgoing() {
for target in outgoing {
self.incoming.write().entry(*target).or_default().push(hash);
}
}
hash
}
fn get(&self, hash: &Hash) -> Option<Atom> {
self.atoms.read().get(hash).cloned()
}
fn get_node(&self, typ: &str, name: &str) -> Option<(Hash, Atom)> {
let key = format!("{}:{}", typ, name);
let hash = *self.by_name.read().get(&key)?;
let atom = self.atoms.read().get(&hash)?.clone();
Some((hash, atom))
}
fn contains(&self, hash: &Hash) -> bool {
self.atoms.read().contains_key(hash)
}
fn iter_by_type<'a>(&'a self, typ: &str) -> Box<dyn Iterator<Item = (Hash, Atom)> + 'a> {
let hashes: Vec<Hash> = self
.by_type
.read()
.get(typ)
.cloned()
.unwrap_or_default();
let atoms = self.atoms.read();
let results: Vec<(Hash, Atom)> = hashes
.into_iter()
.filter_map(|h| atoms.get(&h).map(|a| (h, a.clone())))
.collect();
Box::new(results.into_iter())
}
fn get_incoming(&self, target: &Hash) -> Vec<Hash> {
self.incoming.read().get(target).cloned().unwrap_or_default()
}
fn set_value(&self, hash: &Hash, key: &str, value: &Value) {
let db_key = format!("{:?}:{}", hash, key);
self.values.write().insert(db_key, value.clone());
}
fn get_value(&self, hash: &Hash, key: &str) -> Option<Value> {
let db_key = format!("{:?}:{}", hash, key);
self.values.read().get(&db_key).cloned()
}
fn len(&self) -> usize {
self.atoms.read().len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_memory_store() {
let store = MemoryStore::new();
let dog = Atom::node("ConceptNode", "dog");
let animal = Atom::node("ConceptNode", "animal");
let dog_hash = store.add(&dog);
let animal_hash = store.add(&animal);
// Verify retrieval
assert_eq!(store.get(&dog_hash), Some(dog.clone()));
assert_eq!(store.get_node("ConceptNode", "dog"), Some((dog_hash, dog)));
// Add link
let link = Atom::link("InheritanceLink", vec![dog_hash, animal_hash]);
let link_hash = store.add(&link);
// Verify incoming
let incoming = store.get_incoming(&animal_hash);
assert!(incoming.contains(&link_hash));
// Verify iteration
let concepts: Vec<_> = store.iter_by_type("ConceptNode").collect();
assert_eq!(concepts.len(), 2);
}
}