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