Skip to main content

llmtxt_core/disclosure/
code.rs

1//! Code section parsing for progressive disclosure.
2
3use crate::calculate_tokens;
4use crate::disclosure::types::Section;
5
6/// Pattern types that indicate the start of a code section.
7struct Pattern {
8    prefix: &'static str,
9    section_type: &'static str,
10}
11
12const PATTERNS: &[Pattern] = &[
13    Pattern {
14        prefix: "export async function ",
15        section_type: "function",
16    },
17    Pattern {
18        prefix: "async function ",
19        section_type: "function",
20    },
21    Pattern {
22        prefix: "export function ",
23        section_type: "function",
24    },
25    Pattern {
26        prefix: "function ",
27        section_type: "function",
28    },
29    Pattern {
30        prefix: "export class ",
31        section_type: "class",
32    },
33    Pattern {
34        prefix: "class ",
35        section_type: "class",
36    },
37    Pattern {
38        prefix: "export const ",
39        section_type: "function",
40    },
41    Pattern {
42        prefix: "const ",
43        section_type: "function",
44    },
45    Pattern {
46        prefix: "export let ",
47        section_type: "function",
48    },
49    Pattern {
50        prefix: "let ",
51        section_type: "function",
52    },
53    Pattern {
54        prefix: "export var ",
55        section_type: "function",
56    },
57    Pattern {
58        prefix: "var ",
59        section_type: "function",
60    },
61    Pattern {
62        prefix: "def ",
63        section_type: "function",
64    },
65    Pattern {
66        prefix: "pub fn ",
67        section_type: "function",
68    },
69    Pattern {
70        prefix: "fn ",
71        section_type: "function",
72    },
73    Pattern {
74        prefix: "func ",
75        section_type: "function",
76    },
77];
78
79/// Check whether a line matches a function/class pattern and return the symbol name.
80fn match_pattern(line: &str) -> Option<(&'static str, String)> {
81    for pat in PATTERNS {
82        if let Some(rest) = line.strip_prefix(pat.prefix) {
83            // Extract identifier (stop at whitespace, '(', '<', '=', '{')
84            let name: String = rest
85                .chars()
86                .take_while(|&c| c.is_alphanumeric() || c == '_')
87                .collect();
88            if !name.is_empty() {
89                return Some((pat.section_type, name));
90            }
91        }
92    }
93    None
94}
95
96/// Parse code lines into sections based on function/class declarations.
97pub fn parse_code_sections(lines: &[&str]) -> Vec<Section> {
98    let mut sections: Vec<Section> = Vec::new();
99
100    for (i, line) in lines.iter().enumerate() {
101        if let Some((section_type, name)) = match_pattern(line.trim_start()) {
102            // Find end of this section (next pattern match)
103            let mut end_line = lines.len();
104            for (j, next_line) in lines.iter().enumerate().skip(i + 1) {
105                if match_pattern(next_line.trim_start()).is_some() {
106                    end_line = j;
107                    break;
108                }
109            }
110            let section_content = lines[i..end_line].join("\n");
111            sections.push(Section {
112                title: name,
113                depth: 0,
114                start_line: i + 1,
115                end_line,
116                token_count: calculate_tokens(&section_content),
117                section_type: section_type.to_string(),
118            });
119        }
120    }
121
122    sections
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn parse_function_declaration() {
131        let lines = vec!["function hello() {", "  return 42;", "}"];
132        let sections = parse_code_sections(&lines);
133        assert_eq!(sections.len(), 1);
134        assert_eq!(sections[0].title, "hello");
135        assert_eq!(sections[0].section_type, "function");
136    }
137
138    #[test]
139    fn parse_class_declaration() {
140        let lines = vec!["class MyClass {", "  constructor() {}", "}"];
141        let sections = parse_code_sections(&lines);
142        assert_eq!(sections.len(), 1);
143        assert_eq!(sections[0].title, "MyClass");
144        assert_eq!(sections[0].section_type, "class");
145    }
146
147    #[test]
148    fn parse_export_function() {
149        let lines = vec!["export function foo() {", "}"];
150        let sections = parse_code_sections(&lines);
151        assert_eq!(sections.len(), 1);
152        assert_eq!(sections[0].title, "foo");
153    }
154}