Skip to main content

lean_ctx/core/
signatures.rs

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