From 1a80fe074bfd142c39e6de3e85e52fb079ff797a Mon Sep 17 00:00:00 2001
From: Klesti Selimaj <klesti.selimay@icloud.com>
Date: Sat, 16 May 2026 05:32:10 +0200
Subject: [PATCH] Improved bug implementations

---
 parser/src/error.rs                   |  3 +-
 parser/src/lib.rs                     | 16 +++++++++
 parser/src/parser/common/mod.rs       |  5 ++-
 parser/src/parser/common/statement.rs | 48 +++++++++++----------------
 4 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/parser/src/error.rs b/parser/src/error.rs
index 44984d5..e9fc2dd 100644
--- a/parser/src/error.rs
+++ b/parser/src/error.rs
@@ -20,7 +20,8 @@ pub struct AstError<'a, T> {
 
 #[derive(Debug, Clone)]
 pub enum ErrorCode {
-    InvalidStatement = 200,
+    InvalidStatement,
+    AstGenBug,
 }
 
 impl<T> From<pest::error::Error<Rule>> for ParseError<'_, T> {
diff --git a/parser/src/lib.rs b/parser/src/lib.rs
index b9c126b..df5cbf5 100644
--- a/parser/src/lib.rs
+++ b/parser/src/lib.rs
@@ -32,6 +32,22 @@ pub fn parse<'a>(source: &'a str) -> Result<Vec<TopLevel>, ParseError<'a, Vec<To
     }
 }
 
+#[macro_export]
+macro_rules! ast_ensure {
+    ($pair:expr, $rule:expr => $body:block) => {
+        if $pair.as_rule() == $rule
+            $body
+        else {
+            Err(AstError {
+                span: $pair.as_span(),
+                error_code: crate::error::ErrorCode::AstGenBug,
+                error_message: format!("Possible bug: expected {:?}, got {}", $rule, $pair),
+                recovered: None,
+            })
+        }
+    };
+}
+
 #[macro_export]
 macro_rules! ast_expr {
     ($($item:ident)::+ { $($k:ident: $v:expr),* $(,)? }) => {{
diff --git a/parser/src/parser/common/mod.rs b/parser/src/parser/common/mod.rs
index 59c47cb..4fd4853 100644
--- a/parser/src/parser/common/mod.rs
+++ b/parser/src/parser/common/mod.rs
@@ -6,6 +6,7 @@ pub mod types;
 use crate::{
     Rule,
     ast::*,
+    ast_ensure,
     error::{AstError, IntoErr, collect_recovered},
     parser::consume_rule,
 };
@@ -14,7 +15,9 @@ impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for Identifier {
     type Error = AstError<'a, Self>;
 
     fn try_from(pair: pest::iterators::Pair<'a, Rule>) -> Result<Self, Self::Error> {
-        Ok(Identifier(pair.as_str().to_string()))
+        ast_ensure!(pair, Rule::identifier => {
+            Ok(Identifier(pair.as_str().to_string()))
+        })
     }
 }
 
diff --git a/parser/src/parser/common/statement.rs b/parser/src/parser/common/statement.rs
index 32882a5..3bd8e9c 100644
--- a/parser/src/parser/common/statement.rs
+++ b/parser/src/parser/common/statement.rs
@@ -1,7 +1,7 @@
 use crate::{
     Rule,
     ast::*,
-    ast_expr,
+    ast_ensure, ast_expr,
     error::{AstError, AstResult, ErrorCode, IntoErr, collect_recovered},
     parser::listen_rule,
 };
@@ -10,19 +10,25 @@ impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for Block {
     type Error = AstError<'a, Self>;
 
     fn try_from(pair: pest::iterators::Pair<'a, Rule>) -> Result<Self, Self::Error> {
-        if pair.as_rule() == Rule::block {
+        ast_ensure!(pair, Rule::block => {
             ast_expr!(Block(collect_recovered(pair.into_inner())))
-        } else {
-            Err(AstError {
-                span: pair.as_span(),
-                error_code: ErrorCode::InvalidStatement,
-                error_message: format!(
-                    "BUG: AST requires a block, this isn't a block, it's a {:?}",
-                    pair.as_rule()
-                ),
-                recovered: None,
-            })
-        }
+        })
+    }
+}
+
+impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for StatementBranch {
+    type Error = AstError<'a, Self>;
+
+    fn try_from(pair: pest::iterators::Pair<'a, Rule>) -> Result<Self, Self::Error> {
+        let mut inner = pair.into_inner();
+
+        let condition = inner.next().unwrap().try_into().get()?;
+        let body = inner.next().unwrap().try_into().get()?;
+
+        Ok(StatementBranch {
+            condition,
+            body: Box::new(body),
+        })
     }
 }
 
@@ -114,19 +120,3 @@ impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for Statement {
         })
     }
 }
-
-impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for StatementBranch {
-    type Error = AstError<'a, Self>;
-
-    fn try_from(pair: pest::iterators::Pair<'a, Rule>) -> Result<Self, Self::Error> {
-        let mut inner = pair.into_inner();
-
-        let condition = inner.next().unwrap().try_into().get()?;
-        let body = inner.next().unwrap().try_into().get()?;
-
-        Ok(StatementBranch {
-            condition,
-            body: Box::new(body),
-        })
-    }
-}