Skip to main content

decy_llm/
context_builder.rs

1//! Context builder for LLM prompts.
2//!
3//! Builds structured JSON context from static analysis results.
4
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7
8/// Complete analysis context for LLM consumption.
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct AnalysisContext {
11    /// Functions in the file with their analysis results
12    pub functions: Vec<FunctionContext>,
13}
14
15/// Per-function analysis context.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct FunctionContext {
18    /// Function name
19    pub name: String,
20    /// Original C signature
21    pub c_signature: String,
22    /// Ownership inferences for variables
23    pub ownership: HashMap<String, OwnershipInfo>,
24    /// Lifetime relationships
25    pub lifetimes: Vec<LifetimeInfo>,
26    /// Lock-to-data mappings
27    pub lock_mappings: HashMap<String, Vec<String>>,
28}
29
30/// Serializable ownership information.
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct OwnershipInfo {
33    /// Ownership kind: "owning", "immutable_borrow", "mutable_borrow", "unknown"
34    pub kind: String,
35    /// Confidence score 0.0-1.0
36    pub confidence: f64,
37    /// Human-readable reasoning
38    pub reason: String,
39}
40
41/// Serializable lifetime information.
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct LifetimeInfo {
44    /// Variable name
45    pub variable: String,
46    /// Scope depth
47    pub scope_depth: usize,
48    /// Whether variable escapes its scope
49    pub escapes: bool,
50    /// Dependent lifetimes (other variables this depends on)
51    pub depends_on: Vec<String>,
52}
53
54/// Builder for LLM context from analysis results.
55#[derive(Debug, Default)]
56pub struct ContextBuilder {
57    /// Functions being built
58    functions: Vec<FunctionContext>,
59}
60
61impl ContextBuilder {
62    /// Create a new context builder.
63    pub fn new() -> Self {
64        Self { functions: Vec::new() }
65    }
66
67    /// Build context from ownership and lifetime analysis results.
68    pub fn build(&self) -> AnalysisContext {
69        AnalysisContext { functions: self.functions.clone() }
70    }
71
72    /// Add a function with its analysis results.
73    pub fn add_function(&mut self, name: &str, c_signature: &str) -> &mut Self {
74        self.functions.push(FunctionContext {
75            name: name.to_string(),
76            c_signature: c_signature.to_string(),
77            ownership: HashMap::new(),
78            lifetimes: Vec::new(),
79            lock_mappings: HashMap::new(),
80        });
81        self
82    }
83
84    /// Add ownership inference for a variable.
85    pub fn add_ownership(
86        &mut self,
87        function: &str,
88        variable: &str,
89        kind: &str,
90        confidence: f64,
91        reason: &str,
92    ) -> &mut Self {
93        if let Some(func) = self.functions.iter_mut().find(|f| f.name == function) {
94            func.ownership.insert(
95                variable.to_string(),
96                OwnershipInfo { kind: kind.to_string(), confidence, reason: reason.to_string() },
97            );
98        }
99        self
100    }
101
102    /// Add lifetime information.
103    pub fn add_lifetime(
104        &mut self,
105        function: &str,
106        variable: &str,
107        scope_depth: usize,
108        escapes: bool,
109    ) -> &mut Self {
110        if let Some(func) = self.functions.iter_mut().find(|f| f.name == function) {
111            func.lifetimes.push(LifetimeInfo {
112                variable: variable.to_string(),
113                scope_depth,
114                escapes,
115                depends_on: Vec::new(),
116            });
117        }
118        self
119    }
120
121    /// Add lock-to-data mapping.
122    pub fn add_lock_mapping(
123        &mut self,
124        function: &str,
125        lock: &str,
126        protected_data: Vec<String>,
127    ) -> &mut Self {
128        if let Some(func) = self.functions.iter_mut().find(|f| f.name == function) {
129            func.lock_mappings.insert(lock.to_string(), protected_data);
130        }
131        self
132    }
133
134    /// Serialize context to JSON string.
135    pub fn to_json(&self) -> Result<String, serde_json::Error> {
136        let context = self.build();
137        serde_json::to_string_pretty(&context)
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn builder_new_empty() {
147        let builder = ContextBuilder::new();
148        let ctx = builder.build();
149        assert!(ctx.functions.is_empty());
150    }
151
152    #[test]
153    fn builder_default() {
154        let builder = ContextBuilder::default();
155        let ctx = builder.build();
156        assert!(ctx.functions.is_empty());
157    }
158
159    #[test]
160    fn builder_add_function() {
161        let mut builder = ContextBuilder::new();
162        builder.add_function("process", "void process(int* data, int len)");
163        let ctx = builder.build();
164        assert_eq!(ctx.functions.len(), 1);
165        assert_eq!(ctx.functions[0].name, "process");
166        assert_eq!(ctx.functions[0].c_signature, "void process(int* data, int len)");
167    }
168
169    #[test]
170    fn builder_add_multiple_functions() {
171        let mut builder = ContextBuilder::new();
172        builder.add_function("foo", "void foo()");
173        builder.add_function("bar", "int bar(int x)");
174        let ctx = builder.build();
175        assert_eq!(ctx.functions.len(), 2);
176        assert_eq!(ctx.functions[0].name, "foo");
177        assert_eq!(ctx.functions[1].name, "bar");
178    }
179
180    #[test]
181    fn builder_add_ownership() {
182        let mut builder = ContextBuilder::new();
183        builder.add_function("alloc", "void* alloc()");
184        builder.add_ownership("alloc", "ptr", "owning", 0.95, "malloc detected");
185        let ctx = builder.build();
186        let func = &ctx.functions[0];
187        assert!(func.ownership.contains_key("ptr"));
188        let info = &func.ownership["ptr"];
189        assert_eq!(info.kind, "owning");
190        assert!((info.confidence - 0.95).abs() < 0.01);
191        assert_eq!(info.reason, "malloc detected");
192    }
193
194    #[test]
195    fn builder_add_ownership_nonexistent_function() {
196        let mut builder = ContextBuilder::new();
197        builder.add_function("foo", "void foo()");
198        // Adding ownership to a nonexistent function should be a no-op
199        builder.add_ownership("nonexistent", "ptr", "owning", 0.9, "test");
200        let ctx = builder.build();
201        assert!(ctx.functions[0].ownership.is_empty());
202    }
203
204    #[test]
205    fn builder_add_lifetime() {
206        let mut builder = ContextBuilder::new();
207        builder.add_function("borrow", "int* borrow(int* src)");
208        builder.add_lifetime("borrow", "src", 1, false);
209        let ctx = builder.build();
210        let func = &ctx.functions[0];
211        assert_eq!(func.lifetimes.len(), 1);
212        assert_eq!(func.lifetimes[0].variable, "src");
213        assert_eq!(func.lifetimes[0].scope_depth, 1);
214        assert!(!func.lifetimes[0].escapes);
215        assert!(func.lifetimes[0].depends_on.is_empty());
216    }
217
218    #[test]
219    fn builder_add_lifetime_escaping() {
220        let mut builder = ContextBuilder::new();
221        builder.add_function("escape", "int* escape()");
222        builder.add_lifetime("escape", "result", 0, true);
223        let ctx = builder.build();
224        assert!(ctx.functions[0].lifetimes[0].escapes);
225    }
226
227    #[test]
228    fn builder_add_lifetime_nonexistent_function() {
229        let mut builder = ContextBuilder::new();
230        builder.add_function("foo", "void foo()");
231        builder.add_lifetime("nonexistent", "var", 0, false);
232        let ctx = builder.build();
233        assert!(ctx.functions[0].lifetimes.is_empty());
234    }
235
236    #[test]
237    fn builder_add_lock_mapping() {
238        let mut builder = ContextBuilder::new();
239        builder.add_function("sync", "void sync()");
240        builder.add_lock_mapping(
241            "sync",
242            "mutex_a",
243            vec!["counter".to_string(), "buffer".to_string()],
244        );
245        let ctx = builder.build();
246        let func = &ctx.functions[0];
247        assert!(func.lock_mappings.contains_key("mutex_a"));
248        let protected = &func.lock_mappings["mutex_a"];
249        assert_eq!(protected.len(), 2);
250        assert!(protected.contains(&"counter".to_string()));
251    }
252
253    #[test]
254    fn builder_add_lock_mapping_nonexistent_function() {
255        let mut builder = ContextBuilder::new();
256        builder.add_function("foo", "void foo()");
257        builder.add_lock_mapping("nonexistent", "lock", vec!["data".to_string()]);
258        let ctx = builder.build();
259        assert!(ctx.functions[0].lock_mappings.is_empty());
260    }
261
262    #[test]
263    fn builder_to_json() {
264        let mut builder = ContextBuilder::new();
265        builder.add_function("main", "int main()");
266        builder.add_ownership("main", "buf", "mutable_borrow", 0.85, "mutation");
267        let json = builder.to_json().unwrap();
268        assert!(json.contains("\"name\": \"main\""));
269        assert!(json.contains("\"buf\""));
270        assert!(json.contains("mutable_borrow"));
271    }
272
273    #[test]
274    fn builder_to_json_empty() {
275        let builder = ContextBuilder::new();
276        let json = builder.to_json().unwrap();
277        assert!(json.contains("functions"));
278    }
279
280    #[test]
281    fn builder_chaining() {
282        let mut builder = ContextBuilder::new();
283        builder
284            .add_function("f", "void f(int* p)")
285            .add_ownership("f", "p", "immutable_borrow", 0.9, "read-only")
286            .add_lifetime("f", "p", 1, false)
287            .add_lock_mapping("f", "mtx", vec!["shared".to_string()]);
288        let ctx = builder.build();
289        assert_eq!(ctx.functions.len(), 1);
290        assert_eq!(ctx.functions[0].ownership.len(), 1);
291        assert_eq!(ctx.functions[0].lifetimes.len(), 1);
292        assert_eq!(ctx.functions[0].lock_mappings.len(), 1);
293    }
294
295    #[test]
296    fn analysis_context_serde_roundtrip() {
297        let mut builder = ContextBuilder::new();
298        builder.add_function("test", "void test()");
299        builder.add_ownership("test", "x", "owning", 0.8, "test reason");
300        let ctx = builder.build();
301        let json = serde_json::to_string(&ctx).unwrap();
302        let parsed: AnalysisContext = serde_json::from_str(&json).unwrap();
303        assert_eq!(parsed.functions.len(), 1);
304        assert_eq!(parsed.functions[0].name, "test");
305    }
306
307    #[test]
308    fn function_context_default_fields() {
309        let mut builder = ContextBuilder::new();
310        builder.add_function("empty", "void empty()");
311        let ctx = builder.build();
312        let func = &ctx.functions[0];
313        assert!(func.ownership.is_empty());
314        assert!(func.lifetimes.is_empty());
315        assert!(func.lock_mappings.is_empty());
316    }
317
318    #[test]
319    fn ownership_info_serde() {
320        let info = OwnershipInfo {
321            kind: "owning".to_string(),
322            confidence: 0.99,
323            reason: "test".to_string(),
324        };
325        let json = serde_json::to_string(&info).unwrap();
326        let parsed: OwnershipInfo = serde_json::from_str(&json).unwrap();
327        assert_eq!(parsed.kind, "owning");
328        assert!((parsed.confidence - 0.99).abs() < 0.01);
329    }
330
331    #[test]
332    fn lifetime_info_serde() {
333        let info = LifetimeInfo {
334            variable: "p".to_string(),
335            scope_depth: 2,
336            escapes: true,
337            depends_on: vec!["q".to_string()],
338        };
339        let json = serde_json::to_string(&info).unwrap();
340        let parsed: LifetimeInfo = serde_json::from_str(&json).unwrap();
341        assert_eq!(parsed.variable, "p");
342        assert!(parsed.escapes);
343        assert_eq!(parsed.depends_on, vec!["q".to_string()]);
344    }
345}