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(u32::try_from(i).unwrap_or(0)) {
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 let grammar_name = node.grammar_name();
284 if grammar_name != kind {
285 builder = builder.constraint(&vertex_id, "pre-alias-symbol", grammar_name);
286 }
287
288 if node.named_child_count() == 0 {
290 if let Ok(text) = node.utf8_text(self.source) {
291 builder = builder.constraint(&vertex_id, "literal-value", text);
292 }
293 }
294
295 builder = self.capture_anonymous_field_constraints(node, &vertex_id, builder);
305
306 if self.config.capture_formatting {
308 builder = self.emit_formatting_constraints(node, &vertex_id, builder);
309 }
310
311 let entered_scope = if let Some(leaf) = recorded_named_leaf {
317 id_gen.push_recorded_scope(leaf);
318 true
319 } else if !is_root_wrapper && self.block_kinds.contains(kind) {
320 id_gen.push_anonymous_scope();
321 true
322 } else {
323 false
324 };
325
326 builder = self.walk_children_with_interstitials(node, builder, id_gen, &vertex_id)?;
327
328 if entered_scope {
329 id_gen.pop_scope();
330 }
331
332 Ok(builder)
333 }
334
335 fn walk_children_with_interstitials(
347 &self,
348 node: tree_sitter::Node<'_>,
349 mut builder: SchemaBuilder,
350 id_gen: &mut IdGenerator,
351 vertex_id: &str,
352 ) -> Result<SchemaBuilder, ParseError> {
353 let cursor = &mut node.walk();
354 let children: Vec<_> = node.named_children(cursor).collect();
355 let mut interstitial_idx = 0;
356 let mut prev_end = node.start_byte();
357 let mut fingerprint_parts: Vec<String> = Vec::new();
358 let mut child_kinds: Vec<String> = Vec::new();
359
360 for child in &children {
361 let gap_start = prev_end;
362 let gap_end = child.start_byte();
363 builder = self.capture_interstitial(
364 builder,
365 vertex_id,
366 gap_start,
367 gap_end,
368 &mut interstitial_idx,
369 &mut fingerprint_parts,
370 );
371 let child_kind = child.kind();
381 if !child_kind.starts_with('_') {
382 child_kinds.push(child_kind.to_owned());
383 }
384 builder = self.walk_node(*child, builder, id_gen, Some(vertex_id))?;
385 prev_end = child.end_byte();
386 }
387
388 builder = self.capture_interstitial(
390 builder,
391 vertex_id,
392 prev_end,
393 node.end_byte(),
394 &mut interstitial_idx,
395 &mut fingerprint_parts,
396 );
397
398 if !fingerprint_parts.is_empty() {
399 builder = builder.constraint(
400 vertex_id,
401 "chose-alt-fingerprint",
402 &fingerprint_parts.join(" "),
403 );
404 }
405 if !child_kinds.is_empty() {
406 builder =
407 builder.constraint(vertex_id, "chose-alt-child-kinds", &child_kinds.join(" "));
408 }
409
410 Ok(builder)
411 }
412
413 fn capture_anonymous_field_constraints(
428 &self,
429 node: tree_sitter::Node<'_>,
430 vertex_id: &str,
431 mut builder: SchemaBuilder,
432 ) -> SchemaBuilder {
433 let child_count = node.child_count();
434 for i in 0..child_count {
435 let Some(child) = node.child(u32::try_from(i).unwrap_or(0)) else {
436 continue;
437 };
438 if child.is_named() {
442 continue;
443 }
444 let Some(field_name) = u32::try_from(i)
445 .ok()
446 .and_then(|idx| node.field_name_for_child(idx))
447 else {
448 continue;
449 };
450 let Ok(text) = child.utf8_text(self.source) else {
451 continue;
452 };
453 let sort = format!("field:{field_name}");
454 builder = builder.constraint(vertex_id, &sort, text);
455 }
456 builder
457 }
458
459 fn capture_interstitial(
460 &self,
461 mut builder: SchemaBuilder,
462 vertex_id: &str,
463 gap_start: usize,
464 gap_end: usize,
465 idx: &mut usize,
466 fingerprint: &mut Vec<String>,
467 ) -> SchemaBuilder {
468 if gap_end > gap_start && gap_end <= self.source.len() {
469 if let Ok(gap_text) = std::str::from_utf8(&self.source[gap_start..gap_end]) {
470 if !gap_text.is_empty() {
471 let sort = format!("interstitial-{}", *idx);
472 builder = builder.constraint(vertex_id, &sort, gap_text);
473 builder = builder.constraint(
474 vertex_id,
475 &format!("{sort}-start-byte"),
476 &gap_start.to_string(),
477 );
478 *idx += 1;
479 let trimmed = gap_text.trim();
480 if !trimmed.is_empty() {
481 fingerprint.push(trimmed.to_owned());
482 }
483 }
484 }
485 }
486 builder
487 }
488
489 fn emit_formatting_constraints(
491 &self,
492 node: tree_sitter::Node<'_>,
493 vertex_id: &str,
494 mut builder: SchemaBuilder,
495 ) -> SchemaBuilder {
496 let start = node.start_position();
497
498 if start.column > 0 {
500 let line_start = node.start_byte().saturating_sub(start.column);
502 if line_start < self.source.len() {
503 let indent_end = line_start + start.column.min(self.source.len() - line_start);
504 if let Ok(indent) = std::str::from_utf8(&self.source[line_start..indent_end]) {
505 if !indent.is_empty() && indent.trim().is_empty() {
507 builder = builder.constraint(vertex_id, "indent", indent);
508 }
509 }
510 }
511 }
512
513 if let Some(prev) = node.prev_named_sibling() {
516 let gap_start = prev.end_byte();
517 let gap_end = node.start_byte();
518 if gap_start < gap_end && gap_end <= self.source.len() {
519 let gap = &self.source[gap_start..gap_end];
520 let blank_lines = memchr::memchr_iter(b'\n', gap).count().saturating_sub(1);
521 if blank_lines > 0 {
522 builder = builder.constraint(
523 vertex_id,
524 "blank-lines-before",
525 &blank_lines.to_string(),
526 );
527 }
528 }
529 }
530
531 builder
532 }
533}
534
535#[cfg(test)]
536#[allow(clippy::unwrap_used)]
537mod tests {
538 use super::*;
539
540 fn make_test_protocol() -> Protocol {
541 Protocol {
542 name: "test".into(),
543 schema_theory: "ThTest".into(),
544 instance_theory: "ThTestInst".into(),
545 schema_composition: None,
546 instance_composition: None,
547 obj_kinds: vec![], edge_rules: vec![],
549 constraint_sorts: vec![],
550 has_order: true,
551 has_coproducts: false,
552 has_recursion: false,
553 has_causal: false,
554 nominal_identity: false,
555 has_defaults: false,
556 has_coercions: false,
557 has_mergers: false,
558 has_policies: false,
559 }
560 }
561
562 fn make_test_meta() -> ExtractedTheoryMeta {
563 use panproto_gat::{Sort, Theory};
564 ExtractedTheoryMeta {
565 theory: Theory::new("ThTest", vec![Sort::simple("Vertex")], vec![], vec![]),
566 supertypes: FxHashSet::default(),
567 subtype_map: Vec::new(),
568 optional_fields: FxHashSet::default(),
569 ordered_fields: FxHashSet::default(),
570 vertex_kinds: Vec::new(),
571 edge_kinds: Vec::new(),
572 }
573 }
574
575 #[cfg(feature = "grammars")]
577 fn get_grammar(name: &str) -> panproto_grammars::Grammar {
578 panproto_grammars::grammars()
579 .into_iter()
580 .find(|g| g.name == name)
581 .unwrap_or_else(|| panic!("grammar '{name}' not enabled in features"))
582 }
583
584 #[test]
585 #[cfg(feature = "grammars")]
586 fn walk_simple_typescript() {
587 let source = b"function greet(name: string): string { return name; }";
588 let grammar = get_grammar("typescript");
589
590 let mut parser = tree_sitter::Parser::new();
591 parser.set_language(&grammar.language).unwrap();
592 let tree = parser.parse(source, None).unwrap();
593
594 let protocol = make_test_protocol();
595 let meta = make_test_meta();
596 let mut detector =
597 crate::scope_detector::ScopeDetector::new(&grammar.language, grammar.tags_query, None)
598 .unwrap();
599 let walker = AstWalker::new(
600 source,
601 &meta,
602 &protocol,
603 WalkerConfig::standard(),
604 Some(&mut detector),
605 );
606
607 let schema = walker.walk(&tree, "test.ts").unwrap();
608
609 assert!(
611 schema.vertices.len() > 1,
612 "expected multiple vertices, got {}",
613 schema.vertices.len()
614 );
615
616 let root_name: panproto_gat::Name = "test.ts".into();
618 assert!(
619 schema.vertices.contains_key(&root_name),
620 "missing root vertex"
621 );
622
623 if detector.has_query() {
625 let has_greet = schema
626 .vertices
627 .keys()
628 .any(|n| n.to_string().ends_with("::greet"));
629 assert!(
630 has_greet,
631 "expected a vertex ID ending in ::greet, got: {:?}",
632 schema
633 .vertices
634 .keys()
635 .map(ToString::to_string)
636 .collect::<Vec<_>>()
637 );
638 }
639 }
640
641 #[test]
642 #[cfg(feature = "grammars")]
643 fn walk_simple_python() {
644 let source = b"def add(a, b):\n return a + b\n";
645 let grammar = get_grammar("python");
646
647 let mut parser = tree_sitter::Parser::new();
648 parser.set_language(&grammar.language).unwrap();
649 let tree = parser.parse(source, None).unwrap();
650
651 let protocol = make_test_protocol();
652 let meta = make_test_meta();
653 let mut detector =
654 crate::scope_detector::ScopeDetector::new(&grammar.language, grammar.tags_query, None)
655 .unwrap();
656 let walker = AstWalker::new(
657 source,
658 &meta,
659 &protocol,
660 WalkerConfig::standard(),
661 Some(&mut detector),
662 );
663
664 let schema = walker.walk(&tree, "test.py").unwrap();
665
666 assert!(
667 schema.vertices.len() > 1,
668 "expected multiple vertices, got {}",
669 schema.vertices.len()
670 );
671
672 if detector.has_query() {
673 let has_add = schema
674 .vertices
675 .keys()
676 .any(|n| n.to_string().ends_with("::add"));
677 assert!(has_add, "expected ::add vertex");
678 }
679 }
680
681 #[test]
682 #[cfg(feature = "grammars")]
683 fn walk_simple_rust() {
684 let source = b"fn verify_push() {}\nstruct Foo;\nimpl Foo { fn bar(&self) {} }\n";
685 let grammar = get_grammar("rust");
686
687 let mut parser = tree_sitter::Parser::new();
688 parser.set_language(&grammar.language).unwrap();
689 let tree = parser.parse(source, None).unwrap();
690
691 let protocol = make_test_protocol();
692 let meta = make_test_meta();
693 let mut detector =
694 crate::scope_detector::ScopeDetector::new(&grammar.language, grammar.tags_query, None)
695 .unwrap();
696 let walker = AstWalker::new(
697 source,
698 &meta,
699 &protocol,
700 WalkerConfig::standard(),
701 Some(&mut detector),
702 );
703
704 let schema = walker.walk(&tree, "test.rs").unwrap();
705
706 assert!(
707 schema.vertices.len() > 1,
708 "expected multiple vertices, got {}",
709 schema.vertices.len()
710 );
711
712 if detector.has_query() {
713 let vertex_ids: Vec<String> = schema.vertices.keys().map(ToString::to_string).collect();
714
715 assert!(
718 vertex_ids.iter().any(|id| id.ends_with("::verify_push")),
719 "expected ::verify_push named scope, got: {vertex_ids:?}"
720 );
721 assert!(
722 vertex_ids.iter().any(|id| id.ends_with("::Foo")),
723 "expected ::Foo named scope, got: {vertex_ids:?}"
724 );
725 }
726 }
727
728 #[cfg(feature = "group-data")]
730 fn assert_roundtrip(grammar_name: &str, source: &[u8], file_path: &str) {
731 use crate::registry::AstParser;
732 let grammar = panproto_grammars::grammars()
733 .into_iter()
734 .find(|g| g.name == grammar_name)
735 .unwrap_or_else(|| panic!("grammar '{grammar_name}' not enabled"));
736
737 let config = crate::languages::walker_configs::walker_config_for(grammar_name);
738 let lang_parser = crate::languages::common::LanguageParser::from_language(
739 grammar_name,
740 grammar.extensions.to_vec(),
741 grammar.language,
742 grammar.node_types,
743 grammar.tags_query,
744 config,
745 )
746 .unwrap();
747
748 let schema = lang_parser.parse(source, file_path).unwrap();
749 let emitted = lang_parser.emit(&schema).unwrap();
750
751 assert_eq!(
752 std::str::from_utf8(source).unwrap(),
753 std::str::from_utf8(&emitted).unwrap(),
754 "round-trip failed for {grammar_name}: emitted bytes differ from source"
755 );
756 }
757
758 #[test]
759 #[cfg(feature = "group-data")]
760 fn roundtrip_json_simple() {
761 assert_roundtrip("json", br#"{"name": "test", "value": 42}"#, "test.json");
762 }
763
764 #[test]
765 #[cfg(feature = "group-data")]
766 fn roundtrip_json_formatted() {
767 let source =
768 b"{\n \"name\": \"test\",\n \"value\": 42,\n \"nested\": {\n \"a\": true\n }\n}";
769 assert_roundtrip("json", source, "test.json");
770 }
771
772 #[test]
773 #[cfg(feature = "group-data")]
774 fn roundtrip_json_array() {
775 let source = b"[\n 1,\n 2,\n 3\n]";
776 assert_roundtrip("json", source, "test.json");
777 }
778
779 #[test]
780 #[cfg(feature = "group-data")]
781 fn roundtrip_xml_simple() {
782 let source = b"<root>\n <child attr=\"val\">text</child>\n</root>";
783 assert_roundtrip("xml", source, "test.xml");
784 }
785
786 #[test]
787 #[cfg(feature = "group-data")]
788 fn roundtrip_yaml_simple() {
789 let source = b"name: test\nvalue: 42\nnested:\n a: true\n";
790 assert_roundtrip("yaml", source, "test.yaml");
791 }
792
793 #[test]
794 #[cfg(feature = "group-data")]
795 fn roundtrip_toml_simple() {
796 let source = b"[package]\nname = \"test\"\nversion = \"0.1.0\"\n";
797 assert_roundtrip("toml", source, "test.toml");
798 }
799
800 #[cfg(feature = "group-data")]
801 fn parse_with(grammar_name: &str, source: &[u8], file_path: &str) -> panproto_schema::Schema {
802 use crate::registry::AstParser;
803 let grammar = panproto_grammars::grammars()
804 .into_iter()
805 .find(|g| g.name == grammar_name)
806 .unwrap_or_else(|| panic!("grammar '{grammar_name}' not enabled"));
807 let config = crate::languages::walker_configs::walker_config_for(grammar_name);
808 let lang_parser = crate::languages::common::LanguageParser::from_language(
809 grammar_name,
810 grammar.extensions.to_vec(),
811 grammar.language,
812 grammar.node_types,
813 grammar.tags_query,
814 config,
815 )
816 .unwrap();
817 lang_parser.parse(source, file_path).unwrap()
818 }
819
820 #[test]
821 #[cfg(feature = "group-data")]
822 fn fingerprint_and_child_kinds_emitted_separately() {
823 let schema = parse_with("json", br#"{"a": 1}"#, "test.json");
830
831 let saw_fingerprint = schema.constraints.values().any(|cs| {
832 cs.iter()
833 .any(|c| c.sort.as_ref() == "chose-alt-fingerprint")
834 });
835 let saw_child_kinds = schema.constraints.values().any(|cs| {
836 cs.iter()
837 .any(|c| c.sort.as_ref() == "chose-alt-child-kinds")
838 });
839
840 assert!(
841 saw_fingerprint,
842 "walker must emit chose-alt-fingerprint (literal-token witness)"
843 );
844 assert!(
845 saw_child_kinds,
846 "walker must emit chose-alt-child-kinds (named-kind witness)"
847 );
848 }
849
850 #[test]
851 #[cfg(feature = "group-data")]
852 fn child_kinds_excludes_hidden_rules() {
853 let schema = parse_with("json", br#"{"k": "v"}"#, "test.json");
856
857 for cs in schema.constraints.values() {
858 for c in cs {
859 if c.sort.as_ref() == "chose-alt-child-kinds" {
860 for kind in c.value.split_whitespace() {
861 assert!(
862 !kind.starts_with('_'),
863 "hidden-rule kind '{kind}' must not appear in chose-alt-child-kinds"
864 );
865 }
866 }
867 }
868 }
869 }
870}