use serde::{Deserialize, Deserializer, Serialize, Serializer};
use crate::Provenance;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NodeKind {
Fn,
Struct,
Enum,
Trait,
Module,
File,
Adr,
AdrSection,
Blueprint,
Doc,
Marker,
Other(String),
}
impl NodeKind {
#[must_use]
pub fn as_str(&self) -> &str {
match self {
Self::Fn => "fn",
Self::Struct => "struct",
Self::Enum => "enum",
Self::Trait => "trait",
Self::Module => "module",
Self::File => "file",
Self::Adr => "adr",
Self::AdrSection => "adr_section",
Self::Blueprint => "blueprint",
Self::Doc => "doc",
Self::Marker => "marker",
Self::Other(s) => s,
}
}
#[must_use]
pub fn from_token(s: &str) -> Self {
match s {
"fn" => Self::Fn,
"struct" => Self::Struct,
"enum" => Self::Enum,
"trait" => Self::Trait,
"module" => Self::Module,
"file" => Self::File,
"adr" => Self::Adr,
"adr_section" => Self::AdrSection,
"blueprint" => Self::Blueprint,
"doc" => Self::Doc,
"marker" => Self::Marker,
other => Self::Other(other.to_owned()),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EdgeKind {
Calls,
Imports,
Defines,
Contains,
References,
Supersedes,
AuthoredBy,
InferredFrom,
Related,
Other(String),
}
impl EdgeKind {
#[must_use]
pub fn as_str(&self) -> &str {
match self {
Self::Calls => "calls",
Self::Imports => "imports",
Self::Defines => "defines",
Self::Contains => "contains",
Self::References => "references",
Self::Supersedes => "supersedes",
Self::AuthoredBy => "authored_by",
Self::InferredFrom => "inferred_from",
Self::Related => "related",
Self::Other(s) => s,
}
}
#[must_use]
pub fn from_token(s: &str) -> Self {
match s {
"calls" => Self::Calls,
"imports" => Self::Imports,
"defines" => Self::Defines,
"contains" => Self::Contains,
"references" => Self::References,
"supersedes" => Self::Supersedes,
"authored_by" => Self::AuthoredBy,
"inferred_from" => Self::InferredFrom,
"related" => Self::Related,
other => Self::Other(other.to_owned()),
}
}
}
impl Serialize for NodeKind {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.serialize_str(self.as_str())
}
}
impl<'de> Deserialize<'de> for NodeKind {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let s = String::deserialize(deserializer)?;
Ok(Self::from_token(&s))
}
}
impl Serialize for EdgeKind {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
serializer.serialize_str(self.as_str())
}
}
impl<'de> Deserialize<'de> for EdgeKind {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let s = String::deserialize(deserializer)?;
Ok(Self::from_token(&s))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct Span {
pub start: u32,
pub end: u32,
}
impl Span {
#[must_use]
pub fn new(start: u32, end: u32) -> Self {
Self { start, end }
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Node {
pub key: String,
pub kind: NodeKind,
pub name: String,
pub path: Option<String>,
pub lang: Option<String>,
pub blob_hash: Option<String>,
pub span: Option<Span>,
#[serde(default)]
pub provenance: Provenance,
#[serde(default)]
pub meta: serde_json::Value,
}
impl Node {
#[must_use]
pub fn new(key: impl Into<String>, kind: NodeKind, name: impl Into<String>) -> Self {
Self {
key: key.into(),
kind,
name: name.into(),
path: None,
lang: None,
blob_hash: None,
span: None,
provenance: Provenance::Derived,
meta: serde_json::Value::Null,
}
}
#[must_use]
pub fn with_provenance(mut self, provenance: Provenance) -> Self {
self.provenance = provenance;
self
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Edge {
pub src: String,
pub dst: String,
pub kind: EdgeKind,
pub provenance: Provenance,
pub confidence: Option<f64>,
pub src_ref: Option<String>,
}
impl Edge {
#[must_use]
pub fn derived(src: impl Into<String>, dst: impl Into<String>, kind: EdgeKind) -> Self {
Self {
src: src.into(),
dst: dst.into(),
kind,
provenance: Provenance::Derived,
confidence: None,
src_ref: None,
}
}
#[must_use]
pub fn authored(src: impl Into<String>, dst: impl Into<String>, kind: EdgeKind) -> Self {
Self {
src: src.into(),
dst: dst.into(),
kind,
provenance: Provenance::Authored,
confidence: None,
src_ref: None,
}
}
#[must_use]
pub fn inferred(
src: impl Into<String>,
dst: impl Into<String>,
kind: EdgeKind,
confidence: f64,
) -> Self {
Self {
src: src.into(),
dst: dst.into(),
kind,
provenance: Provenance::Inferred,
confidence: Some(confidence),
src_ref: None,
}
}
#[must_use]
pub fn is_valid(&self) -> bool {
let inferred = matches!(self.provenance, Provenance::Inferred);
match self.confidence {
Some(c) => inferred && (0.0..=1.0).contains(&c),
None => !inferred,
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct FactSet {
pub nodes: Vec<Node>,
pub edges: Vec<Edge>,
}
impl FactSet {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_node(mut self, node: Node) -> Self {
self.nodes.push(node);
self
}
#[must_use]
pub fn with_edge(mut self, edge: Edge) -> Self {
self.edges.push(edge);
self
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.nodes.is_empty() && self.edges.is_empty()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Direction {
Outgoing,
Incoming,
Both,
}
#[cfg(test)]
mod tests {
use super::{Edge, EdgeKind, FactSet, Node, NodeKind};
use crate::Provenance;
#[test]
fn node_kind_tokens_round_trip() {
let kinds = [
NodeKind::Fn,
NodeKind::Struct,
NodeKind::Enum,
NodeKind::Trait,
NodeKind::Module,
NodeKind::File,
NodeKind::Adr,
NodeKind::AdrSection,
NodeKind::Blueprint,
NodeKind::Doc,
NodeKind::Marker,
NodeKind::Other("weird".to_owned()),
];
for k in kinds {
assert_eq!(NodeKind::from_token(k.as_str()), k);
}
}
#[test]
fn edge_kind_tokens_round_trip() {
let kinds = [
EdgeKind::Calls,
EdgeKind::Imports,
EdgeKind::Defines,
EdgeKind::Contains,
EdgeKind::References,
EdgeKind::Supersedes,
EdgeKind::AuthoredBy,
EdgeKind::InferredFrom,
EdgeKind::Other("weird".to_owned()),
];
for k in kinds {
assert_eq!(EdgeKind::from_token(k.as_str()), k);
}
}
#[test]
fn kinds_serialize_as_bare_tokens() {
assert_eq!(
serde_json::to_string(&NodeKind::AdrSection).unwrap(),
"\"adr_section\""
);
assert_eq!(
serde_json::to_string(&EdgeKind::AuthoredBy).unwrap(),
"\"authored_by\""
);
let k: NodeKind = serde_json::from_str("\"struct\"").unwrap();
assert_eq!(k, NodeKind::Struct);
}
#[test]
fn edge_validity_tracks_provenance() {
assert!(Edge::derived("a", "b", EdgeKind::Calls).is_valid());
assert!(Edge::authored("a", "b", EdgeKind::AuthoredBy).is_valid());
assert!(Edge::inferred("a", "b", EdgeKind::References, 0.5).is_valid());
assert!(Edge::inferred("a", "b", EdgeKind::References, 0.0).is_valid());
assert!(Edge::inferred("a", "b", EdgeKind::References, 1.0).is_valid());
let inferred = Edge::inferred("a", "b", EdgeKind::References, 0.5);
let bad = Edge {
provenance: Provenance::Derived,
confidence: Some(0.9),
..Edge::derived("a", "b", EdgeKind::Calls)
};
assert!(!bad.is_valid());
assert!(
!Edge {
confidence: None,
..inferred.clone()
}
.is_valid()
);
for c in [-0.1, 1.1, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
assert!(
!Edge {
confidence: Some(c),
..inferred.clone()
}
.is_valid(),
"confidence {c} should be rejected"
);
}
}
#[test]
fn factset_builders() {
let fs = FactSet::new()
.with_node(Node::new("a", NodeKind::Fn, "a"))
.with_edge(Edge::derived("a", "a", EdgeKind::Calls));
assert_eq!(fs.nodes.len(), 1);
assert_eq!(fs.edges.len(), 1);
assert!(!fs.is_empty());
assert!(FactSet::new().is_empty());
}
#[test]
fn node_provenance_defaults_and_builder() {
let legacy = r#"{"key":"k","kind":"fn","name":"n","path":null,"lang":null,"blob_hash":null,"span":null,"meta":null}"#;
let node: Node = serde_json::from_str(legacy).expect("legacy node deserializes");
assert_eq!(node.provenance, Provenance::Derived);
assert_eq!(
Node::new("k", NodeKind::Fn, "n").provenance,
Provenance::Derived
);
assert_eq!(
Node::new("k", NodeKind::Adr, "n")
.with_provenance(Provenance::Authored)
.provenance,
Provenance::Authored
);
}
}