1use std::ffi::c_void;
2use std::fs;
3use std::path::{Path, PathBuf};
4use std::sync::atomic::{AtomicBool, Ordering};
5
6use rpi_plugin_sdk::{
7 register_entrypoint, FreeStringFn, PluginApiVt, StbString, StbStringRef, StepHandle,
8 StepResult, ToolPartialCb,
9};
10use serde_json::{json, Value};
11
12struct Drive {
13 params: Value,
14 cancelled: AtomicBool,
15 completed: bool,
16}
17
18fn language(path: &Path) -> Option<&'static str> {
19 match path.extension().and_then(|v| v.to_str()).unwrap_or("") {
20 "rs" => Some("rust"),
21 "ts" | "tsx" => Some("typescript"),
22 "js" | "jsx" => Some("javascript"),
23 "py" => Some("python"),
24 "go" => Some("go"),
25 "java" => Some("java"),
26 "c" | "h" | "cpp" | "hpp" => Some("c-family"),
27 _ => None,
28 }
29}
30
31fn ignored(path: &Path) -> bool {
32 path.file_name()
33 .and_then(|v| v.to_str())
34 .map(|v| {
35 matches!(
36 v,
37 ".git" | ".rpi" | "target" | "node_modules" | "dist" | "build" | "vendor"
38 )
39 })
40 .unwrap_or(false)
41}
42
43fn scan(
44 root: &Path,
45 current: &Path,
46 files: &mut Vec<Value>,
47 edges: &mut Vec<Value>,
48 max_files: usize,
49 max_edges: usize,
50) {
51 if files.len() >= max_files || edges.len() >= max_edges || ignored(current) {
52 return;
53 }
54 let Ok(mut entries) = fs::read_dir(current).map(|v| v.flatten().collect::<Vec<_>>()) else {
55 return;
56 };
57 entries.sort_by_key(|e| e.path());
58 for entry in entries {
59 if files.len() >= max_files || edges.len() >= max_edges {
60 break;
61 }
62 let path = entry.path();
63 if ignored(&path) {
64 continue;
65 }
66 if path.is_dir() {
67 scan(root, &path, files, edges, max_files, max_edges);
68 continue;
69 }
70 let Some(lang) = language(&path) else {
71 continue;
72 };
73 let Ok(text) = fs::read_to_string(&path) else {
74 continue;
75 };
76 let rel = path
77 .strip_prefix(root)
78 .unwrap_or(&path)
79 .to_string_lossy()
80 .replace('\\', "/");
81 files.push(json!({"path":rel,"language":lang,"lines":text.lines().count()}));
82 for line in text.lines() {
83 let trimmed = line.trim();
84 let target = if trimmed.starts_with("use ") {
85 trimmed
86 .strip_prefix("use ")
87 .and_then(|v| v.split("::").next())
88 } else if trimmed.starts_with("mod ") {
89 trimmed
90 .strip_prefix("mod ")
91 .and_then(|v| v.split_whitespace().next())
92 } else if trimmed.starts_with("from ") {
93 trimmed
94 .strip_prefix("from ")
95 .and_then(|v| v.split_whitespace().next())
96 } else if trimmed.starts_with("import ") {
97 trimmed
98 .strip_prefix("import ")
99 .and_then(|v| v.split(|c: char| c == ',' || c.is_whitespace()).next())
100 } else {
101 None
102 };
103 if let Some(target) = target.filter(|v| !v.is_empty()) {
104 edges.push(json!({"from":rel,"to":target.trim_matches(|c: char| c == '"' || c == '\'' || c == ';'),"kind":"import"}));
105 if edges.len() >= max_edges {
106 break;
107 }
108 }
109 }
110 }
111}
112
113fn codegraph(params: &Value) -> Result<String, String> {
114 let root = params
115 .get("cwd")
116 .and_then(Value::as_str)
117 .filter(|v| !v.is_empty())
118 .map(PathBuf::from)
119 .unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
120 if !root.is_dir() {
121 return Err(format!(
122 "codegraph root is not a directory: {}",
123 root.display()
124 ));
125 }
126 let max_files = params
127 .get("maxFiles")
128 .and_then(Value::as_u64)
129 .unwrap_or(500)
130 .clamp(1, 2000) as usize;
131 let max_edges = params
132 .get("maxEdges")
133 .and_then(Value::as_u64)
134 .unwrap_or(2000)
135 .clamp(1, 10000) as usize;
136 let mut files = Vec::new();
137 let mut edges = Vec::new();
138 scan(&root, &root, &mut files, &mut edges, max_files, max_edges);
139 Ok(json!({"root":root,"files":files,"edges":edges,"truncated":files.len() >= max_files || edges.len() >= max_edges}).to_string())
140}
141
142extern "C" fn execute(
143 _: StbStringRef,
144 params: StbString,
145 free: Option<FreeStringFn>,
146) -> StepHandle {
147 let text = params.to_string_lossy();
148 params.free_with(free);
149 Box::into_raw(Box::new(Drive {
150 params: serde_json::from_str(&text).unwrap_or(Value::Null),
151 cancelled: AtomicBool::new(false),
152 completed: false,
153 })) as StepHandle
154}
155extern "C" fn poll(handle: StepHandle, _: Option<ToolPartialCb>, _: *mut c_void) -> StepResult {
156 if handle.is_null() {
157 return StepResult::err(StbString::from_string("null codegraph handle".into()));
158 }
159 let drive = unsafe { &mut *(handle as *mut Drive) };
160 if drive.cancelled.load(Ordering::SeqCst) {
161 return StepResult::err(StbString::from_string("codegraph cancelled".into()));
162 }
163 if drive.completed {
164 return StepResult::err(StbString::from_string(
165 "codegraph polled after completion".into(),
166 ));
167 }
168 drive.completed = true;
169 match codegraph(&drive.params) {
170 Ok(text) => StepResult::done(StbString::from_string(
171 json!({"content":[{"type":"text","text":text}]}).to_string(),
172 )),
173 Err(e) => StepResult::err(StbString::from_string(e)),
174 }
175}
176extern "C" fn cancel(handle: StepHandle) {
177 if !handle.is_null() {
178 unsafe {
179 (&*(handle as *mut Drive))
180 .cancelled
181 .store(true, Ordering::SeqCst);
182 }
183 }
184}
185extern "C" fn destroy(handle: StepHandle) {
186 if !handle.is_null() {
187 unsafe {
188 drop(Box::from_raw(handle as *mut Drive));
189 }
190 }
191}
192extern "C" fn free_string(s: StbString) {
193 if !s.is_empty() && !s.ptr.is_null() {
194 unsafe {
195 let slice = std::slice::from_raw_parts(s.ptr as *const u8, s.len);
196 let _ = Box::from_raw(slice as *const [u8] as *mut [u8]);
197 }
198 }
199}
200
201#[no_mangle]
202pub extern "C" fn rpi_plugin_register(api: *const PluginApiVt, abi: u32) -> i32 {
203 register_entrypoint(api, abi, |api| {
204 let Some(register) = api.register_tool else {
205 return 1;
206 };
207 let schema=Box::new(rpi_plugin_sdk::StableToolSchema{name:StbString::from_string("codegraph".into()),description:StbString::from_string("Inspect source files and import/module edges.".into()),parameters:StbString::from_string(r#"{"type":"object","properties":{"cwd":{"type":"string"},"maxFiles":{"type":"integer"},"maxEdges":{"type":"integer"}}}"#.into())});
208 let rc = register(&*schema, execute, poll, cancel, destroy, free_string);
209 drop(schema);
210 rc
211 })
212}
213
214#[cfg(test)]
215mod tests {
216 use super::*;
217 #[test]
218 fn skips_dependency_dirs() {
219 assert!(ignored(Path::new("target")));
220 assert_eq!(language(Path::new("main.rs")), Some("rust"));
221 }
222}