llmtxt_core/disclosure/
code.rs1use crate::calculate_tokens;
4use crate::disclosure::types::Section;
5
6struct 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
79fn match_pattern(line: &str) -> Option<(&'static str, String)> {
81 for pat in PATTERNS {
82 if let Some(rest) = line.strip_prefix(pat.prefix) {
83 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
96pub 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 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(§ion_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}