use petgraph::stable_graph::{NodeIndex, StableDiGraph};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};
#[cfg(not(target_arch = "wasm32"))]
use std::time::Instant;
#[cfg(target_arch = "wasm32")]
use web_time::Instant;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum GitChangeKind {
Modified,
Added,
Deleted,
RenamedFrom,
RenamedTo,
Untracked,
}
impl GitChangeKind {
pub fn label(&self) -> &'static str {
match self {
GitChangeKind::Modified => "Modified",
GitChangeKind::Added => "Added",
GitChangeKind::Deleted => "Deleted",
GitChangeKind::RenamedFrom => "Renamed (from)",
GitChangeKind::RenamedTo => "Renamed (to)",
GitChangeKind::Untracked => "Untracked",
}
}
pub fn symbol(&self) -> &'static str {
match self {
GitChangeKind::Modified => "M",
GitChangeKind::Added => "+",
GitChangeKind::Deleted => "-",
GitChangeKind::RenamedFrom => "R←",
GitChangeKind::RenamedTo => "R→",
GitChangeKind::Untracked => "?",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GitFileChange {
pub path: PathBuf,
pub kind: GitChangeKind,
pub staged: bool,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct GitChangeSnapshot {
pub changes: Vec<GitFileChange>,
#[serde(skip)]
pub captured_at: Option<Instant>,
}
impl GitChangeSnapshot {
pub fn new() -> Self {
Self {
changes: Vec::new(),
captured_at: Some(Instant::now()),
}
}
pub fn has_changes(&self, path: &Path) -> bool {
self.changes.iter().any(|c| c.path == path)
}
pub fn get_change(&self, path: &Path) -> Option<&GitFileChange> {
self.changes.iter().find(|c| c.path == path)
}
pub fn changed_paths(&self) -> impl Iterator<Item = &Path> {
self.changes.iter().map(|c| c.path.as_path())
}
pub fn count_by_kind(&self, kind: GitChangeKind) -> usize {
self.changes.iter().filter(|c| c.kind == kind).count()
}
pub fn is_stale(&self, max_age: Duration) -> bool {
match self.captured_at {
Some(at) => at.elapsed() > max_age,
None => true,
}
}
pub fn age(&self) -> Option<Duration> {
self.captured_at.map(|at| at.elapsed())
}
}
#[derive(Debug, Clone)]
pub struct ChangeIndicatorState {
pub phase: f32,
pub speed: f32,
pub enabled: bool,
}
impl Default for ChangeIndicatorState {
fn default() -> Self {
Self {
phase: 0.0,
speed: 1.0,
enabled: true,
}
}
}
impl ChangeIndicatorState {
pub fn tick(&mut self, dt: f32) {
if self.enabled {
self.phase = (self.phase + dt * self.speed) % 1.0;
}
}
pub fn pulse_scale(&self) -> f32 {
let t = self.phase * std::f32::consts::TAU;
1.0 + 0.15 * (t.sin() * 0.5 + 0.5)
}
pub fn ring_alpha(&self) -> f32 {
let t = self.phase * std::f32::consts::TAU;
0.3 + 0.4 * (t.sin() * 0.5 + 0.5)
}
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct NodeId(pub u64);
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct EdgeId(pub u64);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
pub enum GraphNodeKind {
Module,
File,
Directory,
Service,
Test,
#[default]
Other,
}
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct GraphNode {
pub id: NodeId,
pub name: String,
pub kind: GraphNodeKind,
pub metadata: HashMap<String, String>,
}
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct GraphEdge {
pub id: EdgeId,
pub from: NodeId,
pub to: NodeId,
pub relationship: String,
pub metadata: HashMap<String, String>,
}
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct SourceCodeGraph {
pub nodes: Vec<GraphNode>,
pub edges: Vec<GraphEdge>,
pub metadata: HashMap<String, String>,
}
impl SourceCodeGraph {
pub fn empty() -> Self {
Self::default()
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
pub fn edge_count(&self) -> usize {
self.edges.len()
}
pub fn to_petgraph(&self) -> (StableDiGraph<GraphNode, String>, HashMap<NodeId, NodeIndex>) {
let mut graph = StableDiGraph::new();
let mut id_to_index = HashMap::new();
for node in &self.nodes {
let idx = graph.add_node(node.clone());
id_to_index.insert(node.id, idx);
}
for edge in &self.edges {
if let (Some(&from_idx), Some(&to_idx)) =
(id_to_index.get(&edge.from), id_to_index.get(&edge.to))
{
graph.add_edge(from_idx, to_idx, edge.relationship.clone());
}
}
(graph, id_to_index)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ReferenceKind {
Uses,
Imports,
Implements,
Contains,
}
impl std::fmt::Display for ReferenceKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ReferenceKind::Uses => write!(f, "uses"),
ReferenceKind::Imports => write!(f, "imports"),
ReferenceKind::Implements => write!(f, "implements"),
ReferenceKind::Contains => write!(f, "contains"),
}
}
}
#[derive(Debug, Clone)]
pub struct SourceReference {
pub source_path: PathBuf,
pub kind: ReferenceKind,
pub target_route: PathBuf,
}
#[derive(Debug, Default)]
pub struct SourceCodeGraphBuilder {
nodes: Vec<GraphNode>,
edges: Vec<GraphEdge>,
path_to_node: HashMap<PathBuf, NodeId>,
next_node_id: u64,
next_edge_id: u64,
metadata: HashMap<String, String>,
}
impl SourceCodeGraphBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.metadata.insert(key.into(), value.into());
self
}
pub fn add_directory(&mut self, path: &Path) -> NodeId {
if let Some(&id) = self.path_to_node.get(path) {
return id;
}
let id = NodeId(self.next_node_id);
self.next_node_id += 1;
let name = path
.file_name()
.and_then(|n| n.to_str())
.unwrap_or_else(|| path.to_str().unwrap_or("."))
.to_string();
let mut metadata = HashMap::new();
metadata.insert("path".to_string(), path.to_string_lossy().to_string());
self.nodes.push(GraphNode {
id,
name,
kind: GraphNodeKind::Directory,
metadata,
});
self.path_to_node.insert(path.to_path_buf(), id);
id
}
pub fn add_file(&mut self, path: &Path, relative_path: &str) -> NodeId {
if let Some(&id) = self.path_to_node.get(path) {
return id;
}
let id = NodeId(self.next_node_id);
self.next_node_id += 1;
let name = path
.file_name()
.and_then(|n| n.to_str())
.unwrap_or_else(|| path.to_str().unwrap_or("unknown"))
.to_string();
let kind = match path.extension().and_then(|e| e.to_str()) {
Some("rs") | Some("py") | Some("js") | Some("ts") | Some("tsx") | Some("jsx")
| Some("go") | Some("java") | Some("c") | Some("cpp") | Some("h") | Some("hpp") => {
if relative_path.contains("test") || name.starts_with("test_") {
GraphNodeKind::Test
} else if name == "mod.rs"
|| name == "__init__.py"
|| name == "index.ts"
|| name == "index.js"
{
GraphNodeKind::Module
} else {
GraphNodeKind::File
}
}
_ => GraphNodeKind::File,
};
let mut metadata = HashMap::new();
metadata.insert("path".to_string(), path.to_string_lossy().to_string());
metadata.insert("relative_path".to_string(), relative_path.to_string());
if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
metadata.insert("extension".to_string(), ext.to_string());
metadata.insert(
"language".to_string(),
extension_to_language(ext).to_string(),
);
}
self.nodes.push(GraphNode {
id,
name,
kind,
metadata,
});
self.path_to_node.insert(path.to_path_buf(), id);
id
}
pub fn add_hierarchy_edge(&mut self, parent_path: &Path, child_path: &Path) {
if let (Some(&parent_id), Some(&child_id)) = (
self.path_to_node.get(parent_path),
self.path_to_node.get(child_path),
) {
if parent_id != child_id {
self.add_edge(parent_id, child_id, ReferenceKind::Contains);
}
}
}
pub fn add_edge(&mut self, from: NodeId, to: NodeId, kind: ReferenceKind) {
let id = EdgeId(self.next_edge_id);
self.next_edge_id += 1;
self.edges.push(GraphEdge {
id,
from,
to,
relationship: kind.to_string(),
metadata: HashMap::new(),
});
}
pub fn get_node_id(&self, path: &Path) -> Option<NodeId> {
self.path_to_node.get(path).copied()
}
pub fn find_node_by_path_suffix(&self, route: &Path) -> Option<NodeId> {
let route_str = route.to_string_lossy();
for (path, &node_id) in &self.path_to_node {
let path_str = path.to_string_lossy();
if path_str.ends_with(route_str.as_ref()) {
return Some(node_id);
}
let normalized_path: String = path_str.trim_start_matches("./").replace('\\', "/");
let normalized_route: String = route_str.trim_start_matches("./").replace('\\', "/");
if normalized_path.ends_with(&normalized_route) {
return Some(node_id);
}
let route_parts: Vec<&str> = normalized_route.split('/').collect();
let path_parts: Vec<&str> = normalized_path.split('/').collect();
if route_parts.len() <= path_parts.len() {
for window in path_parts.windows(route_parts.len()) {
if window == route_parts.as_slice() {
return Some(node_id);
}
}
}
if let (Some(file_name), Some(route_name)) = (
path.file_name().and_then(|n| n.to_str()),
route.file_name().and_then(|n| n.to_str()),
) {
if file_name == route_name {
return Some(node_id);
}
}
}
None
}
pub fn set_node_metadata(
&mut self,
node_id: NodeId,
key: impl Into<String>,
value: impl Into<String>,
) {
if let Some(node) = self.nodes.iter_mut().find(|n| n.id == node_id) {
node.metadata.insert(key.into(), value.into());
}
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
pub fn edge_count(&self) -> usize {
self.edges.len()
}
pub fn build(self) -> SourceCodeGraph {
SourceCodeGraph {
nodes: self.nodes,
edges: self.edges,
metadata: self.metadata,
}
}
}
pub fn detect_rust_references(content: &str, source_path: &Path) -> Vec<SourceReference> {
let mut refs = Vec::new();
for line in content.lines() {
let trimmed = line.trim();
if trimmed.starts_with("pub mod ") || trimmed.starts_with("mod ") {
let mod_part = trimmed
.strip_prefix("pub mod ")
.or_else(|| trimmed.strip_prefix("mod "))
.unwrap_or("")
.split(';')
.next()
.unwrap_or("")
.trim();
if !mod_part.is_empty() && !mod_part.contains('{') {
refs.push(SourceReference {
source_path: source_path.to_path_buf(),
kind: ReferenceKind::Uses,
target_route: PathBuf::from(format!("{}.rs", mod_part)),
});
refs.push(SourceReference {
source_path: source_path.to_path_buf(),
kind: ReferenceKind::Uses,
target_route: PathBuf::from(format!("{}/mod.rs", mod_part)),
});
}
}
if !trimmed.starts_with("use ") {
continue;
}
let use_part = trimmed
.strip_prefix("use ")
.unwrap_or("")
.split(';')
.next()
.unwrap_or("")
.split('{')
.next()
.unwrap_or("")
.trim();
if use_part.is_empty() {
continue;
}
let module_path = use_part
.strip_prefix("crate::")
.or_else(|| use_part.strip_prefix("self::"))
.or_else(|| use_part.strip_prefix("super::"))
.unwrap_or(use_part);
let path_str = module_path
.replace("::", "/")
.trim_end_matches(|c: char| !c.is_alphanumeric() && c != '_')
.to_string();
refs.push(SourceReference {
source_path: source_path.to_path_buf(),
kind: ReferenceKind::Uses,
target_route: PathBuf::from(format!("{}.rs", path_str)),
});
}
refs
}
pub fn detect_python_references(content: &str, source_path: &Path) -> Vec<SourceReference> {
let mut refs = Vec::new();
for line in content.lines() {
let trimmed = line.trim();
if trimmed.starts_with("import ") && !trimmed.starts_with("import(") {
let import_part = trimmed
.strip_prefix("import ")
.unwrap_or("")
.split_whitespace()
.next()
.unwrap_or("")
.split(',')
.next()
.unwrap_or("")
.trim();
if !import_part.is_empty() {
let path_str = import_part.replace('.', "/");
refs.push(SourceReference {
source_path: source_path.to_path_buf(),
kind: ReferenceKind::Imports,
target_route: PathBuf::from(format!("{}.py", path_str)),
});
}
}
if let Some(module_part) = trimmed
.strip_prefix("from ")
.and_then(|s| s.split(" import ").next())
{
let module = module_part.trim();
if !module.is_empty() && module != "." && !module.starts_with("..") {
let path_str = module.replace('.', "/");
refs.push(SourceReference {
source_path: source_path.to_path_buf(),
kind: ReferenceKind::Imports,
target_route: PathBuf::from(format!("{}.py", path_str)),
});
}
}
}
refs
}
pub fn detect_ts_references(content: &str, source_path: &Path) -> Vec<SourceReference> {
let mut refs = Vec::new();
for line in content.lines() {
let trimmed = line.trim();
if trimmed.starts_with("import ") {
if let Some(path_start) = trimmed.find(['\'', '"']) {
let quote_char = trimmed.chars().nth(path_start).unwrap();
let rest = &trimmed[path_start + 1..];
if let Some(path_end) = rest.find(quote_char) {
let import_path = &rest[..path_end];
if import_path.starts_with('.') {
refs.push(SourceReference {
source_path: source_path.to_path_buf(),
kind: ReferenceKind::Imports,
target_route: PathBuf::from(import_path),
});
}
}
}
}
}
refs
}
pub fn detect_lean_references(content: &str, source_path: &Path) -> Vec<SourceReference> {
let mut refs = Vec::new();
for line in content.lines() {
let trimmed = line.trim();
if trimmed.is_empty() || trimmed.starts_with("--") || trimmed.starts_with("/-") {
continue;
}
let import_part = trimmed
.strip_prefix("public import ")
.or_else(|| trimmed.strip_prefix("import "));
if let Some(module_path) = import_part {
let module = module_path.split_whitespace().next().unwrap_or("").trim();
if !module.is_empty() {
let file_path = module.replace('.', "/");
refs.push(SourceReference {
source_path: source_path.to_path_buf(),
kind: ReferenceKind::Imports,
target_route: PathBuf::from(format!("{}.lean", file_path)),
});
}
continue;
}
if let Some(rest) = trimmed.strip_prefix("open ") {
let namespaces = rest.split(" in").next().unwrap_or(rest).split_whitespace();
for ns in namespaces {
let ns = ns.trim_end_matches(|c: char| !c.is_alphanumeric() && c != '.');
if ns.is_empty() || ns == "scoped" {
continue;
}
let file_path = ns.replace('.', "/");
refs.push(SourceReference {
source_path: source_path.to_path_buf(),
kind: ReferenceKind::Uses,
target_route: PathBuf::from(format!("{}.lean", file_path)),
});
}
}
}
refs
}
pub fn detect_references(content: &str, source_path: &Path) -> Vec<SourceReference> {
match source_path.extension().and_then(|e| e.to_str()) {
Some("rs") => detect_rust_references(content, source_path),
Some("py") => detect_python_references(content, source_path),
Some("ts") | Some("tsx") | Some("js") | Some("jsx") => {
detect_ts_references(content, source_path)
}
Some("lean") => detect_lean_references(content, source_path),
_ => Vec::new(),
}
}
fn extension_to_language(ext: &str) -> &'static str {
match ext {
"rs" => "rust",
"py" => "python",
"js" => "javascript",
"ts" => "typescript",
"tsx" => "typescript",
"jsx" => "javascript",
"go" => "go",
"java" => "java",
"lean" => "lean",
"c" | "h" => "c",
"cpp" | "hpp" | "cc" | "cxx" => "cpp",
"md" => "markdown",
"json" => "json",
"yaml" | "yml" => "yaml",
"toml" => "toml",
_ => "unknown",
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Vibe {
pub id: String,
pub title: String,
pub description: String,
pub targets: Vec<NodeId>,
pub created_by: String,
pub created_at: SystemTime,
pub metadata: HashMap<String, String>,
}
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct Constitution {
pub name: String,
pub version: String,
pub description: String,
pub policies: Vec<String>,
}
pub type StatePayload = Value;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CellState {
pub node_id: NodeId,
pub payload: StatePayload,
pub activation: f32,
pub last_updated: SystemTime,
pub annotations: HashMap<String, String>,
}
impl CellState {
pub fn new(node_id: NodeId, payload: StatePayload) -> Self {
Self {
node_id,
payload,
activation: 0.0,
last_updated: SystemTime::now(),
annotations: HashMap::new(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Snapshot {
pub id: String,
pub graph: SourceCodeGraph,
pub vibes: Vec<Vibe>,
pub cell_states: Vec<CellState>,
pub constitution: Constitution,
pub created_at: SystemTime,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
pub enum LayoutStrategy {
#[default]
Flat,
Lattice { width: usize, group_by_row: bool },
Direct,
Preserve,
Modular,
}
#[derive(Debug, Clone)]
pub struct SampleContext<'a> {
pub node: &'a GraphNode,
pub neighbors: Vec<NeighborRef<'a>>,
pub content: Option<&'a str>,
pub annotations: &'a HashMap<String, Value>,
pub graph_metadata: &'a HashMap<String, String>,
}
#[derive(Debug, Clone)]
pub struct NeighborRef<'a> {
pub node: &'a GraphNode,
pub edge: &'a GraphEdge,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SampleArtifact {
pub node_id: NodeId,
pub value: Value,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct SampleResult {
pub sampler_id: String,
pub artifacts: Vec<SampleArtifact>,
pub metadata: HashMap<String, Value>,
}
impl SampleResult {
pub fn get(&self, node_id: NodeId) -> Option<&SampleArtifact> {
self.artifacts.iter().find(|a| a.node_id == node_id)
}
pub fn len(&self) -> usize {
self.artifacts.len()
}
pub fn is_empty(&self) -> bool {
self.artifacts.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = (NodeId, &Value)> {
self.artifacts.iter().map(|a| (a.node_id, &a.value))
}
}
#[derive(Default)]
pub enum NodeSelector {
#[default]
All,
ByKind(GraphNodeKind),
Explicit(Vec<NodeId>),
HasMetadata(String),
Predicate(Box<dyn Fn(&GraphNode) -> bool + Send + Sync>),
}
impl std::fmt::Debug for NodeSelector {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
NodeSelector::All => write!(f, "All"),
NodeSelector::ByKind(k) => write!(f, "ByKind({:?})", k),
NodeSelector::Explicit(ids) => write!(f, "Explicit({:?})", ids),
NodeSelector::HasMetadata(key) => write!(f, "HasMetadata({:?})", key),
NodeSelector::Predicate(_) => write!(f, "Predicate(<fn>)"),
}
}
}
impl NodeSelector {
pub fn matches(&self, node: &GraphNode) -> bool {
match self {
NodeSelector::All => true,
NodeSelector::ByKind(kind) => node.kind == *kind,
NodeSelector::Explicit(ids) => ids.contains(&node.id),
NodeSelector::HasMetadata(key) => node.metadata.contains_key(key),
NodeSelector::Predicate(f) => f(node),
}
}
}
pub type AnnotationMap = HashMap<NodeId, HashMap<String, Value>>;
pub trait Sampler: Send + Sync {
fn id(&self) -> &str;
fn selector(&self) -> NodeSelector {
NodeSelector::All
}
fn compute(&self, ctx: &SampleContext<'_>) -> Result<Option<Value>, SamplerError>;
fn sample(
&self,
graph: &SourceCodeGraph,
annotations: &AnnotationMap,
) -> Result<SampleResult, SamplerError> {
let selector = self.selector();
let selected: Vec<&GraphNode> =
graph.nodes.iter().filter(|n| selector.matches(n)).collect();
let mut artifacts = Vec::with_capacity(selected.len());
for node in &selected {
let neighbors = graph.neighbors(node.id);
let empty = HashMap::new();
let node_annotations = annotations.get(&node.id).unwrap_or(&empty);
let ctx = SampleContext {
node,
neighbors,
content: None,
annotations: node_annotations,
graph_metadata: &graph.metadata,
};
if let Some(value) = self.compute(&ctx)? {
artifacts.push(SampleArtifact {
node_id: node.id,
value,
});
}
}
Ok(SampleResult {
sampler_id: self.id().to_string(),
artifacts,
metadata: HashMap::new(),
})
}
}
#[derive(Debug, Clone)]
pub struct SamplerError {
pub sampler_id: String,
pub message: String,
}
impl std::fmt::Display for SamplerError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "sampler '{}': {}", self.sampler_id, self.message)
}
}
impl std::error::Error for SamplerError {}
impl SamplerError {
pub fn new(sampler_id: impl Into<String>, message: impl Into<String>) -> Self {
Self {
sampler_id: sampler_id.into(),
message: message.into(),
}
}
}
pub struct SamplerPipeline {
stages: Vec<Box<dyn Sampler>>,
}
impl SamplerPipeline {
pub fn new() -> Self {
Self { stages: Vec::new() }
}
pub fn with_stage(mut self, sampler: Box<dyn Sampler>) -> Self {
self.stages.push(sampler);
self
}
pub fn run(
&self,
graph: &SourceCodeGraph,
) -> Result<(Vec<SampleResult>, AnnotationMap), SamplerError> {
let mut annotations: AnnotationMap = HashMap::new();
let mut results = Vec::with_capacity(self.stages.len());
for stage in &self.stages {
let result = stage.sample(graph, &annotations)?;
for artifact in &result.artifacts {
annotations
.entry(artifact.node_id)
.or_default()
.insert(result.sampler_id.clone(), artifact.value.clone());
}
results.push(result);
}
Ok((results, annotations))
}
}
impl Default for SamplerPipeline {
fn default() -> Self {
Self::new()
}
}
impl SourceCodeGraph {
pub fn neighbors(&self, node_id: NodeId) -> Vec<NeighborRef<'_>> {
let node_map: HashMap<NodeId, &GraphNode> = self.nodes.iter().map(|n| (n.id, n)).collect();
self.edges
.iter()
.filter_map(|edge| {
let peer_id = if edge.from == node_id {
Some(edge.to)
} else if edge.to == node_id {
Some(edge.from)
} else {
None
};
peer_id.and_then(|pid| {
node_map.get(&pid).map(|peer_node| NeighborRef {
node: peer_node,
edge,
})
})
})
.collect()
}
}
pub struct NoOpSampler;
impl Sampler for NoOpSampler {
fn id(&self) -> &str {
"noop"
}
fn compute(&self, _ctx: &SampleContext<'_>) -> Result<Option<Value>, SamplerError> {
Ok(None)
}
}
pub struct DegreeSampler;
impl Sampler for DegreeSampler {
fn id(&self) -> &str {
"degree"
}
fn selector(&self) -> NodeSelector {
NodeSelector::ByKind(GraphNodeKind::File)
}
fn compute(&self, ctx: &SampleContext<'_>) -> Result<Option<Value>, SamplerError> {
let incoming = ctx
.neighbors
.iter()
.filter(|n| n.edge.to == ctx.node.id)
.count();
let outgoing = ctx
.neighbors
.iter()
.filter(|n| n.edge.from == ctx.node.id)
.count();
Ok(Some(serde_json::json!({
"in": incoming,
"out": outgoing,
"total": incoming + outgoing,
})))
}
}
pub struct MetadataSampler {
keys: Vec<String>,
}
impl MetadataSampler {
pub fn new(keys: Vec<String>) -> Self {
Self { keys }
}
pub fn all() -> Self {
Self { keys: Vec::new() }
}
}
impl Sampler for MetadataSampler {
fn id(&self) -> &str {
"metadata"
}
fn compute(&self, ctx: &SampleContext<'_>) -> Result<Option<Value>, SamplerError> {
let extracted: serde_json::Map<String, Value> = if self.keys.is_empty() {
ctx.node
.metadata
.iter()
.map(|(k, v)| (k.clone(), Value::String(v.clone())))
.collect()
} else {
self.keys
.iter()
.filter_map(|key| {
ctx.node
.metadata
.get(key)
.map(|v| (key.clone(), Value::String(v.clone())))
})
.collect()
};
if extracted.is_empty() {
Ok(None)
} else {
Ok(Some(Value::Object(extracted)))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn test_graph() -> SourceCodeGraph {
let mut meta_a = HashMap::new();
meta_a.insert("relative_path".to_string(), "src/main.rs".to_string());
meta_a.insert("extension".to_string(), "rs".to_string());
meta_a.insert("language".to_string(), "rust".to_string());
let mut meta_b = HashMap::new();
meta_b.insert("relative_path".to_string(), "src/lib.rs".to_string());
meta_b.insert("extension".to_string(), "rs".to_string());
meta_b.insert("language".to_string(), "rust".to_string());
let mut meta_dir = HashMap::new();
meta_dir.insert("relative_path".to_string(), "src".to_string());
SourceCodeGraph {
nodes: vec![
GraphNode {
id: NodeId(0),
name: "src".to_string(),
kind: GraphNodeKind::Directory,
metadata: meta_dir,
},
GraphNode {
id: NodeId(1),
name: "main.rs".to_string(),
kind: GraphNodeKind::File,
metadata: meta_a,
},
GraphNode {
id: NodeId(2),
name: "lib.rs".to_string(),
kind: GraphNodeKind::Module,
metadata: meta_b,
},
],
edges: vec![
GraphEdge {
id: EdgeId(0),
from: NodeId(0),
to: NodeId(1),
relationship: "contains".to_string(),
metadata: HashMap::new(),
},
GraphEdge {
id: EdgeId(1),
from: NodeId(0),
to: NodeId(2),
relationship: "contains".to_string(),
metadata: HashMap::new(),
},
GraphEdge {
id: EdgeId(2),
from: NodeId(1),
to: NodeId(2),
relationship: "uses".to_string(),
metadata: HashMap::new(),
},
],
metadata: {
let mut m = HashMap::new();
m.insert("name".to_string(), "test-project".to_string());
m
},
}
}
#[test]
fn test_neighbors_returns_both_directions() {
let graph = test_graph();
let neighbors = graph.neighbors(NodeId(2));
assert_eq!(neighbors.len(), 2);
let peer_ids: Vec<NodeId> = neighbors.iter().map(|n| n.node.id).collect();
assert!(peer_ids.contains(&NodeId(0))); assert!(peer_ids.contains(&NodeId(1))); }
#[test]
fn test_neighbors_empty_for_unknown_node() {
let graph = test_graph();
let neighbors = graph.neighbors(NodeId(999));
assert!(neighbors.is_empty());
}
#[test]
fn test_selector_all() {
let graph = test_graph();
let sel = NodeSelector::All;
assert!(graph.nodes.iter().all(|n| sel.matches(n)));
}
#[test]
fn test_selector_by_kind() {
let graph = test_graph();
let sel = NodeSelector::ByKind(GraphNodeKind::File);
let matched: Vec<_> = graph.nodes.iter().filter(|n| sel.matches(n)).collect();
assert_eq!(matched.len(), 1);
assert_eq!(matched[0].name, "main.rs");
}
#[test]
fn test_selector_explicit() {
let graph = test_graph();
let sel = NodeSelector::Explicit(vec![NodeId(0), NodeId(2)]);
let matched: Vec<_> = graph.nodes.iter().filter(|n| sel.matches(n)).collect();
assert_eq!(matched.len(), 2);
}
#[test]
fn test_selector_has_metadata() {
let graph = test_graph();
let sel = NodeSelector::HasMetadata("language".to_string());
let matched: Vec<_> = graph.nodes.iter().filter(|n| sel.matches(n)).collect();
assert_eq!(matched.len(), 2); }
#[test]
fn test_selector_predicate() {
let graph = test_graph();
let sel = NodeSelector::Predicate(Box::new(|n| n.name.ends_with(".rs")));
let matched: Vec<_> = graph.nodes.iter().filter(|n| sel.matches(n)).collect();
assert_eq!(matched.len(), 2);
}
#[test]
fn test_noop_sampler_produces_nothing() {
let graph = test_graph();
let sampler = NoOpSampler;
let result = sampler.sample(&graph, &HashMap::new()).unwrap();
assert!(result.is_empty());
assert_eq!(result.sampler_id, "noop");
}
#[test]
fn test_degree_sampler() {
let graph = test_graph();
let sampler = DegreeSampler;
let result = sampler.sample(&graph, &HashMap::new()).unwrap();
assert_eq!(result.sampler_id, "degree");
assert_eq!(result.len(), 1);
let artifact = result.get(NodeId(1)).unwrap();
assert_eq!(artifact.value["in"], 1);
assert_eq!(artifact.value["out"], 1);
assert_eq!(artifact.value["total"], 2);
}
#[test]
fn test_metadata_sampler_specific_keys() {
let graph = test_graph();
let sampler = MetadataSampler::new(vec!["language".to_string()]);
let result = sampler.sample(&graph, &HashMap::new()).unwrap();
assert_eq!(result.len(), 2);
for (_, val) in result.iter() {
assert_eq!(val["language"], "rust");
}
}
#[test]
fn test_metadata_sampler_all_keys() {
let graph = test_graph();
let sampler = MetadataSampler::all();
let result = sampler.sample(&graph, &HashMap::new()).unwrap();
assert_eq!(result.len(), 3); }
#[test]
fn test_sample_result_get_and_iter() {
let result = SampleResult {
sampler_id: "test".to_string(),
artifacts: vec![
SampleArtifact {
node_id: NodeId(1),
value: json!({"score": 0.9}),
},
SampleArtifact {
node_id: NodeId(2),
value: json!({"score": 0.5}),
},
],
metadata: HashMap::new(),
};
assert_eq!(result.len(), 2);
assert!(!result.is_empty());
assert_eq!(result.get(NodeId(1)).unwrap().value["score"], 0.9);
assert!(result.get(NodeId(99)).is_none());
assert_eq!(result.iter().count(), 2);
}
#[test]
fn test_pipeline_threads_annotations() {
let graph = test_graph();
let pipeline = SamplerPipeline::new()
.with_stage(Box::new(MetadataSampler::all()))
.with_stage(Box::new(DegreeSampler));
let (results, annotations) = pipeline.run(&graph).unwrap();
assert_eq!(results.len(), 2);
let main_annot = annotations.get(&NodeId(1)).unwrap();
assert!(main_annot.contains_key("metadata"));
assert!(main_annot.contains_key("degree"));
}
#[test]
fn test_pipeline_empty() {
let graph = test_graph();
let pipeline = SamplerPipeline::new();
let (results, annotations) = pipeline.run(&graph).unwrap();
assert!(results.is_empty());
assert!(annotations.is_empty());
}
#[test]
fn test_sampler_error_display() {
let err = SamplerError::new("embed", "model not loaded");
assert_eq!(err.to_string(), "sampler 'embed': model not loaded");
}
struct FailingSampler;
impl Sampler for FailingSampler {
fn id(&self) -> &str {
"failing"
}
fn compute(&self, _ctx: &SampleContext<'_>) -> Result<Option<Value>, SamplerError> {
Err(SamplerError::new("failing", "intentional test failure"))
}
}
#[test]
fn test_sampler_propagates_error() {
let graph = test_graph();
let sampler = FailingSampler;
let result = sampler.sample(&graph, &HashMap::new());
assert!(result.is_err());
assert_eq!(result.unwrap_err().sampler_id, "failing");
}
#[test]
fn test_detect_lean_references_imports() {
let content = r#"/-
Copyright (c) 2024 Someone. All rights reserved.
-/
module
public import Mathlib.Topology.ContinuousMap.Compact
public import Mathlib.Topology.MetricSpace.Ultra.Basic
import Aesop
/-!
# Some module docs
-/
open Topology Filter in
def someDef := sorry
"#;
let path = std::path::Path::new("Mathlib/Topology/Example.lean");
let refs = detect_lean_references(content, path);
let import_targets: Vec<String> = refs
.iter()
.filter(|r| matches!(r.kind, ReferenceKind::Imports))
.map(|r| r.target_route.to_string_lossy().to_string())
.collect();
assert_eq!(import_targets.len(), 3);
assert!(import_targets.contains(&"Mathlib/Topology/ContinuousMap/Compact.lean".to_string()));
assert!(
import_targets.contains(&"Mathlib/Topology/MetricSpace/Ultra/Basic.lean".to_string())
);
assert!(import_targets.contains(&"Aesop.lean".to_string()));
let uses_targets: Vec<String> = refs
.iter()
.filter(|r| matches!(r.kind, ReferenceKind::Uses))
.map(|r| r.target_route.to_string_lossy().to_string())
.collect();
assert_eq!(uses_targets.len(), 2);
assert!(uses_targets.contains(&"Topology.lean".to_string()));
assert!(uses_targets.contains(&"Filter.lean".to_string()));
}
#[test]
fn test_detect_lean_references_empty() {
let content = "-- just a comment\ndef x := 42\n";
let path = std::path::Path::new("test.lean");
let refs = detect_lean_references(content, path);
assert!(refs.is_empty());
}
}