Skip to main content

vibe_graph_ops/
architect.rs

1use anyhow::Result;
2use std::collections::{HashMap, HashSet};
3use std::ffi::OsStr;
4use std::path::{Path, PathBuf};
5use vibe_graph_core::{
6    GraphNodeKind, LayoutStrategy, NodeId, ReferenceKind, SourceCodeGraph, SourceCodeGraphBuilder,
7};
8
9/// Transforms a raw logical graph into a deployable filesystem graph.
10pub trait GraphArchitect {
11    /// Apply the architecture strategy to the graph.
12    fn architect(&self, logical_graph: &SourceCodeGraph) -> Result<SourceCodeGraph>;
13}
14
15/// Factory for creating architects.
16pub struct ArchitectFactory;
17
18impl ArchitectFactory {
19    pub fn create(strategy: LayoutStrategy, root_dir: &Path) -> Box<dyn GraphArchitect> {
20        match strategy {
21            LayoutStrategy::Flat => Box::new(FlatArchitect {
22                root_dir: root_dir.to_path_buf(),
23            }),
24            LayoutStrategy::Lattice {
25                width,
26                group_by_row,
27            } => Box::new(LatticeArchitect {
28                root_dir: root_dir.to_path_buf(),
29                width,
30                group_by_row,
31            }),
32            LayoutStrategy::Preserve | LayoutStrategy::Direct => Box::new(PreserveArchitect {
33                root_dir: root_dir.to_path_buf(),
34            }),
35            _ => Box::new(FlatArchitect {
36                root_dir: root_dir.to_path_buf(),
37            }), // Default to flat
38        }
39    }
40}
41
42/// Preserves the existing directory structure defined in the graph.
43pub struct PreserveArchitect {
44    pub root_dir: PathBuf,
45}
46
47impl GraphArchitect for PreserveArchitect {
48    fn architect(&self, graph: &SourceCodeGraph) -> Result<SourceCodeGraph> {
49        let mut builder = SourceCodeGraphBuilder::new();
50        let _root_id = builder.add_directory(&self.root_dir);
51
52        let mut id_map: HashMap<NodeId, NodeId> = HashMap::new();
53
54        // Detect common prefix to handle absolute paths
55        // Filter for absolute paths only to avoid mixing relative/absolute
56        let all_paths: Vec<PathBuf> = graph
57            .nodes
58            .iter()
59            .filter_map(|n| n.metadata.get("path").map(PathBuf::from))
60            .filter(|p| p.is_absolute())
61            .collect();
62
63        let common_prefix = if all_paths.is_empty() {
64            PathBuf::new()
65        } else {
66            let mut prefix = all_paths[0].clone();
67            for path in &all_paths[1..] {
68                while !path.starts_with(&prefix) {
69                    if !prefix.pop() {
70                        break;
71                    }
72                }
73            }
74            prefix
75        };
76
77        // 1. Process Nodes
78        for node in &graph.nodes {
79            match node.kind {
80                GraphNodeKind::Directory => {
81                    // Try to use relative path if available
82                    let rel_path = if let Some(p) = node.metadata.get("path") {
83                        let p_buf = PathBuf::from(p);
84                        if p_buf.starts_with(&common_prefix) {
85                            p_buf
86                                .strip_prefix(&common_prefix)
87                                .unwrap_or(&p_buf)
88                                .to_path_buf()
89                        } else {
90                            PathBuf::from(&node.name)
91                        }
92                    } else {
93                        PathBuf::from(&node.name)
94                    };
95
96                    // Skip root directory itself if it matches "" or "."
97                    if rel_path.as_os_str().is_empty() || rel_path == OsStr::new(".") {
98                        // This node corresponds to the root itself, which we already created implicitly?
99                        // We map it to the root dir we passed in, but we didn't save root_id in this implementation
100                        // Let's create it explicitly as the root dir provided
101                        let new_id = builder.add_directory(&self.root_dir);
102                        id_map.insert(node.id, new_id);
103                        continue;
104                    }
105
106                    let path = self.root_dir.join(rel_path);
107                    let new_id = builder.add_directory(&path);
108                    id_map.insert(node.id, new_id);
109                }
110                _ => {
111                    // For files/modules/tests/etc
112                    let file_name = &node.name;
113
114                    let rel_path = if let Some(p) = node.metadata.get("path") {
115                        let p_buf = PathBuf::from(p);
116                        if p_buf.starts_with(&common_prefix) {
117                            p_buf
118                                .strip_prefix(&common_prefix)
119                                .unwrap_or(&p_buf)
120                                .to_path_buf()
121                        } else {
122                            PathBuf::from(file_name)
123                        }
124                    } else {
125                        PathBuf::from(file_name)
126                    };
127
128                    let full_path = self.root_dir.join(rel_path);
129
130                    let new_id = builder.add_file(&full_path, file_name);
131                    id_map.insert(node.id, new_id);
132                }
133            }
134        }
135
136        // 2. Process Edges to reconstruct hierarchy and dependencies
137        for edge in &graph.edges {
138            if let (Some(&from), Some(&to)) = (id_map.get(&edge.from), id_map.get(&edge.to)) {
139                let kind = match edge.relationship.as_str() {
140                    "contains" => ReferenceKind::Contains,
141                    "uses" => ReferenceKind::Uses,
142                    "imports" => ReferenceKind::Imports,
143                    "implements" => ReferenceKind::Implements,
144                    _ => ReferenceKind::Uses,
145                };
146                builder.add_edge(from, to, kind);
147            }
148        }
149
150        Ok(builder.build())
151    }
152}
153
154/// Puts all files in a single flat directory.
155pub struct FlatArchitect {
156    pub root_dir: PathBuf,
157}
158
159impl GraphArchitect for FlatArchitect {
160    fn architect(&self, graph: &SourceCodeGraph) -> Result<SourceCodeGraph> {
161        let mut builder = SourceCodeGraphBuilder::new();
162        let root_id = builder.add_directory(&self.root_dir);
163
164        // Map old NodeId to new NodeId
165        let mut id_map: HashMap<NodeId, NodeId> = HashMap::new();
166        let mut processed_names = HashSet::new();
167
168        for node in &graph.nodes {
169            // In Flat mode, we flatten the hierarchy.
170            // We ignore Directory nodes unless they are modules (have content).
171            // We only care about "Leaf" content nodes.
172
173            match node.kind {
174                GraphNodeKind::Directory => {
175                    // Skip pure directories in flat mode
176                    continue;
177                }
178                _ => {
179                    // Keep original name/extension
180                    let file_name = &node.name;
181
182                    // Handle name collisions in flat namespace
183                    if processed_names.contains(file_name) {
184                        // Skip or rename? For now skip to avoid overwrite/error
185                        continue;
186                    }
187                    processed_names.insert(file_name.clone());
188
189                    let path = self.root_dir.join(file_name);
190                    let new_id = builder.add_file(&path, file_name);
191
192                    // Everything is contained by root
193                    builder.add_edge(root_id, new_id, ReferenceKind::Contains);
194                    id_map.insert(node.id, new_id);
195                }
196            }
197        }
198
199        // Reconnect edges
200        for edge in &graph.edges {
201            if let (Some(&from), Some(&to)) = (id_map.get(&edge.from), id_map.get(&edge.to)) {
202                // Ignore original 'contains' edges since we flattened everything to root
203                // Preserve semantic edges
204                if edge.relationship != "contains" {
205                    let kind = match edge.relationship.as_str() {
206                        "imports" => ReferenceKind::Imports,
207                        "implements" => ReferenceKind::Implements,
208                        _ => ReferenceKind::Uses,
209                    };
210                    builder.add_edge(from, to, kind);
211                }
212            }
213        }
214
215        Ok(builder.build())
216    }
217}
218
219/// Organizes nodes based on spatial metadata (x, y) or index.
220pub struct LatticeArchitect {
221    pub root_dir: PathBuf,
222    pub width: usize,
223    pub group_by_row: bool,
224}
225
226impl GraphArchitect for LatticeArchitect {
227    fn architect(&self, graph: &SourceCodeGraph) -> Result<SourceCodeGraph> {
228        let mut builder = SourceCodeGraphBuilder::new();
229        let root_id = builder.add_directory(&self.root_dir);
230
231        let mut id_map: HashMap<NodeId, NodeId> = HashMap::new();
232        let mut row_dirs: HashMap<i32, NodeId> = HashMap::new();
233
234        // 1. Process Nodes
235        for (idx, node) in graph.nodes.iter().enumerate() {
236            // Try to get coordinates from metadata, or fallback to index
237            let x = node
238                .metadata
239                .get("x")
240                .and_then(|v| v.parse::<i32>().ok())
241                .unwrap_or((idx % self.width) as i32);
242
243            let y = node
244                .metadata
245                .get("y")
246                .and_then(|v| v.parse::<i32>().ok())
247                .unwrap_or((idx / self.width) as i32);
248
249            // Determine Parent
250            let parent_id = if self.group_by_row {
251                *row_dirs.entry(y).or_insert_with(|| {
252                    let dir_name = format!("row_{}", y);
253                    let dir_path = self.root_dir.join(&dir_name);
254                    let dir_id = builder.add_directory(&dir_path);
255                    builder.add_edge(root_id, dir_id, ReferenceKind::Contains);
256                    dir_id
257                })
258            } else {
259                root_id
260            };
261
262            // Determine Path
263            let file_name = format!("cell_{}_{}.rs", x, y);
264            let full_path = if self.group_by_row {
265                self.root_dir.join(format!("row_{}", y)).join(&file_name)
266            } else {
267                self.root_dir.join(&file_name)
268            };
269
270            let new_id = builder.add_file(&full_path, &file_name);
271            builder.add_edge(parent_id, new_id, ReferenceKind::Contains);
272            id_map.insert(node.id, new_id);
273
274            // Port metadata (optional, but good for keeping context)
275            // Note: builder nodes have their own metadata, we might want to copy some over
276            // but the builder creates fresh nodes.
277        }
278
279        // 2. Process Edges
280        for edge in &graph.edges {
281            if let (Some(&from), Some(&to)) = (id_map.get(&edge.from), id_map.get(&edge.to)) {
282                // In lattice, all connections become 'Uses' (neighbors)
283                builder.add_edge(from, to, ReferenceKind::Uses);
284            }
285        }
286
287        Ok(builder.build())
288    }
289}