impl TursoVectorDB {
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "path_exists")]
pub async fn new_local<P: AsRef<Path>>(_path: P) -> Result<Self, String> {
let config = VectorStoreConfig {
dimension: 1536,
..Default::default()
};
Ok(Self {
store: RwLock::new(VectorStore::new(config)),
file_index: RwLock::new(HashMap::new()),
metadata: RwLock::new(HashMap::new()),
next_id: RwLock::new(1),
})
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub async fn insert(&self, entry: &EmbeddingEntry) -> Result<i64, String> {
let embedding_dim = entry.embedding.len();
{
let store = self.store.read().map_err(|e| format!("Lock error: {e}"))?;
if store.config().dimension != embedding_dim {
drop(store);
let mut store = self.store.write().map_err(|e| format!("Lock error: {e}"))?;
*store = VectorStore::new(VectorStoreConfig {
dimension: embedding_dim,
..Default::default()
});
}
}
let id = {
let mut next_id = self
.next_id
.write()
.map_err(|e| format!("Lock error: {e}"))?;
let id = *next_id;
*next_id += 1;
id
};
let doc_id = DocumentId::new();
let content = format!(
"{}:{}:{}",
entry.file_path, entry.chunk_name, entry.chunk_type
);
let mut chunk = Chunk::new(doc_id, content, entry.start_line, entry.end_line);
chunk.set_embedding(entry.embedding.clone());
let chunk_id = chunk.id;
{
let mut store = self.store.write().map_err(|e| format!("Lock error: {e}"))?;
store
.insert(chunk)
.map_err(|e| format!("Insert failed: {e}"))?;
}
{
let mut file_index = self
.file_index
.write()
.map_err(|e| format!("Lock error: {e}"))?;
file_index
.entry(entry.file_path.clone())
.or_default()
.push(chunk_id);
}
{
let mut metadata = self
.metadata
.write()
.map_err(|e| format!("Lock error: {e}"))?;
metadata.insert(
chunk_id,
EmbeddingMetadata {
id,
file_path: entry.file_path.clone(),
chunk_name: entry.chunk_name.clone(),
chunk_type: entry.chunk_type.clone(),
language: entry.language.clone(),
start_line: entry.start_line,
end_line: entry.end_line,
content_checksum: entry.content_checksum.clone(),
model: entry.model.clone(),
},
);
}
Ok(id)
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub async fn batch_insert(&self, entries: &[EmbeddingEntry]) -> Result<Vec<i64>, String> {
let mut ids = Vec::new();
for entry in entries {
let id = self.insert(entry).await?;
ids.push(id);
}
Ok(ids)
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub async fn query_by_file(&self, file_path: &str) -> Result<Vec<SearchResult>, String> {
let file_index = self
.file_index
.read()
.map_err(|e| format!("Lock error: {e}"))?;
let metadata_map = self
.metadata
.read()
.map_err(|e| format!("Lock error: {e}"))?;
let store = self.store.read().map_err(|e| format!("Lock error: {e}"))?;
let chunk_ids = match file_index.get(file_path) {
Some(ids) => ids.clone(),
None => return Ok(Vec::new()),
};
let mut results = Vec::new();
for chunk_id in chunk_ids {
if let Some(chunk) = store.get(chunk_id) {
if let Some(meta) = metadata_map.get(&chunk_id) {
results.push(SearchResult {
id: meta.id,
file_path: meta.file_path.clone(),
chunk_name: meta.chunk_name.clone(),
chunk_type: meta.chunk_type.clone(),
language: meta.language.clone(),
start_line: meta.start_line,
end_line: meta.end_line,
similarity: 1.0, embedding: chunk.embedding.clone().unwrap_or_default(),
});
}
}
}
Ok(results)
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub async fn query_by_language(&self, language: &str) -> Result<Vec<SearchResult>, String> {
let metadata_map = self
.metadata
.read()
.map_err(|e| format!("Lock error: {e}"))?;
let store = self.store.read().map_err(|e| format!("Lock error: {e}"))?;
let mut results = Vec::new();
for (chunk_id, meta) in metadata_map.iter() {
if meta.language == language {
if let Some(chunk) = store.get(*chunk_id) {
results.push(SearchResult {
id: meta.id,
file_path: meta.file_path.clone(),
chunk_name: meta.chunk_name.clone(),
chunk_type: meta.chunk_type.clone(),
language: meta.language.clone(),
start_line: meta.start_line,
end_line: meta.end_line,
similarity: 1.0,
embedding: chunk.embedding.clone().unwrap_or_default(),
});
}
}
}
Ok(results)
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub async fn similarity_search(
&self,
query: &[f32],
limit: usize,
) -> Result<Vec<SearchResult>, String> {
let store = self.store.read().map_err(|e| format!("Lock error: {e}"))?;
let metadata_map = self
.metadata
.read()
.map_err(|e| format!("Lock error: {e}"))?;
let search_results = store
.search(query, limit)
.map_err(|e| format!("Search failed: {e}"))?;
let mut results = Vec::new();
for (chunk_id, score) in search_results {
if let Some(meta) = metadata_map.get(&chunk_id) {
if let Some(chunk) = store.get(chunk_id) {
results.push(SearchResult {
id: meta.id,
file_path: meta.file_path.clone(),
chunk_name: meta.chunk_name.clone(),
chunk_type: meta.chunk_type.clone(),
language: meta.language.clone(),
start_line: meta.start_line,
end_line: meta.end_line,
similarity: score as f64,
embedding: chunk.embedding.clone().unwrap_or_default(),
});
}
}
}
Ok(results)
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub async fn delete_by_file(&self, file_path: &str) -> Result<usize, String> {
let chunk_ids = {
let mut file_index = self
.file_index
.write()
.map_err(|e| format!("Lock error: {e}"))?;
file_index.remove(file_path).unwrap_or_default()
};
let count = chunk_ids.len();
{
let mut store = self.store.write().map_err(|e| format!("Lock error: {e}"))?;
let mut metadata = self
.metadata
.write()
.map_err(|e| format!("Lock error: {e}"))?;
for chunk_id in chunk_ids {
store.remove(chunk_id);
metadata.remove(&chunk_id);
}
}
Ok(count)
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub async fn delete_file_entries(&self, file_path: &str) -> Result<usize, String> {
self.delete_by_file(file_path).await
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub async fn get_entry(
&self,
file_path: &str,
chunk_name: &str,
) -> Result<Option<SearchResult>, String> {
let results = self.query_by_file(file_path).await?;
Ok(results.into_iter().find(|r| r.chunk_name == chunk_name))
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub async fn get_stats(&self) -> Result<DbStats, String> {
let store = self.store.read().map_err(|e| format!("Lock error: {e}"))?;
let file_index = self
.file_index
.read()
.map_err(|e| format!("Lock error: {e}"))?;
Ok(DbStats {
total_entries: store.len(),
unique_files: file_index.len(),
})
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn cosine_similarity(v1: &[f32], v2: &[f32]) -> f64 {
if v1.len() != v2.len() {
return 0.0;
}
let dot_product: f32 = v1.iter().zip(v2.iter()).map(|(a, b)| a * b).sum();
let magnitude1: f32 = v1.iter().map(|x| x * x).sum::<f32>().sqrt();
let magnitude2: f32 = v2.iter().map(|x| x * x).sum::<f32>().sqrt();
if magnitude1 == 0.0 || magnitude2 == 0.0 {
return 0.0;
}
(dot_product / (magnitude1 * magnitude2)) as f64
}
#[inline]
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn cosine_similarity_simd(v1: &[f32], v2: &[f32]) -> f64 {
if v1.len() != v2.len() || v1.is_empty() {
return 0.0;
}
let len = v1.len();
let chunks = len / 4;
let remainder = len % 4;
let mut dot0 = 0.0f32;
let mut dot1 = 0.0f32;
let mut dot2 = 0.0f32;
let mut dot3 = 0.0f32;
let mut norm1_0 = 0.0f32;
let mut norm1_1 = 0.0f32;
let mut norm1_2 = 0.0f32;
let mut norm1_3 = 0.0f32;
let mut norm2_0 = 0.0f32;
let mut norm2_1 = 0.0f32;
let mut norm2_2 = 0.0f32;
let mut norm2_3 = 0.0f32;
for i in 0..chunks {
let base = i * 4;
let a0 = v1[base];
let a1 = v1[base + 1];
let a2 = v1[base + 2];
let a3 = v1[base + 3];
let b0 = v2[base];
let b1 = v2[base + 1];
let b2 = v2[base + 2];
let b3 = v2[base + 3];
dot0 += a0 * b0;
dot1 += a1 * b1;
dot2 += a2 * b2;
dot3 += a3 * b3;
norm1_0 += a0 * a0;
norm1_1 += a1 * a1;
norm1_2 += a2 * a2;
norm1_3 += a3 * a3;
norm2_0 += b0 * b0;
norm2_1 += b1 * b1;
norm2_2 += b2 * b2;
norm2_3 += b3 * b3;
}
let remainder_start = chunks * 4;
for i in 0..remainder {
let idx = remainder_start + i;
let a = v1[idx];
let b = v2[idx];
dot0 += a * b;
norm1_0 += a * a;
norm2_0 += b * b;
}
let dot_product = dot0 + dot1 + dot2 + dot3;
let magnitude1 = (norm1_0 + norm1_1 + norm1_2 + norm1_3).sqrt();
let magnitude2 = (norm2_0 + norm2_1 + norm2_2 + norm2_3).sqrt();
if magnitude1 == 0.0 || magnitude2 == 0.0 {
return 0.0;
}
(dot_product / (magnitude1 * magnitude2)) as f64
}
}