1use anyhow::Result;
2use ck_core::Span;
3use serde::{Deserialize, Serialize};
4
5mod query_chunker;
6
7pub use ck_embed::TokenEstimator;
9
10fn estimate_tokens(text: &str) -> usize {
12 TokenEstimator::estimate_tokens(text)
13}
14
15pub fn get_model_chunk_config(model_name: Option<&str>) -> (usize, usize) {
18 let model = model_name.unwrap_or("nomic-embed-text-v1.5");
19
20 match model {
21 "BAAI/bge-small-en-v1.5" | "sentence-transformers/all-MiniLM-L6-v2" => {
23 (400, 80) }
25
26 "nomic-embed-text-v1" | "nomic-embed-text-v1.5" | "jina-embeddings-v2-base-code" => {
29 (1024, 200) }
31
32 "BAAI/bge-base-en-v1.5" | "BAAI/bge-large-en-v1.5" => {
34 (400, 80) }
36
37 _ => (1024, 200), }
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct StrideInfo {
45 pub original_chunk_id: String,
47 pub stride_index: usize,
49 pub total_strides: usize,
51 pub overlap_start: usize,
53 pub overlap_end: usize,
55}
56
57#[derive(Debug, Clone, Serialize, Deserialize, Default)]
58pub struct ChunkMetadata {
59 pub ancestry: Vec<String>,
60 pub breadcrumb: Option<String>,
61 pub leading_trivia: Vec<String>,
62 pub trailing_trivia: Vec<String>,
63 pub byte_length: usize,
64 pub estimated_tokens: usize,
65}
66
67impl ChunkMetadata {
68 fn from_context(
69 text: &str,
70 ancestry: Vec<String>,
71 leading_trivia: Vec<String>,
72 trailing_trivia: Vec<String>,
73 ) -> Self {
74 let breadcrumb = if ancestry.is_empty() {
75 None
76 } else {
77 Some(ancestry.join("::"))
78 };
79
80 Self {
81 ancestry,
82 breadcrumb,
83 leading_trivia,
84 trailing_trivia,
85 byte_length: text.len(),
86 estimated_tokens: estimate_tokens(text),
87 }
88 }
89
90 fn from_text(text: &str) -> Self {
91 Self {
92 ancestry: Vec::new(),
93 breadcrumb: None,
94 leading_trivia: Vec::new(),
95 trailing_trivia: Vec::new(),
96 byte_length: text.len(),
97 estimated_tokens: estimate_tokens(text),
98 }
99 }
100
101 fn with_updated_text(&self, text: &str) -> Self {
102 let mut cloned = self.clone();
103 cloned.byte_length = text.len();
104 cloned.estimated_tokens = estimate_tokens(text);
105 cloned
106 }
107}
108
109#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct Chunk {
111 pub span: Span,
112 pub text: String,
113 pub chunk_type: ChunkType,
114 pub stride_info: Option<StrideInfo>,
116 pub metadata: ChunkMetadata,
117}
118
119#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
120pub enum ChunkType {
121 Text,
122 Function,
123 Class,
124 Method,
125 Module,
126}
127
128#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
129pub enum ParseableLanguage {
130 Python,
131 TypeScript,
132 JavaScript,
133 Haskell,
134 Rust,
135 Ruby,
136 Go,
137 C,
138 Cpp,
139 CSharp,
140 Zig,
141
142 Dart,
143
144 Elixir,
145
146 Markdown,
147}
148
149impl std::fmt::Display for ParseableLanguage {
150 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
151 let name = match self {
152 ParseableLanguage::Python => "python",
153 ParseableLanguage::TypeScript => "typescript",
154 ParseableLanguage::JavaScript => "javascript",
155 ParseableLanguage::Haskell => "haskell",
156 ParseableLanguage::Rust => "rust",
157 ParseableLanguage::Ruby => "ruby",
158 ParseableLanguage::Go => "go",
159 ParseableLanguage::C => "c",
160 ParseableLanguage::Cpp => "cpp",
161 ParseableLanguage::CSharp => "csharp",
162 ParseableLanguage::Zig => "zig",
163
164 ParseableLanguage::Dart => "dart",
165
166 ParseableLanguage::Elixir => "elixir",
167
168 ParseableLanguage::Markdown => "markdown",
169 };
170 write!(f, "{name}")
171 }
172}
173
174impl TryFrom<ck_core::Language> for ParseableLanguage {
175 type Error = anyhow::Error;
176
177 fn try_from(lang: ck_core::Language) -> Result<Self, Self::Error> {
178 match lang {
179 ck_core::Language::Python => Ok(ParseableLanguage::Python),
180 ck_core::Language::TypeScript => Ok(ParseableLanguage::TypeScript),
181 ck_core::Language::JavaScript => Ok(ParseableLanguage::JavaScript),
182 ck_core::Language::Haskell => Ok(ParseableLanguage::Haskell),
183 ck_core::Language::Rust => Ok(ParseableLanguage::Rust),
184 ck_core::Language::Ruby => Ok(ParseableLanguage::Ruby),
185 ck_core::Language::Go => Ok(ParseableLanguage::Go),
186 ck_core::Language::C => Ok(ParseableLanguage::C),
187 ck_core::Language::Cpp => Ok(ParseableLanguage::Cpp),
188 ck_core::Language::CSharp => Ok(ParseableLanguage::CSharp),
189 ck_core::Language::Zig => Ok(ParseableLanguage::Zig),
190
191 ck_core::Language::Dart => Ok(ParseableLanguage::Dart),
192
193 ck_core::Language::Elixir => Ok(ParseableLanguage::Elixir),
194
195 ck_core::Language::Markdown => Ok(ParseableLanguage::Markdown),
196
197 _ => Err(anyhow::anyhow!(
198 "Language {lang:?} is not supported for parsing"
199 )),
200 }
201 }
202}
203
204pub fn chunk_text(text: &str, language: Option<ck_core::Language>) -> Result<Vec<Chunk>> {
205 chunk_text_with_config(text, language, &ChunkConfig::default())
206}
207
208#[derive(Debug, Clone)]
210pub struct ChunkConfig {
211 pub max_tokens: usize,
213 pub stride_overlap: usize,
215 pub enable_striding: bool,
217}
218
219impl Default for ChunkConfig {
220 fn default() -> Self {
221 Self {
222 max_tokens: 8192, stride_overlap: 1024, enable_striding: true,
225 }
226 }
227}
228
229pub fn chunk_text_with_model(
231 text: &str,
232 language: Option<ck_core::Language>,
233 model_name: Option<&str>,
234) -> Result<Vec<Chunk>> {
235 let (target_tokens, overlap_tokens) = get_model_chunk_config(model_name);
236
237 let config = ChunkConfig {
239 max_tokens: target_tokens,
240 stride_overlap: overlap_tokens,
241 enable_striding: true,
242 };
243
244 chunk_text_with_config_and_model(text, language, &config, model_name)
245}
246
247pub fn chunk_text_with_config(
248 text: &str,
249 language: Option<ck_core::Language>,
250 config: &ChunkConfig,
251) -> Result<Vec<Chunk>> {
252 chunk_text_with_config_and_model(text, language, config, None)
253}
254
255fn chunk_text_with_config_and_model(
256 text: &str,
257 language: Option<ck_core::Language>,
258 config: &ChunkConfig,
259 model_name: Option<&str>,
260) -> Result<Vec<Chunk>> {
261 tracing::debug!(
262 "Chunking text with language: {:?}, length: {} chars, config: {:?}",
263 language,
264 text.len(),
265 config
266 );
267
268 let result = match language.map(ParseableLanguage::try_from) {
269 Some(Ok(lang)) => {
270 tracing::debug!("Using {} tree-sitter parser", lang);
271 chunk_language_with_model(text, lang, model_name)
272 }
273 Some(Err(_)) => {
274 tracing::debug!("Language not supported for parsing, using generic chunking strategy");
275 chunk_generic_with_token_config(text, model_name)
276 }
277 None => {
278 tracing::debug!("Using generic chunking strategy");
279 chunk_generic_with_token_config(text, model_name)
280 }
281 };
282
283 let mut chunks = result?;
284
285 if config.enable_striding {
287 chunks = apply_striding(chunks, config)?;
288 }
289
290 tracing::debug!("Successfully created {} final chunks", chunks.len());
291 Ok(chunks)
292}
293
294fn chunk_generic(text: &str) -> Result<Vec<Chunk>> {
295 chunk_generic_with_token_config(text, None)
296}
297
298fn chunk_generic_with_token_config(text: &str, model_name: Option<&str>) -> Result<Vec<Chunk>> {
299 let mut chunks = Vec::new();
300 let lines: Vec<&str> = text.lines().collect();
301
302 let (target_tokens, overlap_tokens) = get_model_chunk_config(model_name);
304
305 let avg_tokens_per_line = 10.0; let target_lines = ((target_tokens as f32) / avg_tokens_per_line) as usize;
309 let overlap_lines = ((overlap_tokens as f32) / avg_tokens_per_line) as usize;
310
311 let chunk_size = target_lines.max(5); let overlap = overlap_lines.max(1); let mut line_byte_offsets = Vec::with_capacity(lines.len() + 1);
316 line_byte_offsets.push(0);
317 let mut cumulative_offset = 0;
318 let mut byte_pos = 0;
319
320 for line in lines.iter() {
321 cumulative_offset += line.len();
322
323 let line_end_pos = byte_pos + line.len();
325 let newline_len = if line_end_pos < text.len() && text.as_bytes()[line_end_pos] == b'\r' {
326 if line_end_pos + 1 < text.len() && text.as_bytes()[line_end_pos + 1] == b'\n' {
327 2 } else {
329 1 }
331 } else if line_end_pos < text.len() && text.as_bytes()[line_end_pos] == b'\n' {
332 1 } else {
334 0 };
336
337 cumulative_offset += newline_len;
338 byte_pos = cumulative_offset;
339 line_byte_offsets.push(cumulative_offset);
340 }
341
342 let mut i = 0;
343 while i < lines.len() {
344 let end = (i + chunk_size).min(lines.len());
345 let chunk_lines = &lines[i..end];
346 let chunk_text = chunk_lines.join("\n");
347 let byte_start = line_byte_offsets[i];
348 let byte_end = line_byte_offsets[end];
349 let metadata = ChunkMetadata::from_text(&chunk_text);
350
351 chunks.push(Chunk {
352 span: Span {
353 byte_start,
354 byte_end,
355 line_start: i + 1,
356 line_end: end,
357 },
358 text: chunk_text,
359 chunk_type: ChunkType::Text,
360 stride_info: None,
361 metadata,
362 });
363
364 i += chunk_size - overlap;
365 if i >= lines.len() {
366 break;
367 }
368 }
369
370 Ok(chunks)
371}
372
373pub(crate) fn tree_sitter_language(language: ParseableLanguage) -> Result<tree_sitter::Language> {
374 if language == ParseableLanguage::Markdown {
375 return Ok(tree_sitter_md::LANGUAGE.into());
376 }
377
378 let ts_language = match language {
379 ParseableLanguage::Python => tree_sitter_python::LANGUAGE,
380 ParseableLanguage::TypeScript | ParseableLanguage::JavaScript => {
381 tree_sitter_typescript::LANGUAGE_TYPESCRIPT
382 }
383 ParseableLanguage::Haskell => tree_sitter_haskell::LANGUAGE,
384 ParseableLanguage::Rust => tree_sitter_rust::LANGUAGE,
385 ParseableLanguage::Ruby => tree_sitter_ruby::LANGUAGE,
386 ParseableLanguage::Go => tree_sitter_go::LANGUAGE,
387 ParseableLanguage::C => tree_sitter_c::LANGUAGE,
388 ParseableLanguage::Cpp => tree_sitter_cpp::LANGUAGE,
389 ParseableLanguage::CSharp => tree_sitter_c_sharp::LANGUAGE,
390 ParseableLanguage::Zig => tree_sitter_zig::LANGUAGE,
391
392 ParseableLanguage::Dart => tree_sitter_dart::LANGUAGE,
393
394 ParseableLanguage::Elixir => tree_sitter_elixir::LANGUAGE,
395
396 ParseableLanguage::Markdown => unreachable!("Handled above via early return"),
397 };
398
399 Ok(ts_language.into())
400}
401
402fn chunk_language(text: &str, language: ParseableLanguage) -> Result<Vec<Chunk>> {
403 let mut parser = tree_sitter::Parser::new();
404 let ts_language = tree_sitter_language(language)?;
405 parser.set_language(&ts_language)?;
406
407 let tree = parser
408 .parse(text, None)
409 .ok_or_else(|| anyhow::anyhow!("Failed to parse {language} code"))?;
410
411 let mut chunks = match query_chunker::chunk_with_queries(language, ts_language, &tree, text)? {
412 Some(query_chunks) if !query_chunks.is_empty() => query_chunks,
413 _ => {
414 let mut legacy_chunks = Vec::new();
415 let mut cursor = tree.walk();
416 extract_code_chunks(&mut cursor, text, &mut legacy_chunks, language);
417 legacy_chunks
418 }
419 };
420
421 if chunks.is_empty() {
422 return chunk_generic(text);
423 }
424
425 if language == ParseableLanguage::Haskell {
427 chunks = merge_haskell_functions(chunks, text);
428 }
429
430 chunks = fill_gaps(chunks, text);
436
437 if language == ParseableLanguage::Cpp {
439 chunks = merge_cpp_template_prefix_chunks(chunks, text);
440 }
441
442 if matches!(language, ParseableLanguage::C | ParseableLanguage::Cpp) {
444 chunks = suppress_contained_text_chunks(chunks);
445 }
446
447 if language == ParseableLanguage::Markdown {
449 let (target_tokens, _) = get_model_chunk_config(None);
450 chunks = merge_small_chunks(chunks, text, target_tokens);
451 }
452
453 Ok(chunks)
454}
455
456fn suppress_contained_text_chunks(chunks: Vec<Chunk>) -> Vec<Chunk> {
457 if chunks.is_empty() {
458 return chunks;
459 }
460
461 let mut containers: Vec<(usize, usize)> = chunks
462 .iter()
463 .filter(|chunk| {
464 matches!(
465 chunk.chunk_type,
466 ChunkType::Class | ChunkType::Method | ChunkType::Function
467 )
468 })
469 .map(|chunk| (chunk.span.byte_start, chunk.span.byte_end))
470 .collect();
471
472 if containers.is_empty() {
473 return chunks;
474 }
475
476 containers.sort_by_key(|(start, _)| *start);
477
478 chunks
479 .into_iter()
480 .filter(|chunk| {
481 if chunk.chunk_type != ChunkType::Text {
482 return true;
483 }
484
485 let start = chunk.span.byte_start;
486 let end = chunk.span.byte_end;
487 !containers
488 .iter()
489 .any(|(c_start, c_end)| *c_start <= start && end <= *c_end)
490 })
491 .collect()
492}
493
494fn merge_cpp_template_prefix_chunks(chunks: Vec<Chunk>, text: &str) -> Vec<Chunk> {
495 if chunks.len() < 2 {
496 return chunks;
497 }
498
499 let mut merged = Vec::with_capacity(chunks.len());
500 let mut idx = 0;
501
502 while idx < chunks.len() {
503 if idx + 1 < chunks.len() && is_template_prefix_chunk(&chunks[idx]) {
504 let template_chunk = &chunks[idx];
505 let mut next_chunk = chunks[idx + 1].clone();
506
507 let gap_is_whitespace = template_chunk.span.byte_end <= next_chunk.span.byte_start
512 && text
513 .get(template_chunk.span.byte_end..next_chunk.span.byte_start)
514 .is_some_and(|gap| gap.chars().all(char::is_whitespace));
515
516 if gap_is_whitespace
517 && template_chunk.span.byte_start < next_chunk.span.byte_end
518 && next_chunk.span.byte_end <= text.len()
519 {
520 let new_start = template_chunk.span.byte_start;
521 let new_end = next_chunk.span.byte_end;
522
523 if let Some(new_text) = text.get(new_start..new_end) {
524 let (line_start, line_end) = line_range_for_span(text, new_start, new_end);
525
526 next_chunk.span.byte_start = new_start;
527 next_chunk.span.line_start = line_start;
528 next_chunk.span.line_end = line_end;
529 next_chunk.text = new_text.to_string();
530 next_chunk.metadata = next_chunk.metadata.with_updated_text(new_text);
531
532 merged.push(next_chunk);
533 idx += 2;
534 continue;
535 }
536 }
537 }
538
539 merged.push(chunks[idx].clone());
540 idx += 1;
541 }
542
543 merged
544}
545
546fn is_template_prefix_chunk(chunk: &Chunk) -> bool {
547 if chunk.chunk_type != ChunkType::Text {
548 return false;
549 }
550
551 let mut has_template = false;
552 for line in chunk.text.lines() {
553 let trimmed = line.trim();
554 if trimmed.is_empty() {
555 continue;
556 }
557 if trimmed.starts_with("template <") || trimmed.starts_with("template<") {
558 has_template = true;
559 continue;
560 }
561 return false;
562 }
563
564 has_template
565}
566
567fn line_range_for_span(text: &str, byte_start: usize, byte_end: usize) -> (usize, usize) {
568 let line_start = text[..byte_start].matches('\n').count() + 1;
569 let newlines_up_to_end = text[..byte_end].matches('\n').count();
570 let line_end = if newlines_up_to_end >= line_start - 1 {
571 newlines_up_to_end.max(line_start)
572 } else {
573 line_start
574 };
575
576 (line_start, line_end)
577}
578
579fn fill_gaps(mut chunks: Vec<Chunk>, text: &str) -> Vec<Chunk> {
583 if chunks.is_empty() {
584 return chunks;
585 }
586
587 chunks.sort_by_key(|c| c.span.byte_start);
589
590 let mut result = Vec::new();
591 let mut last_end = 0;
592
593 let mut gaps = Vec::new();
595
596 for chunk in &chunks {
597 if last_end < chunk.span.byte_start {
598 let gap_start = last_end;
600 let gap_text = &text[gap_start..chunk.span.byte_start];
601
602 let mut current_byte = gap_start;
604 let mut segment_start = gap_start;
605
606 for line in gap_text.split('\n') {
607 let line_start_in_gap = current_byte - gap_start;
608 let _line_end_in_gap = line_start_in_gap + line.len();
609
610 if line.trim().is_empty() {
611 if segment_start < current_byte {
613 let segment_text = &text[segment_start..current_byte];
614 if !segment_text.trim().is_empty() {
615 gaps.push((segment_start, current_byte));
616 }
617 }
618 segment_start = current_byte + line.len() + 1;
620 }
621
622 current_byte += line.len() + 1; }
624
625 if segment_start < chunk.span.byte_start {
627 let remaining = &text[segment_start..chunk.span.byte_start];
628 if !remaining.trim().is_empty() {
629 gaps.push((segment_start, chunk.span.byte_start));
630 }
631 }
632 }
633 last_end = last_end.max(chunk.span.byte_end);
634 }
635
636 if last_end < text.len() {
638 let gap_text = &text[last_end..];
639 if !gap_text.trim().is_empty() {
640 gaps.push((last_end, text.len()));
641 }
642 }
643
644 let combined_gaps = gaps;
645
646 let mut gap_idx = 0;
648
649 for chunk in chunks {
650 while gap_idx < combined_gaps.len() && combined_gaps[gap_idx].1 <= chunk.span.byte_start {
652 let (gap_start, gap_end) = combined_gaps[gap_idx];
653 let gap_text = &text[gap_start..gap_end];
654
655 let line_start = text[..gap_start].matches('\n').count() + 1;
657 let newlines_up_to_end = text[..gap_end].matches('\n').count();
660 let line_end = if newlines_up_to_end >= line_start - 1 {
661 newlines_up_to_end.max(line_start)
662 } else {
663 line_start
664 };
665
666 let gap_chunk = Chunk {
667 text: gap_text.to_string(),
668 span: Span {
669 byte_start: gap_start,
670 byte_end: gap_end,
671 line_start,
672 line_end,
673 },
674 chunk_type: ChunkType::Text,
675 metadata: ChunkMetadata::from_text(gap_text),
676 stride_info: None,
677 };
678 result.push(gap_chunk);
679 gap_idx += 1;
680 }
681
682 result.push(chunk.clone());
683 }
684
685 while gap_idx < combined_gaps.len() {
687 let (gap_start, gap_end) = combined_gaps[gap_idx];
688 let gap_text = &text[gap_start..gap_end];
689
690 let line_start = text[..gap_start].matches('\n').count() + 1;
692 let newlines_up_to_end = text[..gap_end].matches('\n').count();
694 let line_end = if newlines_up_to_end >= line_start - 1 {
695 newlines_up_to_end.max(line_start)
696 } else {
697 line_start
698 };
699
700 let gap_chunk = Chunk {
701 text: gap_text.to_string(),
702 span: Span {
703 byte_start: gap_start,
704 byte_end: gap_end,
705 line_start,
706 line_end,
707 },
708 chunk_type: ChunkType::Text,
709 metadata: ChunkMetadata::from_text(gap_text),
710 stride_info: None,
711 };
712 result.push(gap_chunk);
713 gap_idx += 1;
714 }
715
716 result
717}
718
719fn merge_haskell_functions(chunks: Vec<Chunk>, source: &str) -> Vec<Chunk> {
721 if chunks.is_empty() {
722 return chunks;
723 }
724
725 let mut merged = Vec::new();
726 let mut i = 0;
727
728 while i < chunks.len() {
729 let chunk = &chunks[i];
730
731 let trimmed = chunk.text.trim();
733 if trimmed.is_empty()
734 || trimmed.starts_with("--")
735 || trimmed.starts_with("{-")
736 || !chunk.text.contains(|c: char| c.is_alphanumeric())
737 {
738 i += 1;
739 continue;
740 }
741
742 let is_signature = chunk.text.contains("::");
745 let function_name = if is_signature {
746 chunk
748 .text
749 .split("::")
750 .next()
751 .and_then(|s| s.split_whitespace().next())
752 .map(std::string::ToString::to_string)
753 } else {
754 extract_haskell_function_name(&chunk.text)
755 };
756
757 if function_name.is_none() {
758 merged.push(chunk.clone());
760 i += 1;
761 continue;
762 }
763
764 let name = function_name.unwrap();
765 let group_start = chunk.span.byte_start;
766 let mut group_end = chunk.span.byte_end;
767 let line_start = chunk.span.line_start;
768 let mut line_end = chunk.span.line_end;
769 let mut trailing_trivia = chunk.metadata.trailing_trivia.clone();
770
771 let mut j = i + 1;
773 while j < chunks.len() {
774 let next_chunk = &chunks[j];
775
776 let next_trimmed = next_chunk.text.trim();
778 if next_trimmed.starts_with("--") || next_trimmed.starts_with("{-") {
779 j += 1;
780 continue;
781 }
782
783 let next_is_signature = next_chunk.text.contains("::");
784 let next_name = if next_is_signature {
785 next_chunk
786 .text
787 .split("::")
788 .next()
789 .and_then(|s| s.split_whitespace().next())
790 .map(std::string::ToString::to_string)
791 } else {
792 extract_haskell_function_name(&next_chunk.text)
793 };
794
795 if next_name == Some(name.clone()) {
796 group_end = next_chunk.span.byte_end;
798 line_end = next_chunk.span.line_end;
799 trailing_trivia = next_chunk.metadata.trailing_trivia.clone();
800 j += 1;
801 } else {
802 break;
803 }
804 }
805
806 let merged_text = source.get(group_start..group_end).unwrap_or("").to_string();
808 let mut metadata = chunk.metadata.with_updated_text(&merged_text);
809 metadata.trailing_trivia = trailing_trivia;
810
811 merged.push(Chunk {
812 span: Span {
813 byte_start: group_start,
814 byte_end: group_end,
815 line_start,
816 line_end,
817 },
818 text: merged_text,
819 chunk_type: ChunkType::Function,
820 stride_info: None,
821 metadata,
822 });
823
824 i = j; }
826
827 merged
828}
829
830fn extract_haskell_function_name(text: &str) -> Option<String> {
832 let trimmed = text.trim();
835
836 let first_word = trimmed
838 .split_whitespace()
839 .next()?
840 .trim_end_matches(|c: char| !c.is_alphanumeric() && c != '_' && c != '\'');
841
842 if first_word.is_empty() {
844 return None;
845 }
846
847 let first_char = first_word.chars().next()?;
848 if first_char.is_lowercase() || first_char == '_' {
849 Some(first_word.to_string())
850 } else {
851 None
852 }
853}
854
855fn chunk_language_with_model(
856 text: &str,
857 language: ParseableLanguage,
858 _model_name: Option<&str>,
859) -> Result<Vec<Chunk>> {
860 chunk_language(text, language)
864}
865
866fn extract_code_chunks(
867 cursor: &mut tree_sitter::TreeCursor,
868 source: &str,
869 chunks: &mut Vec<Chunk>,
870 language: ParseableLanguage,
871) {
872 let node = cursor.node();
873
874 let should_skip = if language == ParseableLanguage::Haskell && node.kind() == "function" {
877 let mut current = node.parent();
879 while let Some(parent) = current {
880 if parent.kind() == "signature" {
881 return; }
883 current = parent.parent();
884 }
885 false
886 } else {
887 false
888 };
889
890 if !should_skip
891 && let Some(initial_chunk_type) = chunk_type_for_node(language, &node)
892 && let Some(chunk) = build_chunk(node, source, initial_chunk_type, language)
893 {
894 let is_duplicate = chunks.iter().any(|existing| {
895 existing.span.byte_start == chunk.span.byte_start
896 && existing.span.byte_end == chunk.span.byte_end
897 });
898
899 if !is_duplicate {
900 chunks.push(chunk);
901 }
902 }
903
904 let should_recurse = !(language == ParseableLanguage::Haskell && node.kind() == "signature");
906
907 if should_recurse && cursor.goto_first_child() {
908 loop {
909 extract_code_chunks(cursor, source, chunks, language);
910 if !cursor.goto_next_sibling() {
911 break;
912 }
913 }
914 cursor.goto_parent();
915 }
916}
917
918fn chunk_type_for_node(
919 language: ParseableLanguage,
920 node: &tree_sitter::Node<'_>,
921) -> Option<ChunkType> {
922 let kind = node.kind();
923
924 let supported = match language {
925 ParseableLanguage::Python => matches!(kind, "function_definition" | "class_definition"),
926 ParseableLanguage::TypeScript | ParseableLanguage::JavaScript => matches!(
927 kind,
928 "function_declaration" | "class_declaration" | "method_definition" | "arrow_function"
929 ),
930 ParseableLanguage::Haskell => matches!(
931 kind,
932 "function" | "signature" | "data_type"
935 | "newtype"
936 | "type_synonym"
937 | "type_family"
938 | "class"
939 | "instance"
940 ),
941 ParseableLanguage::Rust => matches!(
942 kind,
943 "function_item" | "impl_item" | "struct_item" | "enum_item" | "trait_item" | "mod_item"
944 ),
945 ParseableLanguage::Ruby => {
946 matches!(kind, "method" | "class" | "module" | "singleton_method")
947 }
948 ParseableLanguage::Go => matches!(
949 kind,
950 "function_declaration"
951 | "method_declaration"
952 | "type_declaration"
953 | "var_declaration"
954 | "const_declaration"
955 ),
956 ParseableLanguage::C => matches!(
957 kind,
958 "function_definition"
959 | "struct_specifier"
960 | "enum_specifier"
961 | "union_specifier"
962 | "type_definition"
963 | "declaration"
964 | "preproc_function_def"
965 | "preproc_def"
966 ),
967 ParseableLanguage::Cpp => matches!(
968 kind,
969 "function_definition"
970 | "class_specifier"
971 | "struct_specifier"
972 | "enum_specifier"
973 | "union_specifier"
974 | "namespace_definition"
975 | "template_declaration"
976 | "type_definition"
977 | "alias_declaration"
978 | "declaration"
979 | "preproc_function_def"
980 | "preproc_def"
981 ),
982 ParseableLanguage::CSharp => matches!(
983 kind,
984 "method_declaration"
985 | "class_declaration"
986 | "interface_declaration"
987 | "variable_declaration"
988 ),
989 ParseableLanguage::Dart => matches!(
990 kind,
991 "class_definition"
992 | "class_declaration"
993 | "mixin_declaration"
994 | "enum_declaration"
995 | "function_declaration"
996 | "method_declaration"
997 | "constructor_declaration"
998 | "variable_declaration"
999 | "local_variable_declaration"
1000 | "lambda_expression"
1001 | "class_member_definition"
1002 ),
1003 ParseableLanguage::Zig => matches!(
1004 kind,
1005 "function_declaration"
1006 | "test_declaration"
1007 | "variable_declaration"
1008 | "struct_declaration"
1009 | "enum_declaration"
1010 | "union_declaration"
1011 | "opaque_declaration"
1012 | "error_set_declaration"
1013 | "comptime_declaration"
1014 ),
1015 ParseableLanguage::Elixir => matches!(kind, "call" | "do_block"),
1018 ParseableLanguage::Markdown => matches!(
1019 kind,
1020 "atx_heading"
1021 | "setext_heading"
1022 | "heading"
1023 | "section"
1024 | "fenced_code_block"
1025 | "indented_code_block"
1026 | "block_quote"
1027 | "list"
1028 | "list_item"
1029 | "paragraph"
1030 | "thematic_break"
1031 ),
1032 };
1033
1034 if !supported {
1035 return None;
1036 }
1037
1038 match language {
1039 ParseableLanguage::Go
1040 if matches!(node.kind(), "var_declaration" | "const_declaration")
1041 && node.parent().is_some_and(|p| p.kind() == "block") =>
1042 {
1043 return None;
1044 }
1045 ParseableLanguage::CSharp if node.kind() == "variable_declaration" => {
1046 if !is_csharp_field_like(*node) {
1047 return None;
1048 }
1049 }
1050 _ => {}
1051 }
1052
1053 Some(classify_chunk_kind(kind))
1054}
1055
1056fn classify_chunk_kind(kind: &str) -> ChunkType {
1057 match kind {
1058 "function_definition"
1059 | "function_declaration"
1060 | "arrow_function"
1061 | "function"
1062 | "function_item"
1063 | "def"
1064 | "defp"
1065 | "defn"
1066 | "defn-"
1067 | "method"
1068 | "singleton_method"
1069 | "preproc_function_def" => ChunkType::Function,
1070 "signature" => ChunkType::Function, "class_definition"
1072 | "class_declaration"
1073 | "instance_declaration"
1074 | "class"
1075 | "instance"
1076 | "struct_item"
1077 | "enum_item"
1078 | "class_specifier"
1079 | "struct_specifier"
1080 | "enum_specifier"
1081 | "union_specifier"
1082 | "defstruct"
1083 | "defrecord"
1084 | "deftype"
1085 | "type_declaration"
1086 | "struct_declaration"
1087 | "enum_declaration"
1088 | "union_declaration"
1089 | "opaque_declaration"
1090 | "error_set_declaration" => ChunkType::Class,
1091 "method_definition" | "method_declaration" | "defmacro" => ChunkType::Method,
1092 "data_type"
1093 | "newtype"
1094 | "type_synonym"
1095 | "type_family"
1096 | "impl_item"
1097 | "trait_item"
1098 | "mod_item"
1099 | "namespace_definition"
1100 | "defmodule"
1101 | "module"
1102 | "defprotocol"
1103 | "interface_declaration"
1104 | "ns"
1105 | "var_declaration"
1106 | "const_declaration"
1107 | "variable_declaration"
1108 | "test_declaration"
1109 | "comptime_declaration"
1110 | "atx_heading"
1111 | "setext_heading"
1112 | "heading"
1113 | "section" => ChunkType::Module,
1114 _ => ChunkType::Text,
1115 }
1116}
1117
1118pub(crate) fn build_chunk(
1119 node: tree_sitter::Node<'_>,
1120 source: &str,
1121 initial_type: ChunkType,
1122 language: ParseableLanguage,
1123) -> Option<Chunk> {
1124 let target_node = adjust_node_for_language(node, language);
1125
1126 if matches!(language, ParseableLanguage::C | ParseableLanguage::Cpp)
1127 && matches!(initial_type, ChunkType::Class)
1128 && matches!(
1129 target_node.kind(),
1130 "struct_specifier" | "union_specifier" | "enum_specifier"
1131 )
1132 && !c_cpp_type_has_body_node(target_node)
1133 {
1134 return None;
1135 }
1136 let (byte_start, start_row, leading_segments) =
1137 extend_with_leading_trivia(target_node, language, source);
1138 let trailing_segments = collect_trailing_trivia(target_node, language, source);
1139
1140 let byte_end = target_node.end_byte();
1141 let end_pos = target_node.end_position();
1142
1143 if byte_start >= byte_end || byte_end > source.len() {
1144 return None;
1145 }
1146
1147 let chunk_type = adjust_chunk_type_for_context(target_node, initial_type, language);
1148 let mut text = source.get(byte_start..byte_end)?.to_string();
1149 if matches!(language, ParseableLanguage::C | ParseableLanguage::Cpp)
1150 && chunk_type == ChunkType::Class
1151 {
1152 text = strip_method_bodies_in_class_text(target_node, source, byte_start, byte_end);
1153 }
1154
1155 if text.trim().is_empty() {
1156 return None;
1157 }
1158 let ancestry = collect_ancestry(target_node, language, source);
1159 let leading_trivia = segments_to_strings(&leading_segments, source);
1160 let trailing_trivia = segments_to_strings(&trailing_segments, source);
1161 let mut metadata =
1162 ChunkMetadata::from_context(&text, ancestry, leading_trivia, trailing_trivia);
1163 if matches!(language, ParseableLanguage::C | ParseableLanguage::Cpp)
1164 && matches!(chunk_type, ChunkType::Function | ChunkType::Method)
1165 && let Some(full_name) = c_cpp_function_breadcrumb(target_node, language, source)
1166 {
1167 metadata.breadcrumb = Some(full_name);
1168 }
1169
1170 Some(Chunk {
1171 span: Span {
1172 byte_start,
1173 byte_end,
1174 line_start: start_row + 1,
1175 line_end: end_pos.row + 1,
1176 },
1177 text,
1178 chunk_type,
1179 stride_info: None,
1180 metadata,
1181 })
1182}
1183
1184fn c_cpp_type_has_body_node(node: tree_sitter::Node<'_>) -> bool {
1185 let mut cursor = node.walk();
1186
1187 match node.kind() {
1188 "struct_specifier" | "union_specifier" => node
1189 .children(&mut cursor)
1190 .any(|child| child.kind() == "field_declaration_list"),
1191 "enum_specifier" => node
1192 .children(&mut cursor)
1193 .any(|child| child.kind() == "enumerator_list"),
1194 _ => false,
1195 }
1196}
1197
1198fn c_cpp_function_breadcrumb(
1199 node: tree_sitter::Node<'_>,
1200 language: ParseableLanguage,
1201 source: &str,
1202) -> Option<String> {
1203 let name = display_name_for_node(node, language, source, ChunkType::Function)?;
1204 let context = collect_c_cpp_context_names(node, language, source);
1205 let context_path = context.join("::");
1206
1207 if name.contains("::") {
1208 if context_path.is_empty() || name.starts_with(&format!("{}::", context_path)) {
1209 Some(name)
1210 } else {
1211 Some(format!("{}::{}", context_path, name))
1212 }
1213 } else if context_path.is_empty() {
1214 Some(name)
1215 } else {
1216 Some(format!("{}::{}", context_path, name))
1217 }
1218}
1219
1220fn collect_c_cpp_context_names(
1221 mut node: tree_sitter::Node<'_>,
1222 language: ParseableLanguage,
1223 source: &str,
1224) -> Vec<String> {
1225 let mut parts = Vec::new();
1226
1227 while let Some(parent) = node.parent() {
1228 let kind = parent.kind();
1229 let include = match language {
1230 ParseableLanguage::Cpp => matches!(
1231 kind,
1232 "namespace_definition" | "class_specifier" | "struct_specifier"
1233 ),
1234 ParseableLanguage::C => matches!(kind, "struct_specifier"),
1235 _ => false,
1236 };
1237
1238 if include
1239 && let Some(name) = display_name_for_node(parent, language, source, ChunkType::Class)
1240 {
1241 parts.push(name);
1242 }
1243
1244 node = parent;
1245 }
1246
1247 parts.reverse();
1248 parts
1249}
1250
1251fn strip_method_bodies_in_class_text(
1252 class_node: tree_sitter::Node<'_>,
1253 source: &str,
1254 byte_start: usize,
1255 byte_end: usize,
1256) -> String {
1257 let mut replacements: Vec<(usize, usize, String)> = Vec::new();
1258 let mut stack = vec![class_node];
1259
1260 while let Some(node) = stack.pop() {
1261 if is_method_like_node(node.kind())
1262 && let Some(body) = find_method_body_node(node)
1263 {
1264 let start = body.start_byte();
1265 let end = body.end_byte();
1266 if start >= byte_start && end <= byte_end && start < end {
1267 let replacement = method_body_placeholder(body, source);
1268 replacements.push((start, end, replacement));
1269 }
1270 }
1271
1272 let child_count = node.child_count();
1273 for idx in (0..child_count).rev() {
1274 if let Some(child) = node.child(idx) {
1275 stack.push(child);
1276 }
1277 }
1278 }
1279
1280 if replacements.is_empty() {
1281 return source
1282 .get(byte_start..byte_end)
1283 .unwrap_or_default()
1284 .to_string();
1285 }
1286
1287 replacements.sort_by(|a, b| b.0.cmp(&a.0));
1288 let mut text = source
1289 .get(byte_start..byte_end)
1290 .unwrap_or_default()
1291 .to_string();
1292
1293 for (start, end, replacement) in replacements {
1294 if start < byte_start || end > byte_end || end <= start {
1295 continue;
1296 }
1297 let local_start = start - byte_start;
1298 let local_end = end - byte_start;
1299 if local_end <= text.len() {
1300 text.replace_range(local_start..local_end, &replacement);
1301 }
1302 }
1303
1304 text
1305}
1306
1307fn is_method_like_node(kind: &str) -> bool {
1308 matches!(
1309 kind,
1310 "function_definition"
1311 | "method_definition"
1312 | "method_declaration"
1313 | "constructor_declaration"
1314 | "destructor_declaration"
1315 | "function_item"
1316 | "method"
1317 | "singleton_method"
1318 )
1319}
1320
1321fn find_method_body_node(node: tree_sitter::Node<'_>) -> Option<tree_sitter::Node<'_>> {
1322 let body_kinds = [
1323 "compound_statement",
1324 "statement_block",
1325 "block",
1326 "body",
1327 "body_statement",
1328 "declaration_list",
1329 ];
1330
1331 for idx in 0..node.child_count() {
1332 if let Some(child) = node.child(idx)
1333 && body_kinds.contains(&child.kind())
1334 {
1335 return Some(child);
1336 }
1337 }
1338
1339 None
1340}
1341
1342fn method_body_placeholder(_body: tree_sitter::Node<'_>, _source: &str) -> String {
1343 ";".to_string()
1344}
1345
1346fn adjust_node_for_language(
1347 node: tree_sitter::Node<'_>,
1348 language: ParseableLanguage,
1349) -> tree_sitter::Node<'_> {
1350 match language {
1351 ParseableLanguage::TypeScript | ParseableLanguage::JavaScript => {
1352 if node.kind() == "arrow_function" {
1353 return expand_arrow_function_context(node);
1354 }
1355 node
1356 }
1357 _ => node,
1358 }
1359}
1360
1361fn expand_arrow_function_context(mut node: tree_sitter::Node<'_>) -> tree_sitter::Node<'_> {
1362 const PARENTS: &[&str] = &[
1363 "parenthesized_expression",
1364 "variable_declarator",
1365 "variable_declaration",
1366 "lexical_declaration",
1367 "assignment_expression",
1368 "expression_statement",
1369 "public_field_definition",
1370 "export_statement",
1371 ];
1372
1373 while let Some(parent) = node.parent() {
1374 let kind = parent.kind();
1375 if PARENTS.contains(&kind) {
1376 node = parent;
1377 continue;
1378 }
1379 break;
1380 }
1381
1382 node
1383}
1384
1385#[derive(Clone, Copy)]
1386struct TriviaSegment {
1387 start_byte: usize,
1388 end_byte: usize,
1389}
1390
1391fn extend_with_leading_trivia(
1392 node: tree_sitter::Node<'_>,
1393 language: ParseableLanguage,
1394 source: &str,
1395) -> (usize, usize, Vec<TriviaSegment>) {
1396 let mut start_byte = node.start_byte();
1397 let mut start_row = node.start_position().row;
1398 let mut current = node;
1399 let mut segments = Vec::new();
1400
1401 while let Some(prev) = current.prev_sibling() {
1402 if should_attach_leading_trivia(language, &prev)
1403 && only_whitespace_between(source, prev.end_byte(), start_byte)
1404 {
1405 start_byte = prev.start_byte();
1406 start_row = prev.start_position().row;
1407 segments.push(TriviaSegment {
1408 start_byte: prev.start_byte(),
1409 end_byte: prev.end_byte(),
1410 });
1411 current = prev;
1412 continue;
1413 }
1414 break;
1415 }
1416
1417 segments.reverse();
1418 (start_byte, start_row, segments)
1419}
1420
1421fn should_attach_leading_trivia(language: ParseableLanguage, node: &tree_sitter::Node<'_>) -> bool {
1422 let kind = node.kind();
1423 if kind == "comment" {
1424 return true;
1425 }
1426
1427 match language {
1428 ParseableLanguage::Rust => {
1429 matches!(kind, "line_comment" | "block_comment" | "attribute_item")
1430 }
1431 ParseableLanguage::Python => kind == "decorator",
1432 ParseableLanguage::TypeScript | ParseableLanguage::JavaScript => kind == "decorator",
1433 ParseableLanguage::C | ParseableLanguage::Cpp | ParseableLanguage::Markdown => {
1434 kind == "comment"
1435 }
1436 ParseableLanguage::CSharp => matches!(kind, "attribute_list" | "attribute"),
1437 _ => false,
1438 }
1439}
1440
1441fn collect_trailing_trivia(
1442 node: tree_sitter::Node<'_>,
1443 language: ParseableLanguage,
1444 source: &str,
1445) -> Vec<TriviaSegment> {
1446 let mut segments = Vec::new();
1447 let mut current = node;
1448 let mut previous_end = node.end_byte();
1449
1450 while let Some(next) = current.next_sibling() {
1451 if should_attach_trailing_trivia(language, &next)
1452 && only_whitespace_between(source, previous_end, next.start_byte())
1453 {
1454 segments.push(TriviaSegment {
1455 start_byte: next.start_byte(),
1456 end_byte: next.end_byte(),
1457 });
1458 previous_end = next.end_byte();
1459 current = next;
1460 continue;
1461 }
1462 break;
1463 }
1464
1465 segments
1466}
1467
1468fn should_attach_trailing_trivia(
1469 _language: ParseableLanguage,
1470 node: &tree_sitter::Node<'_>,
1471) -> bool {
1472 node.kind() == "comment"
1473}
1474
1475fn segments_to_strings(segments: &[TriviaSegment], source: &str) -> Vec<String> {
1476 let mut result = Vec::new();
1477
1478 for segment in segments {
1479 if let Some(text) = source
1480 .get(segment.start_byte..segment.end_byte)
1481 .map(std::string::ToString::to_string)
1482 {
1483 result.push(text);
1484 }
1485 }
1486
1487 result
1488}
1489
1490fn collect_ancestry(
1491 mut node: tree_sitter::Node<'_>,
1492 language: ParseableLanguage,
1493 source: &str,
1494) -> Vec<String> {
1495 if language == ParseableLanguage::Markdown {
1496 return markdown_heading_ancestry(node, source);
1497 }
1498
1499 let mut parts = Vec::new();
1500
1501 while let Some(parent) = node.parent() {
1502 if let Some(parent_chunk_type) = chunk_type_for_node(language, &parent)
1503 && let Some(name) = display_name_for_node(parent, language, source, parent_chunk_type)
1504 {
1505 parts.push(name);
1506 }
1507 node = parent;
1508 }
1509
1510 parts.reverse();
1511 parts
1512}
1513
1514fn display_name_for_node(
1515 node: tree_sitter::Node<'_>,
1516 language: ParseableLanguage,
1517 source: &str,
1518 chunk_type: ChunkType,
1519) -> Option<String> {
1520 if let Some(name_node) = node.child_by_field_name("name") {
1521 return text_for_node(name_node, source);
1522 }
1523
1524 match language {
1525 ParseableLanguage::Rust => rust_display_name(node, source, chunk_type),
1526 ParseableLanguage::Python => find_identifier(node, source, &["identifier"]),
1527 ParseableLanguage::TypeScript | ParseableLanguage::JavaScript => find_identifier(
1528 node,
1529 source,
1530 &["identifier", "type_identifier", "property_identifier"],
1531 ),
1532 ParseableLanguage::Haskell => {
1533 find_identifier(node, source, &["identifier", "type_identifier", "variable"])
1534 .or_else(|| first_word_of_node(node, source))
1535 }
1536 ParseableLanguage::Ruby => find_identifier(node, source, &["identifier"]),
1537 ParseableLanguage::Go => find_identifier(node, source, &["identifier", "type_identifier"]),
1538 ParseableLanguage::C => c_display_name(node, source, chunk_type),
1539 ParseableLanguage::Cpp => cpp_display_name(node, source, chunk_type),
1540 ParseableLanguage::CSharp => find_identifier(node, source, &["identifier"]),
1541 ParseableLanguage::Zig => find_identifier(node, source, &["identifier"]),
1542
1543 ParseableLanguage::Markdown => markdown_display_name(node, source, chunk_type),
1544
1545 ParseableLanguage::Dart => {
1546 find_identifier(node, source, &["identifier", "type_identifier"])
1547 }
1548 ParseableLanguage::Elixir => {
1549 find_identifier(node, source, &["alias", "identifier", "atom"])
1551 }
1552 }
1553}
1554
1555fn markdown_display_name(
1556 node: tree_sitter::Node<'_>,
1557 source: &str,
1558 _chunk_type: ChunkType,
1559) -> Option<String> {
1560 if node.kind() == "section" {
1561 return markdown_section_heading(node, source);
1562 }
1563
1564 if markdown_heading_kind(node.kind()) {
1565 return markdown_heading_text(node, source);
1566 }
1567
1568 None
1569}
1570
1571fn markdown_section_heading(node: tree_sitter::Node<'_>, source: &str) -> Option<String> {
1572 let mut cursor = node.walk();
1573 for child in node.children(&mut cursor) {
1574 if markdown_heading_kind(child.kind()) {
1575 return markdown_heading_text(child, source);
1576 }
1577 }
1578 None
1579}
1580
1581fn markdown_heading_kind(kind: &str) -> bool {
1582 matches!(kind, "atx_heading" | "setext_heading" | "heading")
1583}
1584
1585fn markdown_heading_text(node: tree_sitter::Node<'_>, source: &str) -> Option<String> {
1586 let text = text_for_node(node, source)?;
1587 let mut lines = text.lines();
1588 let first_line = lines.next().unwrap_or("");
1589
1590 if let Some((_, heading)) = parse_atx_heading_line(first_line) {
1591 return Some(heading);
1592 }
1593
1594 let second_line = lines.next().unwrap_or("");
1595 if parse_setext_level(second_line).is_some() {
1596 let trimmed = first_line.trim();
1597 if !trimmed.is_empty() {
1598 return Some(trimmed.to_string());
1599 }
1600 }
1601
1602 let trimmed = first_line.trim();
1603 if trimmed.is_empty() {
1604 None
1605 } else {
1606 Some(trimmed.to_string())
1607 }
1608}
1609
1610fn markdown_heading_ancestry(node: tree_sitter::Node<'_>, source: &str) -> Vec<String> {
1611 let mut target_row = node.start_position().row;
1612 if node.kind() == "section" || markdown_heading_kind(node.kind()) {
1613 target_row = target_row.saturating_sub(1);
1614 }
1615 let lines: Vec<&str> = source.lines().collect();
1616 let mut stack: Vec<(usize, String)> = Vec::new();
1617 let mut i = 0;
1618
1619 while i < lines.len() && i <= target_row {
1620 let line = lines[i];
1621
1622 if let Some((level, heading)) = parse_atx_heading_line(line) {
1623 update_markdown_heading_stack(&mut stack, level, heading);
1624 i += 1;
1625 continue;
1626 }
1627
1628 if i + 1 < lines.len() && i < target_row {
1629 let underline = lines[i + 1];
1630 if let Some(level) = parse_setext_level(underline) {
1631 let heading_text = line.trim();
1632 if !heading_text.is_empty() {
1633 update_markdown_heading_stack(&mut stack, level, heading_text.to_string());
1634 }
1635 i += 2;
1636 continue;
1637 }
1638 }
1639
1640 i += 1;
1641 }
1642
1643 stack.into_iter().map(|(_, heading)| heading).collect()
1644}
1645
1646fn update_markdown_heading_stack(stack: &mut Vec<(usize, String)>, level: usize, text: String) {
1647 while let Some((existing_level, _)) = stack.last() {
1648 if *existing_level < level {
1649 break;
1650 }
1651 stack.pop();
1652 }
1653 stack.push((level, text));
1654}
1655
1656fn parse_atx_heading_line(line: &str) -> Option<(usize, String)> {
1657 let trimmed = line.trim_start();
1658 if !trimmed.starts_with('#') {
1659 return None;
1660 }
1661
1662 let level = trimmed.chars().take_while(|c| *c == '#').count();
1663 if level == 0 {
1664 return None;
1665 }
1666
1667 let mut text = trimmed[level..].trim();
1668 text = text.trim_end_matches('#').trim();
1669 if text.is_empty() {
1670 return None;
1671 }
1672
1673 Some((level, text.to_string()))
1674}
1675
1676fn parse_setext_level(line: &str) -> Option<usize> {
1677 let trimmed = line.trim();
1678 if trimmed.is_empty() {
1679 return None;
1680 }
1681
1682 if trimmed.chars().all(|c| c == '=') {
1683 Some(1)
1684 } else if trimmed.chars().all(|c| c == '-') {
1685 Some(2)
1686 } else {
1687 None
1688 }
1689}
1690
1691fn rust_display_name(
1692 node: tree_sitter::Node<'_>,
1693 source: &str,
1694 chunk_type: ChunkType,
1695) -> Option<String> {
1696 match node.kind() {
1697 "impl_item" => {
1698 let mut parts = Vec::new();
1699 if let Some(ty) = node.child_by_field_name("type")
1700 && let Some(text) = text_for_node(ty, source)
1701 {
1702 parts.push(text);
1703 }
1704 if let Some(trait_node) = node.child_by_field_name("trait")
1705 && let Some(text) = text_for_node(trait_node, source)
1706 {
1707 if let Some(last) = parts.first() {
1708 parts[0] = format!("{} (impl {})", last, text.trim());
1709 } else {
1710 parts.push(format!("impl {}", text.trim()));
1711 }
1712 }
1713 if parts.is_empty() {
1714 find_identifier(node, source, &["identifier"])
1715 } else {
1716 Some(parts.remove(0))
1717 }
1718 }
1719 "mod_item" if chunk_type == ChunkType::Module => {
1720 find_identifier(node, source, &["identifier"])
1721 }
1722 _ => find_identifier(node, source, &["identifier", "type_identifier"]),
1723 }
1724}
1725
1726fn c_display_name(
1727 node: tree_sitter::Node<'_>,
1728 source: &str,
1729 _chunk_type: ChunkType,
1730) -> Option<String> {
1731 match node.kind() {
1732 "function_definition" => {
1733 if let Some(declarator) = node.child_by_field_name("declarator") {
1735 return find_identifier_recursive(declarator, source, &["identifier"]);
1736 }
1737 None
1738 }
1739 "struct_specifier" | "enum_specifier" | "union_specifier" => {
1740 find_identifier(node, source, &["type_identifier", "identifier"])
1741 }
1742 "type_definition" => find_identifier(node, source, &["type_identifier", "identifier"]),
1743 "preproc_function_def" | "preproc_def" => find_identifier(node, source, &["identifier"]),
1744 _ => find_identifier(node, source, &["identifier", "type_identifier"]),
1745 }
1746}
1747
1748fn cpp_display_name(
1749 node: tree_sitter::Node<'_>,
1750 source: &str,
1751 _chunk_type: ChunkType,
1752) -> Option<String> {
1753 match node.kind() {
1754 "function_definition" => {
1755 if let Some(declarator) = node.child_by_field_name("declarator") {
1756 return find_identifier_recursive(
1757 declarator,
1758 source,
1759 &[
1760 "identifier",
1761 "field_identifier",
1762 "destructor_name",
1763 "qualified_identifier",
1764 ],
1765 );
1766 }
1767 None
1768 }
1769 "declaration" => {
1770 if let Some(declarator) = node.child_by_field_name("declarator") {
1771 return find_identifier_recursive(
1772 declarator,
1773 source,
1774 &[
1775 "identifier",
1776 "field_identifier",
1777 "destructor_name",
1778 "qualified_identifier",
1779 ],
1780 );
1781 }
1782 find_identifier(node, source, &["identifier", "type_identifier"])
1783 }
1784 "class_specifier" | "struct_specifier" | "enum_specifier" | "union_specifier" => {
1785 find_identifier(node, source, &["type_identifier", "identifier"])
1786 }
1787 "namespace_definition" => {
1788 find_identifier(node, source, &["identifier", "namespace_identifier"])
1789 }
1790 "alias_declaration" | "type_definition" => {
1791 find_identifier(node, source, &["type_identifier", "identifier"])
1792 }
1793 "template_declaration" => {
1794 let mut cursor = node.walk();
1795 for child in node.children(&mut cursor) {
1796 if matches!(
1797 child.kind(),
1798 "class_specifier"
1799 | "struct_specifier"
1800 | "enum_specifier"
1801 | "union_specifier"
1802 | "function_definition"
1803 | "declaration"
1804 | "alias_declaration"
1805 | "type_definition"
1806 | "concept_definition"
1807 ) {
1808 return cpp_display_name(child, source, _chunk_type);
1809 }
1810 }
1811 find_identifier(node, source, &["type_identifier", "identifier"])
1812 }
1813 _ => find_identifier(node, source, &["identifier", "type_identifier"]),
1814 }
1815}
1816
1817fn find_identifier_recursive(
1819 node: tree_sitter::Node<'_>,
1820 source: &str,
1821 candidate_kinds: &[&str],
1822) -> Option<String> {
1823 if candidate_kinds.contains(&node.kind()) {
1824 return text_for_node(node, source).map(|s| s.trim().to_string());
1825 }
1826 let mut cursor = node.walk();
1827 for child in node.children(&mut cursor) {
1828 if let Some(result) = find_identifier_recursive(child, source, candidate_kinds) {
1829 return Some(result);
1830 }
1831 }
1832 None
1833}
1834
1835fn find_identifier(
1836 node: tree_sitter::Node<'_>,
1837 source: &str,
1838 candidate_kinds: &[&str],
1839) -> Option<String> {
1840 let mut cursor = node.walk();
1841 for child in node.children(&mut cursor) {
1842 if candidate_kinds.contains(&child.kind())
1843 && let Some(text) = text_for_node(child, source)
1844 {
1845 return Some(text.trim().to_string());
1846 }
1847 }
1848 None
1849}
1850
1851fn first_word_of_node(node: tree_sitter::Node<'_>, source: &str) -> Option<String> {
1852 let text = text_for_node(node, source)?;
1853 text.split_whitespace().next().map(|s| {
1854 s.trim_end_matches(|c: char| !c.is_alphanumeric() && c != '_')
1855 .to_string()
1856 })
1857}
1858
1859fn text_for_node(node: tree_sitter::Node<'_>, source: &str) -> Option<String> {
1860 node.utf8_text(source.as_bytes())
1861 .ok()
1862 .map(std::string::ToString::to_string)
1863}
1864
1865fn only_whitespace_between(source: &str, start: usize, end: usize) -> bool {
1866 if start >= end || end > source.len() {
1867 return true;
1868 }
1869
1870 source[start..end].chars().all(char::is_whitespace)
1871}
1872
1873fn adjust_chunk_type_for_context(
1874 node: tree_sitter::Node<'_>,
1875 chunk_type: ChunkType,
1876 language: ParseableLanguage,
1877) -> ChunkType {
1878 if chunk_type != ChunkType::Function {
1879 return chunk_type;
1880 }
1881
1882 if is_method_context(node, language) {
1883 ChunkType::Method
1884 } else {
1885 chunk_type
1886 }
1887}
1888
1889fn is_method_context(node: tree_sitter::Node<'_>, language: ParseableLanguage) -> bool {
1890 const PYTHON_CONTAINERS: &[&str] = &["class_definition"];
1891 const TYPESCRIPT_CONTAINERS: &[&str] = &["class_body", "class_declaration"];
1892 const RUBY_CONTAINERS: &[&str] = &["class", "module"];
1893 const RUST_CONTAINERS: &[&str] = &["impl_item", "trait_item"];
1894 const DART_CONTAINERS: &[&str] = &[
1895 "class_definition",
1896 "class_declaration",
1897 "mixin_declaration",
1898 "enum_declaration",
1899 ];
1900
1901 match language {
1902 ParseableLanguage::Python => ancestor_has_kind(node, PYTHON_CONTAINERS),
1903 ParseableLanguage::TypeScript | ParseableLanguage::JavaScript => {
1904 ancestor_has_kind(node, TYPESCRIPT_CONTAINERS)
1905 }
1906 ParseableLanguage::Ruby => ancestor_has_kind(node, RUBY_CONTAINERS),
1907 ParseableLanguage::Rust => ancestor_has_kind(node, RUST_CONTAINERS),
1908 ParseableLanguage::Go => false,
1909 ParseableLanguage::C => ancestor_has_kind(node, &["struct_specifier"]),
1910 ParseableLanguage::Cpp => ancestor_has_kind(node, &["class_specifier", "struct_specifier"]),
1911 ParseableLanguage::CSharp => false,
1912 ParseableLanguage::Haskell => false,
1913 ParseableLanguage::Zig => false,
1914
1915 ParseableLanguage::Dart => ancestor_has_kind(node, DART_CONTAINERS),
1916
1917 ParseableLanguage::Elixir => false, ParseableLanguage::Markdown => false,
1919 }
1920}
1921
1922fn ancestor_has_kind(node: tree_sitter::Node<'_>, kinds: &[&str]) -> bool {
1923 let mut current = node;
1924 while let Some(parent) = current.parent() {
1925 if kinds.contains(&parent.kind()) {
1926 return true;
1927 }
1928 current = parent;
1929 }
1930 false
1931}
1932
1933fn is_csharp_field_like(node: tree_sitter::Node<'_>) -> bool {
1934 if let Some(parent) = node.parent() {
1935 return matches!(
1936 parent.kind(),
1937 "field_declaration" | "event_field_declaration"
1938 );
1939 }
1940 false
1941}
1942
1943fn apply_striding(chunks: Vec<Chunk>, config: &ChunkConfig) -> Result<Vec<Chunk>> {
1945 let mut result = Vec::new();
1946
1947 for chunk in chunks {
1948 let estimated_tokens = estimate_tokens(&chunk.text);
1949
1950 if estimated_tokens <= config.max_tokens {
1951 result.push(chunk);
1953 } else {
1954 tracing::debug!(
1956 "Chunk with {} tokens exceeds limit of {}, applying striding",
1957 estimated_tokens,
1958 config.max_tokens
1959 );
1960
1961 let strided_chunks = stride_large_chunk(chunk, config)?;
1962 result.extend(strided_chunks);
1963 }
1964 }
1965
1966 Ok(result)
1967}
1968
1969fn stride_large_chunk(chunk: Chunk, config: &ChunkConfig) -> Result<Vec<Chunk>> {
1971 let text = &chunk.text;
1972
1973 if text.is_empty() {
1975 return Ok(vec![chunk]);
1976 }
1977
1978 let char_count = text.chars().count();
1981 let estimated_tokens = estimate_tokens(text);
1982 let chars_per_token = if estimated_tokens == 0 {
1984 4.5 } else {
1986 char_count as f32 / estimated_tokens as f32
1987 };
1988 let window_chars = ((config.max_tokens as f32 * 0.9) * chars_per_token) as usize; let overlap_chars = (config.stride_overlap as f32 * chars_per_token) as usize;
1990 let stride_chars = window_chars.saturating_sub(overlap_chars);
1991
1992 if stride_chars == 0 {
1993 return Err(anyhow::anyhow!("Stride size is too small"));
1994 }
1995
1996 let char_byte_indices: Vec<(usize, char)> = text.char_indices().collect();
1998 let mut strided_chunks = Vec::new();
2001 let original_chunk_id = format!("{}:{}", chunk.span.byte_start, chunk.span.byte_end);
2002 let mut start_char_idx = 0;
2003 let mut stride_index = 0;
2004
2005 let total_strides = if char_count <= window_chars {
2007 1
2008 } else {
2009 ((char_count - overlap_chars) as f32 / stride_chars as f32).ceil() as usize
2010 };
2011
2012 while start_char_idx < char_count {
2013 let end_char_idx = (start_char_idx + window_chars).min(char_count);
2014
2015 let start_byte_pos = char_byte_indices[start_char_idx].0;
2017 let end_byte_pos = if end_char_idx < char_count {
2018 char_byte_indices[end_char_idx].0
2019 } else {
2020 text.len()
2021 };
2022
2023 let stride_text = &text[start_byte_pos..end_byte_pos];
2024
2025 let overlap_start = if stride_index > 0 { overlap_chars } else { 0 };
2027 let overlap_end = if end_char_idx < char_count {
2028 overlap_chars
2029 } else {
2030 0
2031 };
2032
2033 let byte_offset_start = chunk.span.byte_start + start_byte_pos;
2035 let byte_offset_end = chunk.span.byte_start + end_byte_pos;
2036
2037 let text_before_start = &text[..start_byte_pos];
2039 let line_offset_start = text_before_start.lines().count().saturating_sub(1);
2040 let stride_lines = stride_text.lines().count();
2041 let metadata = chunk.metadata.with_updated_text(stride_text);
2042
2043 let stride_chunk = Chunk {
2044 span: Span {
2045 byte_start: byte_offset_start,
2046 byte_end: byte_offset_end,
2047 line_start: chunk.span.line_start + line_offset_start,
2048 line_end: chunk.span.line_start
2050 + line_offset_start
2051 + stride_lines.saturating_sub(1),
2052 },
2053 text: stride_text.to_string(),
2054 chunk_type: chunk.chunk_type.clone(),
2055 stride_info: Some(StrideInfo {
2056 original_chunk_id: original_chunk_id.clone(),
2057 stride_index,
2058 total_strides,
2059 overlap_start,
2060 overlap_end,
2061 }),
2062 metadata,
2063 };
2064
2065 strided_chunks.push(stride_chunk);
2066
2067 if end_char_idx >= char_count {
2069 break;
2070 }
2071
2072 start_char_idx += stride_chars;
2073 stride_index += 1;
2074 }
2075
2076 tracing::debug!(
2077 "Created {} strides from chunk of {} tokens",
2078 strided_chunks.len(),
2079 estimate_tokens(text)
2080 );
2081
2082 Ok(strided_chunks)
2083}
2084
2085fn merge_small_chunks(chunks: Vec<Chunk>, text: &str, target_tokens: usize) -> Vec<Chunk> {
2098 if chunks.is_empty() {
2099 return chunks;
2100 }
2101
2102 let mut result = Vec::new();
2103 let mut current_group: Vec<Chunk> = Vec::new();
2104 let mut current_tokens = 0;
2105
2106 for chunk in chunks {
2107 let chunk_tokens = chunk.metadata.estimated_tokens;
2108
2109 if current_tokens + chunk_tokens > target_tokens {
2110 if !current_group.is_empty() {
2112 result.push(merge_group(¤t_group, text));
2113 current_group.clear();
2114 current_tokens = 0;
2115 }
2116 }
2117
2118 if chunk_tokens > target_tokens {
2120 if !current_group.is_empty() {
2121 result.push(merge_group(¤t_group, text));
2122 current_group.clear();
2123 current_tokens = 0;
2124 }
2125 result.push(chunk);
2126 continue;
2127 }
2128
2129 current_group.push(chunk);
2130 current_tokens += chunk_tokens;
2131 }
2132
2133 if !current_group.is_empty() {
2135 result.push(merge_group(¤t_group, text));
2136 }
2137
2138 result
2139}
2140
2141fn merge_group(group: &[Chunk], text: &str) -> Chunk {
2142 if group.len() == 1 {
2143 return group[0].clone();
2144 }
2145
2146 let first = &group[0];
2147 let last = &group[group.len() - 1];
2148
2149 let byte_start = first.span.byte_start;
2152 let byte_end = last.span.byte_end;
2153 let line_start = first.span.line_start;
2154 let line_end = last.span.line_end;
2155
2156 let chunk_text = if byte_end <= text.len() {
2158 text[byte_start..byte_end].to_string()
2159 } else {
2160 text.get(byte_start..).unwrap_or("").to_string()
2162 };
2163
2164 let metadata = ChunkMetadata::from_text(&chunk_text);
2165
2166 let chunk_type = if group.iter().all(|c| c.chunk_type == first.chunk_type) {
2173 first.chunk_type.clone()
2174 } else {
2175 ChunkType::Text
2176 };
2177
2178 Chunk {
2179 span: Span {
2180 byte_start,
2181 byte_end,
2182 line_start,
2183 line_end,
2184 },
2185 text: chunk_text,
2186 chunk_type,
2187 stride_info: None,
2188 metadata,
2189 }
2190}
2191
2192#[cfg(test)]
2195mod tests {
2196 use super::*;
2197
2198 fn canonicalize_spans(
2199 mut spans: Vec<(usize, usize, ChunkType)>,
2200 ) -> Vec<(usize, usize, ChunkType)> {
2201 fn chunk_type_order(chunk_type: &ChunkType) -> u8 {
2202 match chunk_type {
2203 ChunkType::Text => 0,
2204 ChunkType::Function => 1,
2205 ChunkType::Class => 2,
2206 ChunkType::Method => 3,
2207 ChunkType::Module => 4,
2208 }
2209 }
2210
2211 spans.sort_by(|a, b| {
2212 let order_a = chunk_type_order(&a.2);
2213 let order_b = chunk_type_order(&b.2);
2214 order_a
2215 .cmp(&order_b)
2216 .then_with(|| a.0.cmp(&b.0))
2217 .then_with(|| a.1.cmp(&b.1))
2218 });
2219
2220 let mut result: Vec<(usize, usize, ChunkType)> = Vec::new();
2221 for (start, end, ty) in spans {
2222 if let Some(last) = result.last_mut()
2223 && last.0 == start
2224 && last.2 == ty
2225 {
2226 if end > last.1 {
2227 last.1 = end;
2228 }
2229 continue;
2230 }
2231 result.push((start, end, ty));
2232 }
2233
2234 result
2235 }
2236
2237 fn assert_query_parity(language: ParseableLanguage, source: &str) {
2238 let mut parser = tree_sitter::Parser::new();
2239 let ts_language = tree_sitter_language(language).expect("language");
2240 parser.set_language(&ts_language).expect("set language");
2241 let tree = parser.parse(source, None).expect("parse source");
2242
2243 let query_chunks = query_chunker::chunk_with_queries(language, ts_language, &tree, source)
2244 .expect("query execution")
2245 .expect("queries available");
2246
2247 let mut legacy_chunks = Vec::new();
2248 let mut cursor = tree.walk();
2249 extract_code_chunks(&mut cursor, source, &mut legacy_chunks, language);
2250
2251 let query_spans = canonicalize_spans(
2252 query_chunks
2253 .iter()
2254 .map(|chunk| {
2255 (
2256 chunk.span.byte_start,
2257 chunk.span.byte_end,
2258 chunk.chunk_type.clone(),
2259 )
2260 })
2261 .collect(),
2262 );
2263 let legacy_spans = canonicalize_spans(
2264 legacy_chunks
2265 .iter()
2266 .map(|chunk| {
2267 (
2268 chunk.span.byte_start,
2269 chunk.span.byte_end,
2270 chunk.chunk_type.clone(),
2271 )
2272 })
2273 .collect(),
2274 );
2275
2276 assert_eq!(query_spans, legacy_spans);
2277 }
2278
2279 #[test]
2280 fn test_chunk_generic_byte_offsets() {
2281 let text = "line 1\nline 2\nline 3\nline 4\nline 5";
2283 let chunks = chunk_generic(text).unwrap();
2284
2285 assert!(!chunks.is_empty());
2286
2287 assert_eq!(chunks[0].span.byte_start, 0);
2289
2290 for chunk in &chunks {
2292 let expected_len = chunk.text.len();
2293 let actual_len = chunk.span.byte_end - chunk.span.byte_start;
2294 assert_eq!(actual_len, expected_len);
2295 }
2296 }
2297
2298 #[test]
2299 fn test_chunk_generic_large_file_performance() {
2300 let lines: Vec<String> = (0..1000)
2302 .map(|i| format!("Line {i}: Some content here"))
2303 .collect();
2304 let text = lines.join("\n");
2305
2306 let start = std::time::Instant::now();
2307 let chunks = chunk_generic(&text).unwrap();
2308 let duration = start.elapsed();
2309
2310 assert!(
2312 duration.as_millis() < 100,
2313 "Chunking took too long: {duration:?}"
2314 );
2315 assert!(!chunks.is_empty());
2316
2317 for chunk in &chunks {
2319 assert!(chunk.span.line_start > 0);
2320 assert!(chunk.span.line_end >= chunk.span.line_start);
2321 }
2322 }
2323
2324 #[test]
2325 fn test_chunk_rust() {
2326 let rust_code = r"
2327pub struct Calculator {
2328 memory: f64,
2329}
2330
2331impl Calculator {
2332 pub fn new() -> Self {
2333 Calculator { memory: 0.0 }
2334 }
2335
2336 pub fn add(&mut self, a: f64, b: f64) -> f64 {
2337 a + b
2338 }
2339}
2340
2341fn main() {
2342 let calc = Calculator::new();
2343}
2344
2345pub mod utils {
2346 pub fn helper() {}
2347}
2348";
2349
2350 let chunks = chunk_language(rust_code, ParseableLanguage::Rust).unwrap();
2351 assert!(!chunks.is_empty());
2352
2353 let chunk_types: Vec<&ChunkType> = chunks.iter().map(|c| &c.chunk_type).collect();
2355 assert!(chunk_types.contains(&&ChunkType::Class)); assert!(chunk_types.contains(&&ChunkType::Module)); assert!(chunk_types.contains(&&ChunkType::Function)); }
2359
2360 #[test]
2361 fn test_rust_doc_comments_attached() {
2362 let rust_code = r"
2363/// Doc comment
2364pub struct Foo {}
2365";
2366 let chunks = chunk_language(rust_code, ParseableLanguage::Rust).unwrap();
2367 let struct_chunk = chunks
2368 .iter()
2369 .find(|c| c.text.contains("struct Foo"))
2370 .unwrap();
2371 assert!(
2372 struct_chunk.text.contains("/// Doc comment"),
2373 "Doc comment should be attached"
2374 );
2375 }
2376
2377 #[test]
2378 fn test_rust_query_matches_legacy() {
2379 let source = r"
2380 mod sample {
2381 struct Thing;
2382
2383 impl Thing {
2384 fn new() -> Self { Self }
2385 fn helper(&self) {}
2386 }
2387 }
2388
2389 fn util() {}
2390 ";
2391
2392 assert_query_parity(ParseableLanguage::Rust, source);
2393 }
2394
2395 #[test]
2396 fn test_python_query_matches_legacy() {
2397 let source = r"
2398class Example:
2399 @classmethod
2400 def build(cls):
2401 return cls()
2402
2403
2404def helper():
2405 return 1
2406
2407
2408async def async_helper():
2409 return 2
2410";
2411
2412 assert_query_parity(ParseableLanguage::Python, source);
2413 }
2414
2415 #[test]
2416 fn test_chunk_ruby() {
2417 let ruby_code = r#"
2418class Calculator
2419 def initialize
2420 @memory = 0.0
2421 end
2422
2423 def add(a, b)
2424 a + b
2425 end
2426
2427 def self.class_method
2428 "class method"
2429 end
2430
2431 private
2432
2433 def private_method
2434 "private"
2435 end
2436end
2437
2438module Utils
2439 def self.helper
2440 "helper"
2441 end
2442end
2443
2444def main
2445 calc = Calculator.new
2446end
2447"#;
2448
2449 let chunks = chunk_language(ruby_code, ParseableLanguage::Ruby).unwrap();
2450 assert!(!chunks.is_empty());
2451
2452 let chunk_types: Vec<&ChunkType> = chunks.iter().map(|c| &c.chunk_type).collect();
2454 assert!(chunk_types.contains(&&ChunkType::Class)); assert!(chunk_types.contains(&&ChunkType::Module)); assert!(chunk_types.contains(&&ChunkType::Function)); }
2458
2459 #[test]
2460 fn test_language_detection_fallback() {
2461 let generic_text = "Some text\nwith multiple lines\nto chunk generically";
2463
2464 let chunks_unknown = chunk_text(generic_text, None).unwrap();
2465 let chunks_generic = chunk_generic(generic_text).unwrap();
2466
2467 assert_eq!(chunks_unknown.len(), chunks_generic.len());
2469 assert_eq!(chunks_unknown[0].text, chunks_generic[0].text);
2470 }
2471
2472 #[test]
2473 fn test_chunk_go() {
2474 let go_code = r#"
2475package main
2476
2477import "fmt"
2478
2479const Pi = 3.14159
2480
2481var memory float64
2482
2483type Calculator struct {
2484 memory float64
2485}
2486
2487type Operation interface {
2488 Calculate(a, b float64) float64
2489}
2490
2491func NewCalculator() *Calculator {
2492 return &Calculator{memory: 0.0}
2493}
2494
2495func (c *Calculator) Add(a, b float64) float64 {
2496 return a + b
2497}
2498
2499func main() {
2500 calc := NewCalculator()
2501}
2502"#;
2503
2504 let chunks = chunk_language(go_code, ParseableLanguage::Go).unwrap();
2505 assert!(!chunks.is_empty());
2506
2507 let chunk_types: Vec<&ChunkType> = chunks.iter().map(|c| &c.chunk_type).collect();
2509 assert!(chunk_types.contains(&&ChunkType::Module)); assert!(chunk_types.contains(&&ChunkType::Class)); assert!(chunk_types.contains(&&ChunkType::Function)); assert!(chunk_types.contains(&&ChunkType::Method)); }
2514
2515 #[test]
2516 #[ignore] fn test_chunk_typescript_arrow_context() {
2518 let ts_code = r"
2519// Utility function
2520export const util = () => {
2521 // comment about util
2522 return 42;
2523};
2524
2525export class Example {
2526 // leading comment for method
2527 constructor() {}
2528
2529 // Another comment
2530 run = () => {
2531 return util();
2532 };
2533}
2534
2535const compute = (x: number) => x * 2;
2536";
2537
2538 let chunks = chunk_language(ts_code, ParseableLanguage::TypeScript).unwrap();
2539
2540 let util_chunk = chunks
2541 .iter()
2542 .find(|chunk| chunk.text.contains("export const util"))
2543 .expect("Expected chunk for util arrow function");
2544 assert_eq!(util_chunk.chunk_type, ChunkType::Function);
2545 assert!(
2546 util_chunk.text.contains("// Utility function"),
2547 "expected leading comment to be included"
2548 );
2549 assert!(util_chunk.text.contains("export const util ="));
2550
2551 let method_chunk = chunks
2553 .iter()
2554 .find(|chunk| {
2555 chunk.chunk_type == ChunkType::Method && chunk.text.contains("run = () =>")
2556 })
2557 .expect("Expected chunk for class field arrow function");
2558
2559 assert_eq!(method_chunk.chunk_type, ChunkType::Method);
2560 assert!(
2561 method_chunk.text.contains("// Another comment"),
2562 "expected inline comment to be included"
2563 );
2564
2565 let compute_chunk = chunks
2566 .iter()
2567 .find(|chunk| chunk.text.contains("const compute"))
2568 .expect("Expected chunk for compute arrow function");
2569 assert_eq!(compute_chunk.chunk_type, ChunkType::Function);
2570 assert!(
2571 compute_chunk
2572 .text
2573 .contains("const compute = (x: number) => x * 2;")
2574 );
2575
2576 assert!(
2578 chunks
2579 .iter()
2580 .all(|chunk| !chunk.text.trim_start().starts_with("() =>"))
2581 );
2582 assert!(
2583 chunks
2584 .iter()
2585 .all(|chunk| !chunk.text.trim_start().starts_with("(x: number) =>"))
2586 );
2587 }
2588
2589 #[test]
2593 #[ignore]
2594 fn test_typescript_query_matches_legacy() {
2595 let source = r"
2596export const util = () => {
2597 return 42;
2598};
2599
2600export class Example {
2601 run = () => {
2602 return util();
2603 };
2604}
2605
2606const compute = (x: number) => x * 2;
2607";
2608
2609 assert_query_parity(ParseableLanguage::TypeScript, source);
2610 }
2611
2612 #[test]
2613 fn test_ruby_query_matches_legacy() {
2614 let source = r#"
2615class Calculator
2616 def initialize
2617 @memory = 0.0
2618 end
2619
2620 def add(a, b)
2621 a + b
2622 end
2623
2624 def self.class_method
2625 "class method"
2626 end
2627end
2628"#;
2629
2630 assert_query_parity(ParseableLanguage::Ruby, source);
2631 }
2632
2633 #[test]
2634 fn test_go_query_matches_legacy() {
2635 let source = r#"
2636package main
2637
2638import "fmt"
2639
2640const Pi = 3.14159
2641
2642var memory float64
2643
2644type Calculator struct {
2645 memory float64
2646}
2647
2648func (c *Calculator) Add(a, b float64) float64 {
2649 return a + b
2650}
2651
2652func Helper() {}
2653"#;
2654
2655 assert_query_parity(ParseableLanguage::Go, source);
2656 }
2657
2658 #[test]
2659 fn test_chunk_c_corner_cases() {
2660 let c_code = r#"
2661#define MAX(a,b) ((a) > (b) ? (a) : (b))
2662#define VERSION 3
2663
2664typedef struct Node {
2665 int value;
2666 struct Node* next;
2667} Node;
2668
2669union Payload {
2670 int i;
2671 float f;
2672};
2673
2674enum Color {
2675 Red,
2676 Green,
2677 Blue,
2678};
2679
2680static inline int add(int a, int b) {
2681 return a + b;
2682}
2683
2684int main(void) {
2685 return MAX(add(1, 2), VERSION);
2686}
2687"#;
2688
2689 let chunks = chunk_language(c_code, ParseableLanguage::C).unwrap();
2690 assert!(!chunks.is_empty());
2691
2692 assert!(
2693 chunks
2694 .iter()
2695 .any(|c| { c.chunk_type == ChunkType::Function && c.text.contains("#define MAX") })
2696 );
2697 assert!(
2698 chunks
2699 .iter()
2700 .any(|c| { c.chunk_type == ChunkType::Text && c.text.contains("#define VERSION") })
2701 );
2702 assert!(
2703 chunks
2704 .iter()
2705 .any(|c| { c.chunk_type == ChunkType::Class && c.text.contains("struct Node") })
2706 );
2707 assert!(
2708 chunks
2709 .iter()
2710 .any(|c| { c.chunk_type == ChunkType::Class && c.text.contains("union Payload") })
2711 );
2712 assert!(
2713 chunks
2714 .iter()
2715 .any(|c| { c.chunk_type == ChunkType::Class && c.text.contains("enum Color") })
2716 );
2717 assert!(chunks.iter().any(|c| {
2718 c.chunk_type == ChunkType::Function && c.text.contains("static inline int add")
2719 }));
2720 assert!(
2721 chunks
2722 .iter()
2723 .any(|c| { c.chunk_type == ChunkType::Function && c.text.contains("int main") })
2724 );
2725 }
2726
2727 #[test]
2728 fn test_chunk_c_struct_declaration_without_body_stays_intact() {
2729 let c_code = r#"
2730#include <stdint.h>
2731
2732struct mtd_info_user meminfo;
2733struct foo forward;
2734"#;
2735
2736 let chunks = chunk_language(c_code, ParseableLanguage::C).unwrap();
2737
2738 assert!(
2739 chunks
2740 .iter()
2741 .any(|c| { c.text.contains("struct mtd_info_user meminfo;") })
2742 );
2743 assert!(
2744 chunks
2745 .iter()
2746 .any(|c| c.text.contains("struct foo forward;"))
2747 );
2748 assert!(
2749 !chunks
2750 .iter()
2751 .any(|c| c.text.trim() == "struct mtd_info_user")
2752 );
2753 assert!(!chunks.iter().any(|c| c.text.trim() == "struct foo"));
2754 }
2755
2756 #[test]
2757 fn test_chunk_cpp_corner_cases() {
2758 let cpp_code = r#"
2759#include <vector>
2760#define SQUARE(x) ((x) * (x))
2761
2762namespace math {
2763template <typename T>
2764T add(T a, T b) {
2765 return a + b;
2766}
2767
2768using Vec = std::vector<int>;
2769typedef unsigned long ulong_t;
2770
2771struct Point {
2772 int x;
2773 int y;
2774};
2775
2776class Calculator {
2777public:
2778 int add(int a, int b) { return a + b; }
2779};
2780
2781enum class Color { Red, Green, Blue };
2782} // namespace math
2783
2784int main() {
2785 return math::add(1, 2);
2786}
2787"#;
2788
2789 let chunks = chunk_language(cpp_code, ParseableLanguage::Cpp).unwrap();
2790 assert!(!chunks.is_empty());
2791
2792 assert!(
2793 chunks
2794 .iter()
2795 .any(|c| c.text.contains("template <typename T>"))
2796 );
2797 assert!(
2798 chunks
2799 .iter()
2800 .any(|c| { c.chunk_type == ChunkType::Text && c.text.contains("using Vec") })
2801 );
2802 assert!(chunks.iter().any(|c| {
2803 c.chunk_type == ChunkType::Text && c.text.contains("typedef unsigned long")
2804 }));
2805 assert!(
2806 chunks.iter().any(|c| {
2807 c.chunk_type == ChunkType::Function && c.text.contains("#define SQUARE")
2808 })
2809 );
2810 let calculator_chunk = chunks
2811 .iter()
2812 .find(|c| c.chunk_type == ChunkType::Class && c.text.contains("class Calculator"));
2813 assert!(calculator_chunk.is_some());
2814 let calculator_chunk = calculator_chunk.unwrap();
2815 assert!(calculator_chunk.text.contains("int add"));
2816 assert!(!calculator_chunk.text.contains("return a + b"));
2817
2818 assert!(
2819 chunks
2820 .iter()
2821 .any(|c| { c.chunk_type == ChunkType::Class && c.text.contains("struct Point") })
2822 );
2823 assert!(
2824 chunks.iter().any(|c| {
2825 c.chunk_type == ChunkType::Class && c.text.contains("enum class Color")
2826 })
2827 );
2828 assert!(
2829 chunks
2830 .iter()
2831 .any(|c| { c.chunk_type == ChunkType::Function && c.text.contains("int main") })
2832 );
2833 assert!(
2834 chunks
2835 .iter()
2836 .any(|c| { c.chunk_type == ChunkType::Function && c.text.contains("T add") })
2837 );
2838 assert!(chunks.iter().any(|c| {
2839 c.chunk_type == ChunkType::Method && c.text.contains("int add(int a, int b)")
2840 }));
2841 }
2842
2843 #[test]
2844 fn test_cpp_suppresses_contained_text_chunks() {
2845 let cpp_code = r#"
2846class Widget {
2847public:
2848 using Alias = int;
2849 int calc() { int local = 1; return local; }
2850};
2851
2852using TopLevel = double;
2853"#;
2854
2855 let chunks = chunk_language(cpp_code, ParseableLanguage::Cpp).unwrap();
2856
2857 assert!(
2858 !chunks
2859 .iter()
2860 .any(|c| { c.chunk_type == ChunkType::Text && c.text.contains("using Alias") })
2861 );
2862 assert!(
2863 !chunks
2864 .iter()
2865 .any(|c| { c.chunk_type == ChunkType::Text && c.text.contains("int local") })
2866 );
2867 assert!(
2868 chunks
2869 .iter()
2870 .any(|c| { c.chunk_type == ChunkType::Text && c.text.contains("using TopLevel") })
2871 );
2872 assert!(
2873 chunks
2874 .iter()
2875 .any(|c| { c.chunk_type == ChunkType::Method && c.text.contains("int calc") })
2876 );
2877 }
2878
2879 #[test]
2880 fn test_cpp_template_prefix_merges_with_definition() {
2881 let cpp_code = r#"
2882template <typename T>
2883struct Box {
2884 static int value;
2885};
2886
2887template <typename T>
2888int Box<T>::value = 0;
2889"#;
2890
2891 let chunks = chunk_language(cpp_code, ParseableLanguage::Cpp).unwrap();
2892
2893 let def_chunk = chunks
2894 .iter()
2895 .find(|c| c.text.contains("int Box<T>::value = 0;"))
2896 .expect("static member definition chunk present");
2897
2898 assert!(def_chunk.text.contains("template <typename T>"));
2899
2900 assert!(!chunks.iter().any(|c| {
2901 c.chunk_type == ChunkType::Text && c.text.trim() == "template <typename T>"
2902 }));
2903 }
2904
2905 #[test]
2906 fn test_cpp_template_method_breadcrumb_in_namespaces() {
2907 let cpp_code = r#"
2908namespace com {
2909namespace ford {
2910
2911template <typename T>
2912class Wrapper {
2913public:
2914 template <typename U>
2915 U convert(U value) { return value; }
2916};
2917
2918} // namespace ford
2919} // namespace com
2920"#;
2921
2922 let chunks = chunk_language(cpp_code, ParseableLanguage::Cpp).unwrap();
2923 let method_chunk = chunks
2924 .iter()
2925 .find(|c| c.chunk_type == ChunkType::Method && c.text.contains("convert"))
2926 .expect("convert method chunk present");
2927
2928 assert_eq!(
2929 method_chunk.metadata.breadcrumb.as_deref(),
2930 Some("com::ford::Wrapper::convert")
2931 );
2932 }
2933
2934 #[test]
2935 fn test_cpp_function_breadcrumb_qualification() {
2936 let cpp_code = r#"
2937namespace outer {
2938class A {
2939public:
2940 void m();
2941};
2942}
2943
2944void outer::A::m() {
2945 // body
2946}
2947"#;
2948
2949 let chunks = chunk_language(cpp_code, ParseableLanguage::Cpp).unwrap();
2950 let method_chunk = chunks
2951 .iter()
2952 .find(|c| c.chunk_type == ChunkType::Function && c.text.contains("outer::A::m"))
2953 .expect("method chunk should exist");
2954 assert_eq!(
2955 method_chunk.metadata.breadcrumb.as_deref(),
2956 Some("outer::A::m")
2957 );
2958 }
2959
2960 #[test]
2961 fn test_haskell_query_matches_legacy() {
2962 let source = r#"
2963module Example where
2964
2965data Shape
2966 = Circle Float
2967 | Square Float
2968
2969type family Area a
2970
2971class Printable a where
2972 printValue :: a -> String
2973
2974instance Printable Shape where
2975 printValue (Circle _) = "circle"
2976 printValue (Square _) = "square"
2977
2978shapeDescription :: Shape -> String
2979shapeDescription (Circle r) = "circle of radius " ++ show r
2980shapeDescription (Square s) = "square of side " ++ show s
2981"#;
2982
2983 assert_query_parity(ParseableLanguage::Haskell, source);
2984 }
2985
2986 #[test]
2987 fn test_markdown_real_file_breadcrumbs() {
2988 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
2989 .join("tests/fixtures/markdown_breadcrumbs.md");
2990 let source = std::fs::read_to_string(&path).expect("read markdown file");
2991
2992 let chunks =
2993 chunk_text(&source, Some(ck_core::Language::Markdown)).expect("chunk markdown");
2994
2995 let all_text: String = chunks
3003 .iter()
3004 .map(|c| c.text.as_str())
3005 .collect::<Vec<_>>()
3006 .join("\n");
3007
3008 assert!(
3009 all_text.contains("Project Overview"),
3010 "expected top-level heading text to be present in merged chunk"
3011 );
3012 assert!(
3013 all_text.contains("## Usage"),
3014 "expected second-level heading text to be present in merged chunk"
3015 );
3016 assert!(
3017 all_text.contains("Setext Section"),
3018 "expected setext heading text to be present in merged chunk"
3019 );
3020 }
3021
3022 #[test]
3023 fn test_markdown_inline_fixtures_cover_blocks() {
3024 let source = r#"
3025# Title
3026
3027Intro paragraph with **bold** text.
3028
3029## Usage
3030
3031```rust
3032fn main() {
3033 println!("hi");
3034}
3035```
3036
3037> Blockquote with _emphasis_.
3038
3039- Item one
3040- Item two
3041
3042Setext Section
3043==============
3044
3045Trailing paragraph.
3046"#;
3047
3048 let chunks = chunk_text(source, Some(ck_core::Language::Markdown)).expect("chunk markdown");
3049
3050 let all_text: String = chunks
3058 .iter()
3059 .map(|c| c.text.as_str())
3060 .collect::<Vec<_>>()
3061 .join("\n");
3062
3063 assert!(
3064 all_text.contains("# Title") || all_text.contains("## Usage"),
3065 "expected heading text to be present after merging"
3066 );
3067 assert!(
3068 all_text.contains("```rust"),
3069 "expected markdown to include fenced code block"
3070 );
3071 assert!(
3072 all_text.contains("> Blockquote"),
3073 "expected markdown to include blockquote text"
3074 );
3075 assert!(
3076 all_text.contains("Setext Section"),
3077 "expected markdown to include Setext heading text"
3078 );
3079 }
3080
3081 #[test]
3094 fn cpp_class_body_gap_text_is_suppressed() {
3095 let source = r"
3096class Counter {
3097public:
3098 int value;
3099 int spare;
3100
3101 void inc() { value++; }
3102 void dec() { value--; }
3103};
3104";
3105 let chunks = chunk_text(source, Some(ck_core::Language::Cpp)).expect("chunk cpp");
3106
3107 assert!(
3109 chunks
3110 .iter()
3111 .any(|c| c.chunk_type == ChunkType::Class && c.text.contains("class Counter")),
3112 "expected Class chunk for Counter"
3113 );
3114
3115 let leaked_field = chunks.iter().any(|c| {
3119 c.chunk_type == ChunkType::Text
3120 && (c.text.contains("int value") || c.text.contains("int spare"))
3121 });
3122 let leaked_access = chunks
3123 .iter()
3124 .any(|c| c.chunk_type == ChunkType::Text && c.text.trim() == "public:");
3125 assert!(
3126 !leaked_field,
3127 "field declarations leaked as standalone Text chunks: {:?}",
3128 chunks
3129 .iter()
3130 .filter(|c| c.chunk_type == ChunkType::Text)
3131 .map(|c| &c.text)
3132 .collect::<Vec<_>>()
3133 );
3134 assert!(!leaked_access, "access specifier leaked as Text chunk");
3135 }
3136
3137 #[test]
3142 fn cpp_multiline_template_prefix_merges_with_definition() {
3143 let source = r"
3145template <
3146 typename T,
3147 typename U
3148>
3149
3150T identity(T value, U /*ignored*/) {
3151 return value;
3152}
3153";
3154 let chunks = chunk_text(source, Some(ck_core::Language::Cpp)).expect("chunk cpp");
3155
3156 let combined = chunks.iter().find(|c| {
3160 matches!(c.chunk_type, ChunkType::Function | ChunkType::Method)
3161 && c.text.contains("template")
3162 && c.text.contains("T identity")
3163 });
3164 assert!(
3165 combined.is_some(),
3166 "expected one Function chunk combining the multiline template clause with its definition; got {:?}",
3167 chunks
3168 .iter()
3169 .map(|c| (&c.chunk_type, c.text.lines().next().unwrap_or("")))
3170 .collect::<Vec<_>>()
3171 );
3172
3173 let orphan_template = chunks.iter().any(|c| {
3176 c.chunk_type == ChunkType::Text
3177 && c.text.contains("template")
3178 && !c.text.contains("identity")
3179 });
3180 assert!(
3181 !orphan_template,
3182 "template clause was emitted as a standalone Text chunk instead of merging"
3183 );
3184 }
3185
3186 #[test]
3191 fn cpp_extern_c_block_functions_are_captured() {
3192 let source = r#"
3193extern "C" {
3194 int add(int a, int b) {
3195 return a + b;
3196 }
3197 int sub(int a, int b) {
3198 return a - b;
3199 }
3200}
3201"#;
3202 let chunks = chunk_text(source, Some(ck_core::Language::Cpp)).expect("chunk cpp");
3203
3204 let has_add = chunks.iter().any(|c| {
3205 matches!(c.chunk_type, ChunkType::Function | ChunkType::Method)
3206 && c.text.contains("int add")
3207 });
3208 let has_sub = chunks.iter().any(|c| {
3209 matches!(c.chunk_type, ChunkType::Function | ChunkType::Method)
3210 && c.text.contains("int sub")
3211 });
3212 assert!(
3213 has_add && has_sub,
3214 "expected both `add` and `sub` to be captured inside extern \"C\"; got {:?}",
3215 chunks
3216 .iter()
3217 .map(|c| (&c.chunk_type, c.text.lines().next().unwrap_or("")))
3218 .collect::<Vec<_>>()
3219 );
3220 }
3221
3222 #[test]
3223 fn test_csharp_query_matches_legacy() {
3224 let source = r"
3225namespace Calculator;
3226
3227public interface ICalculator
3228{
3229 double Add(double x, double y);
3230}
3231
3232public class Calculator
3233{
3234 public static double PI = 3.14159;
3235 private double _memory;
3236
3237 public Calculator()
3238 {
3239 _memory = 0.0;
3240 }
3241
3242 public double Add(double x, double y)
3243 {
3244 return x + y;
3245 }
3246}
3247";
3248
3249 assert_query_parity(ParseableLanguage::CSharp, source);
3250 }
3251
3252 #[test]
3253 fn test_zig_query_matches_legacy() {
3254 let source = r#"
3255const std = @import("std");
3256
3257const Calculator = struct {
3258 memory: f64,
3259
3260 pub fn init() Calculator {
3261 return Calculator{ .memory = 0.0 };
3262 }
3263
3264 pub fn add(self: *Calculator, a: f64, b: f64) f64 {
3265 return a + b;
3266 }
3267};
3268
3269test "calculator addition" {
3270 var calc = Calculator.init();
3271 const result = calc.add(2.0, 3.0);
3272 try std.testing.expect(result == 5.0);
3273}
3274"#;
3275
3276 assert_query_parity(ParseableLanguage::Zig, source);
3277 }
3278
3279 #[test]
3280 fn test_chunk_zig() {
3281 let zig_code = r#"
3282const std = @import("std");
3283
3284const Calculator = struct {
3285 memory: f64,
3286
3287 pub fn init() Calculator {
3288 return Calculator{ .memory = 0.0 };
3289 }
3290
3291 pub fn add(self: *Calculator, a: f64, b: f64) f64 {
3292 const result = a + b;
3293 self.memory = result;
3294 return result;
3295 }
3296};
3297
3298const Color = enum {
3299 Red,
3300 Green,
3301 Blue,
3302};
3303
3304const Value = union(enum) {
3305 int: i32,
3306 float: f64,
3307};
3308
3309const Handle = opaque {};
3310
3311const MathError = error{
3312 DivisionByZero,
3313 Overflow,
3314};
3315
3316pub fn multiply(a: i32, b: i32) i32 {
3317 return a * b;
3318}
3319
3320pub fn divide(a: i32, b: i32) MathError!i32 {
3321 if (b == 0) return error.DivisionByZero;
3322 return @divTrunc(a, b);
3323}
3324
3325comptime {
3326 @compileLog("Compile-time validation");
3327}
3328
3329pub fn main() !void {
3330 var calc = Calculator.init();
3331 const result = calc.add(2.0, 3.0);
3332 std.debug.print("Result: {}\n", .{result});
3333}
3334
3335test "calculator addition" {
3336 var calc = Calculator.init();
3337 const result = calc.add(2.0, 3.0);
3338 try std.testing.expect(result == 5.0);
3339}
3340
3341test "multiply function" {
3342 const result = multiply(3, 4);
3343 try std.testing.expect(result == 12);
3344}
3345"#;
3346
3347 let chunks = chunk_language(zig_code, ParseableLanguage::Zig).unwrap();
3348 assert!(!chunks.is_empty());
3349
3350 let chunk_types: Vec<&ChunkType> = chunks.iter().map(|c| &c.chunk_type).collect();
3351
3352 let class_count = chunk_types
3353 .iter()
3354 .filter(|&&t| t == &ChunkType::Class)
3355 .count();
3356 let function_count = chunk_types
3357 .iter()
3358 .filter(|&&t| t == &ChunkType::Function)
3359 .count();
3360 let module_count = chunk_types
3361 .iter()
3362 .filter(|&&t| t == &ChunkType::Module)
3363 .count();
3364
3365 assert!(
3366 class_count >= 5,
3367 "Expected at least 5 Class chunks (struct, enum, union, opaque, error set), found {class_count}"
3368 );
3369
3370 assert!(
3371 function_count >= 3,
3372 "Expected at least 3 functions (multiply, divide, main), found {function_count}"
3373 );
3374
3375 assert!(
3376 module_count >= 4,
3377 "Expected at least 4 module-type chunks (const std, comptime, 2 tests), found {module_count}"
3378 );
3379
3380 assert!(
3381 chunk_types.contains(&&ChunkType::Class),
3382 "Expected to find Class chunks"
3383 );
3384 assert!(
3385 chunk_types.contains(&&ChunkType::Function),
3386 "Expected to find Function chunks"
3387 );
3388 assert!(
3389 chunk_types.contains(&&ChunkType::Module),
3390 "Expected to find Module chunks"
3391 );
3392 }
3393
3394 #[test]
3395 fn test_chunk_csharp() {
3396 let csharp_code = r"
3397namespace Calculator;
3398
3399public interface ICalculator
3400{
3401 double Add(double x, double y);
3402}
3403
3404public class Calculator
3405{
3406 public static const double PI = 3.14159;
3407 private double _memory;
3408
3409 public Calculator()
3410 {
3411 _memory = 0.0;
3412 }
3413
3414 public double Add(double x, double y)
3415 {
3416 return x + y;
3417 }
3418
3419 public static void Main(string[] args)
3420 {
3421 var calc = new Calculator();
3422 }
3423}
3424";
3425
3426 let chunks = chunk_language(csharp_code, ParseableLanguage::CSharp).unwrap();
3427 assert!(!chunks.is_empty());
3428
3429 let chunk_types: Vec<&ChunkType> = chunks.iter().map(|c| &c.chunk_type).collect();
3431 assert!(chunk_types.contains(&&ChunkType::Module)); assert!(chunk_types.contains(&&ChunkType::Class)); assert!(chunk_types.contains(&&ChunkType::Method)); }
3435
3436 #[test]
3437 fn test_stride_large_chunk_empty_text() {
3438 let empty_chunk = Chunk {
3440 span: Span {
3441 byte_start: 0,
3442 byte_end: 0,
3443 line_start: 1,
3444 line_end: 1,
3445 },
3446 text: String::new(), chunk_type: ChunkType::Text,
3448 stride_info: None,
3449 metadata: ChunkMetadata::from_text(""),
3450 };
3451
3452 let config = ChunkConfig::default();
3453 let result = stride_large_chunk(empty_chunk.clone(), &config);
3454
3455 assert!(result.is_ok());
3457 let chunks = result.unwrap();
3458 assert_eq!(chunks.len(), 1);
3459 assert_eq!(chunks[0].text, "");
3460 }
3461
3462 #[test]
3463 fn test_stride_large_chunk_zero_token_estimate() {
3464 let chunk = Chunk {
3466 span: Span {
3467 byte_start: 0,
3468 byte_end: 5,
3469 line_start: 1,
3470 line_end: 1,
3471 },
3472 text: " ".to_string(), chunk_type: ChunkType::Text,
3474 stride_info: None,
3475 metadata: ChunkMetadata::from_text(" "),
3476 };
3477
3478 let config = ChunkConfig::default();
3479 let result = stride_large_chunk(chunk, &config);
3480
3481 assert!(result.is_ok());
3483 }
3484
3485 #[test]
3486 fn test_strided_chunk_line_calculation() {
3487 let long_text = (1..=50).map(|i| format!("This is a longer line {i} with more content to ensure token count is high enough")).collect::<Vec<_>>().join("\n");
3490
3491 let metadata = ChunkMetadata::from_text(&long_text);
3492 let chunk = Chunk {
3493 span: Span {
3494 byte_start: 0,
3495 byte_end: long_text.len(),
3496 line_start: 1,
3497 line_end: 50,
3498 },
3499 text: long_text,
3500 chunk_type: ChunkType::Text,
3501 stride_info: None,
3502 metadata,
3503 };
3504
3505 let config = ChunkConfig {
3506 max_tokens: 100, stride_overlap: 10, ..Default::default()
3509 };
3510
3511 let result = stride_large_chunk(chunk, &config);
3512 if let Err(e) = &result {
3513 eprintln!("Stride error: {e}");
3514 }
3515 assert!(result.is_ok());
3516
3517 let chunks = result.unwrap();
3518 assert!(
3519 chunks.len() > 1,
3520 "Should create multiple chunks when striding"
3521 );
3522
3523 for chunk in chunks {
3524 assert!(chunk.span.line_end >= chunk.span.line_start);
3527
3528 let line_count = chunk.text.lines().count();
3530 if line_count > 0 {
3531 let calculated_line_span = chunk.span.line_end - chunk.span.line_start + 1;
3532
3533 assert!(
3535 calculated_line_span <= line_count + 1,
3536 "Line span {calculated_line_span} should not exceed content lines {line_count} by more than 1"
3537 );
3538 }
3539 }
3540 }
3541
3542 #[test]
3543 fn test_gap_filling_coverage() {
3544 let test_cases = vec![
3546 (
3547 ParseableLanguage::Rust,
3548 r#"// This is a test file with imports at the top
3549use std::collections::HashMap;
3550use std::sync::Arc;
3551
3552// A comment between imports and code
3553const VERSION: &str = "1.0.0";
3554
3555// Main function
3556fn main() {
3557 println!("Hello, world!");
3558}
3559
3560// Some trailing content
3561// that should be indexed
3562"#,
3563 ),
3564 (
3565 ParseableLanguage::Python,
3566 r#"# Imports at the top
3567import os
3568import sys
3569
3570# Some constant
3571VERSION = "1.0.0"
3572
3573# Main function
3574def main():
3575 print("Hello, world!")
3576
3577# Trailing comment
3578# should be indexed
3579"#,
3580 ),
3581 (
3582 ParseableLanguage::TypeScript,
3583 r#"// Imports at the top
3584import { foo } from 'bar';
3585
3586// Some constant
3587const VERSION = "1.0.0";
3588
3589// Main function
3590function main() {
3591 console.log("Hello, world!");
3592}
3593
3594// Trailing comment
3595// should be indexed
3596"#,
3597 ),
3598 ];
3599
3600 for (language, code) in test_cases {
3601 eprintln!("\n=== Testing {language} ===");
3602 let chunks = chunk_language(code, language).unwrap();
3603
3604 let mut covered_bytes = vec![false; code.len()];
3606 for chunk in &chunks {
3607 for item in covered_bytes
3608 .iter_mut()
3609 .take(chunk.span.byte_end)
3610 .skip(chunk.span.byte_start)
3611 {
3612 *item = true;
3613 }
3614 }
3615
3616 let uncovered_non_ws: Vec<usize> = covered_bytes
3617 .iter()
3618 .enumerate()
3619 .filter(|(i, covered)| !**covered && !code.as_bytes()[*i].is_ascii_whitespace())
3620 .map(|(i, _)| i)
3621 .collect();
3622
3623 if !uncovered_non_ws.is_empty() {
3624 eprintln!("\n=== UNCOVERED NON-WHITESPACE for {language} ===");
3625 eprintln!("Total bytes: {}", code.len());
3626 eprintln!("Uncovered non-whitespace: {}", uncovered_non_ws.len());
3627
3628 for &pos in uncovered_non_ws.iter().take(10) {
3630 let context_start = pos.saturating_sub(20);
3631 let context_end = (pos + 20).min(code.len());
3632 eprintln!(
3633 "Uncovered at byte {}: {:?}",
3634 pos,
3635 &code[context_start..context_end]
3636 );
3637 }
3638
3639 eprintln!("\n=== CHUNKS ===");
3640 for (i, chunk) in chunks.iter().enumerate() {
3641 eprintln!(
3642 "Chunk {}: {:?} bytes {}-{} (len {})",
3643 i,
3644 chunk.chunk_type,
3645 chunk.span.byte_start,
3646 chunk.span.byte_end,
3647 chunk.span.byte_end - chunk.span.byte_start
3648 );
3649 eprintln!(" Text: {:?}", &chunk.text[..chunk.text.len().min(60)]);
3650 }
3651 }
3652
3653 assert!(
3654 uncovered_non_ws.is_empty(),
3655 "{}: Expected all non-whitespace covered but found {} uncovered non-whitespace bytes",
3656 language,
3657 uncovered_non_ws.len()
3658 );
3659 }
3660 }
3661
3662 #[test]
3663 fn test_web_server_file_coverage() {
3664 let code = std::fs::read_to_string("../examples/code/web_server.rs")
3666 .expect("Failed to read web_server.rs");
3667
3668 let chunks = chunk_language(&code, ParseableLanguage::Rust).unwrap();
3669
3670 let mut covered = vec![false; code.len()];
3672 for chunk in &chunks {
3673 for item in covered
3674 .iter_mut()
3675 .take(chunk.span.byte_end)
3676 .skip(chunk.span.byte_start)
3677 {
3678 *item = true;
3679 }
3680 }
3681
3682 let uncovered_non_whitespace: Vec<(usize, char)> = covered
3684 .iter()
3685 .enumerate()
3686 .filter(|(i, covered)| !**covered && !code.as_bytes()[*i].is_ascii_whitespace())
3687 .map(|(i, _)| (i, code.chars().nth(i).unwrap_or('?')))
3688 .collect();
3689
3690 if !uncovered_non_whitespace.is_empty() {
3691 eprintln!("\n=== WEB_SERVER.RS UNCOVERED NON-WHITESPACE ===");
3692 eprintln!("File size: {} bytes", code.len());
3693 eprintln!("Total chunks: {}", chunks.len());
3694 eprintln!(
3695 "Uncovered non-whitespace: {}",
3696 uncovered_non_whitespace.len()
3697 );
3698
3699 for &(pos, ch) in uncovered_non_whitespace.iter().take(10) {
3700 let start = pos.saturating_sub(30);
3701 let end = (pos + 30).min(code.len());
3702 eprintln!(
3703 "\nUncovered '{}' at byte {}: {:?}",
3704 ch,
3705 pos,
3706 &code[start..end]
3707 );
3708 }
3709
3710 eprintln!("\n=== CHUNKS ===");
3711 for (i, chunk) in chunks.iter().enumerate().take(20) {
3712 eprintln!(
3713 "Chunk {}: {:?} bytes {}-{} lines {}-{}",
3714 i,
3715 chunk.chunk_type,
3716 chunk.span.byte_start,
3717 chunk.span.byte_end,
3718 chunk.span.line_start,
3719 chunk.span.line_end
3720 );
3721 }
3722 }
3723
3724 assert!(
3725 uncovered_non_whitespace.is_empty(),
3726 "Expected all non-whitespace content covered but found {} uncovered non-whitespace bytes",
3727 uncovered_non_whitespace.len()
3728 );
3729 }
3730
3731 #[test]
3732 fn test_haskell_function_chunking() {
3733 let haskell_code = r"
3734factorial :: Integer -> Integer
3735factorial 0 = 1
3736factorial n = n * factorial (n - 1)
3737
3738fibonacci :: Integer -> Integer
3739fibonacci 0 = 0
3740fibonacci 1 = 1
3741fibonacci n = fibonacci (n - 1) + fibonacci (n - 2)
3742";
3743
3744 let mut parser = tree_sitter::Parser::new();
3745 parser
3746 .set_language(&tree_sitter_haskell::LANGUAGE.into())
3747 .unwrap();
3748 let tree = parser.parse(haskell_code, None).unwrap();
3749
3750 fn walk(node: tree_sitter::Node, _src: &str, depth: usize) {
3752 let kind = node.kind();
3753 let start = node.start_position();
3754 let end = node.end_position();
3755 eprintln!(
3756 "{}{:30} L{}-{}",
3757 " ".repeat(depth),
3758 kind,
3759 start.row + 1,
3760 end.row + 1
3761 );
3762
3763 let mut cursor = node.walk();
3764 if cursor.goto_first_child() {
3765 loop {
3766 walk(cursor.node(), _src, depth + 1);
3767 if !cursor.goto_next_sibling() {
3768 break;
3769 }
3770 }
3771 }
3772 }
3773
3774 eprintln!("\n=== TREE STRUCTURE ===");
3775 walk(tree.root_node(), haskell_code, 0);
3776 eprintln!("=== END TREE ===\n");
3777
3778 let chunks = chunk_language(haskell_code, ParseableLanguage::Haskell).unwrap();
3779
3780 eprintln!("\n=== CHUNKS ===");
3781 for (i, chunk) in chunks.iter().enumerate() {
3782 eprintln!(
3783 "Chunk {}: {:?} L{}-{}",
3784 i, chunk.chunk_type, chunk.span.line_start, chunk.span.line_end
3785 );
3786 eprintln!(" Text: {:?}", chunk.text);
3787 }
3788 eprintln!("=== END CHUNKS ===\n");
3789
3790 assert!(!chunks.is_empty(), "Should find chunks in Haskell code");
3791
3792 let factorial_chunk = chunks.iter().find(|c| c.text.contains("factorial 0 = 1"));
3794 assert!(
3795 factorial_chunk.is_some(),
3796 "Should find factorial function body"
3797 );
3798
3799 let fac = factorial_chunk.unwrap();
3800 assert!(
3801 fac.text.contains("factorial :: Integer -> Integer"),
3802 "Should include type signature"
3803 );
3804 assert!(
3805 fac.text.contains("factorial 0 = 1"),
3806 "Should include base case"
3807 );
3808 assert!(
3809 fac.text.contains("factorial n = n * factorial (n - 1)"),
3810 "Should include recursive case"
3811 );
3812 }
3813
3814 #[test]
3815 fn test_chunk_elixir_basic() {
3816 let elixir_code = r#"
3817defmodule Calculator do
3818 @moduledoc "A simple calculator module"
3819
3820 def add(a, b) do
3821 a + b
3822 end
3823
3824 defp multiply(a, b) do
3825 a * b
3826 end
3827end
3828"#;
3829
3830 let chunks = chunk_language(elixir_code, ParseableLanguage::Elixir).unwrap();
3831
3832 eprintln!("\n=== ELIXIR CHUNKS ===");
3833 for (i, chunk) in chunks.iter().enumerate() {
3834 eprintln!(
3835 "Chunk {}: {:?} L{}-{}",
3836 i, chunk.chunk_type, chunk.span.line_start, chunk.span.line_end
3837 );
3838 eprintln!(" Text: {:?}", &chunk.text[..chunk.text.len().min(80)]);
3839 }
3840 eprintln!("=== END CHUNKS ===\n");
3841
3842 assert!(!chunks.is_empty(), "Should find chunks in Elixir code");
3843
3844 let has_module = chunks.iter().any(|c| c.chunk_type == ChunkType::Module);
3846 let has_function = chunks.iter().any(|c| c.chunk_type == ChunkType::Function);
3847
3848 assert!(has_module, "Should detect defmodule as Module");
3849 assert!(has_function, "Should detect def/defp as Function");
3850 }
3851
3852 #[test]
3853 fn test_chunk_elixir_protocol() {
3854 let elixir_code = r#"
3855defprotocol Stringable do
3856 @doc "Converts to string"
3857 def to_string(value)
3858end
3859
3860defimpl Stringable, for: Integer do
3861 def to_string(value), do: Integer.to_string(value)
3862end
3863"#;
3864
3865 let chunks = chunk_language(elixir_code, ParseableLanguage::Elixir).unwrap();
3866
3867 eprintln!("\n=== ELIXIR PROTOCOL CHUNKS ===");
3868 for (i, chunk) in chunks.iter().enumerate() {
3869 eprintln!(
3870 "Chunk {}: {:?} L{}-{}",
3871 i, chunk.chunk_type, chunk.span.line_start, chunk.span.line_end
3872 );
3873 eprintln!(" Text: {:?}", &chunk.text[..chunk.text.len().min(80)]);
3874 }
3875 eprintln!("=== END CHUNKS ===\n");
3876
3877 let modules: Vec<_> = chunks
3879 .iter()
3880 .filter(|c| c.chunk_type == ChunkType::Module)
3881 .collect();
3882
3883 assert!(
3884 modules.len() >= 2,
3885 "Should detect defprotocol and defimpl as modules, found {}",
3886 modules.len()
3887 );
3888 }
3889
3890 #[test]
3891 fn test_chunk_elixir_genserver() {
3892 let elixir_code = r"
3893defmodule MyServer do
3894 use GenServer
3895
3896 def start_link(opts) do
3897 GenServer.start_link(__MODULE__, opts, name: __MODULE__)
3898 end
3899
3900 def init(state) do
3901 {:ok, state}
3902 end
3903
3904 def handle_call(:get, _from, state) do
3905 {:reply, state, state}
3906 end
3907
3908 def handle_cast({:set, value}, _state) do
3909 {:noreply, value}
3910 end
3911end
3912";
3913
3914 let chunks = chunk_language(elixir_code, ParseableLanguage::Elixir).unwrap();
3915
3916 let functions: Vec<_> = chunks
3918 .iter()
3919 .filter(|c| c.chunk_type == ChunkType::Function)
3920 .collect();
3921
3922 assert!(
3923 functions.len() >= 4,
3924 "Should detect at least 4 functions (start_link, init, handle_call, handle_cast), found {}",
3925 functions.len()
3926 );
3927 }
3928
3929 #[test]
3930 fn test_elixir_extension_detection() {
3931 use ck_core::Language;
3932
3933 assert_eq!(Language::from_extension("ex"), Some(Language::Elixir));
3934 assert_eq!(Language::from_extension("exs"), Some(Language::Elixir));
3935 assert_eq!(Language::from_extension("EX"), Some(Language::Elixir));
3936 assert_eq!(Language::from_extension("EXS"), Some(Language::Elixir));
3937 }
3938
3939 #[test]
3940 fn test_chunk_elixir_macros() {
3941 let elixir_code = r"
3942defmodule MyMacros do
3943 defmacro unless(condition, do: block) do
3944 quote do
3945 if !unquote(condition), do: unquote(block)
3946 end
3947 end
3948
3949 defmacrop private_macro(x) do
3950 quote do: unquote(x) * 2
3951 end
3952end
3953";
3954
3955 let chunks = chunk_language(elixir_code, ParseableLanguage::Elixir).unwrap();
3956
3957 let functions: Vec<_> = chunks
3958 .iter()
3959 .filter(|c| c.chunk_type == ChunkType::Function)
3960 .collect();
3961
3962 assert!(
3963 functions.len() >= 2,
3964 "Should detect defmacro and defmacrop as functions, found {}",
3965 functions.len()
3966 );
3967 }
3968
3969 #[test]
3970 fn test_chunk_elixir_module_attributes() {
3971 let elixir_code = r#"
3972defmodule Calculator do
3973 @moduledoc "A calculator with type specs"
3974
3975 @behaviour GenServer
3976
3977 @type operation :: :add | :subtract | :multiply | :divide
3978 @typep internal_state :: %{history: list()}
3979 @opaque result :: {:ok, number()} | {:error, atom()}
3980
3981 @callback init(args :: term()) :: {:ok, state :: term()}
3982 @callback handle_call(request :: term(), from :: term(), state :: term()) :: {:reply, term(), term()}
3983
3984 @optional_callbacks [handle_info: 2]
3985
3986 @spec add(number(), number()) :: number()
3987 def add(a, b), do: a + b
3988
3989 @spec subtract(number(), number()) :: number()
3990 def subtract(a, b), do: a - b
3991end
3992"#;
3993
3994 let chunks = chunk_language(elixir_code, ParseableLanguage::Elixir).unwrap();
3995
3996 eprintln!("\n=== ELIXIR MODULE ATTRIBUTES CHUNKS ===");
3997 for (i, chunk) in chunks.iter().enumerate() {
3998 eprintln!(
3999 "Chunk {}: {:?} L{}-{}",
4000 i, chunk.chunk_type, chunk.span.line_start, chunk.span.line_end
4001 );
4002 eprintln!(" Text: {:?}", &chunk.text[..chunk.text.len().min(80)]);
4003 }
4004 eprintln!("=== END CHUNKS ===\n");
4005
4006 let has_behaviour = chunks
4008 .iter()
4009 .any(|c| c.chunk_type == ChunkType::Text && c.text.contains("@behaviour GenServer"));
4010 assert!(has_behaviour, "Should capture @behaviour declaration");
4011
4012 let type_chunks: Vec<_> = chunks
4014 .iter()
4015 .filter(|c| {
4016 c.chunk_type == ChunkType::Text
4017 && (c.text.contains("@type")
4018 || c.text.contains("@typep")
4019 || c.text.contains("@opaque"))
4020 })
4021 .collect();
4022 assert!(
4023 type_chunks.len() >= 3,
4024 "Should capture @type, @typep, and @opaque, found {}",
4025 type_chunks.len()
4026 );
4027
4028 let callback_chunks: Vec<_> = chunks
4030 .iter()
4031 .filter(|c| c.chunk_type == ChunkType::Text && c.text.contains("@callback"))
4032 .collect();
4033 assert!(
4034 callback_chunks.len() >= 2,
4035 "Should capture @callback definitions, found {}",
4036 callback_chunks.len()
4037 );
4038
4039 let spec_chunks: Vec<_> = chunks
4041 .iter()
4042 .filter(|c| c.chunk_type == ChunkType::Text && c.text.contains("@spec"))
4043 .collect();
4044 assert!(
4045 spec_chunks.len() >= 2,
4046 "Should capture @spec definitions, found {}",
4047 spec_chunks.len()
4048 );
4049
4050 let function_chunks: Vec<_> = chunks
4052 .iter()
4053 .filter(|c| c.chunk_type == ChunkType::Function)
4054 .collect();
4055 assert!(
4056 function_chunks.len() >= 2,
4057 "Should still capture def functions, found {}",
4058 function_chunks.len()
4059 );
4060 }
4061
4062 #[test]
4063 fn test_chunk_elixir_behavior_spelling() {
4064 let elixir_code = r"
4066defmodule BritishModule do
4067 @behaviour GenServer
4068end
4069
4070defmodule AmericanModule do
4071 @behavior GenServer
4072end
4073";
4074
4075 let chunks = chunk_language(elixir_code, ParseableLanguage::Elixir).unwrap();
4076
4077 let behaviour_chunks: Vec<_> = chunks
4078 .iter()
4079 .filter(|c| {
4080 c.chunk_type == ChunkType::Text
4081 && (c.text.contains("@behaviour") || c.text.contains("@behavior"))
4082 })
4083 .collect();
4084
4085 assert!(
4086 behaviour_chunks.len() >= 2,
4087 "Should capture both @behaviour and @behavior spellings, found {}",
4088 behaviour_chunks.len()
4089 );
4090 }
4091}