Skip to main content

lean_ctx/core/
signatures.rs

1use regex::Regex;
2use std::sync::OnceLock;
3
4#[derive(Debug, Clone)]
5pub struct Signature {
6    pub kind: &'static str,
7    pub name: String,
8    pub params: String,
9    pub return_type: String,
10    pub is_async: bool,
11    pub is_exported: bool,
12    pub indent: usize,
13}
14
15impl Signature {
16    pub fn to_compact(&self) -> String {
17        let export = if self.is_exported { "⊛ " } else { "" };
18        let async_prefix = if self.is_async { "async " } else { "" };
19
20        match self.kind {
21            "fn" | "method" => {
22                let ret = if self.return_type.is_empty() {
23                    String::new()
24                } else {
25                    format!(" → {}", self.return_type)
26                };
27                let indent = " ".repeat(self.indent);
28                format!(
29                    "{indent}fn {async_prefix}{export}{}({}){}",
30                    self.name, self.params, ret
31                )
32            }
33            "class" | "struct" => format!("cl {export}{}", self.name),
34            "interface" | "trait" => format!("if {export}{}", self.name),
35            "type" => format!("ty {export}{}", self.name),
36            "enum" => format!("en {export}{}", self.name),
37            "const" | "let" | "var" => {
38                let ty = if self.return_type.is_empty() {
39                    String::new()
40                } else {
41                    format!(":{}", self.return_type)
42                };
43                format!("val {export}{}{ty}", self.name)
44            }
45            _ => format!("{} {}", self.kind, self.name),
46        }
47    }
48
49    pub fn to_tdd(&self) -> String {
50        let vis = if self.is_exported { "+" } else { "-" };
51        let a = if self.is_async { "~" } else { "" };
52
53        match self.kind {
54            "fn" | "method" => {
55                let ret = if self.return_type.is_empty() {
56                    String::new()
57                } else {
58                    format!("→{}", compact_type(&self.return_type))
59                };
60                let params = tdd_params(&self.params);
61                let indent = if self.indent > 0 { " " } else { "" };
62                format!("{indent}{a}λ{vis}{}({params}){ret}", self.name)
63            }
64            "class" | "struct" => format!("§{vis}{}", self.name),
65            "interface" | "trait" => format!("∂{vis}{}", self.name),
66            "type" => format!("τ{vis}{}", self.name),
67            "enum" => format!("ε{vis}{}", self.name),
68            "const" | "let" | "var" => {
69                let ty = if self.return_type.is_empty() {
70                    String::new()
71                } else {
72                    format!(":{}", compact_type(&self.return_type))
73                };
74                format!("ν{vis}{}{ty}", self.name)
75            }
76            _ => format!(
77                "{}{vis}{}",
78                self.kind.chars().next().unwrap_or('?'),
79                self.name
80            ),
81        }
82    }
83}
84
85static FN_RE: OnceLock<Regex> = OnceLock::new();
86static CLASS_RE: OnceLock<Regex> = OnceLock::new();
87static IFACE_RE: OnceLock<Regex> = OnceLock::new();
88static TYPE_RE: OnceLock<Regex> = OnceLock::new();
89static CONST_RE: OnceLock<Regex> = OnceLock::new();
90static RUST_FN_RE: OnceLock<Regex> = OnceLock::new();
91static RUST_STRUCT_RE: OnceLock<Regex> = OnceLock::new();
92static RUST_ENUM_RE: OnceLock<Regex> = OnceLock::new();
93static RUST_TRAIT_RE: OnceLock<Regex> = OnceLock::new();
94static RUST_IMPL_RE: OnceLock<Regex> = OnceLock::new();
95
96fn fn_re() -> &'static Regex {
97    FN_RE.get_or_init(|| {
98        Regex::new(r"^(\s*)(export\s+)?(async\s+)?function\s+(\w+)\s*(?:<[^>]*>)?\s*\(([^)]*)\)(?:\s*:\s*([^\{]+))?\s*\{?")
99            .unwrap()
100    })
101}
102
103fn class_re() -> &'static Regex {
104    CLASS_RE.get_or_init(|| Regex::new(r"^(\s*)(export\s+)?(abstract\s+)?class\s+(\w+)").unwrap())
105}
106
107fn iface_re() -> &'static Regex {
108    IFACE_RE.get_or_init(|| Regex::new(r"^(\s*)(export\s+)?interface\s+(\w+)").unwrap())
109}
110
111fn type_re() -> &'static Regex {
112    TYPE_RE.get_or_init(|| Regex::new(r"^(\s*)(export\s+)?type\s+(\w+)").unwrap())
113}
114
115fn const_re() -> &'static Regex {
116    CONST_RE.get_or_init(|| {
117        Regex::new(r"^(\s*)(export\s+)?(const|let|var)\s+(\w+)(?:\s*:\s*(\w+))?").unwrap()
118    })
119}
120
121fn rust_fn_re() -> &'static Regex {
122    RUST_FN_RE.get_or_init(|| {
123        Regex::new(r"^(\s*)(pub\s+)?(async\s+)?fn\s+(\w+)\s*(?:<[^>]*>)?\s*\(([^)]*)\)(?:\s*->\s*([^\{]+))?\s*\{?")
124            .unwrap()
125    })
126}
127
128fn rust_struct_re() -> &'static Regex {
129    RUST_STRUCT_RE.get_or_init(|| Regex::new(r"^(\s*)(pub\s+)?struct\s+(\w+)").unwrap())
130}
131
132fn rust_enum_re() -> &'static Regex {
133    RUST_ENUM_RE.get_or_init(|| Regex::new(r"^(\s*)(pub\s+)?enum\s+(\w+)").unwrap())
134}
135
136fn rust_trait_re() -> &'static Regex {
137    RUST_TRAIT_RE.get_or_init(|| Regex::new(r"^(\s*)(pub\s+)?trait\s+(\w+)").unwrap())
138}
139
140fn rust_impl_re() -> &'static Regex {
141    RUST_IMPL_RE.get_or_init(|| Regex::new(r"^(\s*)impl\s+(?:(\w+)\s+for\s+)?(\w+)").unwrap())
142}
143
144pub fn extract_signatures(content: &str, file_ext: &str) -> Vec<Signature> {
145    #[cfg(feature = "tree-sitter")]
146    {
147        if let Some(sigs) = super::signatures_ts::extract_signatures_ts(content, file_ext) {
148            return sigs;
149        }
150    }
151
152    match file_ext {
153        "rs" => extract_rust_signatures(content),
154        "ts" | "tsx" | "js" | "jsx" | "svelte" | "vue" => extract_ts_signatures(content),
155        "py" => extract_python_signatures(content),
156        "go" => extract_go_signatures(content),
157        _ => extract_generic_signatures(content),
158    }
159}
160
161pub fn extract_file_map(path: &str, content: &str) -> String {
162    let ext = std::path::Path::new(path)
163        .extension()
164        .and_then(|e| e.to_str())
165        .unwrap_or("rs");
166    let dep_info = super::deps::extract_deps(content, ext);
167    let sigs = extract_signatures(content, ext);
168    let mut parts = Vec::new();
169    if !dep_info.imports.is_empty() {
170        parts.push(dep_info.imports.join(","));
171    }
172    let key_sigs: Vec<String> = sigs
173        .iter()
174        .filter(|s| s.is_exported || s.indent == 0)
175        .map(|s| s.to_compact())
176        .collect();
177    if !key_sigs.is_empty() {
178        parts.push(key_sigs.join("\n"));
179    }
180    parts.join("\n")
181}
182
183fn extract_ts_signatures(content: &str) -> Vec<Signature> {
184    let mut sigs = Vec::new();
185
186    for line in content.lines() {
187        let trimmed = line.trim();
188        if trimmed.starts_with("//") || trimmed.starts_with("/*") || trimmed.starts_with('*') {
189            continue;
190        }
191
192        if let Some(caps) = fn_re().captures(line) {
193            let indent = caps.get(1).map_or(0, |m| m.as_str().len());
194            sigs.push(Signature {
195                kind: if indent > 0 { "method" } else { "fn" },
196                name: caps[4].to_string(),
197                params: compact_params(&caps[5]),
198                return_type: caps
199                    .get(6)
200                    .map_or(String::new(), |m| m.as_str().trim().to_string()),
201                is_async: caps.get(3).is_some(),
202                is_exported: caps.get(2).is_some(),
203                indent: if indent > 0 { 2 } else { 0 },
204            });
205        } else if let Some(caps) = class_re().captures(line) {
206            sigs.push(Signature {
207                kind: "class",
208                name: caps[4].to_string(),
209                params: String::new(),
210                return_type: String::new(),
211                is_async: false,
212                is_exported: caps.get(2).is_some(),
213                indent: 0,
214            });
215        } else if let Some(caps) = iface_re().captures(line) {
216            sigs.push(Signature {
217                kind: "interface",
218                name: caps[3].to_string(),
219                params: String::new(),
220                return_type: String::new(),
221                is_async: false,
222                is_exported: caps.get(2).is_some(),
223                indent: 0,
224            });
225        } else if let Some(caps) = type_re().captures(line) {
226            sigs.push(Signature {
227                kind: "type",
228                name: caps[3].to_string(),
229                params: String::new(),
230                return_type: String::new(),
231                is_async: false,
232                is_exported: caps.get(2).is_some(),
233                indent: 0,
234            });
235        } else if let Some(caps) = const_re().captures(line) {
236            if caps.get(2).is_some() {
237                sigs.push(Signature {
238                    kind: "const",
239                    name: caps[4].to_string(),
240                    params: String::new(),
241                    return_type: caps
242                        .get(5)
243                        .map_or(String::new(), |m| m.as_str().to_string()),
244                    is_async: false,
245                    is_exported: true,
246                    indent: 0,
247                });
248            }
249        }
250    }
251
252    sigs
253}
254
255fn extract_rust_signatures(content: &str) -> Vec<Signature> {
256    let mut sigs = Vec::new();
257
258    for line in content.lines() {
259        let trimmed = line.trim();
260        if trimmed.starts_with("//") || trimmed.starts_with("///") {
261            continue;
262        }
263
264        if let Some(caps) = rust_fn_re().captures(line) {
265            let indent = caps.get(1).map_or(0, |m| m.as_str().len());
266            sigs.push(Signature {
267                kind: if indent > 0 { "method" } else { "fn" },
268                name: caps[4].to_string(),
269                params: compact_params(&caps[5]),
270                return_type: caps
271                    .get(6)
272                    .map_or(String::new(), |m| m.as_str().trim().to_string()),
273                is_async: caps.get(3).is_some(),
274                is_exported: caps.get(2).is_some(),
275                indent: if indent > 0 { 2 } else { 0 },
276            });
277        } else if let Some(caps) = rust_struct_re().captures(line) {
278            sigs.push(Signature {
279                kind: "struct",
280                name: caps[3].to_string(),
281                params: String::new(),
282                return_type: String::new(),
283                is_async: false,
284                is_exported: caps.get(2).is_some(),
285                indent: 0,
286            });
287        } else if let Some(caps) = rust_enum_re().captures(line) {
288            sigs.push(Signature {
289                kind: "enum",
290                name: caps[3].to_string(),
291                params: String::new(),
292                return_type: String::new(),
293                is_async: false,
294                is_exported: caps.get(2).is_some(),
295                indent: 0,
296            });
297        } else if let Some(caps) = rust_trait_re().captures(line) {
298            sigs.push(Signature {
299                kind: "trait",
300                name: caps[3].to_string(),
301                params: String::new(),
302                return_type: String::new(),
303                is_async: false,
304                is_exported: caps.get(2).is_some(),
305                indent: 0,
306            });
307        } else if let Some(caps) = rust_impl_re().captures(line) {
308            let trait_name = caps.get(2).map(|m| m.as_str());
309            let type_name = &caps[3];
310            let name = if let Some(t) = trait_name {
311                format!("{t} for {type_name}")
312            } else {
313                type_name.to_string()
314            };
315            sigs.push(Signature {
316                kind: "class",
317                name,
318                params: String::new(),
319                return_type: String::new(),
320                is_async: false,
321                is_exported: false,
322                indent: 0,
323            });
324        }
325    }
326
327    sigs
328}
329
330fn extract_python_signatures(content: &str) -> Vec<Signature> {
331    use std::sync::OnceLock;
332    static PY_FN: OnceLock<Regex> = OnceLock::new();
333    static PY_CLASS: OnceLock<Regex> = OnceLock::new();
334
335    let mut sigs = Vec::new();
336    let py_fn = PY_FN.get_or_init(|| {
337        Regex::new(r"^(\s*)(async\s+)?def\s+(\w+)\s*\(([^)]*)\)(?:\s*->\s*(\w+))?").unwrap()
338    });
339    let py_class = PY_CLASS.get_or_init(|| Regex::new(r"^(\s*)class\s+(\w+)").unwrap());
340
341    for line in content.lines() {
342        if let Some(caps) = py_fn.captures(line) {
343            let indent = caps.get(1).map_or(0, |m| m.as_str().len());
344            sigs.push(Signature {
345                kind: if indent > 0 { "method" } else { "fn" },
346                name: caps[3].to_string(),
347                params: compact_params(&caps[4]),
348                return_type: caps
349                    .get(5)
350                    .map_or(String::new(), |m| m.as_str().to_string()),
351                is_async: caps.get(2).is_some(),
352                is_exported: !caps[3].starts_with('_'),
353                indent: if indent > 0 { 2 } else { 0 },
354            });
355        } else if let Some(caps) = py_class.captures(line) {
356            sigs.push(Signature {
357                kind: "class",
358                name: caps[2].to_string(),
359                params: String::new(),
360                return_type: String::new(),
361                is_async: false,
362                is_exported: !caps[2].starts_with('_'),
363                indent: 0,
364            });
365        }
366    }
367
368    sigs
369}
370
371fn extract_go_signatures(content: &str) -> Vec<Signature> {
372    use std::sync::OnceLock;
373    static GO_FN: OnceLock<Regex> = OnceLock::new();
374    static GO_TYPE: OnceLock<Regex> = OnceLock::new();
375
376    let mut sigs = Vec::new();
377    let go_fn = GO_FN.get_or_init(|| Regex::new(r"^func\s+(?:\((\w+)\s+\*?(\w+)\)\s+)?(\w+)\s*\(([^)]*)\)(?:\s*(?:\(([^)]*)\)|(\w+)))?\s*\{").unwrap());
378    let go_type =
379        GO_TYPE.get_or_init(|| Regex::new(r"^type\s+(\w+)\s+(struct|interface)").unwrap());
380
381    for line in content.lines() {
382        if let Some(caps) = go_fn.captures(line) {
383            let is_method = caps.get(2).is_some();
384            sigs.push(Signature {
385                kind: if is_method { "method" } else { "fn" },
386                name: caps[3].to_string(),
387                params: compact_params(&caps[4]),
388                return_type: caps
389                    .get(5)
390                    .or(caps.get(6))
391                    .map_or(String::new(), |m| m.as_str().to_string()),
392                is_async: false,
393                is_exported: caps[3].starts_with(char::is_uppercase),
394                indent: if is_method { 2 } else { 0 },
395            });
396        } else if let Some(caps) = go_type.captures(line) {
397            sigs.push(Signature {
398                kind: if &caps[2] == "struct" {
399                    "struct"
400                } else {
401                    "interface"
402                },
403                name: caps[1].to_string(),
404                params: String::new(),
405                return_type: String::new(),
406                is_async: false,
407                is_exported: caps[1].starts_with(char::is_uppercase),
408                indent: 0,
409            });
410        }
411    }
412
413    sigs
414}
415
416pub(crate) fn compact_params(params: &str) -> String {
417    if params.trim().is_empty() {
418        return String::new();
419    }
420    params
421        .split(',')
422        .map(|p| {
423            let p = p.trim();
424            if let Some((name, ty)) = p.split_once(':') {
425                let name = name.trim();
426                let ty = ty.trim();
427                let short = match ty {
428                    "string" | "String" | "&str" | "str" => ":s",
429                    "number" | "i32" | "i64" | "u32" | "u64" | "usize" | "f32" | "f64" => ":n",
430                    "boolean" | "bool" => ":b",
431                    _ => return format!("{name}:{ty}"),
432                };
433                format!("{name}{short}")
434            } else {
435                p.to_string()
436            }
437        })
438        .collect::<Vec<_>>()
439        .join(", ")
440}
441
442fn compact_type(ty: &str) -> String {
443    match ty.trim() {
444        "String" | "string" | "&str" | "str" => "s".to_string(),
445        "bool" | "boolean" => "b".to_string(),
446        "i32" | "i64" | "u32" | "u64" | "usize" | "f32" | "f64" | "number" => "n".to_string(),
447        "void" | "()" => "∅".to_string(),
448        other => {
449            if other.starts_with("Vec<") || other.starts_with("Array<") {
450                let inner = other
451                    .trim_start_matches("Vec<")
452                    .trim_start_matches("Array<")
453                    .trim_end_matches('>');
454                format!("[{}]", compact_type(inner))
455            } else if other.starts_with("Option<") || other.starts_with("Maybe<") {
456                let inner = other
457                    .trim_start_matches("Option<")
458                    .trim_start_matches("Maybe<")
459                    .trim_end_matches('>');
460                format!("?{}", compact_type(inner))
461            } else if other.starts_with("Result<") {
462                "R".to_string()
463            } else if other.starts_with("impl ") {
464                other.trim_start_matches("impl ").to_string()
465            } else {
466                other.to_string()
467            }
468        }
469    }
470}
471
472fn tdd_params(params: &str) -> String {
473    if params.trim().is_empty() {
474        return String::new();
475    }
476    params
477        .split(',')
478        .map(|p| {
479            let p = p.trim();
480            if p.starts_with('&') {
481                let rest = p.trim_start_matches("&mut ").trim_start_matches('&');
482                if let Some((name, ty)) = rest.split_once(':') {
483                    format!("&{}:{}", name.trim(), compact_type(ty))
484                } else {
485                    p.to_string()
486                }
487            } else if let Some((name, ty)) = p.split_once(':') {
488                format!("{}:{}", name.trim(), compact_type(ty))
489            } else if p == "self" || p == "&self" || p == "&mut self" {
490                "⊕".to_string()
491            } else {
492                p.to_string()
493            }
494        })
495        .collect::<Vec<_>>()
496        .join(",")
497}
498
499fn extract_generic_signatures(content: &str) -> Vec<Signature> {
500    static RE_FUNC: OnceLock<Regex> = OnceLock::new();
501    static RE_CLASS: OnceLock<Regex> = OnceLock::new();
502
503    let re_func = RE_FUNC.get_or_init(|| {
504        Regex::new(r"^\s*(?:(?:public|private|protected|static|async|abstract|virtual|override|final|def|func|fun|fn)\s+)+(\w+)\s*\(").unwrap()
505    });
506    let re_class = RE_CLASS.get_or_init(|| {
507        Regex::new(r"^\s*(?:(?:public|private|protected|abstract|final|sealed|partial)\s+)*(?:class|struct|enum|interface|trait|module|object|record)\s+(\w+)").unwrap()
508    });
509
510    let mut sigs = Vec::new();
511    for line in content.lines() {
512        let trimmed = line.trim();
513        if trimmed.is_empty()
514            || trimmed.starts_with("//")
515            || trimmed.starts_with('#')
516            || trimmed.starts_with("/*")
517            || trimmed.starts_with('*')
518        {
519            continue;
520        }
521        if let Some(caps) = re_class.captures(trimmed) {
522            sigs.push(Signature {
523                kind: "type",
524                name: caps[1].to_string(),
525                params: String::new(),
526                return_type: String::new(),
527                is_async: false,
528                is_exported: true,
529                indent: 0,
530            });
531        } else if let Some(caps) = re_func.captures(trimmed) {
532            sigs.push(Signature {
533                kind: "fn",
534                name: caps[1].to_string(),
535                params: String::new(),
536                return_type: String::new(),
537                is_async: trimmed.contains("async"),
538                is_exported: true,
539                indent: 0,
540            });
541        }
542    }
543    sigs
544}