1use std::collections::BTreeMap;
17
18use panproto_schema::{Protocol, Schema, SchemaBuilder};
19use rustc_hash::FxHashSet;
20
21use crate::error::ParseError;
22use crate::id_scheme::IdGenerator;
23use crate::scope_detector::{NamedScope, ScopeDetector};
24use crate::theory_extract::ExtractedTheoryMeta;
25
26const BLOCK_KINDS: &[&str] = &[
33 "block",
34 "statement_block",
35 "compound_statement",
36 "declaration_list",
37 "field_declaration_list",
38 "enum_body",
39 "class_body",
40 "interface_body",
41 "module_body",
42];
43
44#[derive(Debug, Clone, Default)]
46pub struct WalkerConfig {
47 pub extra_block_kinds: Vec<String>,
55 pub capture_comments: bool,
57 pub capture_formatting: bool,
59}
60
61impl WalkerConfig {
62 #[must_use]
65 pub const fn standard() -> Self {
66 Self {
67 extra_block_kinds: Vec::new(),
68 capture_comments: true,
69 capture_formatting: true,
70 }
71 }
72}
73
74pub struct AstWalker<'a> {
83 source: &'a [u8],
85 theory_meta: &'a ExtractedTheoryMeta,
89 protocol: &'a Protocol,
91 config: WalkerConfig,
93 block_kinds: FxHashSet<String>,
95 scope_map: BTreeMap<(usize, usize), NamedScope>,
99}
100
101impl<'a> AstWalker<'a> {
102 #[must_use]
113 pub fn new(
114 source: &'a [u8],
115 theory_meta: &'a ExtractedTheoryMeta,
116 protocol: &'a Protocol,
117 config: WalkerConfig,
118 scope_detector: Option<&mut ScopeDetector>,
119 ) -> Self {
120 let mut block_kinds: FxHashSet<String> =
121 BLOCK_KINDS.iter().map(|s| (*s).to_owned()).collect();
122 for kind in &config.extra_block_kinds {
123 block_kinds.insert(kind.clone());
124 }
125
126 let mut scope_map: BTreeMap<(usize, usize), NamedScope> = BTreeMap::new();
127 if let Some(det) = scope_detector {
128 for scope in det.scopes(source) {
129 scope_map.insert((scope.node_range.start, scope.node_range.end), scope);
130 }
131 }
132
133 Self {
134 source,
135 theory_meta,
136 protocol,
137 config,
138 block_kinds,
139 scope_map,
140 }
141 }
142
143 pub fn walk(&self, tree: &tree_sitter::Tree, file_path: &str) -> Result<Schema, ParseError> {
149 let mut id_gen = IdGenerator::new(file_path);
150 let builder = SchemaBuilder::new(self.protocol);
151 let root = tree.root_node();
152
153 let builder = self.walk_node(root, builder, &mut id_gen, None)?;
154
155 builder.build().map_err(|e| ParseError::SchemaConstruction {
156 reason: e.to_string(),
157 })
158 }
159
160 fn scope_for(&self, node: tree_sitter::Node<'_>) -> Option<&NamedScope> {
162 self.scope_map.get(&(node.start_byte(), node.end_byte()))
163 }
164
165 fn walk_node(
167 &self,
168 node: tree_sitter::Node<'_>,
169 mut builder: SchemaBuilder,
170 id_gen: &mut IdGenerator,
171 parent_vertex_id: Option<&str>,
172 ) -> Result<SchemaBuilder, ParseError> {
173 if !node.is_named() {
175 return Ok(builder);
176 }
177
178 let kind = node.kind();
179
180 let is_root_wrapper = parent_vertex_id.is_none()
183 && (kind == "program"
184 || kind == "source_file"
185 || kind == "module"
186 || kind == "translation_unit");
187
188 let named_scope = if is_root_wrapper {
189 None
190 } else {
191 self.scope_for(node)
192 };
193
194 let (vertex_id, recorded_named_leaf) = if is_root_wrapper {
208 (id_gen.current_prefix(), None)
210 } else if let Some(scope) = named_scope {
211 let leaf = id_gen.record_name(&scope.name);
212 let prefix = id_gen.current_prefix();
213 (format!("{prefix}::{leaf}"), Some(leaf))
214 } else {
215 (id_gen.anonymous_id(), None)
217 };
218
219 let effective_kind = if self.protocol.obj_kinds.is_empty() {
223 kind
225 } else if self.protocol.obj_kinds.iter().any(|k| k == kind) {
226 kind
227 } else if !self.theory_meta.vertex_kinds.is_empty()
228 && self.theory_meta.vertex_kinds.iter().any(|k| k == kind)
229 {
230 kind
232 } else {
233 "node"
234 };
235
236 builder = builder
237 .vertex(&vertex_id, effective_kind, None)
238 .map_err(|e| ParseError::SchemaConstruction {
239 reason: format!("vertex '{vertex_id}' ({kind}): {e}"),
240 })?;
241
242 if let Some(parent_id) = parent_vertex_id {
244 let edge_kind = node
247 .parent()
248 .and_then(|p| {
249 for i in 0..p.child_count() {
251 if let Some(child) = p.child(i) {
252 if child.id() == node.id() {
253 return u32::try_from(i)
254 .ok()
255 .and_then(|idx| p.field_name_for_child(idx));
256 }
257 }
258 }
259 None
260 })
261 .unwrap_or("child_of");
262
263 builder = builder
264 .edge(parent_id, &vertex_id, edge_kind, None)
265 .map_err(|e| ParseError::SchemaConstruction {
266 reason: format!("edge {parent_id} -> {vertex_id} ({edge_kind}): {e}"),
267 })?;
268 }
269
270 builder = builder.constraint(&vertex_id, "start-byte", &node.start_byte().to_string());
272 builder = builder.constraint(&vertex_id, "end-byte", &node.end_byte().to_string());
273
274 if node.named_child_count() == 0 {
276 if let Ok(text) = node.utf8_text(self.source) {
277 builder = builder.constraint(&vertex_id, "literal-value", text);
278 }
279 }
280
281 builder = self.capture_anonymous_field_constraints(node, &vertex_id, builder);
291
292 if self.config.capture_formatting {
294 builder = self.emit_formatting_constraints(node, &vertex_id, builder);
295 }
296
297 let entered_scope = if let Some(leaf) = recorded_named_leaf {
303 id_gen.push_recorded_scope(leaf);
304 true
305 } else if !is_root_wrapper && self.block_kinds.contains(kind) {
306 id_gen.push_anonymous_scope();
307 true
308 } else {
309 false
310 };
311
312 builder = self.walk_children_with_interstitials(node, builder, id_gen, &vertex_id)?;
313
314 if entered_scope {
315 id_gen.pop_scope();
316 }
317
318 Ok(builder)
319 }
320
321 fn walk_children_with_interstitials(
333 &self,
334 node: tree_sitter::Node<'_>,
335 mut builder: SchemaBuilder,
336 id_gen: &mut IdGenerator,
337 vertex_id: &str,
338 ) -> Result<SchemaBuilder, ParseError> {
339 let cursor = &mut node.walk();
340 let children: Vec<_> = node.named_children(cursor).collect();
341 let mut interstitial_idx = 0;
342 let mut prev_end = node.start_byte();
343 let mut fingerprint_parts: Vec<String> = Vec::new();
344 let mut child_kinds: Vec<String> = Vec::new();
345
346 for child in &children {
347 let gap_start = prev_end;
348 let gap_end = child.start_byte();
349 builder = self.capture_interstitial(
350 builder,
351 vertex_id,
352 gap_start,
353 gap_end,
354 &mut interstitial_idx,
355 &mut fingerprint_parts,
356 );
357 let child_kind = child.kind();
367 if !child_kind.starts_with('_') {
368 child_kinds.push(child_kind.to_owned());
369 }
370 builder = self.walk_node(*child, builder, id_gen, Some(vertex_id))?;
371 prev_end = child.end_byte();
372 }
373
374 builder = self.capture_interstitial(
376 builder,
377 vertex_id,
378 prev_end,
379 node.end_byte(),
380 &mut interstitial_idx,
381 &mut fingerprint_parts,
382 );
383
384 if !fingerprint_parts.is_empty() {
385 builder = builder.constraint(
386 vertex_id,
387 "chose-alt-fingerprint",
388 &fingerprint_parts.join(" "),
389 );
390 }
391 if !child_kinds.is_empty() {
392 builder =
393 builder.constraint(vertex_id, "chose-alt-child-kinds", &child_kinds.join(" "));
394 }
395
396 Ok(builder)
397 }
398
399 fn capture_anonymous_field_constraints(
414 &self,
415 node: tree_sitter::Node<'_>,
416 vertex_id: &str,
417 mut builder: SchemaBuilder,
418 ) -> SchemaBuilder {
419 let child_count = node.child_count();
420 for i in 0..child_count {
421 let Some(child) = node.child(i) else { continue };
422 if child.is_named() {
426 continue;
427 }
428 let Some(field_name) = u32::try_from(i)
429 .ok()
430 .and_then(|idx| node.field_name_for_child(idx))
431 else {
432 continue;
433 };
434 let Ok(text) = child.utf8_text(self.source) else {
435 continue;
436 };
437 let sort = format!("field:{field_name}");
438 builder = builder.constraint(vertex_id, &sort, text);
439 }
440 builder
441 }
442
443 fn capture_interstitial(
444 &self,
445 mut builder: SchemaBuilder,
446 vertex_id: &str,
447 gap_start: usize,
448 gap_end: usize,
449 idx: &mut usize,
450 fingerprint: &mut Vec<String>,
451 ) -> SchemaBuilder {
452 if gap_end > gap_start && gap_end <= self.source.len() {
453 if let Ok(gap_text) = std::str::from_utf8(&self.source[gap_start..gap_end]) {
454 if !gap_text.is_empty() {
455 let sort = format!("interstitial-{}", *idx);
456 builder = builder.constraint(vertex_id, &sort, gap_text);
457 builder = builder.constraint(
458 vertex_id,
459 &format!("{sort}-start-byte"),
460 &gap_start.to_string(),
461 );
462 *idx += 1;
463 let trimmed = gap_text.trim();
464 if !trimmed.is_empty() {
465 fingerprint.push(trimmed.to_owned());
466 }
467 }
468 }
469 }
470 builder
471 }
472
473 fn emit_formatting_constraints(
475 &self,
476 node: tree_sitter::Node<'_>,
477 vertex_id: &str,
478 mut builder: SchemaBuilder,
479 ) -> SchemaBuilder {
480 let start = node.start_position();
481
482 if start.column > 0 {
484 let line_start = node.start_byte().saturating_sub(start.column);
486 if line_start < self.source.len() {
487 let indent_end = line_start + start.column.min(self.source.len() - line_start);
488 if let Ok(indent) = std::str::from_utf8(&self.source[line_start..indent_end]) {
489 if !indent.is_empty() && indent.trim().is_empty() {
491 builder = builder.constraint(vertex_id, "indent", indent);
492 }
493 }
494 }
495 }
496
497 if let Some(prev) = node.prev_named_sibling() {
500 let gap_start = prev.end_byte();
501 let gap_end = node.start_byte();
502 if gap_start < gap_end && gap_end <= self.source.len() {
503 let gap = &self.source[gap_start..gap_end];
504 let blank_lines = memchr::memchr_iter(b'\n', gap).count().saturating_sub(1);
505 if blank_lines > 0 {
506 builder = builder.constraint(
507 vertex_id,
508 "blank-lines-before",
509 &blank_lines.to_string(),
510 );
511 }
512 }
513 }
514
515 builder
516 }
517}
518
519#[cfg(test)]
520#[allow(clippy::unwrap_used)]
521mod tests {
522 use super::*;
523
524 fn make_test_protocol() -> Protocol {
525 Protocol {
526 name: "test".into(),
527 schema_theory: "ThTest".into(),
528 instance_theory: "ThTestInst".into(),
529 schema_composition: None,
530 instance_composition: None,
531 obj_kinds: vec![], edge_rules: vec![],
533 constraint_sorts: vec![],
534 has_order: true,
535 has_coproducts: false,
536 has_recursion: false,
537 has_causal: false,
538 nominal_identity: false,
539 has_defaults: false,
540 has_coercions: false,
541 has_mergers: false,
542 has_policies: false,
543 }
544 }
545
546 fn make_test_meta() -> ExtractedTheoryMeta {
547 use panproto_gat::{Sort, Theory};
548 ExtractedTheoryMeta {
549 theory: Theory::new("ThTest", vec![Sort::simple("Vertex")], vec![], vec![]),
550 supertypes: FxHashSet::default(),
551 subtype_map: Vec::new(),
552 optional_fields: FxHashSet::default(),
553 ordered_fields: FxHashSet::default(),
554 vertex_kinds: Vec::new(),
555 edge_kinds: Vec::new(),
556 }
557 }
558
559 #[cfg(feature = "grammars")]
561 fn get_grammar(name: &str) -> panproto_grammars::Grammar {
562 panproto_grammars::grammars()
563 .into_iter()
564 .find(|g| g.name == name)
565 .unwrap_or_else(|| panic!("grammar '{name}' not enabled in features"))
566 }
567
568 #[test]
569 #[cfg(feature = "grammars")]
570 fn walk_simple_typescript() {
571 let source = b"function greet(name: string): string { return name; }";
572 let grammar = get_grammar("typescript");
573
574 let mut parser = tree_sitter::Parser::new();
575 parser.set_language(&grammar.language).unwrap();
576 let tree = parser.parse(source, None).unwrap();
577
578 let protocol = make_test_protocol();
579 let meta = make_test_meta();
580 let mut detector =
581 crate::scope_detector::ScopeDetector::new(&grammar.language, grammar.tags_query, None)
582 .unwrap();
583 let walker = AstWalker::new(
584 source,
585 &meta,
586 &protocol,
587 WalkerConfig::standard(),
588 Some(&mut detector),
589 );
590
591 let schema = walker.walk(&tree, "test.ts").unwrap();
592
593 assert!(
595 schema.vertices.len() > 1,
596 "expected multiple vertices, got {}",
597 schema.vertices.len()
598 );
599
600 let root_name: panproto_gat::Name = "test.ts".into();
602 assert!(
603 schema.vertices.contains_key(&root_name),
604 "missing root vertex"
605 );
606
607 if detector.has_query() {
609 let has_greet = schema
610 .vertices
611 .keys()
612 .any(|n| n.to_string().ends_with("::greet"));
613 assert!(
614 has_greet,
615 "expected a vertex ID ending in ::greet, got: {:?}",
616 schema
617 .vertices
618 .keys()
619 .map(ToString::to_string)
620 .collect::<Vec<_>>()
621 );
622 }
623 }
624
625 #[test]
626 #[cfg(feature = "grammars")]
627 fn walk_simple_python() {
628 let source = b"def add(a, b):\n return a + b\n";
629 let grammar = get_grammar("python");
630
631 let mut parser = tree_sitter::Parser::new();
632 parser.set_language(&grammar.language).unwrap();
633 let tree = parser.parse(source, None).unwrap();
634
635 let protocol = make_test_protocol();
636 let meta = make_test_meta();
637 let mut detector =
638 crate::scope_detector::ScopeDetector::new(&grammar.language, grammar.tags_query, None)
639 .unwrap();
640 let walker = AstWalker::new(
641 source,
642 &meta,
643 &protocol,
644 WalkerConfig::standard(),
645 Some(&mut detector),
646 );
647
648 let schema = walker.walk(&tree, "test.py").unwrap();
649
650 assert!(
651 schema.vertices.len() > 1,
652 "expected multiple vertices, got {}",
653 schema.vertices.len()
654 );
655
656 if detector.has_query() {
657 let has_add = schema
658 .vertices
659 .keys()
660 .any(|n| n.to_string().ends_with("::add"));
661 assert!(has_add, "expected ::add vertex");
662 }
663 }
664
665 #[test]
666 #[cfg(feature = "grammars")]
667 fn walk_simple_rust() {
668 let source = b"fn verify_push() {}\nstruct Foo;\nimpl Foo { fn bar(&self) {} }\n";
669 let grammar = get_grammar("rust");
670
671 let mut parser = tree_sitter::Parser::new();
672 parser.set_language(&grammar.language).unwrap();
673 let tree = parser.parse(source, None).unwrap();
674
675 let protocol = make_test_protocol();
676 let meta = make_test_meta();
677 let mut detector =
678 crate::scope_detector::ScopeDetector::new(&grammar.language, grammar.tags_query, None)
679 .unwrap();
680 let walker = AstWalker::new(
681 source,
682 &meta,
683 &protocol,
684 WalkerConfig::standard(),
685 Some(&mut detector),
686 );
687
688 let schema = walker.walk(&tree, "test.rs").unwrap();
689
690 assert!(
691 schema.vertices.len() > 1,
692 "expected multiple vertices, got {}",
693 schema.vertices.len()
694 );
695
696 if detector.has_query() {
697 let vertex_ids: Vec<String> = schema.vertices.keys().map(ToString::to_string).collect();
698
699 assert!(
702 vertex_ids.iter().any(|id| id.ends_with("::verify_push")),
703 "expected ::verify_push named scope, got: {vertex_ids:?}"
704 );
705 assert!(
706 vertex_ids.iter().any(|id| id.ends_with("::Foo")),
707 "expected ::Foo named scope, got: {vertex_ids:?}"
708 );
709 }
710 }
711
712 #[cfg(feature = "group-data")]
714 fn assert_roundtrip(grammar_name: &str, source: &[u8], file_path: &str) {
715 use crate::registry::AstParser;
716 let grammar = panproto_grammars::grammars()
717 .into_iter()
718 .find(|g| g.name == grammar_name)
719 .unwrap_or_else(|| panic!("grammar '{grammar_name}' not enabled"));
720
721 let config = crate::languages::walker_configs::walker_config_for(grammar_name);
722 let lang_parser = crate::languages::common::LanguageParser::from_language(
723 grammar_name,
724 grammar.extensions.to_vec(),
725 grammar.language,
726 grammar.node_types,
727 grammar.tags_query,
728 config,
729 )
730 .unwrap();
731
732 let schema = lang_parser.parse(source, file_path).unwrap();
733 let emitted = lang_parser.emit(&schema).unwrap();
734
735 assert_eq!(
736 std::str::from_utf8(source).unwrap(),
737 std::str::from_utf8(&emitted).unwrap(),
738 "round-trip failed for {grammar_name}: emitted bytes differ from source"
739 );
740 }
741
742 #[test]
743 #[cfg(feature = "group-data")]
744 fn roundtrip_json_simple() {
745 assert_roundtrip("json", br#"{"name": "test", "value": 42}"#, "test.json");
746 }
747
748 #[test]
749 #[cfg(feature = "group-data")]
750 fn roundtrip_json_formatted() {
751 let source =
752 b"{\n \"name\": \"test\",\n \"value\": 42,\n \"nested\": {\n \"a\": true\n }\n}";
753 assert_roundtrip("json", source, "test.json");
754 }
755
756 #[test]
757 #[cfg(feature = "group-data")]
758 fn roundtrip_json_array() {
759 let source = b"[\n 1,\n 2,\n 3\n]";
760 assert_roundtrip("json", source, "test.json");
761 }
762
763 #[test]
764 #[cfg(feature = "group-data")]
765 fn roundtrip_xml_simple() {
766 let source = b"<root>\n <child attr=\"val\">text</child>\n</root>";
767 assert_roundtrip("xml", source, "test.xml");
768 }
769
770 #[test]
771 #[cfg(feature = "group-data")]
772 fn roundtrip_yaml_simple() {
773 let source = b"name: test\nvalue: 42\nnested:\n a: true\n";
774 assert_roundtrip("yaml", source, "test.yaml");
775 }
776
777 #[test]
778 #[cfg(feature = "group-data")]
779 fn roundtrip_toml_simple() {
780 let source = b"[package]\nname = \"test\"\nversion = \"0.1.0\"\n";
781 assert_roundtrip("toml", source, "test.toml");
782 }
783
784 #[cfg(feature = "group-data")]
785 fn parse_with(grammar_name: &str, source: &[u8], file_path: &str) -> panproto_schema::Schema {
786 use crate::registry::AstParser;
787 let grammar = panproto_grammars::grammars()
788 .into_iter()
789 .find(|g| g.name == grammar_name)
790 .unwrap_or_else(|| panic!("grammar '{grammar_name}' not enabled"));
791 let config = crate::languages::walker_configs::walker_config_for(grammar_name);
792 let lang_parser = crate::languages::common::LanguageParser::from_language(
793 grammar_name,
794 grammar.extensions.to_vec(),
795 grammar.language,
796 grammar.node_types,
797 grammar.tags_query,
798 config,
799 )
800 .unwrap();
801 lang_parser.parse(source, file_path).unwrap()
802 }
803
804 #[test]
805 #[cfg(feature = "group-data")]
806 fn fingerprint_and_child_kinds_emitted_separately() {
807 let schema = parse_with("json", br#"{"a": 1}"#, "test.json");
814
815 let saw_fingerprint = schema.constraints.values().any(|cs| {
816 cs.iter()
817 .any(|c| c.sort.as_ref() == "chose-alt-fingerprint")
818 });
819 let saw_child_kinds = schema.constraints.values().any(|cs| {
820 cs.iter()
821 .any(|c| c.sort.as_ref() == "chose-alt-child-kinds")
822 });
823
824 assert!(
825 saw_fingerprint,
826 "walker must emit chose-alt-fingerprint (literal-token witness)"
827 );
828 assert!(
829 saw_child_kinds,
830 "walker must emit chose-alt-child-kinds (named-kind witness)"
831 );
832 }
833
834 #[test]
835 #[cfg(feature = "group-data")]
836 fn child_kinds_excludes_hidden_rules() {
837 let schema = parse_with("json", br#"{"k": "v"}"#, "test.json");
840
841 for cs in schema.constraints.values() {
842 for c in cs {
843 if c.sort.as_ref() == "chose-alt-child-kinds" {
844 for kind in c.value.split_whitespace() {
845 assert!(
846 !kind.starts_with('_'),
847 "hidden-rule kind '{kind}' must not appear in chose-alt-child-kinds"
848 );
849 }
850 }
851 }
852 }
853 }
854}