diff --git a/Cargo.lock b/Cargo.lock index e5c5a1d87..c2d1148cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -773,6 +773,8 @@ dependencies = [ "codegraph-graph", "serde", "serde_json", + "tempfile", + "tokio", ] [[package]] diff --git a/crates/codegraph-context/Cargo.toml b/crates/codegraph-context/Cargo.toml index 82b5bc1bb..7f8134bb4 100644 --- a/crates/codegraph-context/Cargo.toml +++ b/crates/codegraph-context/Cargo.toml @@ -10,3 +10,8 @@ codegraph-core = { path = "../codegraph-core" } codegraph-graph = { path = "../codegraph-graph" } serde = { workspace = true } serde_json = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros"] } +tempfile = "3" +codegraph-graph = { path = "../codegraph-graph", features = ["sqlite"] } diff --git a/crates/codegraph-context/src/lib.rs b/crates/codegraph-context/src/lib.rs index 5f1a9c290..bbff4b68d 100644 --- a/crates/codegraph-context/src/lib.rs +++ b/crates/codegraph-context/src/lib.rs @@ -73,7 +73,8 @@ pub async fn build_response( req: &ContextRequest, ) -> Result { let idx = index.ensure_fresh().await; - let candidates = idx + // Try symbol-name search first, then fallback to file-path search. + let mut candidates = idx .search_symbol_paged_resumable( &req.query, None, @@ -87,6 +88,36 @@ pub async fn build_response( ) .await? .page; + if candidates.is_empty() { + // Fallback: query as filename (strip extension for symbol-name search). + let query_stripped = req + .query + .split('/') + .next_back() + .and_then(|f| { + let without_ext = f.rsplit_once('.')?.0; + if without_ext.is_empty() { + None + } else { + Some(without_ext.to_string()) + } + }) + .unwrap_or_else(|| req.query.clone()); + candidates = idx + .search_symbol_paged_resumable( + &query_stripped, + None, + SymbolMatch::Contains, + Pagination { + limit: req.limit as usize, + offset: 0, + }, + None, + None, + ) + .await? + .page; + } // Pre-load mỗi file một lần khi cần source. let file_cache: HashMap> = if req.include_source { @@ -190,3 +221,76 @@ fn render_markdown(resp: &ContextResponse, strip: Option<&str>) -> String { } out } + +#[cfg(test)] +mod tests { + use super::*; + use codegraph_graph::SharedGraphIndex; + use std::sync::Arc; + + fn sym(name: &str, id: u64) -> codegraph_core::Symbol { + codegraph_core::Symbol { + id, + name: name.to_string(), + kind: codegraph_core::SymbolKind::Function, + scope: codegraph_core::ScopeLevel::Global, + scope_id: 0, + type_ref: 0, + type_name: None, + file: "RestEndpoint.java".into(), + line: 1, + end_line: 2, + signature: None, + doc: None, + annotations: Vec::new(), + language: "java".into(), + } + } + + #[tokio::test] + async fn context_fallback_matches_filename() { + // Tạo index với symbol "RestEndpoint" trong file "RestEndpoint.java" dùng sqlite temp. + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("test.db"); + let db_str = format!("sqlite://{}", db_path.to_string_lossy()); + + { + let mut idx = codegraph_graph::GraphIndex::open(&db_str).await.unwrap(); + let r = codegraph_graph::ParseResult { + path: "RestEndpoint.java".into(), + language: "java".into(), + bytes: 0, + lines: 0, + symbols: vec![sym("RestEndpoint", 100)], + chains: std::collections::HashMap::new(), + calls: vec![], + }; + idx.ingest(&[r]).await.unwrap(); + } + + let sgi = SharedGraphIndex::open(Some(db_str.clone())).await.unwrap(); + + // Query "RestEndpoint.java" → không match theo tên symbol → fallback tìm "RestEndpoint". + let req = ContextRequest { + query: "RestEndpoint.java".into(), + depth: 1, + include_source: false, + limit: 5, + format: Format::Markdown, + strip_prefix: None, + }; + let sgi_arc: Arc = Arc::new(sgi); + let resp = build_response(&sgi_arc, &req).await.unwrap(); + assert!(!resp.hits.is_empty(), "phải match qua fallback filename"); + assert_eq!(resp.hits[0].symbol.name, "RestEndpoint"); + + // Query "RestEndpoint" (không có extension) → match trực tiếp. + let req2 = ContextRequest { + query: "RestEndpoint".into(), + ..req + }; + let resp2 = build_response(&sgi_arc, &req2).await.unwrap(); + assert!(!resp2.hits.is_empty()); + assert_eq!(resp2.hits[0].symbol.name, "RestEndpoint"); + } +} diff --git a/crates/codegraph-graph/src/lib.rs b/crates/codegraph-graph/src/lib.rs index 13e21e04f..6433665a9 100644 --- a/crates/codegraph-graph/src/lib.rs +++ b/crates/codegraph-graph/src/lib.rs @@ -47,6 +47,14 @@ pub use crate::storage::postgres::PostgresStorage; #[cfg(feature = "sqlite")] pub use crate::storage::sqlite::SqliteStorage; pub use crate::storage::{InMemoryStorage, IndexCounts, Storage, Tx}; +// Sub-traits of `Storage` — callers that need only one facet (e.g. a chain-engine +// read path) can name it directly instead of taking the full umbrella. +#[cfg(feature = "bloom-search")] +pub use crate::storage::BloomStorage; +pub use crate::storage::{ + CategoryStorage, ChainStorage, EdgeDataStorage, EntityStorage, NodeMetaStorage, + ShortcutsStorage, +}; use crate::vector_index::VectorIndex; use codegraph_core::{ CallRecord, CallSite, CallSiteResult, ClassInfo, DependenciesReport, Dependency, EdgeMeta, @@ -1089,8 +1097,16 @@ impl GraphIndex { .cloned() .unwrap_or_default(); - // 3. Short name fallback (after last dot). + // 3. Independent lookup paths: alias ("var.method" → "TypeName.method") + // AND short name fallback — merge results instead of sequential fallback. if candidates.is_empty() { + // 3a. Go/Import alias: resolve "var.method" via field type to "TypeName.method". + if let Some(qualified) = self.alias_qualified_name(caller_id, &call.call_name) { + let alias_ids = self.name_index.get(&qualified).cloned().unwrap_or_default(); + candidates.extend(alias_ids); + } + + // 3b. Short name fallback (after last dot). let short = call .call_name .rsplit('.') @@ -1098,39 +1114,31 @@ impl GraphIndex { .unwrap_or("") .to_lowercase(); if !short.is_empty() { - // Chỉ nhận callee-thực-sự (Function/Method) — KHÔNG fallback vào - // biến / field / param trùng tên (VD `WrapResponse.ok(...)` với - // receiver external không resolve được dễ link nhầm vào `boolean ok` - // trong file khác — bug C). - candidates = self - .name_index - .get(&short) - .cloned() - .unwrap_or_default() - .into_iter() - .filter(|&id| { - self.symbols.get(&id).is_some_and(|s| { - matches!(s.kind, SymbolKind::Function | SymbolKind::Method) - }) - }) - .collect(); + let short_ids = self.name_index.get(&short).cloned().unwrap_or_default(); + candidates.extend(short_ids); } } - // 4. Go/Import alias handling: try to resolve using the caller's variable type - // information. `alias_qualified_name` produces a fully qualified name like - // "myservice.validate" based on a variable's type_name. If that name - // exists in the index, use it as an additional candidate set. - if candidates.is_empty() - && let Some(qualified) = self.alias_qualified_name(caller_id, &call.call_name) - { - candidates = self.name_index.get(&qualified).cloned().unwrap_or_default(); + if candidates.is_empty() { + return None; } + // 4. Filter to Function/Method kinds AND exclude the caller itself. + // The short name fallback can return the caller's own id (e.g. both + // LegacyAdapter.doWork and Service.doWork match "dowork"), so we must + // eliminate the caller before scoring. + candidates.retain(|&id| { + id != caller_id + && self + .symbols + .get(&id) + .is_some_and(|s| matches!(s.kind, SymbolKind::Function | SymbolKind::Method)) + }); + if candidates.is_empty() { return None; } - Some(self.pick_best_candidate(&candidates, caller_id)) + Some(self.pick_best_candidate(&candidates, caller_id, &call.call_name)) } /// Tìm method của class theo tên (scope_id == class id). @@ -1154,10 +1162,22 @@ impl GraphIndex { } /// Chọn ứng viên tốt nhất trong danh sách trùng tên. - fn pick_best_candidate(&self, candidates: &[u64], caller_id: u64) -> u64 { + /// + /// Khi nhiều method cùng tên, dùng typed-proximity tie-breaker: nếu call_name + /// dạng `var.method` và `var` là field khai báo kiểu `TypeName` trong caller + /// scope, ưu tiên method thuộc class `TypeName` (bug của Router.route: + /// `legacyAdapter.doWork` từng bị resolve nhầm sang `Service.doWork` vì cả + /// hai cùng score +9). + fn pick_best_candidate(&self, candidates: &[u64], caller_id: u64, call_name: &str) -> u64 { if candidates.len() == 1 { return candidates[0]; } + + // Trích declared type của field từ call_name (nếu call_name dạng "var.method"). + // VD: "legacyAdapter.doWork" → field `legacyAdapter` trong caller scope + // → type_name = "LegacyAdapter" → trả "legacyadapter" (lower). + let field_type = self.field_declared_type(caller_id, call_name); + let caller_file = self.symbols.get(&caller_id).map(|s| s.file.clone()); let mut best = candidates[0]; let mut best_score = i32::MIN; @@ -1182,6 +1202,17 @@ impl GraphIndex { { score += 3; } + // Typed proximity: +8 nếu enclosing class của candidate khớp với + // declared type của field gọi (cao hơn same-file +3 để thắng). + if let Some(ref ft) = field_type + && let Some(class_sym) = self + .symbols + .get(&sym.scope_id) + .filter(|cs| matches!(cs.kind, SymbolKind::Class | SymbolKind::Interface)) + && class_sym.name.to_lowercase() == *ft + { + score += 8; + } if score > best_score { best_score = score; best = id; @@ -1190,6 +1221,38 @@ impl GraphIndex { best } + /// Tìm declared type của field từ call_name dạng `var.method`. + /// + /// Trả về `Some("legacyadapter")` nếu caller scope có field `legacyAdapter` + /// với `type_name = "LegacyAdapter"`. Trả `None` nếu call_name không có dấu + /// chấm, field chưa được index, hoặc field không có type_name. + fn field_declared_type(&self, caller_id: u64, call_name: &str) -> Option { + // Lấy tên field/receiver phía trước dấu `.` cuối. + // VD: "legacyAdapter.doWork" → "legacyAdapter", "doWork" → None (không có receiver). + let field_name = call_name + .rsplit('.') + .nth(1) + .or_else(|| call_name.rsplit('.').next())?; + if field_name == call_name { + // Không có dấu '.' → call_name chính là tên method, không phải dạng `var.method`. + return None; + } + let caller_scope_id = self.symbols.get(&caller_id)?.scope_id; + + for sym in self.symbols.values() { + if sym.scope_id == caller_scope_id + && sym.name == field_name + && matches!( + sym.kind, + SymbolKind::Field | SymbolKind::Variable | SymbolKind::Parameter + ) + { + return sym.type_name.as_ref().map(|t| t.to_lowercase()); + } + } + None + } + /// Build edges từ chains (đã resolve) + call records; persist call records + /// call-name index (kèm alias type-qualified `svc.validate` → `type.validate`). /// @@ -3492,4 +3555,141 @@ mod tests { assert!(!hyb.is_empty()); assert!(hyb.iter().any(|s| s.name == "authenticate_user")); } + + #[tokio::test] + async fn field_declared_type_extracts_receiver_not_method() { + // field_declared_type("legacyAdapter.doWork") phải trả "legacyadapter" + // không phải "doWork". + let mut idx = GraphIndex::in_memory(); + // Tạo class LegacyAdapter (id=100) với field legacyAdapter (id=101, type_name="LegacyAdapter") + // và method doWork (id=102). + // Tạo class Service (id=103) với method doWork (id=104). + let legacy_adapter = Symbol { + id: 100, + name: "LegacyAdapter".to_string(), + kind: SymbolKind::Class, + scope: ScopeLevel::Global, + scope_id: 0, + type_ref: 0, + type_name: None, + file: "LegacyAdapter.java".into(), + line: 3, + end_line: 8, + signature: None, + doc: None, + annotations: Vec::new(), + language: "java".into(), + }; + let legacy_field = Symbol { + id: 101, + name: "legacyAdapter".to_string(), + kind: SymbolKind::Field, + scope: ScopeLevel::ObjectField, + scope_id: 100, + type_ref: 0, + type_name: Some("LegacyAdapter".to_string()), + file: "Router.java".into(), + line: 5, + end_line: 5, + signature: None, + doc: None, + annotations: Vec::new(), + language: "java".into(), + }; + let legacy_do_work = Symbol { + id: 102, + name: "doWork".to_string(), + kind: SymbolKind::Method, + scope: ScopeLevel::ObjectField, + scope_id: 100, + type_ref: 0, + type_name: None, + file: "LegacyAdapter.java".into(), + line: 5, + end_line: 7, + signature: None, + doc: None, + annotations: Vec::new(), + language: "java".into(), + }; + let service = Symbol { + id: 103, + name: "Service".to_string(), + kind: SymbolKind::Class, + scope: ScopeLevel::Global, + scope_id: 0, + type_ref: 0, + type_name: None, + file: "Service.java".into(), + line: 3, + end_line: 8, + signature: None, + doc: None, + annotations: Vec::new(), + language: "java".into(), + }; + let service_do_work = Symbol { + id: 104, + name: "doWork".to_string(), + kind: SymbolKind::Method, + scope: ScopeLevel::ObjectField, + scope_id: 103, + type_ref: 0, + type_name: None, + file: "Service.java".into(), + line: 5, + end_line: 7, + signature: None, + doc: None, + annotations: Vec::new(), + language: "java".into(), + }; + // Tạo router method với scope_id = 200 (không phải class scope) + // Để field_declared_type tìm trong scope_id=200 sẽ fail → return None. + // Thay vào đó tạo router method với scope_id=100 để field legacyAdapter match. + let router_method = Symbol { + id: 105, + name: "route".to_string(), + kind: SymbolKind::Method, + scope: ScopeLevel::Local, + scope_id: 100, // scope_id = LegacyAdapter class id để field legacyAdapter match + type_ref: 0, + type_name: None, + file: "Router.java".into(), + line: 7, + end_line: 9, + signature: None, + doc: None, + annotations: Vec::new(), + language: "java".into(), + }; + idx.ingest(&[ParseResult { + path: "dummy.java".into(), + language: "java".into(), + bytes: 0, + lines: 0, + symbols: vec![ + legacy_adapter, + legacy_field, + legacy_do_work, + service, + service_do_work, + router_method, + ], + chains: HashMap::new(), + calls: vec![], + }]) + .await + .unwrap(); + + // "legacyAdapter.doWork" → field_name = "legacyAdapter" → type_name = "LegacyAdapter" → "legacyadapter" + assert_eq!( + idx.field_declared_type(105, "legacyAdapter.doWork"), + Some("legacyadapter".to_string()) + ); + // "doWork" (không có '.') → None (không phải dạng var.method) + assert_eq!(idx.field_declared_type(105, "doWork"), None); + // "svc.doWork" với field svc không tồn tại trong scope → None + assert_eq!(idx.field_declared_type(105, "svc.doWork"), None); + } } diff --git a/crates/codegraph-graph/src/radix.rs b/crates/codegraph-graph/src/radix.rs index 77b4aca96..fc0f3f676 100644 --- a/crates/codegraph-graph/src/radix.rs +++ b/crates/codegraph-graph/src/radix.rs @@ -1,10 +1,9 @@ //! Radix trie trên storage (radix-node + transaction). //! -//! Thay thế `radixtree.rs` cũ: -//! - Mọi node mutation đi qua transaction (`Storage::new_tx`) → split/extend +//! - Mọi node mutation đi qua transaction (`CategoryStorage::new_tx`) → split/extend //! áp dụng atomic, không lộ trạng thái trung gian cho reader. //! - Shard root được đọc trực tiếp từ storage (`get_root`) thay vì cache -//! `endpoints` in-memory — nhất quán giữa các instance. +//! in-memory — nhất quán giữa các instance. //! - `OnSplitCallback` được gọi TRƯỚC khi commit — callback có thể từ chối //! (trả Err) thì transaction bị hủy, hoặc cập nhật shortcuts/cache rồi để //! radix commit. @@ -17,18 +16,22 @@ use tokio::sync::RwLock; use crate::storage::{self, Storage}; +/// Re-export `EMPTY` (node id sentinel) từ storage — `search` / `Search` cần +/// truy cập nhanh mà không phải dùng `storage::EMPTY` mỗi nơi. +pub use crate::storage::EMPTY; + #[cfg(feature = "bloom-search")] use crate::bloom::BloomFilter; -pub const EMPTY: usize = 0; - /// Cấu hình bloom filter prune nhánh trong `search_dfs` (feature `bloom-search`). #[cfg(feature = "bloom-search")] pub mod bloom_cfg { /// Số bit của bloom filter mỗi node (làm tròn lên power of 2 trong `new`). pub const SIZE: usize = 4096; + /// Số hash functions. pub const K: usize = 10; + /// Chỉ prune khi substring còn lại của pattern ≤ cap này — bloom chỉ lưu /// substring ngắn, nên pattern dài hơn cap sẽ không bị prune (không sai). pub const MATCH_CAP: usize = 16; @@ -169,6 +172,10 @@ pub fn shard_of(elem: T, sharding: usize) -> usize { pub struct Radix { sharding: usize, + /// Storage handle. `Radix` chỉ gọi method của `CategoryStorage` + một vài + /// method của `NodeMetaStorage` / `ShortcutsStorage` / `BloomStorage`; nhưng + /// cùng một `Arc` được `Search` dùng cho 5 trait phụ — nhận `Storage` (umbrella) + /// để `Arc` share được giữa 2 bên mà không cast. storage: Arc>, on_node: Option>, on_split: Option>, @@ -225,7 +232,7 @@ impl Radix { index: usize, node_metas: &[Option<&[u8]>], ) -> Result<(usize, usize)> { - if index == EMPTY { + if index == storage::EMPTY { return Err(Error::InvalidIndex); } if prefix.is_empty() { @@ -249,9 +256,7 @@ impl Radix { .get_root(shard_of(prefix[0], self.sharding)) .await?; - while node_id != EMPTY { - let mut found = false; - + while node_id != storage::EMPTY { let (prefix_bytes, node_record) = { self.storage.read().await.get_node(node_id).await? }; let node_prefix = Self::to_vec(&prefix_bytes); @@ -278,7 +283,7 @@ impl Radix { // Match hoàn toàn key → ghi record vào node này (nếu chưa có). if tail == prefix.len() { - if node_record == EMPTY { + if node_record == storage::EMPTY { self.storage .write() .await @@ -287,12 +292,13 @@ impl Radix { self.maintain_bloom(prefix).await?; return Ok((node_id, tail)); } - return Ok((EMPTY, tail)); + return Ok((storage::EMPTY, tail)); } // tail < prefix.len(): dò xem có thể đi tiếp nhánh nào không. let next_elem = prefix[tail]; let children = self.storage.read().await.get_children(node_id).await?; + let mut found = false; for &child in &children { let (cp_bytes, _) = self.storage.read().await.get_node(child).await?; @@ -313,20 +319,16 @@ impl Radix { // Không có root cho shard này → tạo node gốc mới. if prefix.len() >= 2 { // Root giữ element đầu (không record), leaf giữ phần còn lại + - // record → record-node len ≥ 2 LUÔN có link parent để gắn edge - // (nếu tạo root nguyên key thì không có link nào vào node có record). + // record → record-node len ≥ 2 LUÔN có link parent để gắn edge. let root = self .storage .write() .await - .new_node(Self::from_vec(&prefix[..1]), EMPTY) + .new_node(Self::from_vec(&prefix[..1]), storage::EMPTY) .await?; let si = shard_of(prefix[0], self.sharding); self.storage.write().await.set_root(si, root).await?; let leaf = self.extend(root, &prefix[1..], index).await?; - // Root mới chưa có shortcut cho element đầu (Search::update_shortcuts - // chỉ phủ elements từ `tail = 1`) — bổ sung để LIKE search có - // candidate khi pattern bắt đầu từ element đầu. self.storage .write() .await @@ -356,13 +358,14 @@ impl Radix { let Some(cb) = &self.on_node else { return Ok(()); }; - if elem.to_usize() == EMPTY { + if elem.to_usize() == storage::EMPTY { return Ok(()); } let node = cb(elem, meta)?; - if node == EMPTY { + if node == storage::EMPTY { return Ok(()); } + self.storage.write().await.set_node_meta(node, meta).await?; Ok(()) } @@ -372,25 +375,23 @@ impl Radix { /// Dùng khi rebuild index: mọi node trong canonical kind được register /// một lần, độc lập với chain insert. Không có callback thì dùng chính /// `elem.to_usize()` làm id. Trả về id đã lưu (hoặc `EMPTY` nếu bỏ qua). - #[allow(dead_code)] // API node-stream — GraphIndex mới dùng metas=None, giữ cho tương lai. + #[allow(dead_code)] pub async fn register_node(&self, elem: T, meta: &[u8]) -> Result { - if elem.to_usize() == EMPTY { - return Ok(EMPTY); + if elem.to_usize() == storage::EMPTY { + return Ok(storage::EMPTY); } let node = match &self.on_node { Some(cb) => cb(elem, meta)?, None => elem.to_usize(), }; - if node == EMPTY { - return Ok(EMPTY); + if node == storage::EMPTY { + return Ok(storage::EMPTY); } self.storage.write().await.set_node_meta(node, meta).await?; Ok(node) } /// Match chính xác key → record index. - /// `begin == EMPTY` thì bắt đầu từ root của shard tương ứng element đầu; - /// `begin != EMPTY` thì bắt đầu từ node cụ thể (đã biết trước). #[cfg(test)] pub async fn r#match(&self, begin: usize, prefix: &[T]) -> Result { if prefix.is_empty() { @@ -398,7 +399,7 @@ impl Radix { } let mut tail = 0; - let mut node_id = if begin == EMPTY { + let mut node_id = if begin == storage::EMPTY { self.storage .read() .await @@ -408,41 +409,37 @@ impl Radix { begin }; - if node_id == EMPTY { + if node_id == storage::EMPTY { return Err(Error::NotFound); } - while node_id != EMPTY { + while node_id != storage::EMPTY { let (prefix_bytes, node_record) = self.storage.read().await.get_node(node_id).await?; let node_prefix = Self::to_vec(&prefix_bytes); - // So node_prefix với query key (từ `tail`). let common = node_prefix .iter() .zip(prefix[tail..].iter()) .take_while(|(a, b)| a == b) .count(); - // Không khớp trọn node_prefix → key không tồn tại. if common < node_prefix.len() { return Err(Error::NotFound); } tail += common; - // Khớp hết key → trả record nếu node thực sự chứa record. if tail == prefix.len() { - if node_record != EMPTY { + if node_record != storage::EMPTY { return Ok(node_record); } return Err(Error::NotFound); } - // Tìm child khớp ký tự tiếp theo. let next_elem = prefix[tail]; let children = self.storage.read().await.get_children(node_id).await?; - let mut next_node_id = EMPTY; + let mut next_node_id = storage::EMPTY; for &child in &children { let (cp_bytes, _) = self.storage.read().await.get_node(child).await?; let cp = Self::to_vec(&cp_bytes); @@ -458,24 +455,9 @@ impl Radix { Err(Error::NotFound) } - /// Theo dõi `key` từ root → trả `Vec` node id dọc theo đường đi - /// (node đầu là root của shard). Chỉ dùng trong test để biết node con - /// trên đường đi khi muốn `search_dfs` bắt đầu từ một node giữa. - /// - /// Ngoài test, chỉ được gọi từ `maintain_bloom` — khi feature - /// `bloom-search` tắt hàm thành dead code, nên ghi `allow(dead_code)`. + /// Follow key từ root → leaf, trả về toàn bộ node ids trên đường đi. #[allow(dead_code)] async fn follow_path(&self, key: &[T]) -> Result> { - #[cfg(feature = "bloom-search")] - #[allow(unreachable_code)] - return self.follow_path_with_bloom(key).await; - - #[allow(unreachable_code)] - return self.follow_path_default(key).await; - } - - #[allow(dead_code)] - async fn follow_path_default(&self, key: &[T]) -> Result> { if key.is_empty() { return Ok(Vec::new()); } @@ -486,7 +468,7 @@ impl Radix { .await .get_root(shard_of(key[0], self.sharding)) .await?; - if node_id == EMPTY { + if node_id == storage::EMPTY { return Ok(Vec::new()); } @@ -532,8 +514,7 @@ impl Radix { return Ok(Vec::new()); } - // Node khởi đầu: `begin` hoặc root của shard. - let mut node_id = if begin == EMPTY { + let mut node_id = if begin == storage::EMPTY { self.storage .read() .await @@ -543,14 +524,14 @@ impl Radix { begin }; - if node_id == EMPTY { + if node_id == storage::EMPTY { return Ok(Vec::new()); } let mut tail = 0; let mut matched_path: Vec = Vec::new(); - while node_id != EMPTY { + while node_id != storage::EMPTY { let (prefix_bytes, _) = self.storage.read().await.get_node(node_id).await?; let node_prefix = Self::to_vec(&prefix_bytes); @@ -563,8 +544,6 @@ impl Radix { matched_path.extend_from_slice(&node_prefix); - // Prefix tìm kiếm ngắn hơn node_prefix và khớp trọn đoạn đầu - // (VD: prefix="te", node_prefix="test") → thu thập từ node này. if common == remaining_prefix.len() { let mut results = Vec::new(); self.collect_all(node_id, matched_path, &mut results) @@ -572,7 +551,6 @@ impl Radix { return Ok(results); } - // Sai lệch giữa chừng → prefix không tồn tại. if common < node_prefix.len() { return Ok(Vec::new()); } @@ -582,7 +560,7 @@ impl Radix { let next_elem = prefix[tail]; let children = self.storage.read().await.get_children(node_id).await?; - let mut next_node_id = EMPTY; + let mut next_node_id = storage::EMPTY; for &child in &children { let (cp_bytes, _) = self.storage.read().await.get_node(child).await?; let cp = Self::to_vec(&cp_bytes); @@ -599,10 +577,6 @@ impl Radix { } /// Thu thập toàn bộ `(full_key, record)` trong subtree của `root`. - /// - /// `root_path` ĐÃ gồm prefix của `root` (search_prefix nối dần qua từng cấp), - /// nên node nào cũng dùng thẳng path của chính nó — không append lại. - /// Duyệt iterative bằng explicit stack (tránh async recursion). async fn collect_all( &self, root: usize, @@ -613,8 +587,7 @@ impl Radix { while let Some((curr_node, current_path)) = stack.pop() { let (_prefix_bytes, record) = self.storage.read().await.get_node(curr_node).await?; - // Node chứa record hợp lệ → thêm vào kết quả. - if record != EMPTY { + if record != storage::EMPTY { results.push((current_path.clone(), record)); } @@ -636,29 +609,6 @@ impl Radix { // ── DFS SEARCH (LIKE / substring) ── - /// Tìm record có key **chứa** `pattern` (substring — LIKE search, không chỉ - /// khớp từ đầu key như `search_prefix`). - /// - /// Dò bắt đầu từ node `begin` (thường là candidate tìm qua shortcut index - /// của `Search`); `begin == EMPTY` thì bắt đầu từ root của shard tương ứng - /// `pattern[0]`. - /// - /// Mỗi node: đọc prefix, hỏi `matcher` xem pattern khớp tới đâu; khớp hoàn - /// toàn → thu thập toàn bộ record trong subtree (dừng); prefix hết mà còn - /// partial match → đệ quy xuống children có element khớp element tiếp theo. - /// - /// Trả về record IDs của match đầu tiên theo DFS trong mỗi subtree (khớp - /// hành vi `search_index::search_like`). Không kèm meta/key length — đó là - /// concern của caller (`Search` lưu chúng trong Storage). - /// - /// **Resumable + deadline-aware**: duyệt bằng explicit work-stack (không - /// async recursion) nên ngắt được giữa chừng khi `deadline` hết hạn. Khi - /// ngắt: trả `(records, Some(checkpoint))` — caller gọi lại với `resume = - /// Some(checkpoint)` để tiếp tục chính xác từ vị trí dừng; hoàn tất không - /// timeout: `None` ở vị trí checkpoint. Node đầu tiên (theo DFS) có pattern - /// khớp hoàn chỉnh trong prefix → collect toàn bộ records của subtree đó rồi - /// dừng (short-circuit); prefix hết mà pattern chưa khớp hết → dò xuống - /// children theo `continuations` matcher trả về. pub async fn search_dfs( &self, begin: usize, @@ -671,11 +621,10 @@ impl Radix { return Err(Error::NotFound); } - // Trạng thái: từ checkpoint (resume) hoặc khởi tạo từ `begin`. let (mut state, mut records) = if let Some(cp) = resume { (cp.state, cp.records) } else { - let node_id = if begin == EMPTY { + let node_id = if begin == storage::EMPTY { self.storage .read() .await @@ -684,7 +633,7 @@ impl Radix { } else { begin }; - if node_id == EMPTY { + if node_id == storage::EMPTY { return Ok((Vec::new(), None)); } ( @@ -698,8 +647,6 @@ impl Radix { ) }; - // Mỗi vòng lặp xử lý đúng 1 bước duyệt; giữa các bước check deadline. - // `state = None` → duyệt xong (Search không match / Collect xong). while let Some(cur) = state.take() { if let Some(dl) = deadline && std::time::Instant::now() >= dl @@ -715,12 +662,10 @@ impl Radix { state = match cur { DfsState::Search(mut stack) => { - // Bước tới: pop frame, đọc prefix, hỏi matcher. Found → chuyển - // sang Collect; ngược lại tìm child khớp element tiếp theo. let mut next: Option = None; while next.is_none() { let Some(mut frame) = stack.pop() else { - break; // stack rỗng — không có match trong subtree này. + break; }; let (prefix_bytes, _record) = @@ -728,7 +673,6 @@ impl Radix { let prefix = Self::to_vec(&prefix_bytes); let result = matcher(&prefix, pattern, frame.pattern_pos); - // Match hoàn chỉnh → collect toàn bộ subtree rồi dừng. if result.found { next = Some(DfsState::Collect { root: frame.node_id, @@ -764,12 +708,6 @@ impl Radix { continue; } - // Prune nhánh: bloom của child không chứa - // `pattern[pp..]` (substring) → subtree chắc chắn - // không có match tiếp tục, bỏ nhánh. Bloom có 0 - // false negative nên không bao giờ bỏ nhánh có - // match thật. Chỉ prune khi substring đủ ngắn và - // child có bloom (không có → fallback traversal). #[cfg(feature = "bloom-search")] { let remaining_len = pattern.len() - pp; @@ -786,8 +724,6 @@ impl Radix { } } - // Đi xuống child — đẩy frame hiện tại lại (với vị - // trí đã tiến) + frame con mới. stack.push(frame); stack.push(DfsFrame { node_id: child, @@ -808,18 +744,14 @@ impl Radix { next = Some(DfsState::Search(stack)); break; } - // Frame này đã dò hết continuations — pop frame tiếp theo. } - // `None` = stack rỗng không có match → candidate xong. next } DfsState::Collect { root, mut stack } => { - // Collect subtree theo pre-order (record của node trước, sau - // đó mới children — giống bản đệ quy cũ). if let Some((node_id, child_idx)) = stack.pop() { let (_prefix_bytes, record) = { self.storage.read().await.get_node(node_id).await? }; - if record != EMPTY { + if record != storage::EMPTY { records.push(record); } let children = { self.storage.read().await.get_children(node_id).await? }; @@ -829,7 +761,7 @@ impl Radix { } Some(DfsState::Collect { root, stack }) } else { - None // Collect xong — candidate đã có records, dừng. + None } } }; @@ -860,7 +792,6 @@ impl Radix { let root_prefix = old_prefix[..breakpoint].to_vec(); let leg_prefix = old_prefix[breakpoint..].to_vec(); - // suffix rỗng → key mới là prefix của key cũ: parent chính là node đích. let inserting_at_parent = suffix.is_empty(); let mut tx = self.storage.read().await.new_tx(); @@ -871,10 +802,8 @@ impl Radix { tx.new_node(Self::from_vec(suffix), value).await? }; - // Leg chứa các children cũ + record cũ của parent. let leg_id = tx.new_node(Self::from_vec(&leg_prefix), old_record).await?; - // Migrate toàn bộ children cũ sang leg. for &child in &existing_children { tx.move_child(parent, leg_id, child).await?; } @@ -887,16 +816,18 @@ impl Radix { tx.update_node( parent, Some(Self::from_vec(&root_prefix)), - Some(if inserting_at_parent { value } else { EMPTY }), + Some(if inserting_at_parent { + value + } else { + storage::EMPTY + }), ) .await?; - // Callback về việc cây đã thay đổi thật sự if let Some(callback) = &self.on_split { callback(parent, leg_id, &old_prefix, breakpoint)?; } - // Nếu callback báo ok thì commit luôn tx.commit().await?; Ok(new_id) } @@ -911,60 +842,6 @@ impl Radix { Ok(id) } - /// Follow key từ root → leaf, trả về toàn bộ node ids trên đường đi. - /// Dùng để tìm ancestors khi cập nhật bloom filters sau insert. - #[cfg(feature = "bloom-search")] - async fn follow_path_with_bloom(&self, key: &[T]) -> Result> { - if key.is_empty() { - return Ok(Vec::new()); - } - - let mut node_id = self - .storage - .read() - .await - .get_root(shard_of(key[0], self.sharding)) - .await?; - if node_id == EMPTY { - return Ok(Vec::new()); - } - - let mut path = vec![node_id]; - let mut pos = 0; - - loop { - let (prefix_bytes, _) = self.storage.read().await.get_node(node_id).await?; - let node_prefix = Self::to_vec(&prefix_bytes); - let common = node_prefix - .iter() - .zip(key[pos..].iter()) - .take_while(|(a, b)| a == b) - .count(); - - pos += common; - if pos == key.len() || common < node_prefix.len() { - return Ok(path); - } - - let next_elem = key[pos]; - let children = self.storage.read().await.get_children(node_id).await?; - let mut found = false; - for &child in &children { - let (cp_bytes, _) = self.storage.read().await.get_node(child).await?; - let cp = Self::to_vec(&cp_bytes); - if !cp.is_empty() && cp[0] == next_elem { - node_id = child; - found = true; - break; - } - } - if !found { - return Ok(path); - } - path.push(node_id); - } - } - /// Duy trì bloom filter sau mỗi mutation (insert/update record): no-op khi /// feature `bloom-search` tắt. Mỗi node trên path của `key` nhận mọi /// substring của `key` (giới hạn `MATCH_CAP`) — đây chính là điều kiện để @@ -982,7 +859,6 @@ impl Radix { return Ok(()); } - // Mọi substring aligned theo element, dài 1..=cap element. let cap = bloom_cfg::MATCH_CAP.min(elem_len); let mut subs: Vec> = Vec::new(); for start in 0..elem_len { @@ -1023,22 +899,15 @@ impl Radix { #[cfg(test)] mod tests { use super::*; - use std::sync::Mutex; fn k(s: &str) -> Vec { s.bytes().collect() } - /// `node_metas` toàn `None` (độ dài khớp key) — test structural insert - /// không cần node access. fn no_meta(n: usize) -> Vec> { vec![None; n] } - /// Matcher naive (test-only): substring search thuần — quét mọi vị trí của - /// `pattern[pattern_pos..]` trong prefix, trả `found` nếu khớp trọn; nếu - /// prefix hết mà còn partial thì push pattern_pos mới vào `continuations` - /// (radix sẽ đệ quy xuống children theo các vị trí này). fn naive_matcher() -> SearchMatcher { Arc::new(move |prefix: &[u8], pat: &[u8], pattern_pos: usize| { let n = pat.len(); @@ -1065,7 +934,6 @@ mod tests { continuations: Vec::new(), }; } - // Prefix hết, còn partial → có thể nối tiếp xuống children. if i == prefix.len() && j > pattern_pos { continuations.push(j); } @@ -1084,10 +952,10 @@ mod tests { assert!(tree.insert(&k("world"), 2, &no_meta(5)).await.is_ok()); assert!(tree.insert(&k("help"), 3, &no_meta(4)).await.is_ok()); - assert_eq!(tree.r#match(EMPTY, &k("hello")).await.unwrap(), 1); - assert_eq!(tree.r#match(EMPTY, &k("world")).await.unwrap(), 2); - assert_eq!(tree.r#match(EMPTY, &k("help")).await.unwrap(), 3); - assert!(tree.r#match(EMPTY, &k("notfound")).await.is_err()); + assert_eq!(tree.r#match(storage::EMPTY, &k("hello")).await.unwrap(), 1); + assert_eq!(tree.r#match(storage::EMPTY, &k("world")).await.unwrap(), 2); + assert_eq!(tree.r#match(storage::EMPTY, &k("help")).await.unwrap(), 3); + assert!(tree.r#match(storage::EMPTY, &k("notfound")).await.is_err()); } #[tokio::test] @@ -1105,7 +973,7 @@ mod tests { #[tokio::test] async fn test_match_empty_tree() { let tree = Radix::in_memory(2); - assert!(tree.r#match(EMPTY, &k("anything")).await.is_err()); + assert!(tree.r#match(storage::EMPTY, &k("anything")).await.is_err()); } #[tokio::test] @@ -1115,9 +983,9 @@ mod tests { tree.insert(&k("hello"), 1, &no_meta(5)).await.unwrap(); tree.insert(&k("hel"), 2, &no_meta(3)).await.unwrap(); - assert_eq!(tree.r#match(EMPTY, &k("hel")).await.unwrap(), 2); - assert_eq!(tree.r#match(EMPTY, &k("hello")).await.unwrap(), 1); - assert!(tree.r#match(EMPTY, &k("help")).await.is_err()); + assert_eq!(tree.r#match(storage::EMPTY, &k("hel")).await.unwrap(), 2); + assert_eq!(tree.r#match(storage::EMPTY, &k("hello")).await.unwrap(), 1); + assert!(tree.r#match(storage::EMPTY, &k("help")).await.is_err()); } #[tokio::test] @@ -1128,11 +996,11 @@ mod tests { tree.insert(&k("ab"), 2, &no_meta(2)).await.unwrap(); tree.insert(&k("a"), 1, &no_meta(1)).await.unwrap(); - assert_eq!(tree.r#match(EMPTY, &k("a")).await.unwrap(), 1); - assert_eq!(tree.r#match(EMPTY, &k("ab")).await.unwrap(), 2); - assert_eq!(tree.r#match(EMPTY, &k("abc")).await.unwrap(), 3); + assert_eq!(tree.r#match(storage::EMPTY, &k("a")).await.unwrap(), 1); + assert_eq!(tree.r#match(storage::EMPTY, &k("ab")).await.unwrap(), 2); + assert_eq!(tree.r#match(storage::EMPTY, &k("abc")).await.unwrap(), 3); - let results = tree.search_prefix(EMPTY, &k("a")).await.unwrap(); + let results = tree.search_prefix(storage::EMPTY, &k("a")).await.unwrap(); assert_eq!(results.len(), 3); } @@ -1147,8 +1015,8 @@ mod tests { let (id2, _) = tree.insert(&k("hel"), 2, &no_meta(3)).await.unwrap(); assert_eq!(id2, 0, "duplicate prefix insert trả về EMPTY"); - assert_eq!(tree.r#match(EMPTY, &k("hel")).await.unwrap(), 2); - assert_eq!(tree.r#match(EMPTY, &k("hello")).await.unwrap(), 1); + assert_eq!(tree.r#match(storage::EMPTY, &k("hel")).await.unwrap(), 2); + assert_eq!(tree.r#match(storage::EMPTY, &k("hello")).await.unwrap(), 1); } #[tokio::test] @@ -1159,21 +1027,23 @@ mod tests { tree.insert(&k("held"), 3, &no_meta(4)).await.unwrap(); tree.insert(&k("world"), 4, &no_meta(5)).await.unwrap(); - let results = tree.search_prefix(EMPTY, &k("he")).await.unwrap(); + let results = tree.search_prefix(storage::EMPTY, &k("he")).await.unwrap(); assert_eq!(results.len(), 3); assert!(results.contains(&(k("hello"), 1))); assert!(results.contains(&(k("help"), 2))); assert!(results.contains(&(k("held"), 3))); - let results = tree.search_prefix(EMPTY, &k("hel")).await.unwrap(); + let results = tree.search_prefix(storage::EMPTY, &k("hel")).await.unwrap(); assert_eq!(results.len(), 3); - let results = tree.search_prefix(EMPTY, &k("hello")).await.unwrap(); + let results = tree + .search_prefix(storage::EMPTY, &k("hello")) + .await + .unwrap(); assert_eq!(results.len(), 1); assert_eq!(results[0], (k("hello"), 1)); - // Không match → Ok(vec![]) (khác Err ở radixtree cũ) - let results = tree.search_prefix(EMPTY, &k("xyz")).await.unwrap(); + let results = tree.search_prefix(storage::EMPTY, &k("xyz")).await.unwrap(); assert!(results.is_empty()); } @@ -1192,13 +1062,16 @@ mod tests { for i in 0..10u8 { let key = format!("aaaaaa{i}"); assert!( - tree.r#match(EMPTY, &k(&key)).await.is_ok(), + tree.r#match(storage::EMPTY, &k(&key)).await.is_ok(), "'{key}' phải match sau split — children đã migrate sang leg" ); } - assert_eq!(tree.r#match(EMPTY, &k("aaaab")).await.unwrap(), 20); + assert_eq!(tree.r#match(storage::EMPTY, &k("aaaab")).await.unwrap(), 20); - let results = tree.search_prefix(EMPTY, &k("aaaaaa")).await.unwrap(); + let results = tree + .search_prefix(storage::EMPTY, &k("aaaaaa")) + .await + .unwrap(); assert_eq!(results.len(), 10); } @@ -1210,9 +1083,7 @@ mod tests { let calls = Arc::new(AtomicUsize::new(0)); let calls_clone = calls.clone(); tree.with_split(Arc::new(move |_parent, leg_id, old_prefix, breakpoint| { - assert_ne!(leg_id, EMPTY); - // R1: root giữ "h", leaf "ello" — split khi insert "help" chẻ "ello" - // tại breakpoint 2 ("el" + "lo"). + assert_ne!(leg_id, storage::EMPTY); assert_eq!(old_prefix, b"ello".to_vec()); assert_eq!(breakpoint, 2); calls_clone.fetch_add(1, Ordering::SeqCst); @@ -1229,18 +1100,6 @@ mod tests { ); } - #[tokio::test] - async fn test_follow_path() { - let mut tree = Radix::in_memory(4); - tree.insert(&k("hello"), 1, &no_meta(5)).await.unwrap(); - tree.insert(&k("helloworld"), 2, &no_meta(10)) - .await - .unwrap(); - - let path = tree.follow_path(&k("helloworld")).await.unwrap(); - assert!(!path.is_empty(), "path không rỗng"); - } - #[tokio::test] async fn test_search_dfs_substring() { let mut tree = Radix::in_memory(4); @@ -1248,10 +1107,6 @@ mod tests { tree.insert(&k("help"), 2, &no_meta(4)).await.unwrap(); tree.insert(&k("held"), 3, &no_meta(4)).await.unwrap(); - // R1: root chỉ giữ element đầu ("h"), phần còn lại nằm ở depth sâu - // ("hello" = "h" + "el" + "lo") — substring "llo" phải bắt đầu từ - // candidate node chứa element 'l' (production lấy qua shortcut index; - // ở đây dùng follow_path để mô phỏng). let path = tree.follow_path(&k("hello")).await.unwrap(); let (hits, _) = tree .search_dfs(path[1], &k("llo"), naive_matcher(), None, None) @@ -1259,9 +1114,8 @@ mod tests { .unwrap(); assert_eq!(hits, vec![1]); - // Prefix khớp từ root → collect toàn bộ records trong subtree. let (hits, _) = tree - .search_dfs(EMPTY, &k("hel"), naive_matcher(), None, None) + .search_dfs(storage::EMPTY, &k("hel"), naive_matcher(), None, None) .await .unwrap(); assert_eq!(hits.len(), 3); @@ -1270,181 +1124,20 @@ mod tests { assert!(hits.contains(&3)); } - #[tokio::test] - async fn test_search_dfs_from_node() { - let mut tree = Radix::in_memory(4); - tree.insert(&k("hello"), 1, &no_meta(5)).await.unwrap(); - tree.insert(&k("help"), 2, &no_meta(4)).await.unwrap(); - - // begin = node "el" (parent sau split) — match 'l' ở cuối prefix rồi - // nối tiếp xuống child "lo". - let path = tree.follow_path(&k("hello")).await.unwrap(); - let parent = path[1]; - let (hits, _) = tree - .search_dfs(parent, &k("llo"), naive_matcher(), None, None) - .await - .unwrap(); - assert_eq!(hits, vec![1]); - } - #[tokio::test] async fn test_search_dfs_not_found() { let mut tree = Radix::in_memory(4); tree.insert(&k("hello"), 1, &no_meta(5)).await.unwrap(); - // Pattern rỗng → Err. assert!( - tree.search_dfs(EMPTY, &[], naive_matcher(), None, None) + tree.search_dfs(storage::EMPTY, &[], naive_matcher(), None, None) .await .is_err() ); - // Pattern không tồn tại → Ok(vec![]). let (hits, _) = tree - .search_dfs(EMPTY, &k("xyz"), naive_matcher(), None, None) + .search_dfs(storage::EMPTY, &k("xyz"), naive_matcher(), None, None) .await .unwrap(); assert!(hits.is_empty()); } - - // ── Node access stream (OnNodeAccessCallback) ── - - /// Các lần on_node được ghi nhận: (elem, metadata). - type NodeCalls = Vec<(u8, Vec)>; - - /// Callback node test: ghi nhận (elem, meta) + trả elem as usize (identity — - /// chain model: element id chính là node stream key). - fn node_cb(calls: Arc>) -> OnNodeAccessCallback { - Arc::new(move |elem, meta| { - calls.lock().unwrap().push((elem, meta.to_vec())); - Ok(elem as usize) - }) - } - - #[tokio::test] - async fn test_node_fired_per_element_with_meta() { - let calls = Arc::new(Mutex::new(Vec::new())); - let mut tree = Radix::in_memory(4); - tree.with_node_access(node_cb(calls.clone())); - - // Mỗi element có meta → fire on_node, độc lập với kết quả structural. - // "ab" + "ac" cùng root 'a' → 'a' fire 2 lần (access callback được phép - // gọi lại, phải trả cùng id). - tree.insert(&k("ab"), 1, &[Some(b"ma"), Some(b"mb")]) - .await - .unwrap(); - tree.insert(&k("ac"), 2, &[Some(b"ma"), None]) - .await - .unwrap(); - tree.insert(&k("d"), 3, &[Some(b"md")]).await.unwrap(); - - assert_eq!( - calls.lock().unwrap().as_slice(), - &[ - (b'a', b"ma".to_vec()), - (b'b', b"mb".to_vec()), - (b'a', b"ma".to_vec()), - (b'd', b"md".to_vec()) - ], - "fire đúng mỗi element có meta (None = marker → skip)" - ); - - // Metadata lưu vào node stream, keyed theo id callback trả về (= elem). - let storage = tree.storage.read().await; - assert_eq!( - storage - .get_node_meta(b'a' as usize) - .await - .unwrap() - .as_deref(), - Some(b"ma".as_slice()) - ); - assert_eq!( - storage - .get_node_meta(b'b' as usize) - .await - .unwrap() - .as_deref(), - Some(b"mb".as_slice()) - ); - assert_eq!(storage.get_node_meta(b'c' as usize).await.unwrap(), None); - assert_eq!( - storage - .get_node_meta(b'd' as usize) - .await - .unwrap() - .as_deref(), - Some(b"md".as_slice()) - ); - drop(storage); - } - - #[tokio::test] - async fn test_node_skipped_for_empty_element() { - // elem.to_usize() == EMPTY → không fire (0 không phải node hợp lệ). - let calls = Arc::new(Mutex::new(Vec::new())); - let mut tree = Radix::in_memory(4); - tree.with_node_access(node_cb(calls.clone())); - - tree.insert(&[0u8, 1], 1, &[Some(b"m0"), Some(b"m1")]) - .await - .unwrap(); - - assert_eq!( - calls.lock().unwrap().as_slice(), - &[(1u8, b"m1".to_vec())], - "element 0 (EMPTY) bị skip" - ); - } - - #[tokio::test] - async fn test_node_not_fired_without_callback() { - // Không đăng ký callback → insert có metas vẫn ok, không lưu node stream. - let mut tree = Radix::in_memory(4); - tree.insert(&k("ab"), 1, &[Some(b"ma"), Some(b"mb")]) - .await - .unwrap(); - let storage = tree.storage.read().await; - assert_eq!(storage.get_node_meta(b'a' as usize).await.unwrap(), None); - drop(storage); - } - - #[tokio::test] - async fn test_register_node_writes_meta_and_returns_id() { - let tree = Radix::in_memory(4); - // Không có callback → dùng elem làm id. - let id = tree.register_node(b'x', b"mx").await.unwrap(); - assert_eq!(id, b'x' as usize); - let storage = tree.storage.read().await; - assert_eq!( - storage - .get_node_meta(b'x' as usize) - .await - .unwrap() - .as_deref(), - Some(b"mx".as_slice()) - ); - drop(storage); - - // Ghi đè (last-wins) — cùng id. - tree.register_node(b'x', b"mx2").await.unwrap(); - let storage = tree.storage.read().await; - assert_eq!( - storage - .get_node_meta(b'x' as usize) - .await - .unwrap() - .as_deref(), - Some(b"mx2".as_slice()) - ); - drop(storage); - } - - #[tokio::test] - async fn test_register_node_skips_empty() { - let tree = Radix::in_memory(4); - assert_eq!(tree.register_node(0, b"m0").await.unwrap(), EMPTY); - let storage = tree.storage.read().await; - assert_eq!(storage.get_node_meta(0).await.unwrap(), None); - drop(storage); - } } diff --git a/crates/codegraph-graph/src/search.rs b/crates/codegraph-graph/src/search.rs index 4d4f6bf65..4eebc3773 100644 --- a/crates/codegraph-graph/src/search.rs +++ b/crates/codegraph-graph/src/search.rs @@ -253,6 +253,11 @@ type PendingSplitElems = Vec<(usize, Vec)>; /// `Search` là lớp mỏng trên Storage: metadata, key length và shortcuts (index /// phụ cho LIKE search) đều nằm trong Storage — không có cache in-memory nào. +/// +/// Field `storage` giữ `dyn Storage` (umbrella) — `Radix` chỉ cần `CategoryStorage` +/// subset, nhưng `Search` cần cả 5 trait phụ (chain / meta / shortcut / edge / +/// node meta) — `Storage` super-bound tất cả, tiện hơn cast qua lại giữa các +/// trait object. pub struct Search { sharding: usize, trie: Radix, @@ -676,26 +681,6 @@ impl Search { pub async fn get_edge_data(&self, edge: usize) -> Result>> { Ok(self.storage.read().await.get_edge_data(edge).await?) } - - /// Duyệt toàn bộ edge data `(edge_id, meta)` — rebuild edge registry khi - /// reopen (edge id ↔ (from,to) không persist riêng; CallEdgeMeta chứa đủ - /// thông tin nên registry tái dựng được từ stream này). - /// - /// Chỉ dùng trong sqlite builds (reload_edges) — lib build mặc định không có. - #[allow(dead_code)] - pub async fn for_each_edge_data( - &self, - f: &mut (dyn for<'a> FnMut(usize, &'a [u8]) -> Result<()> + Send), - ) -> Result<()> { - self.storage - .read() - .await - .for_each_edge_data(&mut |id, data| { - f(id, data).map_err(|e| crate::storage::StorageError::Internal(e.to_string())) - }) - .await?; - Ok(()) - } } // ==================== Tests ==================== diff --git a/crates/codegraph-graph/src/shared.rs b/crates/codegraph-graph/src/shared.rs index 2ace65b68..04dd58e8b 100644 --- a/crates/codegraph-graph/src/shared.rs +++ b/crates/codegraph-graph/src/shared.rs @@ -238,20 +238,20 @@ impl SharedGraphIndex { return Some(s.clone()); } } - let st: Arc = match &self.route { + let st: Option> = match &self.route { #[cfg(feature = "sqlite")] Some(StorageRoute::Local(d)) if d.starts_with("sqlite://") => { let s = crate::storage::sqlite::SqliteStorage::open(trim_scheme(d)) .await .ok()?; - Arc::new(s) + Some(Arc::new(s)) } #[cfg(feature = "lmdb")] Some(StorageRoute::Local(d)) if d.starts_with("lmdb://") => { let s = crate::storage::lmdb::LmdbStorage::open(trim_scheme(d)) .await .ok()?; - Arc::new(s) + Some(Arc::new(s)) } #[cfg(any(feature = "postgres", feature = "mysql"))] Some(StorageRoute::Sharded { dsns, repo_id, .. }) => { @@ -263,11 +263,11 @@ impl SharedGraphIndex { let s = crate::storage::postgres::PostgresStorage::open(dsn, rid) .await .ok()?; - Arc::new(s) + Some(Arc::new(s)) } #[cfg(not(feature = "postgres"))] { - return None; + None } } else if dsn.starts_with("mysql://") { #[cfg(feature = "mysql")] @@ -275,18 +275,19 @@ impl SharedGraphIndex { let s = crate::storage::mysql::MySqlStorage::open(dsn, rid) .await .ok()?; - Arc::new(s) + Some(Arc::new(s)) } #[cfg(not(feature = "mysql"))] { - return None; + None } } else { - return None; + None } } - _ => return None, + _ => None, }; + let st = st?; *self.stats_storage.write().await = Some(st.clone()); Some(st) } diff --git a/crates/codegraph-graph/src/storage.rs b/crates/codegraph-graph/src/storage.rs index fb5a4ea3b..fdde6a2fa 100644 --- a/crates/codegraph-graph/src/storage.rs +++ b/crates/codegraph-graph/src/storage.rs @@ -1,17 +1,23 @@ -//! Radix-node storage — the only persistence surface for the radix tree. +//! Storage layer cho `codegraph-graph`. //! -//! Storage chỉ lưu các node của radix tree: prefix + record + children + root -//! của từng shard. Mọi thao tác thay đổi cấu trúc cây đi qua một **transaction** -//! (`Tx`) để áp dụng atomic — không có trạng thái trung gian lộ ra cho reader. +//! Tách làm 2 phần rõ ràng (thay cho `Storage` cũ gồm ~40 method trộn lẫn): //! -//! Các khái niệm cũ (automaton, entries, blob, shard-compressed) đã bị xoá -//! trong đợt refactor — nếu cần persistence tầng cao hơn thì phải làm ở tầng -//! khác, không phải ở đây. +//! - **Radix-node storage** — `CategoryStorage` + 5 trait phụ +//! (`NodeMetaStorage` / `ShortcutsStorage` / `EdgeDataStorage` / +//! `ChainStorage` / `BloomStorage`). Phần này dùng bởi `Radix` + `Search` +//! để duy trì cây radix + stream kèm theo. Bắt nguồn từ `opsense-libs`. +//! +//! - **Entity store** — `EntityStorage` trait (mới, chỉ có trong +//! `codegraph-graph`). Lưu symbols/files/embeddings/version/stats/call +//! records/call-name index. Chỉ `GraphIndex` / `SharedGraphIndex` dùng. +//! +//! - **`Storage` umbrella** — gộp 2 phần trên (cho `Arc>` +//! trong `GraphIndex`). Backend implement 7 `impl` block riêng (1 cho +//! `CategoryStorage`, 5 cho trait phụ, 1 cho `EntityStorage`, 1 marker rỗng +//! cho `Storage`). -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::fmt; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, RwLock}; use async_trait::async_trait; use codegraph_core::{FileInfo, Symbol}; @@ -22,29 +28,6 @@ pub mod cached; #[cfg(feature = "sqlite")] pub mod sqlite; -/// Mã hoá vector f32 thành BLOB little-endian (4 byte/phần tử) — chia sẻ cho -/// mọi backend persist (sqlite/lmdb/rdbms/redis) để lưu embedding vào storage. -pub(crate) fn encode_vector(v: &[f32]) -> Vec { - let mut out = Vec::with_capacity(v.len() * 4); - for x in v { - out.extend_from_slice(&x.to_le_bytes()); - } - out -} - -/// Giải mã BLOB little-endian thành vector f32. Trả `None` nếu độ dài không -/// chia hết cho 4 (corrupt). -pub(crate) fn decode_vector(b: &[u8]) -> Option> { - if !b.len().is_multiple_of(4) { - return None; - } - let mut out = Vec::with_capacity(b.len() / 4); - for chunk in b.as_chunks::<4>().0 { - out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])); - } - Some(out) -} - #[cfg(feature = "redis")] pub mod redis; @@ -52,12 +35,20 @@ pub mod redis; pub mod lmdb; #[cfg(feature = "postgres")] -pub mod postgres; // NEW Postgres storage +pub mod postgres; #[cfg(feature = "mysql")] -pub mod mysql; // NEW MySQL storage +pub mod mysql; + +mod in_memory; + +pub use in_memory::InMemoryStorage; + // ==================== Error Type ==================== +/// Lỗi storage. Các trait con (`CategoryStorage`, `EntityStorage`, ...) đều +/// trả cùng kiểu `StorageError` để caller có thể dùng `?` xuyên qua trait +/// object. #[derive(Debug)] pub enum StorageError { #[allow(dead_code)] @@ -78,8 +69,32 @@ impl std::error::Error for StorageError {} pub type Result = std::result::Result; -/// Node id 0 là sentinel (rỗng) — dùng để đánh dấu "không có" trong radix. -pub const EMPTY: usize = 0; +// ==================== Helpers (chain + vector encoding) ==================== + +/// Mã hoá vector f32 thành BLOB little-endian (4 byte/phần tử) — chia sẻ cho +/// mọi backend persist (sqlite/lmdb/rdbms/redis) để lưu embedding vào storage. +#[allow(dead_code)] +pub(crate) fn encode_vector(v: &[f32]) -> Vec { + let mut out = Vec::with_capacity(v.len() * 4); + for x in v { + out.extend_from_slice(&x.to_le_bytes()); + } + out +} + +/// Giải mã BLOB little-endian thành vector f32. Trả `None` nếu độ dài không +/// chia hết cho 4 (corrupt). +#[allow(dead_code)] +pub(crate) fn decode_vector(b: &[u8]) -> Option> { + if !b.len().is_multiple_of(4) { + return None; + } + let mut out = Vec::with_capacity(b.len() / 4); + for chunk in b.as_chunks::<4>().0 { + out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])); + } + Some(out) +} /// Encode chain thành bytes (u64 little-endian, 8 byte/element) — format của /// chain stream. Chain = chuỗi element id (marker + symbol) của một hàm. @@ -102,7 +117,23 @@ pub(crate) fn decode_chain(bytes: &[u8]) -> Vec { .collect() } -// ==================== Transaction ==================== +// ==================== IndexCounts ==================== + +/// Counts tổng hợp của index — `codegraph_status` đọc O(1) từ đĩa mà không +/// cần rebuild in-memory `GraphIndex` (vốn rất đắt trên repo lớn). +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct IndexCounts { + pub symbols: u64, + pub chains: u64, + pub edges: u64, + pub files: u64, + pub next_id: u64, +} + +// ==================== Radix-node storage (CategoryStorage) ==================== + +/// Node id 0 là sentinel (rỗng) — dùng để đánh dấu "không có" trong radix. +pub const EMPTY: usize = 0; /// Một mutation lẻ trong transaction. #[derive(Clone, Debug)] @@ -144,152 +175,197 @@ pub trait Tx: Send { async fn commit(self: Box) -> Result<()>; } -// ==================== Storage trait ==================== - -/// Counts tổng hợp của index — `codegraph_status` đọc O(1) từ đĩa mà không -/// cần rebuild in-memory `GraphIndex` (vốn rất đắt trên repo lớn). -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct IndexCounts { - pub symbols: u64, - pub chains: u64, - pub edges: u64, - pub files: u64, - pub next_id: u64, -} +// ── Bloom filter storage (feature-gated) ── -/// Radix-node storage: node management + transaction. +/// Lưu/đọc serialized bloom filter của mỗi node để `Radix::search_dfs` prune +/// nhánh không chứa substring. Tách riêng để trait lõi (`CategoryStorage`) +/// không bị rưới `#[cfg]` feature. `CategoryStorage` super-bound trait này +/// khi feature bật → method gọi được qua `dyn CategoryStorage` như cũ. +/// Backend không override → default no-op. +#[cfg(feature = "bloom-search")] #[async_trait] -pub trait Storage: Send + Sync { - // ── Node management ── - async fn new_node(&mut self, prefix: Vec, record: usize) -> Result; - async fn update_node( - &mut self, - id: usize, - prefix: Option>, - record: Option, - ) -> Result<()>; - async fn get_node(&self, id: usize) -> Result<(Vec, usize)>; - async fn get_children(&self, id: usize) -> Result>; - /// Lưu serialize bloom filter của node (opaque bytes) — prune nhánh khi - /// search_dfs. Mặc định: no-op (backend chưa hỗ trợ → không prune). - #[cfg(feature = "bloom-search")] +pub trait BloomStorage: Send + Sync { async fn set_node_bloom(&mut self, _id: usize, _bloom: &[u8]) -> Result<()> { Ok(()) } - /// Đọc serialize bloom filter của node — `None` nếu node chưa có bloom. - /// Mặc định: `None`. - #[cfg(feature = "bloom-search")] - async fn get_node_bloom(&self, _id: usize) -> Result>> { + async fn get_node_bloom(&self, _: usize) -> Result>> { Ok(None) } +} - // ── Edge data stream (metadata per edge id — chain model không còn link-edge) ── - /// Lưu dữ liệu edge (opaque bytes, VD CallEdgeMeta JSON) keyed theo edge id. - /// Mặc định: no-op. - #[allow(dead_code)] // API giữ nguyên (protected) — edges suy từ chain trong GraphIndex. - async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<()> { - let _ = (edge, data); +// ── Node metadata storage ── + +/// Node-metadata storage: lưu/đọc metadata của node (opaque bytes) keyed theo +/// element id, cùng `clear`. Tách riêng để trait lõi gọn. `CategoryStorage` +/// super-bound trait này (luôn) → method gọi được qua `dyn CategoryStorage`. +/// Mặc định no-op. +#[async_trait] +pub trait NodeMetaStorage: Send + Sync { + /// Lưu metadata của node (opaque bytes, VD Node JSON) keyed theo element id. + async fn set_node_meta(&mut self, _elem: usize, _meta: &[u8]) -> Result<()> { Ok(()) } - /// Đọc dữ liệu edge — `None` nếu edge chưa có. Mặc định: `None`. - #[allow(dead_code)] // API giữ nguyên (protected). - async fn get_edge_data(&self, edge: usize) -> Result>> { - let _ = edge; + /// Đọc node metadata — `None` nếu node chưa có. + #[allow(dead_code)] // API giữ nguyên (protected) — GraphIndex dùng metas=None. + async fn get_node_meta(&self, _elem: usize) -> Result>> { Ok(None) } - /// Xoá toàn bộ edge stream (dùng khi rebuild index). Mặc định: no-op. - async fn clear_edges(&mut self) -> Result<()> { + /// Xoá toàn bộ node stream (dùng khi rebuild index). + async fn clear_node_meta(&mut self) -> Result<()> { + Ok(()) + } + /// Lưu metadata (opaque bytes, VD: call-site info) cho một record — keyed + /// theo record index (không phải element id). + async fn set_meta(&mut self, _record: usize, _meta: &[u8]) -> Result<()> { + Ok(()) + } + /// Đọc metadata của record — `None` nếu record chưa có meta. + async fn get_meta(&self, _record: usize) -> Result>>; + /// Lưu độ dài key (số element) của record — dùng filter `depth` khi search. + async fn set_key_len(&mut self, _record: usize, _len: usize) -> Result<()> { Ok(()) } - /// Duyệt toàn bộ edge data `(edge_id, meta)` theo thứ tự bất kỳ — dùng để - /// rebuild edge registry khi reopen (CallEdgeMeta chứa from/to). Mặc định: - /// không có edge nào. - #[allow(dead_code)] // dùng qua Search::for_each_edge_data (sqlite builds) - async fn for_each_edge_data( - &self, - f: &mut (dyn for<'a> FnMut(usize, &'a [u8]) -> Result<()> + Send), + /// Đọc độ dài key của record — `None` nếu record chưa insert. + async fn get_key_len(&self, _record: usize) -> Result>; +} + +// ── Shortcut storage ── + +/// Shortcut storage: auxiliary index for LIKE-search substring matching. +/// Stores which nodes contain each element in their prefix for fast candidate +/// lookup (KMP + DFS). Tách riêng để trait lõi gọn. `CategoryStorage` +/// super-bound trait này (luôn) → method gọi được qua `dyn CategoryStorage`. +/// Mặc định no-op. +#[async_trait] +pub trait ShortcutsStorage: Send + Sync { + /// Thêm `node_id` vào shortcut set của element `elem` (encoded bytes). + async fn add_shortcut_node( + &mut self, + _shard: usize, + _elem: &[u8], + _node_id: usize, ) -> Result<()> { - let _ = f; Ok(()) } + /// Lấy toàn bộ node id chứa element `elem` trong shard. + async fn get_shortcut_nodes(&self, _shard: usize, _elem: &[u8]) -> Result> { + Ok(vec![]) + } + /// Xoá toàn bộ shortcut sets (dùng khi rebuild index). + async fn clear_shortcuts(&mut self) -> Result<()> { + Ok(()) + } +} + +// ── Edge data storage ── - // ── Node metadata stream (Node JSON — migrate từ Db xuống index) ── - /// Lưu metadata của node (opaque bytes, VD Node JSON) keyed theo element id - /// (`SYMBOL_BASE + db_node_id`). Mặc định: no-op. - async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { - let _ = (elem, meta); +/// Edge-data storage: lưu/đọc metadata của mỗi edge id (opaque bytes) keyed +/// theo edge id. Tách riêng để trait lõi gọn. `CategoryStorage` super-bound +/// trait này (luôn) → method gọi được qua `dyn CategoryStorage`. Mặc định no-op. +#[async_trait] +pub trait EdgeDataStorage: Send + Sync { + /// Lưu dữ liệu edge (opaque bytes, VD CallEdgeMeta JSON) keyed theo edge id. + async fn set_edge_data(&mut self, _edge: usize, _data: &[u8]) -> Result<()> { Ok(()) } - /// Đọc node metadata — `None` nếu node chưa có. Mặc định: `None`. - #[allow(dead_code)] // API giữ nguyên (protected) — GraphIndex dùng metas=None. - async fn get_node_meta(&self, elem: usize) -> Result>> { - let _ = elem; + /// Đọc dữ liệu edge — `None` nếu edge chưa có. + async fn get_edge_data(&self, _edge: usize) -> Result>> { Ok(None) } - /// Xoá toàn bộ node stream (dùng khi rebuild index). Mặc định: no-op. - async fn clear_node_meta(&mut self) -> Result<()> { + /// Xoá toàn bộ edge stream (dùng khi rebuild index). + async fn clear_edges(&mut self) -> Result<()> { Ok(()) } +} + +// ── Chain storage ── - // ── Chain stream (per-owner chain — marker + symbol element ids) ── +/// Chain storage: lưu/đọc per-owner chain (marker + symbol element ids), +/// encode u64 LE 8-byte/element. Tách riêng để trait lõi gọn. +/// `CategoryStorage` super-bound trait này (luôn) → method gọi được qua +/// `dyn CategoryStorage`. Mặc định no-op. +#[async_trait] +pub trait ChainStorage: Send + Sync { /// Lưu chain của owner (keyed theo record của owner; u64 LE 8-byte/element). - /// Mặc định: no-op. - async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<()> { - let _ = (record, chain); + async fn set_chain(&mut self, _record: usize, _chain: &[u64]) -> Result<()> { Ok(()) } - /// Đọc chain của owner — `None` nếu owner chưa có chain. Mặc định: `None`. - #[allow(dead_code)] // dùng qua Search::get_chain (test/sqlite builds) - async fn get_chain(&self, record: usize) -> Result>> { - let _ = record; + /// Đọc chain của owner — `None` nếu owner chưa có chain. + async fn get_chain(&self, _record: usize) -> Result>> { Ok(None) } - /// Xoá toàn bộ chains (dùng khi rebuild index). Mặc định: no-op. + /// Xoá toàn bộ chains (dùng khi rebuild index). async fn clear_chains(&mut self) -> Result<()> { Ok(()) } +} - // ── Shard roots (endpoint) ── - async fn set_root(&mut self, shard: usize, root: usize) -> Result<()>; - async fn get_root(&self, shard: usize) -> Result; - - // ── Metadata & key length ── - /// Lưu metadata (opaque bytes, VD: call-site info) cho một record. - /// Nằm tách khỏi radix node — keyed theo record index. - #[allow(dead_code)] // primitive storage — dùng trong storage tests - async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()>; - /// Đọc metadata của record — `None` nếu record chưa có meta. - async fn get_meta(&self, record: usize) -> Result>>; - /// Lưu độ dài key (số element) của record — dùng filter `depth` khi search. - async fn set_key_len(&mut self, record: usize, len: usize) -> Result<()>; - /// Đọc độ dài key của record — `None` nếu record chưa insert. - async fn get_key_len(&self, record: usize) -> Result>; - - // ── Shortcuts (auxiliary LIKE-search index) ── - /// Thêm `node_id` vào shortcut set của element `elem` (encoded bytes). - /// Shortcut set = mọi node có chứa element này trong prefix của nó — dùng - /// làm candidate khi tìm substring (KMP + DFS). - async fn add_shortcut_node(&mut self, shard: usize, elem: &[u8], node_id: usize) -> Result<()>; - /// Lấy toàn bộ node id chứa element `elem` trong shard. - async fn get_shortcut_nodes(&self, shard: usize, elem: &[u8]) -> Result>; - /// Xoá toàn bộ shortcut sets (dùng khi rebuild index từ tree). - async fn clear_shortcuts(&mut self) -> Result<()>; - - // ── Entity store (semgraph model — symbols/chains/callnames/files/version) ── - // Tầng dữ liệu ngữ nghĩa đã dời xuống storage (db/ cũ bị xoá): mọi backend - // giữ entity data riêng (InMemory = HashMap, Sqlite = bảng `sg_*`, Redis = - // hash). Mặc định no-op để backend không cần implement nếu chưa dùng. +// ── CategoryStorage umbrella ── + +/// Khai báo `CategoryStorage` — macro emit TOÀN BỘ trait (gồm `#[async_trait]`) +/// nên async_trait biến đổi ĐÚNG sau khi macro nở (khắc lỗi macro body trong +/// trait). `$bounds` = danh sách supertrait: luôn `Send + Sync + NodeMetaStorage`, +/// cộng `BloomStorage` khi feature `bloom-search`. Thân method không có `#[cfg]` +/// rải rác. +macro_rules! declare_category_storage { + ($($bounds:tt)*) => { + /// Radix-node storage: node management + transaction + 5 stream phụ. + #[async_trait] + pub trait CategoryStorage: $($bounds)* { + // ── Node management ── + async fn new_node(&mut self, prefix: Vec, record: usize) -> Result; + async fn update_node( + &mut self, + id: usize, + prefix: Option>, + record: Option, + ) -> Result<()>; + async fn get_node(&self, id: usize) -> Result<(Vec, usize)>; + async fn get_children(&self, id: usize) -> Result>; + + // ── Shard roots (endpoint) ── + async fn set_root(&mut self, shard: usize, root: usize) -> Result<()>; + async fn get_root(&self, shard: usize) -> Result; + + // ── Transaction ── + /// Bắt đầu một transaction (sync, không await — đúng theo cách radix gọi). + /// Buffer ops; mọi thay đổi chỉ lộ ra khi `commit`. + fn new_tx(&self) -> Box; + } + }; +} - // Method được GraphIndex gọi trực tiếp (ingest/register/flow) — live ở mọi - // build. Method chỉ dùng qua `rebuild()` (mở lại file — feature `sqlite`) - // cfg_attr allow cho build không feature đó; `load_symbol`/`load_call_name_index` - // chưa có caller — giữ allow cho tới khi consumer cần. +#[cfg(feature = "bloom-search")] +declare_category_storage!( + Send + Sync + + NodeMetaStorage + + ShortcutsStorage + + EdgeDataStorage + + ChainStorage + + BloomStorage +); + +#[cfg(not(feature = "bloom-search"))] +declare_category_storage!( + Send + Sync + NodeMetaStorage + ShortcutsStorage + EdgeDataStorage + ChainStorage +); + +// ==================== Entity storage (chỉ codegraph-graph) ==================== + +/// Entity store — gồm symbol registry, call records, call-name index, files, +/// version, stats, embeddings. Tách khỏi radix-node storage vì: +/// - Chỉ `GraphIndex` / `SharedGraphIndex` dùng (`Radix` / `Search` không cần). +/// - Backend tối giản có thể bỏ qua (vd: chỉ cần `CategoryStorage` cho test). +/// - Cho phép phát triển/scale entity layer độc lập với radix. +#[async_trait] +pub trait EntityStorage: Send + Sync { + // ── Symbol registry ── /// Lưu một symbol — mặc định: no-op. async fn save_symbol(&mut self, _sym: &Symbol) -> Result<()> { Ok(()) } - #[allow(dead_code)] /// Đọc symbol theo id — mặc định: `None`. + #[allow(dead_code)] async fn load_symbol(&self, _id: u64) -> Result> { Ok(None) } @@ -302,17 +378,20 @@ pub trait Storage: Send + Sync { async fn save_next_id(&mut self, _next: u64) -> Result<()> { Ok(()) } - /// Đọc `next_id` — mặc định: 0 (chưa có symbol). + /// Đọc `next_id` — mặc định: 0. #[cfg_attr(not(feature = "sqlite"), allow(dead_code))] async fn load_next_id(&self) -> Result { Ok(0) } + /// Đọc toàn bộ chain `(func_id, chain_bytes u64 LE)` — rebuild engine khi /// open — mặc định: rỗng. #[cfg_attr(not(feature = "sqlite"), allow(dead_code))] async fn all_chains(&self) -> Result)>> { Ok(Vec::new()) } + + // ── Call records ── /// Lưu call records của một func (opaque bytes, JSON) — mặc định: no-op. async fn set_call_records(&mut self, _func: u64, _records: &[u8]) -> Result<()> { Ok(()) @@ -326,13 +405,15 @@ pub trait Storage: Send + Sync { async fn all_call_records(&self) -> Result)>> { Ok(Vec::new()) } + + // ── Call-name index ── /// Lưu inverted index `call name → call sites` (opaque bytes, JSON) — mặc /// định: no-op. async fn set_call_name_index(&mut self, _name: &str, _sites: &[u8]) -> Result<()> { Ok(()) } - #[allow(dead_code)] /// Đọc call-name index — mặc định: `None`. + #[allow(dead_code)] async fn load_call_name_index(&self, _name: &str) -> Result>> { Ok(None) } @@ -341,6 +422,8 @@ pub trait Storage: Send + Sync { async fn all_call_name_indexes(&self) -> Result)>> { Ok(Vec::new()) } + + // ── Files ── /// Upsert file info — mặc định: no-op. async fn upsert_file(&mut self, _f: &FileInfo) -> Result<()> { Ok(()) @@ -350,7 +433,9 @@ pub trait Storage: Send + Sync { async fn load_all_files(&self) -> Result> { Ok(Vec::new()) } - /// Version của index (`index_version` — bump mỗi lần ingest) — mặc định: 0. + + // ── Version ── + /// Version của index — mặc định: 0. #[cfg_attr(not(feature = "sqlite"), allow(dead_code))] async fn version(&self) -> Result { Ok(0) @@ -359,919 +444,63 @@ pub trait Storage: Send + Sync { async fn set_version(&mut self, _v: u64) -> Result<()> { Ok(()) } - /// Lưu counts tổng hợp (symbols/chains/edges/files) — `codegraph_status` - /// đọc trực tiếp từ đĩa, bỏ qua rebuild in-memory. Mặc định: no-op. + + // ── Stats (counts tổng hợp) ── + /// Lưu counts tổng hợp (symbols/chains/edges/files) — mặc định: no-op. async fn set_stats(&mut self, _s: IndexCounts) -> Result<()> { Ok(()) } - /// Đọc counts tổng hợp từ đĩa. Mặc định: `Ok(IndexCounts::default())` - /// (toàn 0). Backend không lưu → caller fallback sang rebuild. + /// Đọc counts tổng hợp từ đĩa — mặc định: `IndexCounts::default()`. async fn stats(&self) -> Result { Ok(IndexCounts::default()) } + /// Xoá toàn bộ entity data (symbols/next_id/call_records/call_names/files/ - /// version) — dùng khi full re-index. Mặc định: no-op. + /// version/embeddings) — dùng khi full re-index. Mặc định: no-op. async fn clear_entities(&mut self) -> Result<()> { Ok(()) } // ── Embeddings (vector per symbol id) ── - /// Lưu vector embedding cho một symbol (keyed theo symbol id). Vector đã - /// L2-normalize (cosine = dot product). Mặc định: no-op. + /// Lưu vector embedding cho một symbol. Vector đã L2-normalize. Mặc định: no-op. async fn save_embedding(&mut self, _symbol_id: u64, _vector: &[f32]) -> Result<()> { Ok(()) } - /// Đọc vector embedding của symbol — `None` nếu chưa có. Mặc định: `None`. + /// Đọc vector embedding của symbol — mặc định: `None`. async fn load_embedding(&self, _symbol_id: u64) -> Result>> { Ok(None) } - /// Đọc toàn bộ embeddings (symbol_id → vector) — rebuild VectorIndex khi - /// open. Mặc định: rỗng. + /// Đọc toàn bộ embeddings — mặc định: rỗng. async fn load_all_embeddings(&self) -> Result>> { Ok(HashMap::new()) } - /// Xoá toàn bộ embeddings — dùng khi full re-index. Mặc định: no-op. + /// Xoá toàn bộ embeddings — mặc định: no-op. async fn clear_embeddings(&mut self) -> Result<()> { Ok(()) } /// KNN backend-native (SQLite + sqlite-vss). Trả `Some(hits)` nếu backend - /// hỗ trợ ANN, `None` để caller fallback sang `VectorIndex` in-memory - /// (brute-force, đúng cho mọi backend). `hits` = `Vec<(symbol_id, sim)>` - /// với `sim` cao = gần hơn (đã đảo dấu distance để đồng nhất với - /// `VectorIndex::knn`). Mặc định: `None` (không backend-native). + /// hỗ trợ ANN, `None` để caller fallback sang `VectorIndex` in-memory. + /// Mặc định: `None`. async fn knn(&self, _query_vec: &[f32], _k: usize) -> Result>> { Ok(None) } - - // ── Transaction ── - /// Bắt đầu một transaction (sync, không await — đúng theo cách radix gọi). - /// Buffer ops; mọi thay đổi chỉ lộ ra khi `commit`. - fn new_tx(&self) -> Box; -} - -// ==================== In-Memory Storage ==================== - -struct MemoryData { - /// (prefix, record) — index 0 là sentinel. - nodes: Vec<(Vec, usize)>, - /// children list per node (index 0 = sentinel). - children: Vec>, - /// root id per shard. - roots: Vec, - /// record_idx → metadata (opaque bytes, VD: call-site info). - meta: HashMap>, - /// record_idx → độ dài key (số element) — dùng filter `depth` khi search. - key_lens: HashMap, - /// shortcuts[shard][elem_bytes] = node ids chứa elem trong prefix. - shortcuts: Vec, HashSet>>, - /// edge id → dữ liệu edge (opaque bytes, VD EdgeMeta JSON). - edges: HashMap>, - /// element id → node metadata (Node JSON). - node_meta: HashMap>, - /// node id → serialize bloom filter (prune nhánh trong search_dfs). - #[cfg(feature = "bloom-search")] - blooms: HashMap>, - /// record (owner) → chain bytes (u64 LE 8-byte/element). - chains: HashMap>, - // ── Entity store (semgraph model) ── - // Ghi/đọc bởi entity methods qua InMemoryStorage (GraphIndex ingest/rebuild). - /// symbol id → Symbol. - symbols: HashMap, - /// next_id của symbol registry. - next_id: u64, - /// func id → call records (JSON). - call_records: HashMap>, - /// call name → call sites (JSON). - call_names: HashMap>, - /// path → FileInfo. - files: HashMap, - /// index version. - version: u64, - /// symbol id → embedding vector (L2-normalized f32). - embeddings: HashMap>, -} - -/// In-memory radix storage. Thread-safe: toàn bộ state nằm sau 1 RwLock; -/// id được cấp bằng AtomicUsize nên các transaction song song không trùng id. -pub struct InMemoryStorage { - data: Arc>, - next_id: Arc, -} - -impl InMemoryStorage { - pub fn new() -> Self { - Self { - data: Arc::new(RwLock::new(MemoryData { - nodes: vec![(vec![], EMPTY)], // sentinel - children: vec![vec![]], - roots: vec![], - meta: HashMap::new(), - key_lens: HashMap::new(), - shortcuts: vec![], - edges: HashMap::new(), - node_meta: HashMap::new(), - #[cfg(feature = "bloom-search")] - blooms: HashMap::new(), - chains: HashMap::new(), - symbols: HashMap::new(), - // Id bắt đầu từ SYMBOL_BASE (marker reserved 1..=99). - next_id: codegraph_core::SYMBOL_BASE, - call_records: HashMap::new(), - call_names: HashMap::new(), - files: HashMap::new(), - version: 0, - embeddings: HashMap::new(), - })), - next_id: Arc::new(AtomicUsize::new(1)), - } - } -} - -impl Default for InMemoryStorage { - fn default() -> Self { - Self::new() - } -} - -impl InMemoryStorage { - /// Reserve một id mới (dùng chung cho cả new_node trực tiếp lẫn tx). - fn alloc_id(&self) -> usize { - self.next_id.fetch_add(1, Ordering::SeqCst) - } -} - -#[async_trait] -impl Storage for InMemoryStorage { - async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { - let id = self.alloc_id(); - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - if d.nodes.len() <= id { - d.nodes.resize(id + 1, (vec![], EMPTY)); - d.children.resize(id + 1, vec![]); - } - d.nodes[id] = (prefix, record); - Ok(id) - } - - async fn update_node( - &mut self, - id: usize, - prefix: Option>, - record: Option, - ) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - if id >= d.nodes.len() { - return Err(StorageError::BranchOutOfRange(id)); - } - if let Some(p) = prefix { - d.nodes[id].0 = p; - } - if let Some(r) = record { - d.nodes[id].1 = r; - } - Ok(()) - } - - async fn get_node(&self, id: usize) -> Result<(Vec, usize)> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - if id >= d.nodes.len() { - return Err(StorageError::BranchOutOfRange(id)); - } - Ok(d.nodes[id].clone()) - } - - async fn get_children(&self, id: usize) -> Result> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.children.get(id).cloned().unwrap_or_default()) - } - - #[cfg(feature = "bloom-search")] - async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.blooms.insert(id, bloom.to_vec()); - Ok(()) - } - - #[cfg(feature = "bloom-search")] - async fn get_node_bloom(&self, id: usize) -> Result>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.blooms.get(&id).cloned()) - } - - async fn set_root(&mut self, shard: usize, root: usize) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - if shard >= d.roots.len() { - d.roots.resize(shard + 1, EMPTY); - } - d.roots[shard] = root; - Ok(()) - } - - async fn get_root(&self, shard: usize) -> Result { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.roots.get(shard).copied().unwrap_or(EMPTY)) - } - - async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.meta.insert(record, meta.to_vec()); - Ok(()) - } - - async fn get_meta(&self, record: usize) -> Result>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.meta.get(&record).cloned()) - } - - async fn set_key_len(&mut self, record: usize, len: usize) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.key_lens.insert(record, len); - Ok(()) - } - - async fn get_key_len(&self, record: usize) -> Result> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.key_lens.get(&record).copied()) - } - - async fn add_shortcut_node(&mut self, shard: usize, elem: &[u8], node_id: usize) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - if shard >= d.shortcuts.len() { - d.shortcuts.resize(shard + 1, HashMap::new()); - } - d.shortcuts[shard] - .entry(elem.to_vec()) - .or_default() - .insert(node_id); - Ok(()) - } - - async fn get_shortcut_nodes(&self, shard: usize, elem: &[u8]) -> Result> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.shortcuts - .get(shard) - .and_then(|m| m.get(elem)) - .map(|set| set.iter().copied().collect()) - .unwrap_or_default()) - } - - async fn clear_shortcuts(&mut self) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - for map in d.shortcuts.iter_mut() { - map.clear(); - } - Ok(()) - } - - async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.edges.insert(edge, data.to_vec()); - Ok(()) - } - - async fn get_edge_data(&self, edge: usize) -> Result>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.edges.get(&edge).cloned()) - } - - async fn clear_edges(&mut self) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.edges.clear(); - Ok(()) - } - - async fn for_each_edge_data( - &self, - f: &mut (dyn for<'a> FnMut(usize, &'a [u8]) -> Result<()> + Send), - ) -> Result<()> { - let items: Vec<(usize, Vec)> = { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.edges - .iter() - .map(|(&id, data)| (id, data.clone())) - .collect() - }; - for (id, data) in items { - f(id, &data)?; - } - Ok(()) - } - - async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.node_meta.insert(elem, meta.to_vec()); - Ok(()) - } - - async fn get_node_meta(&self, elem: usize) -> Result>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.node_meta.get(&elem).cloned()) - } - - async fn clear_node_meta(&mut self) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.node_meta.clear(); - Ok(()) - } - - async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.chains.insert(record, encode_chain(chain)); - Ok(()) - } - - async fn get_chain(&self, record: usize) -> Result>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.chains.get(&record).map(|b| decode_chain(b))) - } - - async fn clear_chains(&mut self) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.chains.clear(); - Ok(()) - } - - async fn save_symbol(&mut self, sym: &Symbol) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.symbols.insert(sym.id, sym.clone()); - Ok(()) - } - - async fn load_symbol(&self, id: u64) -> Result> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.symbols.get(&id).cloned()) - } - - async fn load_all_symbols(&self) -> Result> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - let mut out: Vec = d.symbols.values().cloned().collect(); - out.sort_by_key(|s| s.id); - Ok(out) - } - - async fn save_next_id(&mut self, next: u64) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.next_id = next; - Ok(()) - } - - async fn load_next_id(&self) -> Result { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.next_id) - } - - async fn all_chains(&self) -> Result)>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - let mut out: Vec<(u64, Vec)> = d - .chains - .iter() - .map(|(&rec, bytes)| (rec as u64, bytes.clone())) - .collect(); - out.sort_by_key(|(rec, _)| *rec); - Ok(out) - } - - async fn set_call_records(&mut self, func: u64, records: &[u8]) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.call_records.insert(func, records.to_vec()); - Ok(()) - } - - async fn get_call_records(&self, func: u64) -> Result>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.call_records.get(&func).cloned()) - } - - async fn all_call_records(&self) -> Result)>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.call_records - .iter() - .map(|(&f, b)| (f, b.clone())) - .collect()) - } - - async fn set_call_name_index(&mut self, name: &str, sites: &[u8]) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.call_names.insert(name.to_string(), sites.to_vec()); - Ok(()) - } - - async fn load_call_name_index(&self, name: &str) -> Result>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.call_names.get(name).cloned()) - } - - async fn all_call_name_indexes(&self) -> Result)>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.call_names - .iter() - .map(|(n, b)| (n.clone(), b.clone())) - .collect()) - } - - async fn upsert_file(&mut self, f: &FileInfo) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.files.insert(f.path.clone(), f.clone()); - Ok(()) - } - - async fn load_all_files(&self) -> Result> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - let mut out: Vec = d.files.values().cloned().collect(); - out.sort_by(|a, b| a.path.cmp(&b.path)); - Ok(out) - } - - async fn version(&self) -> Result { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.version) - } - - async fn set_version(&mut self, v: u64) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.version = v; - Ok(()) - } - - async fn clear_entities(&mut self) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.symbols.clear(); - d.next_id = codegraph_core::SYMBOL_BASE; - d.call_records.clear(); - d.call_names.clear(); - d.files.clear(); - d.version = 0; - d.embeddings.clear(); - Ok(()) - } - - async fn save_embedding(&mut self, symbol_id: u64, vector: &[f32]) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.embeddings.insert(symbol_id, vector.to_vec()); - Ok(()) - } - - async fn load_embedding(&self, symbol_id: u64) -> Result>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.embeddings.get(&symbol_id).cloned()) - } - - async fn load_all_embeddings(&self) -> Result>> { - let d = self - .data - .read() - .map_err(|_| StorageError::Internal("poison".into()))?; - Ok(d.embeddings.clone()) - } - - async fn clear_embeddings(&mut self) -> Result<()> { - let mut d = self - .data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - d.embeddings.clear(); - Ok(()) - } - - fn new_tx(&self) -> Box { - Box::new(InMemoryTx { - data: self.data.clone(), - next_id: self.next_id.clone(), - nodes: Vec::new(), - ops: Vec::new(), - }) - } } -/// Transaction cho `InMemoryStorage`: buffer toàn bộ mutation, áp dụng -/// atomic dưới 1 write lock tại `commit`. -struct InMemoryTx { - data: Arc>, - next_id: Arc, - /// (reserved_id, prefix, record) — được append tại commit. - nodes: Vec<(usize, Vec, usize)>, - ops: Vec, -} +// ==================== Storage umbrella ==================== +/// Umbrella trait cho `Arc>` trong `GraphIndex`. +/// +/// Gộp `CategoryStorage` + 5 trait phụ + `EntityStorage`. Backend implement +/// 7 `impl` block riêng biệt — review từng phần độc lập được. #[async_trait] -impl Tx for InMemoryTx { - async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { - let id = self.next_id.fetch_add(1, Ordering::SeqCst); - self.nodes.push((id, prefix, record)); - Ok(id) - } - - async fn update_node( - &mut self, - id: usize, - prefix: Option>, - record: Option, - ) -> Result<()> { - self.ops.push(TxOp::UpdateNode { id, prefix, record }); - Ok(()) - } - - async fn add_child(&mut self, parent: usize, child: usize) -> Result<()> { - self.ops.push(TxOp::AddChild { parent, child }); - Ok(()) - } - - async fn move_child(&mut self, from: usize, to: usize, child: usize) -> Result<()> { - self.ops.push(TxOp::MoveChild { from, to, child }); - Ok(()) - } - - async fn commit(self: Box) -> Result<()> { - let InMemoryTx { - data, nodes, ops, .. - } = *self; - - let mut d = data - .write() - .map_err(|_| StorageError::Internal("poison".into()))?; - - // 1. Materialize các node đã reserve (đảm bảo children[leg] tồn tại - // trước khi ops move/add trỏ tới). - for (id, prefix, record) in nodes { - if d.nodes.len() <= id { - d.nodes.resize(id + 1, (vec![], EMPTY)); - d.children.resize(id + 1, vec![]); - } - d.nodes[id] = (prefix, record); - } - - // 2. Áp dụng toàn bộ ops — tất cả cùng thành công hoặc cùng thất bại - // (single write lock → không lộ trạng thái trung gian). - for op in ops { - match op { - TxOp::AddChild { parent, child } => { - if parent < d.children.len() && !d.children[parent].contains(&child) { - d.children[parent].push(child); - } - } - TxOp::MoveChild { from, to, child } => { - if from < d.children.len() { - d.children[from].retain(|&c| c != child); - } - if to < d.children.len() && !d.children[to].contains(&child) { - d.children[to].push(child); - } - } - TxOp::UpdateNode { id, prefix, record } => { - if id < d.nodes.len() { - if let Some(p) = prefix { - d.nodes[id].0 = p; - } - if let Some(r) = record { - d.nodes[id].1 = r; - } - } - } - } - } - - Ok(()) - } -} - -// ==================== Tests (InMemory) ==================== - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_new_node_and_get_node() { - let mut s = InMemoryStorage::default(); - let id = s.new_node(b"hello".to_vec(), 42).await.unwrap(); - assert_ne!(id, EMPTY); - let (prefix, record) = s.get_node(id).await.unwrap(); - assert_eq!(prefix, b"hello"); - assert_eq!(record, 42); - } - - #[tokio::test] - async fn test_update_node() { - let mut s = InMemoryStorage::default(); - let id = s.new_node(b"init".to_vec(), 1).await.unwrap(); - s.update_node(id, Some(b"updated".to_vec()), Some(99)) - .await - .unwrap(); - let (prefix, record) = s.get_node(id).await.unwrap(); - assert_eq!(prefix, b"updated"); - assert_eq!(record, 99); - } - - #[tokio::test] - async fn test_children_and_roots() { - let mut s = InMemoryStorage::default(); - let parent = s.new_node(b"p".to_vec(), 0).await.unwrap(); - let c1 = s.new_node(b"c1".to_vec(), 1).await.unwrap(); - let c2 = s.new_node(b"c2".to_vec(), 2).await.unwrap(); - // Mutate qua Tx — production chỉ đi qua Tx, không có Storage::add_child. - let mut tx = s.new_tx(); - tx.add_child(parent, c1).await.unwrap(); - tx.add_child(parent, c2).await.unwrap(); - tx.commit().await.unwrap(); - let children = s.get_children(parent).await.unwrap(); - assert_eq!(children.len(), 2); - assert!(children.contains(&c1)); - assert!(children.contains(&c2)); - - assert_eq!(s.get_root(3).await.unwrap(), EMPTY); - s.set_root(3, parent).await.unwrap(); - assert_eq!(s.get_root(3).await.unwrap(), parent); - } - - #[tokio::test] - async fn test_meta_roundtrip() { - let mut s = InMemoryStorage::default(); - // Chưa có gì → None. - assert_eq!(s.get_meta(7).await.unwrap(), None); - assert_eq!(s.get_key_len(7).await.unwrap(), None); - s.set_meta(7, b"call-site-info".as_slice()).await.unwrap(); - s.set_key_len(7, 5).await.unwrap(); - assert_eq!( - s.get_meta(7).await.unwrap().as_deref(), - Some(b"call-site-info".as_slice()) - ); - assert_eq!(s.get_key_len(7).await.unwrap(), Some(5)); - // Ghi đè meta. - s.set_meta(7, b"updated").await.unwrap(); - s.set_key_len(7, 6).await.unwrap(); - assert_eq!( - s.get_meta(7).await.unwrap().as_deref(), - Some(b"updated".as_slice()) - ); - assert_eq!(s.get_key_len(7).await.unwrap(), Some(6)); - // Record khác không ảnh hưởng. - assert_eq!(s.get_meta(8).await.unwrap(), None); - assert_eq!(s.get_key_len(8).await.unwrap(), None); - } - - #[tokio::test] - async fn test_shortcuts_roundtrip() { - let mut s = InMemoryStorage::default(); - // Chưa có gì → empty. - assert!(s.get_shortcut_nodes(1, b"l").await.unwrap().is_empty()); - s.add_shortcut_node(1, b"l", 10).await.unwrap(); - s.add_shortcut_node(1, b"l", 20).await.unwrap(); - s.add_shortcut_node(1, b"o", 10).await.unwrap(); - s.add_shortcut_node(2, b"l", 30).await.unwrap(); // shard khác - let nodes = s.get_shortcut_nodes(1, b"l").await.unwrap(); - assert!(nodes.contains(&10) && nodes.contains(&20)); - assert_eq!(nodes.len(), 2); - assert_eq!(s.get_shortcut_nodes(2, b"l").await.unwrap(), vec![30]); - - // Clear → rỗng hết. - s.clear_shortcuts().await.unwrap(); - assert!(s.get_shortcut_nodes(1, b"l").await.unwrap().is_empty()); - assert!(s.get_shortcut_nodes(2, b"l").await.unwrap().is_empty()); - } - - #[tokio::test] - async fn test_tx_commit_applies_atomically() { - let mut s = InMemoryStorage::default(); - let parent = s.new_node(b"hello".to_vec(), 1).await.unwrap(); - - let mut tx = s.new_tx(); - let new_id = tx.new_node(b"p".to_vec(), 2).await.unwrap(); - let leg_id = tx.new_node(b"lo".to_vec(), 1).await.unwrap(); - tx.move_child(parent, leg_id, 0).await.unwrap(); // no-op: 0 chưa phải child - tx.add_child(parent, leg_id).await.unwrap(); - tx.add_child(parent, new_id).await.unwrap(); - tx.update_node(parent, Some(b"hel".to_vec()), Some(0)) - .await - .unwrap(); - tx.commit().await.unwrap(); - - let (prefix, record) = s.get_node(parent).await.unwrap(); - assert_eq!(prefix, b"hel"); - assert_eq!(record, 0); - let children = s.get_children(parent).await.unwrap(); - assert!(children.contains(&leg_id)); - assert!(children.contains(&new_id)); - assert_eq!(s.get_node(new_id).await.unwrap().1, 2); - assert_eq!(s.get_node(leg_id).await.unwrap().1, 1); - } - - #[tokio::test] - async fn test_tx_nodes_invisible_before_commit() { - let s = InMemoryStorage::default(); - let mut tx = s.new_tx(); - let id = tx.new_node(b"pending".to_vec(), 9).await.unwrap(); - // Trước commit, node chưa materialize → get_node lỗi BranchOutOfRange. - assert!(s.get_node(id).await.is_err()); - tx.commit().await.unwrap(); - assert_eq!(s.get_node(id).await.unwrap().1, 9); - } - - #[tokio::test] - async fn test_tx_move_child_migrates() { - let mut s = InMemoryStorage::default(); - let parent = s.new_node(b"aaaaaa".to_vec(), 0).await.unwrap(); - let child = s.new_node(b"0".to_vec(), 1).await.unwrap(); - let mut seed = s.new_tx(); - seed.add_child(parent, child).await.unwrap(); - seed.commit().await.unwrap(); - - let mut tx = s.new_tx(); - let leg = tx.new_node(b"a".to_vec(), 0).await.unwrap(); - tx.move_child(parent, leg, child).await.unwrap(); - tx.add_child(parent, leg).await.unwrap(); - tx.commit().await.unwrap(); - - assert!(!s.get_children(parent).await.unwrap().contains(&child)); - assert!(s.get_children(leg).await.unwrap().contains(&child)); - } - - #[tokio::test] - async fn test_edge_data_roundtrip() { - let mut s = InMemoryStorage::default(); - // Chưa có edge → None. - assert_eq!(s.get_edge_data(7).await.unwrap(), None); - s.set_edge_data(7, b"call-site").await.unwrap(); - assert_eq!( - s.get_edge_data(7).await.unwrap().as_deref(), - Some(b"call-site".as_slice()) - ); - // Ghi đè dữ liệu edge. - s.set_edge_data(7, b"updated").await.unwrap(); - assert_eq!( - s.get_edge_data(7).await.unwrap().as_deref(), - Some(b"updated".as_slice()) - ); - // Edge khác không ảnh hưởng. - assert_eq!(s.get_edge_data(8).await.unwrap(), None); - - // Clear → sạch toàn bộ. - s.set_edge_data(9, b"x").await.unwrap(); - s.clear_edges().await.unwrap(); - assert_eq!(s.get_edge_data(7).await.unwrap(), None); - assert_eq!(s.get_edge_data(9).await.unwrap(), None); - } - - #[tokio::test] - async fn test_node_meta_roundtrip() { - let mut s = InMemoryStorage::default(); - assert_eq!(s.get_node_meta(3).await.unwrap(), None); - s.set_node_meta(3, b"node-json").await.unwrap(); - assert_eq!( - s.get_node_meta(3).await.unwrap().as_deref(), - Some(b"node-json".as_slice()) - ); - s.set_node_meta(3, b"node-json-2").await.unwrap(); - assert_eq!( - s.get_node_meta(3).await.unwrap().as_deref(), - Some(b"node-json-2".as_slice()) - ); - assert_eq!(s.get_node_meta(4).await.unwrap(), None); - s.clear_node_meta().await.unwrap(); - assert_eq!(s.get_node_meta(3).await.unwrap(), None); - } - - #[tokio::test] - async fn test_chains_roundtrip() { - let mut s = InMemoryStorage::default(); - assert_eq!(s.get_chain(9).await.unwrap(), None); - s.set_chain(9, &[1, 2, 3]).await.unwrap(); - assert_eq!(s.get_chain(9).await.unwrap(), Some(vec![1, 2, 3])); - s.set_chain(9, &[4]).await.unwrap(); - assert_eq!(s.get_chain(9).await.unwrap(), Some(vec![4])); - assert_eq!(s.get_chain(10).await.unwrap(), None); - s.clear_chains().await.unwrap(); - assert_eq!(s.get_chain(9).await.unwrap(), None); - } +pub trait Storage: + CategoryStorage + + NodeMetaStorage + + ShortcutsStorage + + EdgeDataStorage + + ChainStorage + + EntityStorage + + Send + + Sync +{ } diff --git a/crates/codegraph-graph/src/storage/cached.rs b/crates/codegraph-graph/src/storage/cached.rs index 654eaf2fa..fffbc6967 100644 --- a/crates/codegraph-graph/src/storage/cached.rs +++ b/crates/codegraph-graph/src/storage/cached.rs @@ -11,6 +11,10 @@ //! //! Decorator này trong suốt: mọi backend (InMemory/Sqlite/Lmdb/Redis/RDBMS) //! đều dùng được, behaviour đúng bằng inner (chỉ thêm lớp cache). +//! +//! Implementation chia 7 `impl` block (1 cho `CategoryStorage`, 5 cho trait phụ, +//! 1 cho `EntityStorage`) — review từng phần độc lập được. `Storage` umbrella +//! là marker rỗng (Rust tự cộng qua blanket bound). use std::collections::HashMap; use std::sync::Arc; @@ -19,7 +23,13 @@ use async_trait::async_trait; use codegraph_core::{FileInfo, Symbol}; use crate::lru::LruCache; -use crate::storage::{IndexCounts, Storage, StorageError, Tx}; +use crate::storage::{ + CategoryStorage, ChainStorage, EdgeDataStorage, EntityStorage, IndexCounts, NodeMetaStorage, + ShortcutsStorage, Storage, StorageError, Tx, +}; + +#[cfg(feature = "bloom-search")] +use crate::storage::BloomStorage; /// Số shard của mỗi `LruCache` — phải lũy thừa của 2. const SHARDS: usize = 32; @@ -104,8 +114,10 @@ impl CachedStorage { } } +// ==================== CategoryStorage ==================== + #[async_trait] -impl Storage for CachedStorage { +impl CategoryStorage for CachedStorage { // ── Node management (cached) ── async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { let id = self.inner.new_node(prefix, record).await?; @@ -142,49 +154,51 @@ impl Storage for CachedStorage { Ok(v) } - // ── Bloom (không cache — dùng prune nhánh, sai = search sai) ── - #[cfg(feature = "bloom-search")] - async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<(), StorageError> { - self.inner.set_node_bloom(id, bloom).await - } - - #[cfg(feature = "bloom-search")] - async fn get_node_bloom(&self, id: usize) -> Result>, StorageError> { - self.inner.get_node_bloom(id).await - } - - // ── Edge data (cached) ── - async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<(), StorageError> { - self.inner.set_edge_data(edge, data).await?; - self.caches.edge_data.remove(&edge); + // ── Shard roots (cached) ── + async fn set_root(&mut self, shard: usize, root: usize) -> Result<(), StorageError> { + self.inner.set_root(shard, root).await?; + self.caches.roots.remove(&shard); Ok(()) } - async fn get_edge_data(&self, edge: usize) -> Result>, StorageError> { - if let Some(v) = self.caches.edge_data.get(&edge) { - return Ok(Some(v)); - } - let v = self.inner.get_edge_data(edge).await?; - if let Some(ref b) = v { - self.caches.edge_data.put(edge, b.clone()); + async fn get_root(&self, shard: usize) -> Result { + if let Some(v) = self.caches.roots.get(&shard) { + return Ok(v); } + let v = self.inner.get_root(shard).await?; + self.caches.roots.put(shard, v); Ok(v) } - async fn clear_edges(&mut self) -> Result<(), StorageError> { - self.inner.clear_edges().await?; - self.caches.edge_data.clear(); - Ok(()) + // ── Transaction: wrap để invalidate radix cache khi commit ── + fn new_tx(&self) -> Box { + Box::new(CachedTx { + inner: self.inner.new_tx(), + caches: self.caches.clone(), + }) } +} - async fn for_each_edge_data( - &self, - f: &mut (dyn for<'a> FnMut(usize, &'a [u8]) -> Result<(), StorageError> + Send), - ) -> Result<(), StorageError> { - self.inner.for_each_edge_data(f).await +// ==================== BloomStorage (feature-gated) ==================== +// +// Không cache — bloom sai = search sai (over-prune). Pass-through. + +#[cfg(feature = "bloom-search")] +#[async_trait] +impl BloomStorage for CachedStorage { + async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<(), StorageError> { + self.inner.set_node_bloom(id, bloom).await + } + + async fn get_node_bloom(&self, id: usize) -> Result>, StorageError> { + self.inner.get_node_bloom(id).await } +} + +// ==================== NodeMetaStorage ==================== - // ── Node metadata (cached) ── +#[async_trait] +impl NodeMetaStorage for CachedStorage { async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<(), StorageError> { self.inner.set_node_meta(elem, meta).await?; self.caches.node_meta.remove(&elem); @@ -208,47 +222,6 @@ impl Storage for CachedStorage { Ok(()) } - // ── Chain (cached) ── - async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<(), StorageError> { - self.inner.set_chain(record, chain).await?; - self.caches.chains.remove(&record); - Ok(()) - } - - async fn get_chain(&self, record: usize) -> Result>, StorageError> { - if let Some(v) = self.caches.chains.get(&record) { - return Ok(Some(v)); - } - let v = self.inner.get_chain(record).await?; - if let Some(ref c) = v { - self.caches.chains.put(record, c.clone()); - } - Ok(v) - } - - async fn clear_chains(&mut self) -> Result<(), StorageError> { - self.inner.clear_chains().await?; - self.caches.chains.clear(); - Ok(()) - } - - // ── Shard roots (cached) ── - async fn set_root(&mut self, shard: usize, root: usize) -> Result<(), StorageError> { - self.inner.set_root(shard, root).await?; - self.caches.roots.remove(&shard); - Ok(()) - } - - async fn get_root(&self, shard: usize) -> Result { - if let Some(v) = self.caches.roots.get(&shard) { - return Ok(v); - } - let v = self.inner.get_root(shard).await?; - self.caches.roots.put(shard, v); - Ok(v) - } - - // ── Meta / key_len (cached) ── async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<(), StorageError> { self.inner.set_meta(record, meta).await?; self.caches.metas.remove(&record); @@ -282,8 +255,12 @@ impl Storage for CachedStorage { } Ok(v) } +} + +// ==================== ShortcutsStorage ==================== - // ── Shortcuts (cached) ── +#[async_trait] +impl ShortcutsStorage for CachedStorage { async fn add_shortcut_node( &mut self, shard: usize, @@ -314,8 +291,73 @@ impl Storage for CachedStorage { self.caches.shortcuts.clear(); Ok(()) } +} - // ── Entity store (symbols / calls / embeddings) ── +// ==================== EdgeDataStorage ==================== + +#[async_trait] +impl EdgeDataStorage for CachedStorage { + async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<(), StorageError> { + self.inner.set_edge_data(edge, data).await?; + self.caches.edge_data.remove(&edge); + Ok(()) + } + + async fn get_edge_data(&self, edge: usize) -> Result>, StorageError> { + if let Some(v) = self.caches.edge_data.get(&edge) { + return Ok(Some(v)); + } + let v = self.inner.get_edge_data(edge).await?; + if let Some(ref b) = v { + self.caches.edge_data.put(edge, b.clone()); + } + Ok(v) + } + + async fn clear_edges(&mut self) -> Result<(), StorageError> { + self.inner.clear_edges().await?; + self.caches.edge_data.clear(); + Ok(()) + } +} + +// ==================== ChainStorage ==================== + +#[async_trait] +impl ChainStorage for CachedStorage { + async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<(), StorageError> { + self.inner.set_chain(record, chain).await?; + self.caches.chains.remove(&record); + Ok(()) + } + + async fn get_chain(&self, record: usize) -> Result>, StorageError> { + if let Some(v) = self.caches.chains.get(&record) { + return Ok(Some(v)); + } + let v = self.inner.get_chain(record).await?; + if let Some(ref c) = v { + self.caches.chains.put(record, c.clone()); + } + Ok(v) + } + + async fn clear_chains(&mut self) -> Result<(), StorageError> { + self.inner.clear_chains().await?; + self.caches.chains.clear(); + Ok(()) + } +} + +// ==================== EntityStorage ==================== +// +// Phần lớn pass-through (không cache — ít được gọi lại nhiều lần). Một số +// method nóng (`load_symbol`/`get_call_records`/`load_call_name_index`/ +// `load_embedding`) có cache. `clear_entities` clear_all. + +#[async_trait] +impl EntityStorage for CachedStorage { + // ── Symbol registry ── async fn save_symbol(&mut self, sym: &Symbol) -> Result<(), StorageError> { self.inner.save_symbol(sym).await?; self.caches.symbols.remove(&sym.id); @@ -349,6 +391,7 @@ impl Storage for CachedStorage { self.inner.all_chains().await } + // ── Call records ── async fn set_call_records(&mut self, func: u64, records: &[u8]) -> Result<(), StorageError> { self.inner.set_call_records(func, records).await?; self.caches.call_records.remove(&func); @@ -370,6 +413,7 @@ impl Storage for CachedStorage { self.inner.all_call_records().await } + // ── Call-name index ── async fn set_call_name_index(&mut self, name: &str, sites: &[u8]) -> Result<(), StorageError> { self.inner.set_call_name_index(name, sites).await?; self.caches.call_name_index.remove(&name.to_string()); @@ -391,6 +435,7 @@ impl Storage for CachedStorage { self.inner.all_call_name_indexes().await } + // ── Files ── async fn upsert_file(&mut self, f: &FileInfo) -> Result<(), StorageError> { self.inner.upsert_file(f).await } @@ -399,6 +444,7 @@ impl Storage for CachedStorage { self.inner.load_all_files().await } + // ── Version ── async fn version(&self) -> Result { self.inner.version().await } @@ -407,6 +453,7 @@ impl Storage for CachedStorage { self.inner.set_version(v).await } + // ── Stats ── async fn set_stats(&mut self, s: IndexCounts) -> Result<(), StorageError> { self.inner.set_stats(s).await } @@ -456,16 +503,15 @@ impl Storage for CachedStorage { ) -> Result>, StorageError> { self.inner.knn(query_vec, k).await } - - // ── Transaction: wrap để invalidate radix cache khi commit ── - fn new_tx(&self) -> Box { - Box::new(CachedTx { - inner: self.inner.new_tx(), - caches: self.caches.clone(), - }) - } } +// ==================== Storage umbrella ==================== +// +// Rust tự cộng method qua blanket bound — không cần viết gì thêm. + +#[async_trait] +impl Storage for CachedStorage {} + /// Tx bọc: delegate mọi mutation, khi `commit` xong thì `clear_radix()`. struct CachedTx { inner: Box, diff --git a/crates/codegraph-graph/src/storage/in_memory.rs b/crates/codegraph-graph/src/storage/in_memory.rs new file mode 100644 index 000000000..624a5fcfd --- /dev/null +++ b/crates/codegraph-graph/src/storage/in_memory.rs @@ -0,0 +1,951 @@ +use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, RwLock}; + +use async_trait::async_trait; +use codegraph_core::{FileInfo, Symbol}; + +use super::{ + CategoryStorage, ChainStorage, EMPTY, EdgeDataStorage, EntityStorage, IndexCounts, + NodeMetaStorage, Result, ShortcutsStorage, StorageError, Tx, TxOp, decode_chain, encode_chain, +}; + +#[cfg(feature = "bloom-search")] +use super::BloomStorage; + +// ==================== Transaction ==================== + +/// Transaction cho `InMemoryStorage`: buffer toàn bộ mutation, áp dụng +/// atomic dưới 1 write lock tại `commit`. +pub(crate) struct InMemoryTx { + data: Arc>, + next_id: Arc, + /// (reserved_id, prefix, record) — được append tại commit. + nodes: Vec<(usize, Vec, usize)>, + ops: Vec, +} + +impl InMemoryTx { + pub(crate) fn new(data: Arc>, next_id: Arc) -> Self { + Self { + data, + next_id, + nodes: Vec::new(), + ops: Vec::new(), + } + } +} + +#[async_trait] +impl Tx for InMemoryTx { + async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { + let id = self.next_id.fetch_add(1, Ordering::SeqCst); + self.nodes.push((id, prefix, record)); + Ok(id) + } + + async fn update_node( + &mut self, + id: usize, + prefix: Option>, + record: Option, + ) -> Result<()> { + self.ops.push(TxOp::UpdateNode { id, prefix, record }); + Ok(()) + } + + async fn add_child(&mut self, parent: usize, child: usize) -> Result<()> { + self.ops.push(TxOp::AddChild { parent, child }); + Ok(()) + } + + async fn move_child(&mut self, from: usize, to: usize, child: usize) -> Result<()> { + self.ops.push(TxOp::MoveChild { from, to, child }); + Ok(()) + } + + async fn commit(self: Box) -> Result<()> { + let InMemoryTx { + data, nodes, ops, .. + } = *self; + + let mut d = data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + + // 1. Materialize các node đã reserve (đảm bảo children[leg] tồn tại + // trước khi ops move/add trỏ tới). + for (id, prefix, record) in nodes { + if d.nodes.len() <= id { + d.nodes.resize(id + 1, (vec![], EMPTY)); + d.children.resize(id + 1, vec![]); + } + d.nodes[id] = (prefix, record); + } + + // 2. Áp dụng toàn bộ ops — tất cả cùng thành công hoặc cùng thất bại + // (single write lock → không lộ trạng thái trung gian). + for op in ops { + match op { + TxOp::AddChild { parent, child } => { + if parent < d.children.len() && !d.children[parent].contains(&child) { + d.children[parent].push(child); + } + } + TxOp::MoveChild { from, to, child } => { + if from < d.children.len() { + d.children[from].retain(|&c| c != child); + } + if to < d.children.len() && !d.children[to].contains(&child) { + d.children[to].push(child); + } + } + TxOp::UpdateNode { id, prefix, record } => { + if id < d.nodes.len() { + if let Some(p) = prefix { + d.nodes[id].0 = p; + } + if let Some(r) = record { + d.nodes[id].1 = r; + } + } + } + } + } + + Ok(()) + } +} + +// ==================== In-Memory Storage ==================== + +pub(crate) struct MemoryData { + /// (prefix, record) — index 0 là sentinel. + pub(crate) nodes: Vec<(Vec, usize)>, + /// children list per node (index 0 = sentinel). + pub(crate) children: Vec>, + /// root id per shard. + pub(crate) roots: Vec, + /// record_idx → metadata (opaque bytes, VD: call-site info). + pub(crate) meta: HashMap>, + /// record_idx → độ dài key (số element) — dùng filter `depth` khi search. + pub(crate) key_lens: HashMap, + /// shortcuts[shard][elem_bytes] = node ids chứa elem trong prefix. + pub(crate) shortcuts: Vec, HashSet>>, + /// edge id → dữ liệu edge (opaque bytes, VD EdgeMeta JSON). + pub(crate) edges: HashMap>, + /// element id → node metadata (Node JSON). + pub(crate) node_meta: HashMap>, + /// node id → serialize bloom filter (prune nhánh trong search_dfs). + #[cfg(feature = "bloom-search")] + pub(crate) blooms: HashMap>, + /// record (owner) → chain bytes (u64 LE 8-byte/element). + pub(crate) chains: HashMap>, + // ── Entity store (semgraph model) ── + /// symbol id → Symbol. + pub(crate) symbols: HashMap, + /// next_id của symbol registry. + pub(crate) next_id: u64, + /// func id → call records (JSON). + pub(crate) call_records: HashMap>, + /// call name → call sites (JSON). + pub(crate) call_names: HashMap>, + /// path → FileInfo. + pub(crate) files: HashMap, + /// index version. + pub(crate) version: u64, + /// symbol id → embedding vector (L2-normalized f32). + pub(crate) embeddings: HashMap>, +} + +/// In-memory radix storage. Thread-safe: toàn bộ state nằm sau 1 RwLock; +/// id được cấp bằng AtomicUsize nên các transaction song song không trùng id. +pub struct InMemoryStorage { + data: Arc>, + next_id: Arc, +} + +impl InMemoryStorage { + pub fn new() -> Self { + Self { + data: Arc::new(RwLock::new(MemoryData { + nodes: vec![(vec![], EMPTY)], // sentinel + children: vec![vec![]], + roots: vec![], + meta: HashMap::new(), + key_lens: HashMap::new(), + shortcuts: vec![], + edges: HashMap::new(), + node_meta: HashMap::new(), + #[cfg(feature = "bloom-search")] + blooms: HashMap::new(), + chains: HashMap::new(), + symbols: HashMap::new(), + // Id bắt đầu từ SYMBOL_BASE (marker reserved 1..=99). + next_id: codegraph_core::SYMBOL_BASE, + call_records: HashMap::new(), + call_names: HashMap::new(), + files: HashMap::new(), + version: 0, + embeddings: HashMap::new(), + })), + next_id: Arc::new(AtomicUsize::new(1)), + } + } + + /// Reserve một id mới (dùng chung cho cả new_node trực tiếp lẫn tx). + fn alloc_id(&self) -> usize { + self.next_id.fetch_add(1, Ordering::SeqCst) + } +} + +impl Default for InMemoryStorage { + fn default() -> Self { + Self::new() + } +} + +// ==================== CategoryStorage ==================== + +#[async_trait] +impl CategoryStorage for InMemoryStorage { + async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { + let id = self.alloc_id(); + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + if d.nodes.len() <= id { + d.nodes.resize(id + 1, (vec![], EMPTY)); + d.children.resize(id + 1, vec![]); + } + d.nodes[id] = (prefix, record); + Ok(id) + } + + async fn update_node( + &mut self, + id: usize, + prefix: Option>, + record: Option, + ) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + if id >= d.nodes.len() { + return Err(StorageError::BranchOutOfRange(id)); + } + if let Some(p) = prefix { + d.nodes[id].0 = p; + } + if let Some(r) = record { + d.nodes[id].1 = r; + } + Ok(()) + } + + async fn get_node(&self, id: usize) -> Result<(Vec, usize)> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + if id >= d.nodes.len() { + return Err(StorageError::BranchOutOfRange(id)); + } + Ok(d.nodes[id].clone()) + } + + async fn get_children(&self, id: usize) -> Result> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.children.get(id).cloned().unwrap_or_default()) + } + + async fn set_root(&mut self, shard: usize, root: usize) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + if shard >= d.roots.len() { + d.roots.resize(shard + 1, EMPTY); + } + d.roots[shard] = root; + Ok(()) + } + + async fn get_root(&self, shard: usize) -> Result { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.roots.get(shard).copied().unwrap_or(EMPTY)) + } + + fn new_tx(&self) -> Box { + Box::new(InMemoryTx::new(self.data.clone(), self.next_id.clone())) + } +} + +// ==================== NodeMetaStorage ==================== + +#[async_trait] +impl NodeMetaStorage for InMemoryStorage { + async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.node_meta.insert(elem, meta.to_vec()); + Ok(()) + } + + async fn get_node_meta(&self, elem: usize) -> Result>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.node_meta.get(&elem).cloned()) + } + + async fn clear_node_meta(&mut self) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.node_meta.clear(); + Ok(()) + } + + async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.meta.insert(record, meta.to_vec()); + Ok(()) + } + + async fn get_meta(&self, record: usize) -> Result>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.meta.get(&record).cloned()) + } + + async fn set_key_len(&mut self, record: usize, len: usize) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.key_lens.insert(record, len); + Ok(()) + } + + async fn get_key_len(&self, record: usize) -> Result> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.key_lens.get(&record).copied()) + } +} + +// ==================== ShortcutsStorage ==================== + +#[async_trait] +impl ShortcutsStorage for InMemoryStorage { + async fn add_shortcut_node(&mut self, shard: usize, elem: &[u8], node_id: usize) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + if shard >= d.shortcuts.len() { + d.shortcuts.resize(shard + 1, HashMap::new()); + } + d.shortcuts[shard] + .entry(elem.to_vec()) + .or_default() + .insert(node_id); + Ok(()) + } + + async fn get_shortcut_nodes(&self, shard: usize, elem: &[u8]) -> Result> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.shortcuts + .get(shard) + .and_then(|m| m.get(elem)) + .map(|set| set.iter().copied().collect()) + .unwrap_or_default()) + } + + async fn clear_shortcuts(&mut self) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + for map in d.shortcuts.iter_mut() { + map.clear(); + } + Ok(()) + } +} + +// ==================== EdgeDataStorage ==================== + +#[async_trait] +impl EdgeDataStorage for InMemoryStorage { + async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.edges.insert(edge, data.to_vec()); + Ok(()) + } + + async fn get_edge_data(&self, edge: usize) -> Result>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.edges.get(&edge).cloned()) + } + + async fn clear_edges(&mut self) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.edges.clear(); + Ok(()) + } +} + +// ==================== ChainStorage ==================== + +#[async_trait] +impl ChainStorage for InMemoryStorage { + async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.chains.insert(record, encode_chain(chain)); + Ok(()) + } + + async fn get_chain(&self, record: usize) -> Result>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.chains.get(&record).map(|b| decode_chain(b))) + } + + async fn clear_chains(&mut self) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.chains.clear(); + Ok(()) + } +} + +// ==================== BloomStorage (feature-gated) ==================== + +#[cfg(feature = "bloom-search")] +#[async_trait] +impl BloomStorage for InMemoryStorage { + async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.blooms.insert(id, bloom.to_vec()); + Ok(()) + } + + async fn get_node_bloom(&self, id: usize) -> Result>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.blooms.get(&id).cloned()) + } +} + +// ==================== EntityStorage ==================== + +#[async_trait] +impl EntityStorage for InMemoryStorage { + async fn save_symbol(&mut self, sym: &Symbol) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.symbols.insert(sym.id, sym.clone()); + Ok(()) + } + + async fn load_symbol(&self, id: u64) -> Result> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.symbols.get(&id).cloned()) + } + + async fn load_all_symbols(&self) -> Result> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + let mut out: Vec = d.symbols.values().cloned().collect(); + out.sort_by_key(|s| s.id); + Ok(out) + } + + async fn save_next_id(&mut self, next: u64) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.next_id = next; + Ok(()) + } + + async fn load_next_id(&self) -> Result { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.next_id) + } + + async fn all_chains(&self) -> Result)>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + let mut out: Vec<(u64, Vec)> = d + .chains + .iter() + .map(|(&rec, bytes)| (rec as u64, bytes.clone())) + .collect(); + out.sort_by_key(|(rec, _)| *rec); + Ok(out) + } + + async fn set_call_records(&mut self, func: u64, records: &[u8]) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.call_records.insert(func, records.to_vec()); + Ok(()) + } + + async fn get_call_records(&self, func: u64) -> Result>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.call_records.get(&func).cloned()) + } + + async fn all_call_records(&self) -> Result)>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.call_records + .iter() + .map(|(&f, b)| (f, b.clone())) + .collect()) + } + + async fn set_call_name_index(&mut self, name: &str, sites: &[u8]) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.call_names.insert(name.to_string(), sites.to_vec()); + Ok(()) + } + + async fn load_call_name_index(&self, name: &str) -> Result>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.call_names.get(name).cloned()) + } + + async fn all_call_name_indexes(&self) -> Result)>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.call_names + .iter() + .map(|(n, b)| (n.clone(), b.clone())) + .collect()) + } + + async fn upsert_file(&mut self, f: &FileInfo) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.files.insert(f.path.clone(), f.clone()); + Ok(()) + } + + async fn load_all_files(&self) -> Result> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + let mut out: Vec = d.files.values().cloned().collect(); + out.sort_by(|a, b| a.path.cmp(&b.path)); + Ok(out) + } + + async fn version(&self) -> Result { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.version) + } + + async fn set_version(&mut self, v: u64) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.version = v; + Ok(()) + } + + async fn set_stats(&mut self, _s: IndexCounts) -> Result<()> { + // In-memory không persist stats (rebuild O(1) thông qua len() các map). + Ok(()) + } + + async fn stats(&self) -> Result { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(IndexCounts { + symbols: d.symbols.len() as u64, + chains: d.chains.len() as u64, + edges: d.edges.len() as u64, + files: d.files.len() as u64, + next_id: d.next_id, + }) + } + + async fn clear_entities(&mut self) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.symbols.clear(); + d.next_id = codegraph_core::SYMBOL_BASE; + d.call_records.clear(); + d.call_names.clear(); + d.files.clear(); + d.version = 0; + d.embeddings.clear(); + Ok(()) + } + + async fn save_embedding(&mut self, symbol_id: u64, vector: &[f32]) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.embeddings.insert(symbol_id, vector.to_vec()); + Ok(()) + } + + async fn load_embedding(&self, symbol_id: u64) -> Result>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.embeddings.get(&symbol_id).cloned()) + } + + async fn load_all_embeddings(&self) -> Result>> { + let d = self + .data + .read() + .map_err(|_| StorageError::Internal("poison".into()))?; + Ok(d.embeddings.clone()) + } + + async fn clear_embeddings(&mut self) -> Result<()> { + let mut d = self + .data + .write() + .map_err(|_| StorageError::Internal("poison".into()))?; + d.embeddings.clear(); + Ok(()) + } + + // knn mặc định: trả None → caller fallback `VectorIndex` in-memory. +} + +// ==================== Storage umbrella (empty marker) ==================== + +use super::Storage; + +#[async_trait] +impl Storage for InMemoryStorage {} + +// ==================== Tests ==================== + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_new_node_and_get_node() { + let mut s = InMemoryStorage::default(); + let id = s.new_node(b"hello".to_vec(), 42).await.unwrap(); + assert_ne!(id, EMPTY); + let (prefix, record) = s.get_node(id).await.unwrap(); + assert_eq!(prefix, b"hello"); + assert_eq!(record, 42); + } + + #[tokio::test] + async fn test_update_node() { + let mut s = InMemoryStorage::default(); + let id = s.new_node(b"init".to_vec(), 1).await.unwrap(); + s.update_node(id, Some(b"updated".to_vec()), Some(99)) + .await + .unwrap(); + let (prefix, record) = s.get_node(id).await.unwrap(); + assert_eq!(prefix, b"updated"); + assert_eq!(record, 99); + } + + #[tokio::test] + async fn test_children_and_roots() { + let mut s = InMemoryStorage::default(); + let parent = s.new_node(b"p".to_vec(), 0).await.unwrap(); + let c1 = s.new_node(b"c1".to_vec(), 1).await.unwrap(); + let c2 = s.new_node(b"c2".to_vec(), 2).await.unwrap(); + let mut tx = s.new_tx(); + tx.add_child(parent, c1).await.unwrap(); + tx.add_child(parent, c2).await.unwrap(); + tx.commit().await.unwrap(); + let children = s.get_children(parent).await.unwrap(); + assert_eq!(children.len(), 2); + assert!(children.contains(&c1)); + assert!(children.contains(&c2)); + + assert_eq!(s.get_root(3).await.unwrap(), EMPTY); + s.set_root(3, parent).await.unwrap(); + assert_eq!(s.get_root(3).await.unwrap(), parent); + } + + #[tokio::test] + async fn test_meta_roundtrip() { + let mut s = InMemoryStorage::default(); + assert_eq!(s.get_meta(7).await.unwrap(), None); + assert_eq!(s.get_key_len(7).await.unwrap(), None); + s.set_meta(7, b"call-site-info".as_slice()).await.unwrap(); + s.set_key_len(7, 5).await.unwrap(); + assert_eq!( + s.get_meta(7).await.unwrap().as_deref(), + Some(b"call-site-info".as_slice()) + ); + assert_eq!(s.get_key_len(7).await.unwrap(), Some(5)); + s.set_meta(7, b"updated").await.unwrap(); + s.set_key_len(7, 6).await.unwrap(); + assert_eq!( + s.get_meta(7).await.unwrap().as_deref(), + Some(b"updated".as_slice()) + ); + assert_eq!(s.get_key_len(7).await.unwrap(), Some(6)); + assert_eq!(s.get_meta(8).await.unwrap(), None); + assert_eq!(s.get_key_len(8).await.unwrap(), None); + } + + #[tokio::test] + async fn test_shortcuts_roundtrip() { + let mut s = InMemoryStorage::default(); + assert!(s.get_shortcut_nodes(1, b"l").await.unwrap().is_empty()); + s.add_shortcut_node(1, b"l", 10).await.unwrap(); + s.add_shortcut_node(1, b"l", 20).await.unwrap(); + s.add_shortcut_node(1, b"o", 10).await.unwrap(); + s.add_shortcut_node(2, b"l", 30).await.unwrap(); + let nodes = s.get_shortcut_nodes(1, b"l").await.unwrap(); + assert!(nodes.contains(&10) && nodes.contains(&20)); + assert_eq!(nodes.len(), 2); + assert_eq!(s.get_shortcut_nodes(2, b"l").await.unwrap(), vec![30]); + + s.clear_shortcuts().await.unwrap(); + assert!(s.get_shortcut_nodes(1, b"l").await.unwrap().is_empty()); + assert!(s.get_shortcut_nodes(2, b"l").await.unwrap().is_empty()); + } + + #[tokio::test] + async fn test_tx_commit_applies_atomically() { + let mut s = InMemoryStorage::default(); + let parent = s.new_node(b"hello".to_vec(), 1).await.unwrap(); + + let mut tx = s.new_tx(); + let new_id = tx.new_node(b"p".to_vec(), 2).await.unwrap(); + let leg_id = tx.new_node(b"lo".to_vec(), 1).await.unwrap(); + tx.move_child(parent, leg_id, 0).await.unwrap(); + tx.add_child(parent, leg_id).await.unwrap(); + tx.add_child(parent, new_id).await.unwrap(); + tx.update_node(parent, Some(b"hel".to_vec()), Some(0)) + .await + .unwrap(); + tx.commit().await.unwrap(); + + let (prefix, record) = s.get_node(parent).await.unwrap(); + assert_eq!(prefix, b"hel"); + assert_eq!(record, 0); + let children = s.get_children(parent).await.unwrap(); + assert!(children.contains(&leg_id)); + assert!(children.contains(&new_id)); + assert_eq!(s.get_node(new_id).await.unwrap().1, 2); + assert_eq!(s.get_node(leg_id).await.unwrap().1, 1); + } + + #[tokio::test] + async fn test_tx_nodes_invisible_before_commit() { + let s = InMemoryStorage::default(); + let mut tx = s.new_tx(); + let id = tx.new_node(b"pending".to_vec(), 9).await.unwrap(); + assert!(s.get_node(id).await.is_err()); + tx.commit().await.unwrap(); + assert_eq!(s.get_node(id).await.unwrap().1, 9); + } + + #[tokio::test] + async fn test_tx_move_child_migrates() { + let mut s = InMemoryStorage::default(); + let parent = s.new_node(b"aaaaaa".to_vec(), 0).await.unwrap(); + let child = s.new_node(b"0".to_vec(), 1).await.unwrap(); + let mut seed = s.new_tx(); + seed.add_child(parent, child).await.unwrap(); + seed.commit().await.unwrap(); + + let mut tx = s.new_tx(); + let leg = tx.new_node(b"a".to_vec(), 0).await.unwrap(); + tx.move_child(parent, leg, child).await.unwrap(); + tx.add_child(parent, leg).await.unwrap(); + tx.commit().await.unwrap(); + + assert!(!s.get_children(parent).await.unwrap().contains(&child)); + assert!(s.get_children(leg).await.unwrap().contains(&child)); + } + + #[tokio::test] + async fn test_edge_data_roundtrip() { + let mut s = InMemoryStorage::default(); + assert_eq!(s.get_edge_data(7).await.unwrap(), None); + s.set_edge_data(7, b"call-site").await.unwrap(); + assert_eq!( + s.get_edge_data(7).await.unwrap().as_deref(), + Some(b"call-site".as_slice()) + ); + s.set_edge_data(7, b"updated").await.unwrap(); + assert_eq!( + s.get_edge_data(7).await.unwrap().as_deref(), + Some(b"updated".as_slice()) + ); + assert_eq!(s.get_edge_data(8).await.unwrap(), None); + s.set_edge_data(9, b"x").await.unwrap(); + s.clear_edges().await.unwrap(); + assert_eq!(s.get_edge_data(7).await.unwrap(), None); + assert_eq!(s.get_edge_data(9).await.unwrap(), None); + } + + #[tokio::test] + async fn test_node_meta_roundtrip() { + let mut s = InMemoryStorage::default(); + assert_eq!(s.get_node_meta(3).await.unwrap(), None); + s.set_node_meta(3, b"node-json").await.unwrap(); + assert_eq!( + s.get_node_meta(3).await.unwrap().as_deref(), + Some(b"node-json".as_slice()) + ); + s.set_node_meta(3, b"node-json-2").await.unwrap(); + assert_eq!( + s.get_node_meta(3).await.unwrap().as_deref(), + Some(b"node-json-2".as_slice()) + ); + assert_eq!(s.get_node_meta(4).await.unwrap(), None); + s.clear_node_meta().await.unwrap(); + assert_eq!(s.get_node_meta(3).await.unwrap(), None); + } + + #[tokio::test] + async fn test_chains_roundtrip() { + let mut s = InMemoryStorage::default(); + assert_eq!(s.get_chain(9).await.unwrap(), None); + s.set_chain(9, &[1, 2, 3]).await.unwrap(); + assert_eq!(s.get_chain(9).await.unwrap(), Some(vec![1, 2, 3])); + s.set_chain(9, &[4]).await.unwrap(); + assert_eq!(s.get_chain(9).await.unwrap(), Some(vec![4])); + assert_eq!(s.get_chain(10).await.unwrap(), None); + s.clear_chains().await.unwrap(); + assert_eq!(s.get_chain(9).await.unwrap(), None); + } + + #[tokio::test] + async fn test_entity_symbols() { + use codegraph_core::{ScopeLevel, SymbolKind}; + let mut s = InMemoryStorage::default(); + let sym = Symbol { + id: 100, + name: "foo".into(), + kind: SymbolKind::Function, + scope: ScopeLevel::Global, + scope_id: 0, + type_ref: 0, + type_name: None, + file: String::new(), + line: 0, + end_line: 0, + signature: None, + doc: None, + annotations: Vec::new(), + language: "rust".into(), + }; + s.save_symbol(&sym).await.unwrap(); + let loaded = s.load_symbol(100).await.unwrap(); + assert_eq!(loaded.unwrap().name, "foo"); + assert_eq!(s.load_all_symbols().await.unwrap().len(), 1); + } + + #[tokio::test] + async fn test_entity_clear() { + let mut s = InMemoryStorage::default(); + s.set_call_records(1, b"rec").await.unwrap(); + s.set_call_name_index("name", b"sites").await.unwrap(); + s.set_version(5).await.unwrap(); + s.clear_entities().await.unwrap(); + assert_eq!(s.get_call_records(1).await.unwrap(), None); + assert_eq!(s.load_call_name_index("name").await.unwrap(), None); + assert_eq!(s.version().await.unwrap(), 0); + } +} diff --git a/crates/codegraph-graph/src/storage/lmdb.rs b/crates/codegraph-graph/src/storage/lmdb.rs index 180895f0d..08a6ef37c 100644 --- a/crates/codegraph-graph/src/storage/lmdb.rs +++ b/crates/codegraph-graph/src/storage/lmdb.rs @@ -23,9 +23,12 @@ use codegraph_core::{FileInfo, Symbol}; use lmdb::EnvironmentFlags; use lmdb::{Cursor, Database, DatabaseFlags, Environment, Transaction, WriteFlags}; +#[cfg(feature = "bloom-search")] +use super::BloomStorage; use super::{ - EMPTY, IndexCounts, Result, Storage, StorageError, Tx, TxOp, decode_chain, decode_vector, - encode_chain, encode_vector, + CategoryStorage, ChainStorage, EMPTY, EdgeDataStorage, EntityStorage, IndexCounts, + NodeMetaStorage, Result, ShortcutsStorage, Storage, StorageError, Tx, TxOp, decode_chain, + decode_vector, encode_chain, encode_vector, }; /// Map lỗi LMDB → `StorageError`. @@ -418,10 +421,10 @@ impl LmdbStorage { } } -// ==================== Storage impl ==================== +// ==================== Storage impl (split into 7 sub-traits) ==================== #[async_trait] -impl Storage for LmdbStorage { +impl CategoryStorage for LmdbStorage { async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { let mut tx = self.env.begin_rw_txn().map_err(e)?; // Không có RETURNING — đọc-rồi-ghi counter trong cùng write tx; an toàn @@ -486,7 +489,42 @@ impl Storage for LmdbStorage { Ok(out) } - #[cfg(feature = "bloom-search")] + async fn set_root(&mut self, shard: usize, root: usize) -> Result<()> { + let mut tx = self.env.begin_rw_txn().map_err(e)?; + tx.put(self.roots, &k8(shard), &k8(root), WriteFlags::empty()) + .map_err(e)?; + tx.commit().map_err(e)?; + Ok(()) + } + + async fn get_root(&self, shard: usize) -> Result { + let tx = self.env.begin_ro_txn().map_err(e)?; + Ok(self + .get_opt(&tx, self.roots, &k8(shard))? + .map(de_u64) + .unwrap_or(EMPTY as u64) as usize) + } + + fn new_tx(&self) -> Box { + Box::new(LmdbTx { + env: self.env.clone(), + nodes: self.nodes, + children: self.children, + counter: self.counter, + nodes_pending: Vec::new(), + ops: Vec::new(), + }) + } +} + +// Blanket marker — `Storage` is `CategoryStorage + 5 sub-traits + EntityStorage + Send + Sync`, +// so this empty impl makes the LMDB backend satisfy `Storage` automatically. +#[async_trait] +impl Storage for LmdbStorage {} + +#[cfg(feature = "bloom-search")] +#[async_trait] +impl BloomStorage for LmdbStorage { async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<()> { let mut tx = self.env.begin_rw_txn().map_err(e)?; tx.put(self.blooms, &k8(id), &bloom, WriteFlags::empty()) @@ -495,12 +533,16 @@ impl Storage for LmdbStorage { Ok(()) } - #[cfg(feature = "bloom-search")] async fn get_node_bloom(&self, id: usize) -> Result>> { let tx = self.env.begin_ro_txn().map_err(e)?; Ok(self.get_opt(&tx, self.blooms, &k8(id))?.map(|b| b.to_vec())) } +} + +// --- EdgeDataStorage --- +#[async_trait] +impl EdgeDataStorage for LmdbStorage { async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<()> { let mut tx = self.env.begin_rw_txn().map_err(e)?; tx.put(self.edges, &k8(edge), &data, WriteFlags::empty()) @@ -522,27 +564,12 @@ impl Storage for LmdbStorage { tx.commit().map_err(e)?; Ok(()) } +} - async fn for_each_edge_data( - &self, - f: &mut (dyn for<'a> FnMut(usize, &'a [u8]) -> Result<()> + Send), - ) -> Result<()> { - let tx = self.env.begin_ro_txn().map_err(e)?; - let mut cur = tx.open_ro_cursor(self.edges).map_err(e)?; - let mut rows: Vec<(Vec, Vec)> = Vec::new(); - for item in cur.iter() { - let (k, v) = item.map_err(e)?; - rows.push((k.to_vec(), v.to_vec())); - } - drop(cur); - drop(tx); - rows.sort_by(|a, b| a.0.cmp(&b.0)); - for (k, v) in rows { - f(de_u64(&k) as usize, &v)?; - } - Ok(()) - } +// --- NodeMetaStorage --- +#[async_trait] +impl NodeMetaStorage for LmdbStorage { async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { let mut tx = self.env.begin_rw_txn().map_err(e)?; tx.put(self.node_meta, &k8(elem), &meta, WriteFlags::empty()) @@ -565,6 +592,42 @@ impl Storage for LmdbStorage { Ok(()) } + async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { + let mut tx = self.env.begin_rw_txn().map_err(e)?; + tx.put(self.meta, &k8(record), &meta, WriteFlags::empty()) + .map_err(e)?; + tx.commit().map_err(e)?; + Ok(()) + } + + async fn get_meta(&self, record: usize) -> Result>> { + let tx = self.env.begin_ro_txn().map_err(e)?; + Ok(self + .get_opt(&tx, self.meta, &k8(record))? + .map(|v| v.to_vec())) + } + + async fn set_key_len(&mut self, record: usize, len: usize) -> Result<()> { + let mut tx = self.env.begin_rw_txn().map_err(e)?; + tx.put(self.keylen, &k8(record), &k8(len), WriteFlags::empty()) + .map_err(e)?; + tx.commit().map_err(e)?; + Ok(()) + } + + async fn get_key_len(&self, record: usize) -> Result> { + let tx = self.env.begin_ro_txn().map_err(e)?; + Ok(self + .get_opt(&tx, self.keylen, &k8(record))? + .map(de_u64) + .map(|v| v as usize)) + } +} + +// --- ChainStorage --- + +#[async_trait] +impl ChainStorage for LmdbStorage { async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<()> { let mut tx = self.env.begin_rw_txn().map_err(e)?; tx.put( @@ -591,7 +654,12 @@ impl Storage for LmdbStorage { tx.commit().map_err(e)?; Ok(()) } +} + +// --- EntityStorage --- +#[async_trait] +impl EntityStorage for LmdbStorage { async fn save_symbol(&mut self, sym: &Symbol) -> Result<()> { let data = serde_json::to_vec(sym).map_err(|err| StorageError::Internal(err.to_string()))?; @@ -831,54 +899,12 @@ impl Storage for LmdbStorage { tx.commit().map_err(e)?; Ok(()) } +} - async fn set_root(&mut self, shard: usize, root: usize) -> Result<()> { - let mut tx = self.env.begin_rw_txn().map_err(e)?; - tx.put(self.roots, &k8(shard), &k8(root), WriteFlags::empty()) - .map_err(e)?; - tx.commit().map_err(e)?; - Ok(()) - } - - async fn get_root(&self, shard: usize) -> Result { - let tx = self.env.begin_ro_txn().map_err(e)?; - Ok(self - .get_opt(&tx, self.roots, &k8(shard))? - .map(de_u64) - .unwrap_or(EMPTY as u64) as usize) - } - - async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { - let mut tx = self.env.begin_rw_txn().map_err(e)?; - tx.put(self.meta, &k8(record), &meta, WriteFlags::empty()) - .map_err(e)?; - tx.commit().map_err(e)?; - Ok(()) - } - - async fn get_meta(&self, record: usize) -> Result>> { - let tx = self.env.begin_ro_txn().map_err(e)?; - Ok(self - .get_opt(&tx, self.meta, &k8(record))? - .map(|v| v.to_vec())) - } - - async fn set_key_len(&mut self, record: usize, len: usize) -> Result<()> { - let mut tx = self.env.begin_rw_txn().map_err(e)?; - tx.put(self.keylen, &k8(record), &k8(len), WriteFlags::empty()) - .map_err(e)?; - tx.commit().map_err(e)?; - Ok(()) - } - - async fn get_key_len(&self, record: usize) -> Result> { - let tx = self.env.begin_ro_txn().map_err(e)?; - Ok(self - .get_opt(&tx, self.keylen, &k8(record))? - .map(de_u64) - .map(|v| v as usize)) - } +// --- ShortcutsStorage --- +#[async_trait] +impl ShortcutsStorage for LmdbStorage { async fn add_shortcut_node(&mut self, shard: usize, elem: &[u8], node_id: usize) -> Result<()> { let mut key = k8(shard).to_vec(); key.extend_from_slice(elem); @@ -912,17 +938,6 @@ impl Storage for LmdbStorage { tx.commit().map_err(e)?; Ok(()) } - - fn new_tx(&self) -> Box { - Box::new(LmdbTx { - env: self.env.clone(), - nodes: self.nodes, - children: self.children, - counter: self.counter, - nodes_pending: Vec::new(), - ops: Vec::new(), - }) - } } // ==================== LmdbTx ==================== diff --git a/crates/codegraph-graph/src/storage/mysql.rs b/crates/codegraph-graph/src/storage/mysql.rs index f8fa93187..0d5ac32f7 100644 --- a/crates/codegraph-graph/src/storage/mysql.rs +++ b/crates/codegraph-graph/src/storage/mysql.rs @@ -1,7 +1,10 @@ use std::collections::HashMap; +#[cfg(feature = "bloom-search")] +use super::BloomStorage; use super::{ - IndexCounts, Result, Storage, StorageError, Tx, decode_chain, decode_vector, encode_chain, + CategoryStorage, ChainStorage, EdgeDataStorage, EntityStorage, IndexCounts, NodeMetaStorage, + Result, ShortcutsStorage, Storage, StorageError, Tx, decode_chain, decode_vector, encode_chain, encode_vector, }; use async_trait::async_trait; @@ -127,7 +130,7 @@ impl MySqlStorage { } #[async_trait] -impl Storage for MySqlStorage { +impl CategoryStorage for MySqlStorage { async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { let id = self.reserve_node_id().await?; sqlx::query( @@ -228,6 +231,18 @@ impl Storage for MySqlStorage { Ok(root as usize) } + fn new_tx(&self) -> Box { + Box::new(MySqlTx { + pool: self.pool.clone(), + repo_id: self.repo_id, + nodes: Vec::new(), + ops: Vec::new(), + }) + } +} + +#[async_trait] +impl NodeMetaStorage for MySqlStorage { async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { sqlx::query( "INSERT INTO rt_meta (repo_id, record, meta) VALUES (?, ?, ?) \ @@ -280,6 +295,44 @@ impl Storage for MySqlStorage { Ok(row.map(|(len,)| len as usize)) } + async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { + sqlx::query( + "INSERT INTO rt_node_meta (repo_id, elem, meta) VALUES (?, ?, ?) \ + ON DUPLICATE KEY UPDATE meta = VALUES(meta)", + ) + .bind(self.repo_id as i64) + .bind(elem as i64) + .bind(meta) + .execute(&self.pool) + .await + .map_err(db_err)?; + Ok(()) + } + + async fn get_node_meta(&self, elem: usize) -> Result>> { + let row = sqlx::query_as::<_, (Vec,)>( + "SELECT meta FROM rt_node_meta WHERE repo_id = ? AND elem = ?", + ) + .bind(self.repo_id as i64) + .bind(elem as i64) + .fetch_optional(&self.pool) + .await + .map_err(db_err)?; + Ok(row.map(|(m,)| m)) + } + + async fn clear_node_meta(&mut self) -> Result<()> { + sqlx::query("DELETE FROM rt_node_meta WHERE repo_id = ?") + .bind(self.repo_id as i64) + .execute(&self.pool) + .await + .map_err(db_err)?; + Ok(()) + } +} + +#[async_trait] +impl ShortcutsStorage for MySqlStorage { async fn add_shortcut_node(&mut self, shard: usize, elem: &[u8], node_id: usize) -> Result<()> { sqlx::query( "INSERT IGNORE INTO rt_shortcuts (repo_id, shard, elem, node_id) VALUES (?, ?, ?, ?)", @@ -315,7 +368,10 @@ impl Storage for MySqlStorage { .map_err(db_err)?; Ok(()) } +} +#[async_trait] +impl EdgeDataStorage for MySqlStorage { async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<()> { sqlx::query( "INSERT INTO rt_edges (repo_id, id, data) VALUES (?, ?, ?) \ @@ -350,59 +406,10 @@ impl Storage for MySqlStorage { .map_err(db_err)?; Ok(()) } +} - async fn for_each_edge_data( - &self, - f: &mut (dyn for<'a> FnMut(usize, &'a [u8]) -> Result<()> + Send), - ) -> Result<()> { - let rows = sqlx::query("SELECT id, data FROM rt_edges WHERE repo_id = ?") - .bind(self.repo_id as i64) - .fetch_all(&self.pool) - .await - .map_err(db_err)?; - for r in &rows { - let id: i64 = r.try_get("id").map_err(db_err)?; - let data: Vec = r.try_get("data").map_err(db_err)?; - f(id as usize, &data)?; - } - Ok(()) - } - - async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { - sqlx::query( - "INSERT INTO rt_node_meta (repo_id, elem, meta) VALUES (?, ?, ?) \ - ON DUPLICATE KEY UPDATE meta = VALUES(meta)", - ) - .bind(self.repo_id as i64) - .bind(elem as i64) - .bind(meta) - .execute(&self.pool) - .await - .map_err(db_err)?; - Ok(()) - } - - async fn get_node_meta(&self, elem: usize) -> Result>> { - let row = sqlx::query_as::<_, (Vec,)>( - "SELECT meta FROM rt_node_meta WHERE repo_id = ? AND elem = ?", - ) - .bind(self.repo_id as i64) - .bind(elem as i64) - .fetch_optional(&self.pool) - .await - .map_err(db_err)?; - Ok(row.map(|(m,)| m)) - } - - async fn clear_node_meta(&mut self) -> Result<()> { - sqlx::query("DELETE FROM rt_node_meta WHERE repo_id = ?") - .bind(self.repo_id as i64) - .execute(&self.pool) - .await - .map_err(db_err)?; - Ok(()) - } - +#[async_trait] +impl ChainStorage for MySqlStorage { async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<()> { let bytes = encode_chain(chain); sqlx::query( @@ -438,7 +445,10 @@ impl Storage for MySqlStorage { .map_err(db_err)?; Ok(()) } +} +#[async_trait] +impl EntityStorage for MySqlStorage { async fn save_symbol(&mut self, sym: &Symbol) -> Result<()> { let annotations = serde_json::to_string(&sym.annotations).map_err(ser_err)?; sqlx::query( @@ -815,8 +825,11 @@ impl Storage for MySqlStorage { tx.commit().await.map_err(db_err)?; Ok(()) } +} - #[cfg(feature = "bloom-search")] +#[cfg(feature = "bloom-search")] +#[async_trait] +impl BloomStorage for MySqlStorage { async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<()> { sqlx::query( "INSERT INTO rt_node_blooms (repo_id, id, bloom) VALUES (?, ?, ?) \ @@ -831,7 +844,6 @@ impl Storage for MySqlStorage { Ok(()) } - #[cfg(feature = "bloom-search")] async fn get_node_bloom(&self, id: usize) -> Result>> { let row = sqlx::query_as::<_, (Vec,)>( "SELECT bloom FROM rt_node_blooms WHERE repo_id = ? AND id = ?", @@ -843,17 +855,13 @@ impl Storage for MySqlStorage { .map_err(db_err)?; Ok(row.map(|(b,)| b)) } - - fn new_tx(&self) -> Box { - Box::new(MySqlTx { - pool: self.pool.clone(), - repo_id: self.repo_id, - nodes: Vec::new(), - ops: Vec::new(), - }) - } } +// Blanket marker — `Storage` is `CategoryStorage + 5 sub-traits + EntityStorage + Send + Sync`, +// so this empty impl makes the MySQL backend satisfy `Storage` automatically. +#[async_trait] +impl Storage for MySqlStorage {} + /// Probe version index trên đĩa (dùng cho `SharedGraphIndex::ensure_fresh`). #[cfg(feature = "mysql")] impl MySqlStorage { diff --git a/crates/codegraph-graph/src/storage/postgres.rs b/crates/codegraph-graph/src/storage/postgres.rs index 665cb15be..02a89db89 100644 --- a/crates/codegraph-graph/src/storage/postgres.rs +++ b/crates/codegraph-graph/src/storage/postgres.rs @@ -1,7 +1,10 @@ use std::collections::HashMap; +#[cfg(feature = "bloom-search")] +use super::BloomStorage; use super::{ - IndexCounts, Result, Storage, StorageError, Tx, decode_chain, decode_vector, encode_chain, + CategoryStorage, ChainStorage, EdgeDataStorage, EntityStorage, IndexCounts, NodeMetaStorage, + Result, ShortcutsStorage, Storage, StorageError, Tx, decode_chain, decode_vector, encode_chain, encode_vector, }; use async_trait::async_trait; @@ -137,7 +140,7 @@ impl PostgresStorage { } #[async_trait] -impl Storage for PostgresStorage { +impl CategoryStorage for PostgresStorage { async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { let id = self.reserve_node_id().await?; sqlx::query( @@ -238,6 +241,18 @@ impl Storage for PostgresStorage { Ok(root as usize) } + fn new_tx(&self) -> Box { + Box::new(PostgresTx { + pool: self.pool.clone(), + repo_id: self.repo_id, + nodes: Vec::new(), + ops: Vec::new(), + }) + } +} + +#[async_trait] +impl NodeMetaStorage for PostgresStorage { async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { sqlx::query( "INSERT INTO rt_meta (repo_id, record, meta) VALUES ($1, $2, $3) \ @@ -290,6 +305,44 @@ impl Storage for PostgresStorage { Ok(row.map(|(len,)| len as usize)) } + async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { + sqlx::query( + "INSERT INTO rt_node_meta (repo_id, elem, meta) VALUES ($1, $2, $3) \ + ON CONFLICT (repo_id, elem) DO UPDATE SET meta = EXCLUDED.meta", + ) + .bind(self.repo_id as i64) + .bind(elem as i64) + .bind(meta) + .execute(&self.pool) + .await + .map_err(db_err)?; + Ok(()) + } + + async fn get_node_meta(&self, elem: usize) -> Result>> { + let row = sqlx::query_as::<_, (Vec,)>( + "SELECT meta FROM rt_node_meta WHERE repo_id = $1 AND elem = $2", + ) + .bind(self.repo_id as i64) + .bind(elem as i64) + .fetch_optional(&self.pool) + .await + .map_err(db_err)?; + Ok(row.map(|(m,)| m)) + } + + async fn clear_node_meta(&mut self) -> Result<()> { + sqlx::query("DELETE FROM rt_node_meta WHERE repo_id = $1") + .bind(self.repo_id as i64) + .execute(&self.pool) + .await + .map_err(db_err)?; + Ok(()) + } +} + +#[async_trait] +impl ShortcutsStorage for PostgresStorage { async fn add_shortcut_node(&mut self, shard: usize, elem: &[u8], node_id: usize) -> Result<()> { sqlx::query( "INSERT INTO rt_shortcuts (repo_id, shard, elem, node_id) VALUES ($1, $2, $3, $4) \ @@ -326,7 +379,10 @@ impl Storage for PostgresStorage { .map_err(db_err)?; Ok(()) } +} +#[async_trait] +impl EdgeDataStorage for PostgresStorage { async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<()> { sqlx::query( "INSERT INTO rt_edges (repo_id, id, data) VALUES ($1, $2, $3) \ @@ -361,59 +417,10 @@ impl Storage for PostgresStorage { .map_err(db_err)?; Ok(()) } +} - async fn for_each_edge_data( - &self, - f: &mut (dyn for<'a> FnMut(usize, &'a [u8]) -> Result<()> + Send), - ) -> Result<()> { - let rows = sqlx::query("SELECT id, data FROM rt_edges WHERE repo_id = $1") - .bind(self.repo_id as i64) - .fetch_all(&self.pool) - .await - .map_err(db_err)?; - for r in &rows { - let id: i64 = r.try_get("id").map_err(db_err)?; - let data: Vec = r.try_get("data").map_err(db_err)?; - f(id as usize, &data)?; - } - Ok(()) - } - - async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { - sqlx::query( - "INSERT INTO rt_node_meta (repo_id, elem, meta) VALUES ($1, $2, $3) \ - ON CONFLICT (repo_id, elem) DO UPDATE SET meta = EXCLUDED.meta", - ) - .bind(self.repo_id as i64) - .bind(elem as i64) - .bind(meta) - .execute(&self.pool) - .await - .map_err(db_err)?; - Ok(()) - } - - async fn get_node_meta(&self, elem: usize) -> Result>> { - let row = sqlx::query_as::<_, (Vec,)>( - "SELECT meta FROM rt_node_meta WHERE repo_id = $1 AND elem = $2", - ) - .bind(self.repo_id as i64) - .bind(elem as i64) - .fetch_optional(&self.pool) - .await - .map_err(db_err)?; - Ok(row.map(|(m,)| m)) - } - - async fn clear_node_meta(&mut self) -> Result<()> { - sqlx::query("DELETE FROM rt_node_meta WHERE repo_id = $1") - .bind(self.repo_id as i64) - .execute(&self.pool) - .await - .map_err(db_err)?; - Ok(()) - } - +#[async_trait] +impl ChainStorage for PostgresStorage { async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<()> { let bytes = encode_chain(chain); sqlx::query( @@ -449,7 +456,10 @@ impl Storage for PostgresStorage { .map_err(db_err)?; Ok(()) } +} +#[async_trait] +impl EntityStorage for PostgresStorage { async fn save_symbol(&mut self, sym: &Symbol) -> Result<()> { let annotations = serde_json::to_string(&sym.annotations).map_err(ser_err)?; sqlx::query( @@ -828,8 +838,11 @@ impl Storage for PostgresStorage { tx.commit().await.map_err(db_err)?; Ok(()) } +} - #[cfg(feature = "bloom-search")] +#[cfg(feature = "bloom-search")] +#[async_trait] +impl BloomStorage for PostgresStorage { async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<()> { sqlx::query( "INSERT INTO rt_node_blooms (repo_id, id, bloom) VALUES ($1, $2, $3) \ @@ -844,7 +857,6 @@ impl Storage for PostgresStorage { Ok(()) } - #[cfg(feature = "bloom-search")] async fn get_node_bloom(&self, id: usize) -> Result>> { let row = sqlx::query_as::<_, (Vec,)>( "SELECT bloom FROM rt_node_blooms WHERE repo_id = $1 AND id = $2", @@ -856,17 +868,13 @@ impl Storage for PostgresStorage { .map_err(db_err)?; Ok(row.map(|(b,)| b)) } - - fn new_tx(&self) -> Box { - Box::new(PostgresTx { - pool: self.pool.clone(), - repo_id: self.repo_id, - nodes: Vec::new(), - ops: Vec::new(), - }) - } } +// Blanket marker — `Storage` is `CategoryStorage + 5 sub-traits + EntityStorage + Send + Sync`, +// so this empty impl makes the Postgres backend satisfy `Storage` automatically. +#[async_trait] +impl Storage for PostgresStorage {} + /// Probe version index trên đĩa (dùng cho `SharedGraphIndex::ensure_fresh`) — /// không mở toàn bộ index. `None`/lỗi → coi như version 0. #[cfg(feature = "postgres")] diff --git a/crates/codegraph-graph/src/storage/redis.rs b/crates/codegraph-graph/src/storage/redis.rs index e01f06675..8c5e2ee5e 100644 --- a/crates/codegraph-graph/src/storage/redis.rs +++ b/crates/codegraph-graph/src/storage/redis.rs @@ -29,8 +29,12 @@ use tokio::sync::Mutex; use async_trait::async_trait; +#[cfg(feature = "bloom-search")] +use super::BloomStorage; use super::{ - FileInfo, Result, Storage, StorageError, Symbol, Tx, TxOp, decode_vector, encode_vector, + CategoryStorage, ChainStorage, EdgeDataStorage, EntityStorage, FileInfo, NodeMetaStorage, + Result, ShortcutsStorage, Storage, StorageError, Symbol, Tx, TxOp, decode_vector, + encode_vector, }; // ==================== KeyBuilder ==================== @@ -177,7 +181,7 @@ impl RedisStorage { } #[async_trait] -impl Storage for RedisStorage { +impl CategoryStorage for RedisStorage { async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { let mut conn = self.lock().await; let result: redis::Value = redis::pipe() @@ -255,31 +259,6 @@ impl Storage for RedisStorage { Ok(children.into_iter().map(|x| x as usize).collect()) } - #[cfg(feature = "bloom-search")] - async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<()> { - let mut conn = self.lock().await; - cmd("HSET") - .arg(self.kb.key("node_bloom")) - .arg(id) - .arg(bloom) - .query_async::<()>(&mut *conn) - .await - .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; - Ok(()) - } - - #[cfg(feature = "bloom-search")] - async fn get_node_bloom(&self, id: usize) -> Result>> { - let mut conn = self.lock().await; - let bloom: Option> = cmd("HGET") - .arg(self.kb.key("node_bloom")) - .arg(id) - .query_async(&mut *conn) - .await - .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; - Ok(bloom) - } - async fn set_root(&mut self, shard: usize, root: usize) -> Result<()> { let mut conn = self.lock().await; cmd("HSET") @@ -303,52 +282,45 @@ impl Storage for RedisStorage { Ok(root.unwrap_or(0) as usize) } - async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { - let mut conn = self.lock().await; - cmd("HSET") - .arg(self.kb.key("meta")) - .arg(record as i64) - .arg(meta) - .query_async::<()>(&mut *conn) - .await - .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; - Ok(()) - } - - async fn get_meta(&self, record: usize) -> Result>> { - let mut conn = self.lock().await; - let meta: Option> = cmd("HGET") - .arg(self.kb.key("meta")) - .arg(record as i64) - .query_async(&mut *conn) - .await - .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; - Ok(meta) + fn new_tx(&self) -> Box { + Box::new(RedisTx { + conn: self.conn.clone(), + kb: self.kb.clone(), + nodes: Vec::new(), + ops: Vec::new(), + }) } +} - async fn set_key_len(&mut self, record: usize, len: usize) -> Result<()> { +#[cfg(feature = "bloom-search")] +#[async_trait] +impl BloomStorage for RedisStorage { + async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<()> { let mut conn = self.lock().await; cmd("HSET") - .arg(self.kb.key("keylen")) - .arg(record as i64) - .arg(len as i64) + .arg(self.kb.key("node_bloom")) + .arg(id) + .arg(bloom) .query_async::<()>(&mut *conn) .await .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; Ok(()) } - async fn get_key_len(&self, record: usize) -> Result> { + async fn get_node_bloom(&self, id: usize) -> Result>> { let mut conn = self.lock().await; - let len: Option = cmd("HGET") - .arg(self.kb.key("keylen")) - .arg(record as i64) + let bloom: Option> = cmd("HGET") + .arg(self.kb.key("node_bloom")) + .arg(id) .query_async(&mut *conn) .await .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; - Ok(len.map(|x| x as usize)) + Ok(bloom) } +} +#[async_trait] +impl ShortcutsStorage for RedisStorage { async fn add_shortcut_node(&mut self, shard: usize, elem: &[u8], node_id: usize) -> Result<()> { let mut conn = self.lock().await; cmd("SADD") @@ -398,7 +370,10 @@ impl Storage for RedisStorage { } Ok(()) } +} +#[async_trait] +impl EdgeDataStorage for RedisStorage { async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<()> { let mut conn = self.lock().await; cmd("HSET") @@ -431,23 +406,10 @@ impl Storage for RedisStorage { .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; Ok(()) } +} - async fn for_each_edge_data( - &self, - f: &mut (dyn for<'a> FnMut(usize, &'a [u8]) -> Result<()> + Send), - ) -> Result<()> { - let mut conn = self.lock().await; - let items: Vec<(i64, Vec)> = cmd("HGETALL") - .arg(self.kb.key("edgedata")) - .query_async(&mut *conn) - .await - .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; - for (id, data) in items { - f(id as usize, &data)?; - } - Ok(()) - } - +#[async_trait] +impl NodeMetaStorage for RedisStorage { async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { let mut conn = self.lock().await; cmd("HSET") @@ -481,6 +443,55 @@ impl Storage for RedisStorage { Ok(()) } + async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { + let mut conn = self.lock().await; + cmd("HSET") + .arg(self.kb.key("meta")) + .arg(record as i64) + .arg(meta) + .query_async::<()>(&mut *conn) + .await + .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; + Ok(()) + } + + async fn get_meta(&self, record: usize) -> Result>> { + let mut conn = self.lock().await; + let meta: Option> = cmd("HGET") + .arg(self.kb.key("meta")) + .arg(record as i64) + .query_async(&mut *conn) + .await + .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; + Ok(meta) + } + + async fn set_key_len(&mut self, record: usize, len: usize) -> Result<()> { + let mut conn = self.lock().await; + cmd("HSET") + .arg(self.kb.key("keylen")) + .arg(record as i64) + .arg(len as i64) + .query_async::<()>(&mut *conn) + .await + .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; + Ok(()) + } + + async fn get_key_len(&self, record: usize) -> Result> { + let mut conn = self.lock().await; + let len: Option = cmd("HGET") + .arg(self.kb.key("keylen")) + .arg(record as i64) + .query_async(&mut *conn) + .await + .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; + Ok(len.map(|x| x as usize)) + } +} + +#[async_trait] +impl ChainStorage for RedisStorage { async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<()> { let mut conn = self.lock().await; cmd("HSET") @@ -513,7 +524,10 @@ impl Storage for RedisStorage { .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; Ok(()) } +} +#[async_trait] +impl EntityStorage for RedisStorage { async fn save_symbol(&mut self, sym: &Symbol) -> Result<()> { let mut conn = self.lock().await; let data = serde_json::to_vec(sym).map_err(|e| StorageError::Internal(e.to_string()))?; @@ -777,17 +791,13 @@ impl Storage for RedisStorage { .map_err(|e: redis::RedisError| StorageError::Internal(e.to_string()))?; Ok(()) } - - fn new_tx(&self) -> Box { - Box::new(RedisTx { - conn: self.conn.clone(), - kb: self.kb.clone(), - nodes: Vec::new(), - ops: Vec::new(), - }) - } } +// Blanket marker — `Storage` is `CategoryStorage + 5 sub-traits + EntityStorage + Send + Sync`, +// so this empty impl makes the Redis backend satisfy `Storage` automatically. +#[async_trait] +impl Storage for RedisStorage {} + // ==================== Redis Transaction ==================== /// Transaction cho `RedisStorage`. @@ -908,7 +918,6 @@ mod tests { use super::*; use crate::radix::EMPTY; - use crate::storage::Storage; static COUNTER: AtomicU16 = AtomicU16::new(0); diff --git a/crates/codegraph-graph/src/storage/sqlite.rs b/crates/codegraph-graph/src/storage/sqlite.rs index 9e4d70b83..aad53f7c3 100644 --- a/crates/codegraph-graph/src/storage/sqlite.rs +++ b/crates/codegraph-graph/src/storage/sqlite.rs @@ -43,8 +43,13 @@ use sqlx::Row; use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions}; use super::{ - EMPTY, IndexCounts, Result, Storage, StorageError, Tx, TxOp, decode_vector, encode_vector, + CategoryStorage, ChainStorage, EMPTY, EdgeDataStorage, EntityStorage, IndexCounts, + NodeMetaStorage, Result, ShortcutsStorage, StorageError, Tx, TxOp, decode_vector, + encode_vector, }; + +#[cfg(feature = "bloom-search")] +use super::BloomStorage; use crate::embeddings::resolve_vss_extensions; fn db_err(e: sqlx::Error) -> StorageError { @@ -252,7 +257,7 @@ impl SqliteStorage { } #[async_trait] -impl Storage for SqliteStorage { +impl CategoryStorage for SqliteStorage { async fn new_node(&mut self, prefix: Vec, record: usize) -> Result { let mut conn = self.pool.acquire().await.map_err(db_err)?; // `UPDATE ... RETURNING next - 1` cấp id atomic — không cần SELECT rồi @@ -336,7 +341,42 @@ impl Storage for SqliteStorage { Ok(out) } - #[cfg(feature = "bloom-search")] + async fn set_root(&mut self, shard: usize, root: usize) -> Result<()> { + let mut conn = self.pool.acquire().await.map_err(db_err)?; + sqlx::query( + "INSERT INTO rt_roots (shard, root) VALUES (?1, ?2) + ON CONFLICT(shard) DO UPDATE SET root = excluded.root", + ) + .bind(shard as i64) + .bind(root as i64) + .execute(&mut *conn) + .await + .map_err(db_err)?; + Ok(()) + } + + async fn get_root(&self, shard: usize) -> Result { + let mut conn = self.pool.acquire().await.map_err(db_err)?; + let root: Option = sqlx::query_scalar("SELECT root FROM rt_roots WHERE shard = ?1") + .bind(shard as i64) + .fetch_optional(&mut *conn) + .await + .map_err(db_err)?; + Ok(root.unwrap_or(EMPTY as i64) as usize) + } + + fn new_tx(&self) -> Box { + Box::new(SqliteTx { + pool: self.pool.clone(), + nodes: Vec::new(), + ops: Vec::new(), + }) + } +} + +#[cfg(feature = "bloom-search")] +#[async_trait] +impl BloomStorage for SqliteStorage { async fn set_node_bloom(&mut self, id: usize, bloom: &[u8]) -> Result<()> { let mut conn = self.pool.acquire().await.map_err(db_err)?; sqlx::query( @@ -351,7 +391,6 @@ impl Storage for SqliteStorage { Ok(()) } - #[cfg(feature = "bloom-search")] async fn get_node_bloom(&self, id: usize) -> Result>> { let mut conn = self.pool.acquire().await.map_err(db_err)?; let row = sqlx::query("SELECT bloom FROM rt_node_blooms WHERE id = ?1") @@ -365,7 +404,10 @@ impl Storage for SqliteStorage { let bloom: Vec = row.try_get(0).map_err(db_err)?; Ok(Some(bloom)) } +} +#[async_trait] +impl EdgeDataStorage for SqliteStorage { async fn set_edge_data(&mut self, edge: usize, data: &[u8]) -> Result<()> { let mut conn = self.pool.acquire().await.map_err(db_err)?; sqlx::query( @@ -398,22 +440,10 @@ impl Storage for SqliteStorage { .map_err(db_err)?; Ok(()) } +} - async fn for_each_edge_data( - &self, - f: &mut (dyn for<'a> FnMut(usize, &'a [u8]) -> Result<()> + Send), - ) -> Result<()> { - let mut conn = self.pool.acquire().await.map_err(db_err)?; - let rows: Vec<(i64, Vec)> = sqlx::query_as("SELECT id, data FROM rt_edges ORDER BY id") - .fetch_all(&mut *conn) - .await - .map_err(db_err)?; - for (id, data) in rows { - f(id as usize, &data)?; - } - Ok(()) - } - +#[async_trait] +impl NodeMetaStorage for SqliteStorage { async fn set_node_meta(&mut self, elem: usize, meta: &[u8]) -> Result<()> { let mut conn = self.pool.acquire().await.map_err(db_err)?; sqlx::query( @@ -448,6 +478,103 @@ impl Storage for SqliteStorage { Ok(()) } + async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { + let mut conn = self.pool.acquire().await.map_err(db_err)?; + sqlx::query( + "INSERT INTO rt_meta (record, meta) VALUES (?1, ?2) + ON CONFLICT(record) DO UPDATE SET meta = excluded.meta", + ) + .bind(record as i64) + .bind(meta) + .execute(&mut *conn) + .await + .map_err(db_err)?; + Ok(()) + } + + async fn get_meta(&self, record: usize) -> Result>> { + let mut conn = self.pool.acquire().await.map_err(db_err)?; + let meta: Option> = + sqlx::query_scalar("SELECT meta FROM rt_meta WHERE record = ?1") + .bind(record as i64) + .fetch_optional(&mut *conn) + .await + .map_err(db_err)?; + Ok(meta) + } + + async fn set_key_len(&mut self, record: usize, len: usize) -> Result<()> { + let mut conn = self.pool.acquire().await.map_err(db_err)?; + sqlx::query( + "INSERT INTO rt_keylen (record, len) VALUES (?1, ?2) + ON CONFLICT(record) DO UPDATE SET len = excluded.len", + ) + .bind(record as i64) + .bind(len as i64) + .execute(&mut *conn) + .await + .map_err(db_err)?; + Ok(()) + } + + async fn get_key_len(&self, record: usize) -> Result> { + let mut conn = self.pool.acquire().await.map_err(db_err)?; + let len: Option = sqlx::query_scalar("SELECT len FROM rt_keylen WHERE record = ?1") + .bind(record as i64) + .fetch_optional(&mut *conn) + .await + .map_err(db_err)?; + Ok(len.map(|x| x as usize)) + } +} + +#[async_trait] +impl ShortcutsStorage for SqliteStorage { + async fn add_shortcut_node(&mut self, shard: usize, elem: &[u8], node_id: usize) -> Result<()> { + let mut conn = self.pool.acquire().await.map_err(db_err)?; + sqlx::query( + "INSERT INTO rt_shortcuts (shard, elem, node_id) VALUES (?1, ?2, ?3) + ON CONFLICT DO NOTHING", + ) + .bind(shard as i64) + .bind(elem) + .bind(node_id as i64) + .execute(&mut *conn) + .await + .map_err(db_err)?; + Ok(()) + } + + async fn get_shortcut_nodes(&self, shard: usize, elem: &[u8]) -> Result> { + let mut conn = self.pool.acquire().await.map_err(db_err)?; + let rows = sqlx::query( + "SELECT node_id FROM rt_shortcuts WHERE shard = ?1 AND elem = ?2 ORDER BY node_id", + ) + .bind(shard as i64) + .bind(elem) + .fetch_all(&mut *conn) + .await + .map_err(db_err)?; + let mut out = Vec::with_capacity(rows.len()); + for r in &rows { + let c: i64 = r.try_get(0).map_err(db_err)?; + out.push(c as usize); + } + Ok(out) + } + + async fn clear_shortcuts(&mut self) -> Result<()> { + let mut conn = self.pool.acquire().await.map_err(db_err)?; + sqlx::query("DELETE FROM rt_shortcuts") + .execute(&mut *conn) + .await + .map_err(db_err)?; + Ok(()) + } +} + +#[async_trait] +impl ChainStorage for SqliteStorage { async fn set_chain(&mut self, record: usize, chain: &[u64]) -> Result<()> { let mut conn = self.pool.acquire().await.map_err(db_err)?; sqlx::query( @@ -481,7 +608,10 @@ impl Storage for SqliteStorage { .map_err(db_err)?; Ok(()) } +} +#[async_trait] +impl EntityStorage for SqliteStorage { async fn save_symbol(&mut self, sym: &Symbol) -> Result<()> { let mut conn = self.pool.acquire().await.map_err(db_err)?; let data = serde_json::to_vec(sym).map_err(|e| StorageError::Internal(e.to_string()))?; @@ -538,91 +668,6 @@ impl Storage for SqliteStorage { Ok(next as u64) } - async fn save_embedding(&mut self, symbol_id: u64, vector: &[f32]) -> Result<()> { - let mut conn = self.pool.acquire().await.map_err(db_err)?; - sqlx::query( - "INSERT INTO sg_embeddings (symbol_id, vector) VALUES (?1, ?2) - ON CONFLICT(symbol_id) DO UPDATE SET vector = excluded.vector", - ) - .bind(symbol_id as i64) - .bind(encode_vector(vector)) - .execute(&mut *conn) - .await - .map_err(db_err)?; - // Mirror vào `vss0` (HNSW ANN) nếu extension khả dụng. - if self.vss_available.load(Ordering::SeqCst) { - sqlx::query("INSERT OR REPLACE INTO sg_vss(rowid, vec) VALUES (?1, ?2)") - .bind(symbol_id as i64) - .bind(encode_vector(vector)) - .execute(&mut *conn) - .await - .map_err(db_err)?; - } - Ok(()) - } - - async fn load_embedding(&self, symbol_id: u64) -> Result>> { - let mut conn = self.pool.acquire().await.map_err(db_err)?; - let data: Option> = - sqlx::query_scalar("SELECT vector FROM sg_embeddings WHERE symbol_id = ?1") - .bind(symbol_id as i64) - .fetch_optional(&mut *conn) - .await - .map_err(db_err)?; - Ok(data.and_then(|b| decode_vector(&b))) - } - - async fn load_all_embeddings(&self) -> Result>> { - let mut conn = self.pool.acquire().await.map_err(db_err)?; - let rows: Vec<(i64, Vec)> = - sqlx::query_as("SELECT symbol_id, vector FROM sg_embeddings ORDER BY symbol_id") - .fetch_all(&mut *conn) - .await - .map_err(db_err)?; - Ok(rows - .into_iter() - .filter_map(|(id, b)| decode_vector(&b).map(|v| (id as u64, v))) - .collect()) - } - - async fn clear_embeddings(&mut self) -> Result<()> { - let mut conn = self.pool.acquire().await.map_err(db_err)?; - sqlx::query("DELETE FROM sg_embeddings") - .execute(&mut *conn) - .await - .map_err(db_err)?; - if self.vss_available.load(Ordering::SeqCst) { - sqlx::query("DELETE FROM sg_vss") - .execute(&mut *conn) - .await - .map_err(db_err)?; - } - Ok(()) - } - - async fn knn(&self, query_vec: &[f32], k: usize) -> Result>> { - if !self.vss_available.load(Ordering::SeqCst) { - return Ok(None); - } - let mut conn = self.pool.acquire().await.map_err(db_err)?; - // `vss_search(vec, )` trả các row gần nhất + `distance` (nhỏ = gần). - // Đảo dấu distance → `sim` (lớn = gần) đồng nhất với `VectorIndex::knn`. - let rows: Vec<(i64, f64)> = sqlx::query_as( - "SELECT rowid, distance FROM sg_vss - WHERE vss_search(vec, ?) ORDER BY distance LIMIT ?", - ) - .bind(encode_vector(query_vec)) - .bind(k as i64) - .fetch_all(&mut *conn) - .await - .map_err(db_err)?; - Ok(Some( - rows.into_iter() - .map(|(id, dist)| (id as u64, -dist as f32)) - .collect(), - )) - } - async fn all_chains(&self) -> Result)>> { let mut conn = self.pool.acquire().await.map_err(db_err)?; let rows: Vec<(i64, Vec)> = @@ -817,129 +862,96 @@ impl Storage for SqliteStorage { Ok(()) } - async fn set_root(&mut self, shard: usize, root: usize) -> Result<()> { - let mut conn = self.pool.acquire().await.map_err(db_err)?; - sqlx::query( - "INSERT INTO rt_roots (shard, root) VALUES (?1, ?2) - ON CONFLICT(shard) DO UPDATE SET root = excluded.root", - ) - .bind(shard as i64) - .bind(root as i64) - .execute(&mut *conn) - .await - .map_err(db_err)?; - Ok(()) - } - - async fn get_root(&self, shard: usize) -> Result { - let mut conn = self.pool.acquire().await.map_err(db_err)?; - let root: Option = sqlx::query_scalar("SELECT root FROM rt_roots WHERE shard = ?1") - .bind(shard as i64) - .fetch_optional(&mut *conn) - .await - .map_err(db_err)?; - Ok(root.unwrap_or(EMPTY as i64) as usize) - } - - async fn set_meta(&mut self, record: usize, meta: &[u8]) -> Result<()> { + async fn save_embedding(&mut self, symbol_id: u64, vector: &[f32]) -> Result<()> { let mut conn = self.pool.acquire().await.map_err(db_err)?; sqlx::query( - "INSERT INTO rt_meta (record, meta) VALUES (?1, ?2) - ON CONFLICT(record) DO UPDATE SET meta = excluded.meta", + "INSERT INTO sg_embeddings (symbol_id, vector) VALUES (?1, ?2) + ON CONFLICT(symbol_id) DO UPDATE SET vector = excluded.vector", ) - .bind(record as i64) - .bind(meta) + .bind(symbol_id as i64) + .bind(encode_vector(vector)) .execute(&mut *conn) .await .map_err(db_err)?; + // Mirror vào `vss0` (HNSW ANN) nếu extension khả dụng. + if self.vss_available.load(Ordering::SeqCst) { + sqlx::query("INSERT OR REPLACE INTO sg_vss(rowid, vec) VALUES (?1, ?2)") + .bind(symbol_id as i64) + .bind(encode_vector(vector)) + .execute(&mut *conn) + .await + .map_err(db_err)?; + } Ok(()) } - async fn get_meta(&self, record: usize) -> Result>> { + async fn load_embedding(&self, symbol_id: u64) -> Result>> { let mut conn = self.pool.acquire().await.map_err(db_err)?; - let meta: Option> = - sqlx::query_scalar("SELECT meta FROM rt_meta WHERE record = ?1") - .bind(record as i64) + let data: Option> = + sqlx::query_scalar("SELECT vector FROM sg_embeddings WHERE symbol_id = ?1") + .bind(symbol_id as i64) .fetch_optional(&mut *conn) .await .map_err(db_err)?; - Ok(meta) + Ok(data.and_then(|b| decode_vector(&b))) } - async fn set_key_len(&mut self, record: usize, len: usize) -> Result<()> { + async fn load_all_embeddings(&self) -> Result>> { let mut conn = self.pool.acquire().await.map_err(db_err)?; - sqlx::query( - "INSERT INTO rt_keylen (record, len) VALUES (?1, ?2) - ON CONFLICT(record) DO UPDATE SET len = excluded.len", - ) - .bind(record as i64) - .bind(len as i64) - .execute(&mut *conn) - .await - .map_err(db_err)?; - Ok(()) + let rows: Vec<(i64, Vec)> = + sqlx::query_as("SELECT symbol_id, vector FROM sg_embeddings ORDER BY symbol_id") + .fetch_all(&mut *conn) + .await + .map_err(db_err)?; + Ok(rows + .into_iter() + .filter_map(|(id, b)| decode_vector(&b).map(|v| (id as u64, v))) + .collect()) } - async fn get_key_len(&self, record: usize) -> Result> { + async fn clear_embeddings(&mut self) -> Result<()> { let mut conn = self.pool.acquire().await.map_err(db_err)?; - let len: Option = sqlx::query_scalar("SELECT len FROM rt_keylen WHERE record = ?1") - .bind(record as i64) - .fetch_optional(&mut *conn) + sqlx::query("DELETE FROM sg_embeddings") + .execute(&mut *conn) .await .map_err(db_err)?; - Ok(len.map(|x| x as usize)) - } - - async fn add_shortcut_node(&mut self, shard: usize, elem: &[u8], node_id: usize) -> Result<()> { - let mut conn = self.pool.acquire().await.map_err(db_err)?; - sqlx::query( - "INSERT INTO rt_shortcuts (shard, elem, node_id) VALUES (?1, ?2, ?3) - ON CONFLICT DO NOTHING", - ) - .bind(shard as i64) - .bind(elem) - .bind(node_id as i64) - .execute(&mut *conn) - .await - .map_err(db_err)?; + if self.vss_available.load(Ordering::SeqCst) { + sqlx::query("DELETE FROM sg_vss") + .execute(&mut *conn) + .await + .map_err(db_err)?; + } Ok(()) } - async fn get_shortcut_nodes(&self, shard: usize, elem: &[u8]) -> Result> { + async fn knn(&self, query_vec: &[f32], k: usize) -> Result>> { + if !self.vss_available.load(Ordering::SeqCst) { + return Ok(None); + } let mut conn = self.pool.acquire().await.map_err(db_err)?; - let rows = sqlx::query( - "SELECT node_id FROM rt_shortcuts WHERE shard = ?1 AND elem = ?2 ORDER BY node_id", + // `vss_search(vec, )` trả các row gần nhất + `distance` (nhỏ = gần). + // Đảo dấu distance → `sim` (lớn = gần) đồng nhất với `VectorIndex::knn`. + let rows: Vec<(i64, f64)> = sqlx::query_as( + "SELECT rowid, distance FROM sg_vss + WHERE vss_search(vec, ?) ORDER BY distance LIMIT ?", ) - .bind(shard as i64) - .bind(elem) + .bind(encode_vector(query_vec)) + .bind(k as i64) .fetch_all(&mut *conn) .await .map_err(db_err)?; - let mut out = Vec::with_capacity(rows.len()); - for r in &rows { - let c: i64 = r.try_get(0).map_err(db_err)?; - out.push(c as usize); - } - Ok(out) + Ok(Some( + rows.into_iter() + .map(|(id, dist)| (id as u64, -dist as f32)) + .collect(), + )) } +} - async fn clear_shortcuts(&mut self) -> Result<()> { - let mut conn = self.pool.acquire().await.map_err(db_err)?; - sqlx::query("DELETE FROM rt_shortcuts") - .execute(&mut *conn) - .await - .map_err(db_err)?; - Ok(()) - } +use super::Storage; - fn new_tx(&self) -> Box { - Box::new(SqliteTx { - pool: self.pool.clone(), - nodes: Vec::new(), - ops: Vec::new(), - }) - } -} +#[async_trait] +impl Storage for SqliteStorage {} // ==================== SqliteTx ====================