use petgraph::Direction;
use petgraph::graph::{DiGraph, NodeIndex};
use petgraph::visit::EdgeRef;
use std::collections::HashMap;
use std::path::PathBuf;
use crate::error::SaraError;
use crate::model::{Item, ItemId, ItemType, RelationshipType};
#[derive(Debug)]
pub struct KnowledgeGraph {
graph: DiGraph<Item, RelationshipType>,
index: HashMap<ItemId, NodeIndex>,
}
impl KnowledgeGraph {
pub fn new() -> Self {
Self {
graph: DiGraph::new(),
index: HashMap::new(),
}
}
pub fn item_count(&self) -> usize {
self.graph.node_count()
}
pub fn relationship_count(&self) -> usize {
self.graph.edge_count()
}
fn add_item(&mut self, item: Item) -> NodeIndex {
let id = item.id.clone();
let idx = self.graph.add_node(item);
self.index.insert(id, idx);
idx
}
fn add_relationship(&mut self, from: &ItemId, to: &ItemId, rel_type: RelationshipType) {
if let (Some(from_idx), Some(to_idx)) = (self.index.get(from), self.index.get(to)) {
self.graph.add_edge(*from_idx, *to_idx, rel_type);
}
}
pub fn get(&self, id: &ItemId) -> Option<&Item> {
let idx = self.index.get(id)?;
self.graph.node_weight(*idx)
}
pub fn get_mut(&mut self, id: &ItemId) -> Option<&mut Item> {
let idx = self.index.get(id)?;
self.graph.node_weight_mut(*idx)
}
pub fn contains(&self, id: &ItemId) -> bool {
self.index.contains_key(id)
}
pub fn items(&self) -> impl Iterator<Item = &Item> {
self.graph.node_weights()
}
pub fn item_ids(&self) -> impl Iterator<Item = &ItemId> {
self.index.keys()
}
pub fn items_by_type(&self, item_type: ItemType) -> Vec<&Item> {
self.graph
.node_weights()
.filter(|item| item.item_type == item_type)
.collect()
}
pub fn count_by_type(&self) -> HashMap<ItemType, usize> {
let mut counts = HashMap::new();
for item in self.graph.node_weights() {
*counts.entry(item.item_type).or_insert(0) += 1;
}
counts
}
pub fn parents(&self, id: &ItemId) -> Vec<&Item> {
let Some(idx) = self.index.get(id) else {
return Vec::new();
};
self.graph
.edges_directed(*idx, Direction::Outgoing)
.filter(|edge| edge.weight().is_upstream())
.filter_map(|edge| self.graph.node_weight(edge.target()))
.collect()
}
pub fn children(&self, id: &ItemId) -> Vec<&Item> {
let Some(idx) = self.index.get(id) else {
return Vec::new();
};
self.graph
.edges_directed(*idx, Direction::Incoming)
.filter(|edge| edge.weight().is_upstream())
.filter_map(|edge| self.graph.node_weight(edge.source()))
.collect()
}
pub fn orphans(&self) -> Vec<&Item> {
self.graph
.node_weights()
.filter(|item| {
if item.item_type.is_root() {
return false;
}
item.upstream.is_empty()
})
.collect()
}
pub fn inner(&self) -> &DiGraph<Item, RelationshipType> {
&self.graph
}
pub fn inner_mut(&mut self) -> &mut DiGraph<Item, RelationshipType> {
&mut self.graph
}
pub fn node_index(&self, id: &ItemId) -> Option<NodeIndex> {
self.index.get(id).copied()
}
pub fn has_cycles(&self) -> bool {
petgraph::algo::is_cyclic_directed(&self.graph)
}
pub fn relationships(&self) -> Vec<(ItemId, ItemId, RelationshipType)> {
self.graph
.edge_references()
.filter_map(|edge| {
let from = self.graph.node_weight(edge.source())?;
let to = self.graph.node_weight(edge.target())?;
Some((from.id.clone(), to.id.clone(), *edge.weight()))
})
.collect()
}
}
impl Default for KnowledgeGraph {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Default)]
pub struct KnowledgeGraphBuilder {
items: Vec<Item>,
repositories: Vec<PathBuf>,
}
impl KnowledgeGraphBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn add_repository(mut self, path: impl Into<PathBuf>) -> Self {
self.repositories.push(path.into());
self
}
pub fn add_item(mut self, item: Item) -> Self {
self.items.push(item);
self
}
pub fn add_items(mut self, items: impl IntoIterator<Item = Item>) -> Self {
self.items.extend(items);
self
}
pub fn build(self) -> Result<KnowledgeGraph, SaraError> {
let mut graph = KnowledgeGraph::new();
for item in &self.items {
graph.add_item(item.clone());
}
for item in &self.items {
self.add_relationships_for_item(&mut graph, item);
}
Ok(graph)
}
fn add_relationships_for_item(&self, graph: &mut KnowledgeGraph, item: &Item) {
for target_id in &item.upstream.refines {
graph.add_relationship(&item.id, target_id, RelationshipType::Refines);
}
for target_id in &item.upstream.derives_from {
graph.add_relationship(&item.id, target_id, RelationshipType::DerivesFrom);
}
for target_id in &item.upstream.satisfies {
graph.add_relationship(&item.id, target_id, RelationshipType::Satisfies);
}
for target_id in &item.upstream.justifies {
graph.add_relationship(&item.id, target_id, RelationshipType::Justifies);
graph.add_relationship(target_id, &item.id, RelationshipType::IsJustifiedBy);
}
for target_id in &item.downstream.is_refined_by {
graph.add_relationship(&item.id, target_id, RelationshipType::IsRefinedBy);
graph.add_relationship(target_id, &item.id, RelationshipType::Refines);
}
for target_id in &item.downstream.derives {
graph.add_relationship(&item.id, target_id, RelationshipType::Derives);
graph.add_relationship(target_id, &item.id, RelationshipType::DerivesFrom);
}
for target_id in &item.downstream.is_satisfied_by {
graph.add_relationship(&item.id, target_id, RelationshipType::IsSatisfiedBy);
graph.add_relationship(target_id, &item.id, RelationshipType::Satisfies);
}
for adr_id in &item.downstream.justified_by {
graph.add_relationship(&item.id, adr_id, RelationshipType::IsJustifiedBy);
graph.add_relationship(adr_id, &item.id, RelationshipType::Justifies);
}
for target_id in item.attributes.depends_on() {
graph.add_relationship(&item.id, target_id, RelationshipType::DependsOn);
}
for target_id in item.attributes.supersedes() {
graph.add_relationship(&item.id, target_id, RelationshipType::Supersedes);
graph.add_relationship(target_id, &item.id, RelationshipType::IsSupersededBy);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::UpstreamRefs;
use crate::test_utils::{create_test_adr, create_test_item, create_test_item_with_upstream};
#[test]
fn test_add_and_get_item() {
let graph = KnowledgeGraphBuilder::new()
.add_item(create_test_item("SOL-001", ItemType::Solution))
.build()
.unwrap();
let id = ItemId::new_unchecked("SOL-001");
assert!(graph.contains(&id));
assert_eq!(graph.get(&id).unwrap().name, "Test SOL-001");
}
#[test]
fn test_items_by_type() {
let graph = KnowledgeGraphBuilder::new()
.add_item(create_test_item("SOL-001", ItemType::Solution))
.add_item(create_test_item("UC-001", ItemType::UseCase))
.add_item(create_test_item("UC-002", ItemType::UseCase))
.build()
.unwrap();
let solutions = graph.items_by_type(ItemType::Solution);
assert_eq!(solutions.len(), 1);
let use_cases = graph.items_by_type(ItemType::UseCase);
assert_eq!(use_cases.len(), 2);
}
#[test]
fn test_item_count() {
let graph = KnowledgeGraphBuilder::new().build().unwrap();
assert_eq!(graph.item_count(), 0);
let graph = KnowledgeGraphBuilder::new()
.add_item(create_test_item("SOL-001", ItemType::Solution))
.build()
.unwrap();
assert_eq!(graph.item_count(), 1);
let graph = KnowledgeGraphBuilder::new()
.add_item(create_test_item("SOL-001", ItemType::Solution))
.add_item(create_test_item("UC-001", ItemType::UseCase))
.build()
.unwrap();
assert_eq!(graph.item_count(), 2);
}
#[test]
fn test_build_simple_graph() {
let graph = KnowledgeGraphBuilder::new()
.add_item(create_test_item("SOL-001", ItemType::Solution))
.build()
.unwrap();
assert_eq!(graph.item_count(), 1);
}
#[test]
fn test_build_graph_with_relationships() {
let sol = create_test_item("SOL-001", ItemType::Solution);
let uc = create_test_item_with_upstream(
"UC-001",
ItemType::UseCase,
UpstreamRefs {
refines: vec![ItemId::new_unchecked("SOL-001")],
..Default::default()
},
);
let graph = KnowledgeGraphBuilder::new()
.add_item(sol)
.add_item(uc)
.build()
.unwrap();
assert_eq!(graph.item_count(), 2);
assert_eq!(graph.relationship_count(), 1);
}
#[test]
fn test_adr_justifies_relationship() {
let sysarch = create_test_item("SYSARCH-001", ItemType::SystemArchitecture);
let adr = create_test_adr("ADR-001", &["SYSARCH-001"], &[]);
let graph = KnowledgeGraphBuilder::new()
.add_item(sysarch)
.add_item(adr)
.build()
.unwrap();
assert_eq!(graph.item_count(), 2);
assert_eq!(graph.relationship_count(), 2);
}
#[test]
fn test_adr_supersession_relationship() {
let adr_old = create_test_adr("ADR-001", &[], &[]);
let adr_new = create_test_adr("ADR-002", &[], &["ADR-001"]);
let graph = KnowledgeGraphBuilder::new()
.add_item(adr_old)
.add_item(adr_new)
.build()
.unwrap();
assert_eq!(graph.item_count(), 2);
assert_eq!(graph.relationship_count(), 2);
}
}