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
161fn extract_ts_signatures(content: &str) -> Vec<Signature> {
162    let mut sigs = Vec::new();
163
164    for line in content.lines() {
165        let trimmed = line.trim();
166        if trimmed.starts_with("//") || trimmed.starts_with("/*") || trimmed.starts_with('*') {
167            continue;
168        }
169
170        if let Some(caps) = fn_re().captures(line) {
171            let indent = caps.get(1).map_or(0, |m| m.as_str().len());
172            sigs.push(Signature {
173                kind: if indent > 0 { "method" } else { "fn" },
174                name: caps[4].to_string(),
175                params: compact_params(&caps[5]),
176                return_type: caps
177                    .get(6)
178                    .map_or(String::new(), |m| m.as_str().trim().to_string()),
179                is_async: caps.get(3).is_some(),
180                is_exported: caps.get(2).is_some(),
181                indent: if indent > 0 { 2 } else { 0 },
182            });
183        } else if let Some(caps) = class_re().captures(line) {
184            sigs.push(Signature {
185                kind: "class",
186                name: caps[4].to_string(),
187                params: String::new(),
188                return_type: String::new(),
189                is_async: false,
190                is_exported: caps.get(2).is_some(),
191                indent: 0,
192            });
193        } else if let Some(caps) = iface_re().captures(line) {
194            sigs.push(Signature {
195                kind: "interface",
196                name: caps[3].to_string(),
197                params: String::new(),
198                return_type: String::new(),
199                is_async: false,
200                is_exported: caps.get(2).is_some(),
201                indent: 0,
202            });
203        } else if let Some(caps) = type_re().captures(line) {
204            sigs.push(Signature {
205                kind: "type",
206                name: caps[3].to_string(),
207                params: String::new(),
208                return_type: String::new(),
209                is_async: false,
210                is_exported: caps.get(2).is_some(),
211                indent: 0,
212            });
213        } else if let Some(caps) = const_re().captures(line) {
214            if caps.get(2).is_some() {
215                sigs.push(Signature {
216                    kind: "const",
217                    name: caps[4].to_string(),
218                    params: String::new(),
219                    return_type: caps
220                        .get(5)
221                        .map_or(String::new(), |m| m.as_str().to_string()),
222                    is_async: false,
223                    is_exported: true,
224                    indent: 0,
225                });
226            }
227        }
228    }
229
230    sigs
231}
232
233fn extract_rust_signatures(content: &str) -> Vec<Signature> {
234    let mut sigs = Vec::new();
235
236    for line in content.lines() {
237        let trimmed = line.trim();
238        if trimmed.starts_with("//") || trimmed.starts_with("///") {
239            continue;
240        }
241
242        if let Some(caps) = rust_fn_re().captures(line) {
243            let indent = caps.get(1).map_or(0, |m| m.as_str().len());
244            sigs.push(Signature {
245                kind: if indent > 0 { "method" } else { "fn" },
246                name: caps[4].to_string(),
247                params: compact_params(&caps[5]),
248                return_type: caps
249                    .get(6)
250                    .map_or(String::new(), |m| m.as_str().trim().to_string()),
251                is_async: caps.get(3).is_some(),
252                is_exported: caps.get(2).is_some(),
253                indent: if indent > 0 { 2 } else { 0 },
254            });
255        } else if let Some(caps) = rust_struct_re().captures(line) {
256            sigs.push(Signature {
257                kind: "struct",
258                name: caps[3].to_string(),
259                params: String::new(),
260                return_type: String::new(),
261                is_async: false,
262                is_exported: caps.get(2).is_some(),
263                indent: 0,
264            });
265        } else if let Some(caps) = rust_enum_re().captures(line) {
266            sigs.push(Signature {
267                kind: "enum",
268                name: caps[3].to_string(),
269                params: String::new(),
270                return_type: String::new(),
271                is_async: false,
272                is_exported: caps.get(2).is_some(),
273                indent: 0,
274            });
275        } else if let Some(caps) = rust_trait_re().captures(line) {
276            sigs.push(Signature {
277                kind: "trait",
278                name: caps[3].to_string(),
279                params: String::new(),
280                return_type: String::new(),
281                is_async: false,
282                is_exported: caps.get(2).is_some(),
283                indent: 0,
284            });
285        } else if let Some(caps) = rust_impl_re().captures(line) {
286            let trait_name = caps.get(2).map(|m| m.as_str());
287            let type_name = &caps[3];
288            let name = if let Some(t) = trait_name {
289                format!("{t} for {type_name}")
290            } else {
291                type_name.to_string()
292            };
293            sigs.push(Signature {
294                kind: "class",
295                name,
296                params: String::new(),
297                return_type: String::new(),
298                is_async: false,
299                is_exported: false,
300                indent: 0,
301            });
302        }
303    }
304
305    sigs
306}
307
308fn extract_python_signatures(content: &str) -> Vec<Signature> {
309    use std::sync::OnceLock;
310    static PY_FN: OnceLock<Regex> = OnceLock::new();
311    static PY_CLASS: OnceLock<Regex> = OnceLock::new();
312
313    let mut sigs = Vec::new();
314    let py_fn = PY_FN.get_or_init(|| {
315        Regex::new(r"^(\s*)(async\s+)?def\s+(\w+)\s*\(([^)]*)\)(?:\s*->\s*(\w+))?").unwrap()
316    });
317    let py_class = PY_CLASS.get_or_init(|| Regex::new(r"^(\s*)class\s+(\w+)").unwrap());
318
319    for line in content.lines() {
320        if let Some(caps) = py_fn.captures(line) {
321            let indent = caps.get(1).map_or(0, |m| m.as_str().len());
322            sigs.push(Signature {
323                kind: if indent > 0 { "method" } else { "fn" },
324                name: caps[3].to_string(),
325                params: compact_params(&caps[4]),
326                return_type: caps
327                    .get(5)
328                    .map_or(String::new(), |m| m.as_str().to_string()),
329                is_async: caps.get(2).is_some(),
330                is_exported: !caps[3].starts_with('_'),
331                indent: if indent > 0 { 2 } else { 0 },
332            });
333        } else if let Some(caps) = py_class.captures(line) {
334            sigs.push(Signature {
335                kind: "class",
336                name: caps[2].to_string(),
337                params: String::new(),
338                return_type: String::new(),
339                is_async: false,
340                is_exported: !caps[2].starts_with('_'),
341                indent: 0,
342            });
343        }
344    }
345
346    sigs
347}
348
349fn extract_go_signatures(content: &str) -> Vec<Signature> {
350    use std::sync::OnceLock;
351    static GO_FN: OnceLock<Regex> = OnceLock::new();
352    static GO_TYPE: OnceLock<Regex> = OnceLock::new();
353
354    let mut sigs = Vec::new();
355    let go_fn = GO_FN.get_or_init(|| Regex::new(r"^func\s+(?:\((\w+)\s+\*?(\w+)\)\s+)?(\w+)\s*\(([^)]*)\)(?:\s*(?:\(([^)]*)\)|(\w+)))?\s*\{").unwrap());
356    let go_type =
357        GO_TYPE.get_or_init(|| Regex::new(r"^type\s+(\w+)\s+(struct|interface)").unwrap());
358
359    for line in content.lines() {
360        if let Some(caps) = go_fn.captures(line) {
361            let is_method = caps.get(2).is_some();
362            sigs.push(Signature {
363                kind: if is_method { "method" } else { "fn" },
364                name: caps[3].to_string(),
365                params: compact_params(&caps[4]),
366                return_type: caps
367                    .get(5)
368                    .or(caps.get(6))
369                    .map_or(String::new(), |m| m.as_str().to_string()),
370                is_async: false,
371                is_exported: caps[3].starts_with(char::is_uppercase),
372                indent: if is_method { 2 } else { 0 },
373            });
374        } else if let Some(caps) = go_type.captures(line) {
375            sigs.push(Signature {
376                kind: if &caps[2] == "struct" {
377                    "struct"
378                } else {
379                    "interface"
380                },
381                name: caps[1].to_string(),
382                params: String::new(),
383                return_type: String::new(),
384                is_async: false,
385                is_exported: caps[1].starts_with(char::is_uppercase),
386                indent: 0,
387            });
388        }
389    }
390
391    sigs
392}
393
394pub(crate) fn compact_params(params: &str) -> String {
395    if params.trim().is_empty() {
396        return String::new();
397    }
398    params
399        .split(',')
400        .map(|p| {
401            let p = p.trim();
402            if let Some((name, ty)) = p.split_once(':') {
403                let name = name.trim();
404                let ty = ty.trim();
405                let short = match ty {
406                    "string" | "String" | "&str" | "str" => ":s",
407                    "number" | "i32" | "i64" | "u32" | "u64" | "usize" | "f32" | "f64" => ":n",
408                    "boolean" | "bool" => ":b",
409                    _ => return format!("{name}:{ty}"),
410                };
411                format!("{name}{short}")
412            } else {
413                p.to_string()
414            }
415        })
416        .collect::<Vec<_>>()
417        .join(", ")
418}
419
420fn compact_type(ty: &str) -> String {
421    match ty.trim() {
422        "String" | "string" | "&str" | "str" => "s".to_string(),
423        "bool" | "boolean" => "b".to_string(),
424        "i32" | "i64" | "u32" | "u64" | "usize" | "f32" | "f64" | "number" => "n".to_string(),
425        "void" | "()" => "∅".to_string(),
426        other => {
427            if other.starts_with("Vec<") || other.starts_with("Array<") {
428                let inner = other
429                    .trim_start_matches("Vec<")
430                    .trim_start_matches("Array<")
431                    .trim_end_matches('>');
432                format!("[{}]", compact_type(inner))
433            } else if other.starts_with("Option<") || other.starts_with("Maybe<") {
434                let inner = other
435                    .trim_start_matches("Option<")
436                    .trim_start_matches("Maybe<")
437                    .trim_end_matches('>');
438                format!("?{}", compact_type(inner))
439            } else if other.starts_with("Result<") {
440                "R".to_string()
441            } else if other.starts_with("impl ") {
442                other.trim_start_matches("impl ").to_string()
443            } else {
444                other.to_string()
445            }
446        }
447    }
448}
449
450fn tdd_params(params: &str) -> String {
451    if params.trim().is_empty() {
452        return String::new();
453    }
454    params
455        .split(',')
456        .map(|p| {
457            let p = p.trim();
458            if p.starts_with('&') {
459                let rest = p.trim_start_matches("&mut ").trim_start_matches('&');
460                if let Some((name, ty)) = rest.split_once(':') {
461                    format!("&{}:{}", name.trim(), compact_type(ty))
462                } else {
463                    p.to_string()
464                }
465            } else if let Some((name, ty)) = p.split_once(':') {
466                format!("{}:{}", name.trim(), compact_type(ty))
467            } else if p == "self" || p == "&self" || p == "&mut self" {
468                "⊕".to_string()
469            } else {
470                p.to_string()
471            }
472        })
473        .collect::<Vec<_>>()
474        .join(",")
475}
476
477fn extract_generic_signatures(content: &str) -> Vec<Signature> {
478    static RE_FUNC: OnceLock<Regex> = OnceLock::new();
479    static RE_CLASS: OnceLock<Regex> = OnceLock::new();
480
481    let re_func = RE_FUNC.get_or_init(|| {
482        Regex::new(r"^\s*(?:(?:public|private|protected|static|async|abstract|virtual|override|final|def|func|fun|fn)\s+)+(\w+)\s*\(").unwrap()
483    });
484    let re_class = RE_CLASS.get_or_init(|| {
485        Regex::new(r"^\s*(?:(?:public|private|protected|abstract|final|sealed|partial)\s+)*(?:class|struct|enum|interface|trait|module|object|record)\s+(\w+)").unwrap()
486    });
487
488    let mut sigs = Vec::new();
489    for line in content.lines() {
490        let trimmed = line.trim();
491        if trimmed.is_empty()
492            || trimmed.starts_with("//")
493            || trimmed.starts_with('#')
494            || trimmed.starts_with("/*")
495            || trimmed.starts_with('*')
496        {
497            continue;
498        }
499        if let Some(caps) = re_class.captures(trimmed) {
500            sigs.push(Signature {
501                kind: "type",
502                name: caps[1].to_string(),
503                params: String::new(),
504                return_type: String::new(),
505                is_async: false,
506                is_exported: true,
507                indent: 0,
508            });
509        } else if let Some(caps) = re_func.captures(trimmed) {
510            sigs.push(Signature {
511                kind: "fn",
512                name: caps[1].to_string(),
513                params: String::new(),
514                return_type: String::new(),
515                is_async: trimmed.contains("async"),
516                is_exported: true,
517                indent: 0,
518            });
519        }
520    }
521    sigs
522}