use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt;
use std::str::FromStr;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum RelationshipDirection {
Outgoing,
Incoming,
Both,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RelationType {
Implements,
Extends,
Configures,
ArchitecturalDependency,
Imports,
Calls,
Uses,
FactoryCreates,
ObserverPattern,
StrategyPattern,
AdapterPattern,
SiblingModule,
ParentModule,
ChildModule,
}
impl RelationType {
pub fn importance_weight(&self) -> f32 {
match self {
Self::Implements | Self::Extends | Self::Configures => 1.0,
Self::ArchitecturalDependency => 0.9,
Self::Imports | Self::Calls | Self::Uses => 0.7,
Self::FactoryCreates
| Self::ObserverPattern
| Self::StrategyPattern
| Self::AdapterPattern => 0.8,
Self::SiblingModule | Self::ParentModule | Self::ChildModule => 0.3,
}
}
pub fn as_str(&self) -> &'static str {
match self {
Self::Implements => "implements",
Self::Extends => "extends",
Self::Configures => "configures",
Self::ArchitecturalDependency => "architectural_dependency",
Self::Imports => "imports",
Self::Calls => "calls",
Self::Uses => "uses",
Self::FactoryCreates => "factory_creates",
Self::ObserverPattern => "observer_pattern",
Self::StrategyPattern => "strategy_pattern",
Self::AdapterPattern => "adapter_pattern",
Self::SiblingModule => "sibling_module",
Self::ParentModule => "parent_module",
Self::ChildModule => "child_module",
}
}
}
impl FromStr for RelationType {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
Ok(match s {
"implements" => Self::Implements,
"extends" => Self::Extends,
"configures" => Self::Configures,
"architectural_dependency" => Self::ArchitecturalDependency,
"imports" => Self::Imports,
"calls" => Self::Calls,
"uses" => Self::Uses,
"factory_creates" => Self::FactoryCreates,
"observer_pattern" => Self::ObserverPattern,
"strategy_pattern" => Self::StrategyPattern,
"adapter_pattern" => Self::AdapterPattern,
"sibling_module" => Self::SiblingModule,
"parent_module" => Self::ParentModule,
"child_module" => Self::ChildModule,
_ => Self::Imports,
})
}
}
impl fmt::Display for RelationType {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.as_str())
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeNode {
pub id: String, pub name: String, pub kind: String, pub path: String, pub description: String, pub symbols: Vec<String>, pub hash: String, pub embedding: Vec<f32>, pub imports: Vec<String>, pub exports: Vec<String>, pub functions: Vec<FunctionInfo>, pub size_lines: u32, pub language: String, }
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionInfo {
pub name: String, pub signature: String, pub start_line: u32, pub end_line: u32, pub calls: Vec<String>, pub called_by: Vec<String>, pub parameters: Vec<String>, pub return_type: Option<String>, }
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeRelationship {
pub source: String, pub target: String, pub relation_type: RelationType, pub description: String, pub confidence: f32, pub weight: f32, }
impl CodeRelationship {
pub fn new(
source: String,
target: String,
relation_type: RelationType,
description: String,
) -> Self {
Self {
source,
target,
relation_type,
description,
confidence: 0.9, weight: 1.0, }
}
pub fn from_legacy(
source: String,
target: String,
relation_type_str: &str,
description: String,
confidence: f32,
weight: f32,
) -> Self {
Self {
source,
target,
relation_type: relation_type_str.parse().unwrap_or(RelationType::Imports),
description,
confidence,
weight,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct CodeGraph {
pub nodes: HashMap<String, CodeNode>,
pub relationships: Vec<CodeRelationship>,
}
#[allow(dead_code)]
#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct BatchRelationshipResult {
pub source_id: String,
pub target_id: String,
pub relation_type: String,
pub description: String,
pub confidence: f32,
pub exists: bool,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_relationship_type_importance_weights() {
assert_eq!(RelationType::Implements.importance_weight(), 1.0);
assert_eq!(RelationType::Extends.importance_weight(), 1.0);
assert_eq!(RelationType::Configures.importance_weight(), 1.0);
assert_eq!(
RelationType::ArchitecturalDependency.importance_weight(),
0.9
);
assert_eq!(RelationType::Imports.importance_weight(), 0.7);
assert_eq!(RelationType::Calls.importance_weight(), 0.7);
assert_eq!(RelationType::Uses.importance_weight(), 0.7);
assert_eq!(RelationType::FactoryCreates.importance_weight(), 0.8);
assert_eq!(RelationType::SiblingModule.importance_weight(), 0.3);
assert_eq!(RelationType::ParentModule.importance_weight(), 0.3);
assert_eq!(RelationType::ChildModule.importance_weight(), 0.3);
}
#[test]
fn test_relationship_type_ordering() {
assert!(
RelationType::Implements.importance_weight()
> RelationType::Imports.importance_weight()
);
assert!(
RelationType::Extends.importance_weight() > RelationType::Calls.importance_weight()
);
assert!(
RelationType::Imports.importance_weight()
> RelationType::SiblingModule.importance_weight()
);
assert!(
RelationType::Calls.importance_weight()
> RelationType::ParentModule.importance_weight()
);
}
#[test]
fn test_relationship_type_from_str() {
assert_eq!(
"implements".parse::<RelationType>().unwrap(),
RelationType::Implements
);
assert_eq!(
"extends".parse::<RelationType>().unwrap(),
RelationType::Extends
);
assert_eq!(
"imports".parse::<RelationType>().unwrap(),
RelationType::Imports
);
assert_eq!(
"calls".parse::<RelationType>().unwrap(),
RelationType::Calls
);
assert_eq!(
"sibling_module".parse::<RelationType>().unwrap(),
RelationType::SiblingModule
);
assert_eq!(
"unknown_type".parse::<RelationType>().unwrap(),
RelationType::Imports
);
assert_eq!("".parse::<RelationType>().unwrap(), RelationType::Imports);
}
#[test]
fn test_relationship_type_as_str() {
assert_eq!(RelationType::Implements.as_str(), "implements");
assert_eq!(RelationType::Extends.as_str(), "extends");
assert_eq!(RelationType::Imports.as_str(), "imports");
assert_eq!(RelationType::Calls.as_str(), "calls");
assert_eq!(RelationType::SiblingModule.as_str(), "sibling_module");
}
#[test]
fn test_relationship_type_roundtrip() {
let types = vec![
RelationType::Implements,
RelationType::Extends,
RelationType::Imports,
RelationType::Calls,
RelationType::Uses,
RelationType::SiblingModule,
];
for rel_type in types {
let as_string = rel_type.as_str();
let parsed = as_string.parse::<RelationType>().unwrap();
assert_eq!(parsed, rel_type, "Roundtrip failed for {:?}", rel_type);
}
}
#[test]
fn test_relationship_type_display() {
assert_eq!(format!("{}", RelationType::Implements), "implements");
assert_eq!(format!("{}", RelationType::Extends), "extends");
assert_eq!(format!("{}", RelationType::Imports), "imports");
}
#[test]
fn test_relationship_type_serialization() {
let rel_type = RelationType::Implements;
let serialized = serde_json::to_string(&rel_type).unwrap();
assert_eq!(serialized, "\"implements\"");
let deserialized: RelationType = serde_json::from_str(&serialized).unwrap();
assert_eq!(deserialized, RelationType::Implements);
}
#[test]
fn test_code_relationship_new() {
let rel = CodeRelationship::new(
"src/main.rs".to_string(),
"src/lib.rs".to_string(),
RelationType::Imports,
"Imports lib module".to_string(),
);
assert_eq!(rel.source, "src/main.rs");
assert_eq!(rel.target, "src/lib.rs");
assert_eq!(rel.relation_type, RelationType::Imports);
assert_eq!(rel.description, "Imports lib module");
assert_eq!(rel.confidence, 0.9); assert_eq!(rel.weight, 1.0); }
#[test]
fn test_code_relationship_from_legacy() {
let rel = CodeRelationship::from_legacy(
"src/main.rs".to_string(),
"src/lib.rs".to_string(),
"implements",
"Implements trait".to_string(),
0.95,
2.5,
);
assert_eq!(rel.source, "src/main.rs");
assert_eq!(rel.target, "src/lib.rs");
assert_eq!(rel.relation_type, RelationType::Implements);
assert_eq!(rel.description, "Implements trait");
assert_eq!(rel.confidence, 0.95);
assert_eq!(rel.weight, 2.5);
}
#[test]
fn test_code_relationship_legacy_unknown_type() {
let rel = CodeRelationship::from_legacy(
"src/a.rs".to_string(),
"src/b.rs".to_string(),
"some_old_type",
"Description".to_string(),
0.8,
1.0,
);
assert_eq!(rel.relation_type, RelationType::Imports);
}
#[test]
fn test_relationship_direction_enum() {
let outgoing = RelationshipDirection::Outgoing;
let incoming = RelationshipDirection::Incoming;
let both = RelationshipDirection::Both;
assert_ne!(outgoing, incoming);
assert_ne!(outgoing, both);
assert_ne!(incoming, both);
}
#[test]
fn test_code_relationship_serialization() {
let rel = CodeRelationship::new(
"src/main.rs".to_string(),
"src/lib.rs".to_string(),
RelationType::Imports,
"Test relationship".to_string(),
);
let json = serde_json::to_string(&rel).unwrap();
let deserialized: CodeRelationship = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.source, rel.source);
assert_eq!(deserialized.target, rel.target);
assert_eq!(deserialized.relation_type, rel.relation_type);
assert_eq!(deserialized.description, rel.description);
}
}