From 79e385f12eeefab68181d3e8a49d3e41be331760 Mon Sep 17 00:00:00 2001
From: SpontanCombust <61706594+SpontanCombust@users.noreply.github.com>
Date: Thu, 25 Apr 2024 00:05:37 +0200
Subject: [PATCH] Disambiguated `StateBlock` from `ClassBlock`

---
 crates/core/src/ast/classes.rs | 10 +++----
 crates/core/src/ast/states.rs  | 49 +++++++++++++++++++++++++++++++---
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/crates/core/src/ast/classes.rs b/crates/core/src/ast/classes.rs
index 54616987..c4da06c3 100644
--- a/crates/core/src/ast/classes.rs
+++ b/crates/core/src/ast/classes.rs
@@ -64,14 +64,14 @@ impl DeclarationTraversal for ClassDeclarationNode<'_> {
     fn accept<V: DeclarationVisitor>(&self, visitor: &mut V, _: Self::TraversalCtx) {
         let tp = visitor.visit_class_decl(self);
         if tp.traverse_definition {
-            self.definition().accept(visitor, PropertyTraversalContext::ClassDefinition);
+            self.definition().accept(visitor, ());
         }
         visitor.exit_class_decl(self);
     }
 }
 
 
-//TODO make seperate StateBlockNode type with the same NODE_KIND to disambiguate traversal with state 
+
 pub type ClassBlockNode<'script> = SyntaxNode<'script, tags::ClassBlock>;
 
 impl NamedSyntaxNode for ClassBlockNode<'_> {
@@ -106,10 +106,10 @@ impl<'script> TryFrom<AnyNode<'script>> for ClassBlockNode<'script> {
 }
 
 impl DeclarationTraversal for ClassBlockNode<'_> {
-    type TraversalCtx = PropertyTraversalContext;
+    type TraversalCtx = ();
 
-    fn accept<V: DeclarationVisitor>(&self, visitor: &mut V, ctx: Self::TraversalCtx) {
-        self.iter().for_each(|s| s.accept(visitor, ctx));
+    fn accept<V: DeclarationVisitor>(&self, visitor: &mut V, _: Self::TraversalCtx) {
+        self.iter().for_each(|s| s.accept(visitor, PropertyTraversalContext::ClassDefinition));
     }
 }
 
diff --git a/crates/core/src/ast/states.rs b/crates/core/src/ast/states.rs
index c30c194b..5235d638 100644
--- a/crates/core/src/ast/states.rs
+++ b/crates/core/src/ast/states.rs
@@ -1,10 +1,11 @@
 use std::fmt::Debug;
-use crate::{attribs::StateSpecifierNode, tokens::IdentifierNode, AnyNode, DebugRange, NamedSyntaxNode, SyntaxNode};
+use crate::{attribs::StateSpecifierNode, tokens::IdentifierNode, AnyNode, DebugMaybeAlternate, DebugRange, NamedSyntaxNode, SyntaxNode};
 use super::*;
 
 
 mod tags {
     pub struct StateDeclaration;
+    pub struct StateBlock;
 }
 
 
@@ -31,7 +32,7 @@ impl<'script> StateDeclarationNode<'script> {
         self.field_child("base").map(|n| n.into())
     }
 
-    pub fn definition(&self) -> ClassBlockNode<'script> {
+    pub fn definition(&self) -> StateBlockNode<'script> {
         self.field_child("definition").unwrap().into()
     }
 }
@@ -66,8 +67,50 @@ impl DeclarationTraversal for StateDeclarationNode<'_> {
     fn accept<V: DeclarationVisitor>(&self, visitor: &mut V, _: Self::TraversalCtx) {
         let tp = visitor.visit_state_decl(self);
         if tp.traverse_definition {
-            self.definition().accept(visitor, PropertyTraversalContext::StateDefinition);
+            self.definition().accept(visitor, ());
         }
         visitor.exit_state_decl(self);
     }
+}
+
+
+pub type StateBlockNode<'script> = SyntaxNode<'script, tags::StateBlock>;
+
+impl NamedSyntaxNode for StateBlockNode<'_> {
+    const NODE_KIND: &'static str = "class_block";
+}
+
+impl<'script> StateBlockNode<'script> {
+    pub fn iter(&self) -> impl Iterator<Item = ClassStatementNode<'script>> {
+        self.named_children().map(|n| n.into())
+    }
+}
+
+impl Debug for StateBlockNode<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_maybe_alternate_named(
+            &format!("StateBlock {}", self.range().debug()), 
+            &self.iter().collect::<Vec<_>>()
+        )
+    }
+}
+
+impl<'script> TryFrom<AnyNode<'script>> for StateBlockNode<'script> {
+    type Error = ();
+
+    fn try_from(value: AnyNode<'script>) -> Result<Self, Self::Error> {
+        if value.tree_node.kind() == Self::NODE_KIND {
+            Ok(value.into())
+        } else {
+            Err(())
+        }
+    }
+}
+
+impl DeclarationTraversal for StateBlockNode<'_> {
+    type TraversalCtx = ();
+
+    fn accept<V: DeclarationVisitor>(&self, visitor: &mut V, _: Self::TraversalCtx) {
+        self.iter().for_each(|s| s.accept(visitor, PropertyTraversalContext::StateDefinition));
+    }
 }
\ No newline at end of file