Added TypeAnnotationNode traversal

SpontanCombust · Aug 15, 2024 · 3ff05ed · 3ff05ed
1 parent 2690fd7
commit 3ff05ed
Show file tree

Hide file tree

Showing 10 changed files with 100 additions and 0 deletions.
diff --git a/crates/analysis/src/jobs/scan_symbols.rs b/crates/analysis/src/jobs/scan_symbols.rs
@@ -545,6 +545,7 @@ impl SyntaxNodeVisitor for SymbolScannerVisitor<'_> {
         FunctionDeclarationTraversalPolicy { 
             traverse_annotation: false,
             traverse_params: traverse,
+            traverse_return_type: false,
             traverse_definition: traverse,
             traverse_errors: false
         }
@@ -625,6 +626,7 @@ impl SyntaxNodeVisitor for SymbolScannerVisitor<'_> {
         FunctionDeclarationTraversalPolicy {
             traverse_annotation: false,
             traverse_params: traverse,
+            traverse_return_type: false,
             traverse_definition: traverse,
             traverse_errors: false
         }
@@ -660,6 +662,7 @@ impl SyntaxNodeVisitor for SymbolScannerVisitor<'_> {
 
         EventDeclarationTraversalPolicy { 
             traverse_params: traverse,
+            traverse_return_type: false,
             traverse_definition: traverse,
             traverse_errors: false
         }
@@ -797,6 +800,7 @@ impl SyntaxNodeVisitor for SymbolScannerVisitor<'_> {
         }
 
         VarDeclarationTraversalPolicy {
+            traverse_type: false,
             traverse_init_value: false,
             traverse_errors: false
         }

diff --git a/crates/analysis/src/jobs/syntax_analysis/contextual_analysis.rs b/crates/analysis/src/jobs/syntax_analysis/contextual_analysis.rs
@@ -456,6 +456,7 @@ impl SyntaxNodeVisitor for ContextualSyntaxAnalysis<'_> {
 
         VarDeclarationTraversalPolicy {
             traverse_init_value: false,
+            traverse_type: false,
             traverse_errors: false
         }
     }

diff --git a/crates/analysis/src/jobs/syntax_analysis/syntax_error_visitor.rs b/crates/analysis/src/jobs/syntax_analysis/syntax_error_visitor.rs
@@ -356,6 +356,7 @@ impl SyntaxNodeVisitor for SyntaxErrorVisitor<'_> {
 
         FunctionDeclarationTraversalPolicy { 
             traverse_annotation: false,
+            traverse_return_type: false,
             traverse_params,
             traverse_definition,
             traverse_errors: false
@@ -395,6 +396,7 @@ impl SyntaxNodeVisitor for SyntaxErrorVisitor<'_> {
 
         FunctionDeclarationTraversalPolicy {
             traverse_annotation: false, 
+            traverse_return_type: false,
             traverse_params,
             traverse_definition,
             traverse_errors: false
@@ -430,6 +432,7 @@ impl SyntaxNodeVisitor for SyntaxErrorVisitor<'_> {
 
         EventDeclarationTraversalPolicy { 
             traverse_params,
+            traverse_return_type: false,
             traverse_definition,
             traverse_errors: false
         }
@@ -492,6 +495,7 @@ impl SyntaxNodeVisitor for SyntaxErrorVisitor<'_> {
         }
 
         VarDeclarationTraversalPolicy {
+            traverse_type: false,
             traverse_init_value,
             traverse_errors: false
         }

diff --git a/crates/core/src/ast/annotation.rs b/crates/core/src/ast/annotation.rs
@@ -84,6 +84,8 @@ impl SyntaxNodeTraversal for AnnotationNode<'_> {
         let tp = visitor.visit_annotation(self, ctx);
 
         if tp.any() {
+            ctx.push(TraversalContext::Annotation);
+
             for ch in self.children_detailed().must_be_named(true) {
                 match ch {
                     Err(e) if tp.traverse_errors => {
@@ -92,6 +94,8 @@ impl SyntaxNodeTraversal for AnnotationNode<'_> {
                     _ => {}
                 }
             }
+
+            ctx.pop();
         }
 
         visitor.exit_annotation(self, ctx);

diff --git a/crates/core/src/ast/functions.rs b/crates/core/src/ast/functions.rs
@@ -73,6 +73,11 @@ impl SyntaxNodeTraversal for EventDeclarationNode<'_> {
 
                         params.accept_with_policy(visitor, ctx, tp.traverse_errors);
                     },
+                    Ok((rt, Some("return_type"))) if tp.traverse_return_type => {
+                        let rt: TypeAnnotationNode = rt.unsafe_into();
+
+                        rt.accept(visitor, ctx);
+                    },
                     Ok((def, Some("definition"))) if tp.traverse_definition => {
                         let def: FunctionDefinitionNode = def.unsafe_into();
 
@@ -172,6 +177,11 @@ impl SyntaxNodeTraversal for FunctionDeclarationNode<'_> {
 
                         params.accept_with_policy(visitor, ctx, tp.traverse_errors);
                     },
+                    Ok((rt, Some("return_type"))) if tp.traverse_return_type => {
+                        let rt: TypeAnnotationNode = rt.unsafe_into();
+
+                        rt.accept(visitor, ctx);
+                    },
                     Ok((def, Some("definition"))) if tp.traverse_definition => {
                         let def: FunctionDefinitionNode = def.unsafe_into();
 
@@ -450,6 +460,11 @@ impl SyntaxNodeTraversal for FunctionParameterGroupNode<'_> {
         if tp.any() {
             for ch in self.children_detailed().must_be_named(true) {
                 match ch {
+                    Ok((typ, Some("param_type"))) if tp.traverse_type => {
+                        let typ: TypeAnnotationNode = typ.unsafe_into();
+
+                        typ.accept(visitor, ctx);
+                    },
                     Err(e) if tp.traverse_errors => {
                         e.accept(visitor, ctx);
                     },

diff --git a/crates/core/src/ast/traversal/contexts.rs b/crates/core/src/ast/traversal/contexts.rs
@@ -61,6 +61,9 @@ pub enum TraversalContext {
     TernaryConditionalExpressionAlt,
     ArrayInitializerExpression,
 
+    TypeAnnotation,
+    Annotation,
+
     Error
 }
 

diff --git a/crates/core/src/ast/traversal/policies.rs b/crates/core/src/ast/traversal/policies.rs
@@ -146,6 +146,7 @@ traversal_policy!(EnumVariantDeclarationTraversalPolicy,
 
 traversal_policy!(MemberVarDeclarationTraversalPolicy,
     traverse_annotation,
+    traverse_type,
     traverse_errors
 );
 
@@ -164,22 +165,26 @@ traversal_policy!(MemberHintTraversalPolicy,
 );
 
 traversal_policy!(AutobindDeclarationTraversalPolicy,
+    traverse_type,
     traverse_errors
 );
 
 traversal_policy!(FunctionParameterGroupTraversalPolicy,
+    traverse_type,
     traverse_errors
 );
 
 traversal_policy!(FunctionDeclarationTraversalPolicy,
     traverse_annotation,
     traverse_params,
+    traverse_return_type,
     traverse_definition,
     traverse_errors
 );
 
 traversal_policy!(EventDeclarationTraversalPolicy,
     traverse_params,
+    traverse_return_type,
     traverse_definition,
     traverse_errors
 );
@@ -242,6 +247,7 @@ traversal_policy!(CompoundStatementTraversalPolicy,
 );
 
 traversal_policy!(VarDeclarationTraversalPolicy,
+    traverse_type,
     traverse_init_value,
     traverse_errors
 );
@@ -263,6 +269,11 @@ traversal_policy!(DeleteStatementTraversalPolicy,
 
 
 
+traversal_policy!(TypeAnnotationTraversalPolicy,
+    traverse_type_arg,
+    traverse_errors
+);
+
 traversal_policy!(AnnotationTraversalPolicy,
     traverse_errors
 );

diff --git a/crates/core/src/ast/traversal/visitor.rs b/crates/core/src/ast/traversal/visitor.rs
@@ -266,6 +266,11 @@ pub trait SyntaxNodeVisitor {
     fn visit_nop_stmt(&mut self, n: &NopNode, ctx: &TraversalContextStack) {}
 
 
+    /// Called when visiting a type annotation node
+    fn visit_type_annotation(&mut self, n: &TypeAnnotationNode, ctx: &TraversalContextStack) -> TypeAnnotationTraversalPolicy { TraversalPolicy::default_to(self.traversal_policy_default()) }
+    /// Called after visiting a type annotation node and possibly also children nodes specified in traversal policy.
+    fn exit_type_annotation(&mut self, n: &TypeAnnotationNode, ctx: &TraversalContextStack) {}
+
     /// Called when visiting an annotation node (e.g. @wrapMethod)
     fn visit_annotation(&mut self, n: &AnnotationNode, ctx: &TraversalContextStack) -> AnnotationTraversalPolicy { TraversalPolicy::default_to(self.traversal_policy_default()) }
     /// Called after visiting an annotation node and possibly also children nodes specified in traversal policy.

diff --git a/crates/core/src/ast/traversal/visitor_chain.rs b/crates/core/src/ast/traversal/visitor_chain.rs
@@ -560,6 +560,16 @@ impl<'a> SyntaxNodeVisitor for SyntaxNodeVisitorChain<'a> {
     }
 
 
+    #[inline]
+    fn visit_type_annotation(&mut self, n: &TypeAnnotationNode, ctx: &TraversalContextStack) -> TypeAnnotationTraversalPolicy {
+        self.chain_visit_traversable(move |link| link.visit_type_annotation(n, ctx))
+    }
+
+    #[inline]
+    fn exit_type_annotation(&mut self, n: &TypeAnnotationNode, ctx: &TraversalContextStack) {
+        self.chain_exit(move |link| link.exit_type_annotation(n, ctx))
+    }
+
     #[inline]
     fn visit_annotation(&mut self, n: &AnnotationNode, ctx: &TraversalContextStack) -> AnnotationTraversalPolicy {
         self.chain_visit_traversable(move |link| link.visit_annotation(n, ctx))

diff --git a/crates/core/src/ast/vars.rs b/crates/core/src/ast/vars.rs
@@ -49,6 +49,34 @@ impl<'script> TryFrom<AnyNode<'script>> for TypeAnnotationNode<'script> {
     }
 }
 
+impl SyntaxNodeTraversal for TypeAnnotationNode<'_> {
+    fn accept<V: SyntaxNodeVisitor>(&self, visitor: &mut V, ctx: &mut TraversalContextStack) {
+        let tp = visitor.visit_type_annotation(self, ctx);
+
+        if tp.any() {
+            ctx.push(TraversalContext::TypeAnnotation);
+
+            for ch in self.children_detailed().must_be_named(true) {
+                match ch {
+                    Ok((type_arg, Some("type_arg"))) if tp.traverse_type_arg => {
+                        let type_arg: TypeAnnotationNode = type_arg.unsafe_into();
+
+                        type_arg.accept(visitor, ctx);
+                    },
+                    Err(e) if tp.traverse_errors => {
+                        e.accept(visitor, ctx);
+                    },
+                    _ => {}
+                }
+            }
+
+            ctx.pop();
+        }
+
+        visitor.exit_type_annotation(self, ctx);
+    }
+}
+
 
 
 pub type LocalVarDeclarationNode<'script> = SyntaxNode<'script, tags::LocalVarDeclaration>;
@@ -100,6 +128,11 @@ impl SyntaxNodeTraversal for LocalVarDeclarationNode<'_> {
         if tp.any() {
             for ch in self.children_detailed().must_be_named(true) {
                 match ch {
+                    Ok((typ, Some("var_type"))) if tp.traverse_type => {
+                        let typ: TypeAnnotationNode = typ.unsafe_into();
+
+                        typ.accept(visitor, ctx);
+                    },
                     Ok((init_value, Some("init_value"))) if tp.traverse_init_value => {
                         let init_value: ExpressionNode = init_value.unsafe_into();
 
@@ -179,6 +212,11 @@ impl SyntaxNodeTraversal for MemberVarDeclarationNode<'_> {
 
                         annot.accept(visitor, ctx);
                     },
+                    Ok((typ, Some("var_type"))) if tp.traverse_type => {
+                        let typ: TypeAnnotationNode = typ.unsafe_into();
+
+                        typ.accept(visitor, ctx);
+                    },
                     Err(e) if tp.traverse_errors => {
                         e.accept(visitor, ctx);
                     },
@@ -265,6 +303,11 @@ impl SyntaxNodeTraversal for AutobindDeclarationNode<'_> {
         if tp.any() {
             for ch in self.children_detailed().must_be_named(true) {
                 match ch {
+                    Ok((typ, Some("autobind_type"))) if tp.traverse_type => {
+                        let typ: TypeAnnotationNode = typ.unsafe_into();
+
+                        typ.accept(visitor, ctx);
+                    },
                     Err(e) => {
                         e.accept(visitor, ctx)
                     },