use lxdb_core::ids::TokenId;
use lxdb_storage::BinaryDataset;
use crate::{BinaryDatasetExt, BinaryRelationIter, BinaryToken, EngineError, RelationRecordIter};
#[derive(Debug, Clone, Copy)]
pub struct DatasetQuery<'a> {
dataset: &'a BinaryDataset,
}
impl<'a> DatasetQuery<'a> {
pub const fn new(dataset: &'a BinaryDataset) -> Self {
Self { dataset }
}
pub const fn dataset(&self) -> &'a BinaryDataset {
self.dataset
}
pub fn token_by_text(&self, text: &str) -> Result<Option<BinaryToken<'a>>, EngineError> {
for token in self.dataset.resolved_tokens() {
let token = token?;
if token.text() == text {
return Ok(Some(token));
}
}
Ok(None)
}
pub fn token_by_id(&self, token_id: TokenId) -> Result<Option<BinaryToken<'a>>, EngineError> {
for token in self.dataset.resolved_tokens() {
let token = token?;
if token.id() == token_id {
return Ok(Some(token));
}
}
Ok(None)
}
pub fn outgoing(&self, token_id: TokenId) -> Result<RelationRecordIter<'a>, EngineError> {
self.dataset.outgoing(token_id)
}
pub fn resolved_outgoing(
&self,
token_id: TokenId,
) -> Result<BinaryRelationIter<'a>, EngineError> {
let records = self.dataset.outgoing(token_id)?;
Ok(BinaryRelationIter::new(self.dataset, records))
}
pub fn related_to(&self, text: &str) -> Result<Option<BinaryRelationIter<'a>>, EngineError> {
let Some(token) = self.token_by_text(text)? else {
return Ok(None);
};
let relations = self.resolved_outgoing(token.id())?;
Ok(Some(relations))
}
}
#[cfg(test)]
mod tests {
use lxdb_core::ids::TokenId;
use lxdb_format::{AdjacencyRecord, RelationRecord};
use crate::{BinaryDatasetExt, DatasetQuery, EngineError, test_support};
#[test]
fn finds_token_by_exact_text() {
let dataset = test_support::dataset(
&["rust", "language", "compiler"],
&[],
&[AdjacencyRecord::new(0, 0), AdjacencyRecord::new(0, 0), AdjacencyRecord::new(0, 0)],
);
let query = DatasetQuery::new(&dataset);
let token = query
.token_by_text("language")
.expect("token lookup should succeed")
.expect("token should exist");
assert_eq!(token.id().value(), 1);
assert_eq!(token.text(), "language");
}
#[test]
fn returns_none_for_unknown_token_text() {
let dataset = test_support::dataset(&["rust"], &[], &[AdjacencyRecord::new(0, 0)]);
let token = dataset.query().token_by_text("python").expect("token lookup should succeed");
assert!(token.is_none());
}
#[test]
fn finds_token_by_id() {
let dataset = test_support::dataset(
&["rust", "memory"],
&[],
&[AdjacencyRecord::new(0, 0), AdjacencyRecord::new(0, 0)],
);
let token = dataset
.query()
.token_by_id(TokenId::new(1))
.expect("token lookup should succeed")
.expect("token should exist");
assert_eq!(token.text(), "memory");
}
#[test]
fn exposes_outgoing_relations() {
let relation = RelationRecord::new(0, 0, 1, 0.9);
let dataset = test_support::dataset(
&["rust", "language"],
&[relation],
&[AdjacencyRecord::new(0, 1), AdjacencyRecord::new(1, 0)],
);
let mut outgoing =
dataset.query().outgoing(TokenId::new(0)).expect("outgoing lookup should succeed");
let relation =
outgoing.next().expect("relation should exist").expect("relation should decode");
assert_eq!(relation.source(), 0);
assert_eq!(relation.target(), 1);
assert_eq!(outgoing.len(), 0);
}
#[test]
fn rejects_relations_with_missing_target_tokens() {
let relation = RelationRecord::new(7, 0, 99, 1.0);
let dataset = test_support::dataset(&["rust"], &[relation], &[AdjacencyRecord::new(0, 1)]);
let mut relations = dataset
.query()
.related_to("rust")
.expect("related token query should initialize")
.expect("source token should exist");
let error = relations
.next()
.expect("relation should exist")
.expect_err("missing target should fail");
assert!(matches!(
error,
EngineError::MissingRelationTarget { relation_id: 7, token_id: 99 }
));
}
#[test]
fn resolves_empty_relation_set() {
let dataset = test_support::dataset(&["rust"], &[], &[AdjacencyRecord::new(0, 0)]);
let relations = dataset
.query()
.related_to("rust")
.expect("related token query should succeed")
.expect("source token should exist");
assert_eq!(relations.len(), 0);
}
#[test]
fn returns_none_when_related_source_does_not_exist() {
let dataset = test_support::dataset(&["rust"], &[], &[AdjacencyRecord::new(0, 0)]);
let relations =
dataset.query().related_to("python").expect("related token query should succeed");
assert!(relations.is_none());
}
#[test]
fn resolves_related_tokens_by_text() {
let relation_to_language = RelationRecord::new(0, 0, 1, 0.9);
let relation_to_compiler = RelationRecord::new(1, 0, 2, 0.7);
let dataset = test_support::dataset(
&["rust", "language", "compiler"],
&[relation_to_language, relation_to_compiler],
&[AdjacencyRecord::new(0, 2), AdjacencyRecord::new(2, 0), AdjacencyRecord::new(2, 0)],
);
let mut relations = dataset
.query()
.related_to("rust")
.expect("related token query should succeed")
.expect("source token should exist");
assert_eq!(relations.len(), 2);
let first = relations
.next()
.expect("first relation should exist")
.expect("first relation should resolve");
assert_eq!(first.id(), 0);
assert_eq!(first.source().text(), "rust");
assert_eq!(first.target().text(), "language");
assert_eq!(first.weight().to_bits(), 0.9_f32.to_bits(),);
let second = relations
.next()
.expect("second relation should exist")
.expect("second relation should resolve");
assert_eq!(second.id(), 1);
assert_eq!(second.source().text(), "rust");
assert_eq!(second.target().text(), "compiler");
assert_eq!(second.weight().to_bits(), 0.7_f32.to_bits(),);
assert!(relations.next().is_none());
}
}