statement body parsing

crates/parser/src/ast/statement.rs

@@ -8,27 +8,33 @@ pub struct Block(
     pub Option<Spanned<Expression>>,
 );
 
+#[derive(Debug, Clone, Serialize)]
+pub enum StatementBody {
+    Statement(Box<Statement>),
+    Expression(Expression),
+}
+
 #[derive(Debug, Clone, Serialize)]
 pub enum Statement {
     Block(Block),
     If {
         initial: StatementBranch,
         else_if: Vec<StatementBranch>,
-        else_branch: Option<Expression>,
+        else_branch: Option<StatementBody>,
     },
-    Loop(Expression),
+    Loop(StatementBody),
     While(StatementBranch),
     CStyleFor {
         init: Expression,
         condition: Expression,
         update: Expression,
-        body: Expression,
+        body: StatementBody,
     },
     For {
         mutable: bool,
         pattern: Pattern,
         iterator: Expression,
-        body: Expression,
+        body: StatementBody,
     },
     Match(Expression, Vec<(Vec<Pattern>, Expression)>),
 
@@ -54,7 +60,7 @@ pub struct VarDeclStmt {
 #[derive(Debug, Clone, Serialize)]
 pub struct StatementBranch {
     pub condition: Expression,
-    pub body: Box<Expression>,
+    pub body: Box<StatementBody>,
 }
 
 impl Statement {
@@ -84,3 +90,12 @@ impl Statement {
         }
     }
 }
+
+impl StatementBody {
+    pub fn is_block(&self) -> bool {
+        match self {
+            Self::Expression(_) => false,
+            _ => true,
+        }
+    }
+}

crates/parser/src/grammar.pest

@@ -343,6 +343,10 @@ statement = _{
   | (expr ~ semicolon)
 }
 
+statement_wrapper = { statement }
+
+statement_body = { statement_wrapper | expr }
+
 // ------------------------------------------------------
 // BASIC STATEMENTS
 // ------------------------------------------------------
@@ -384,7 +388,7 @@ control_flow = {
   | block
 }
 
-statement_branch = { "(" ~ expr ~ ")" ~ statement }
+statement_branch = { "(" ~ expr ~ ")" ~ statement_body }
 
 else_if = {
     "else" ~ "if" ~ statement_branch
@@ -395,7 +399,7 @@ else_if_list = {
 }
 
 if_stmt = {
-    "if" ~ statement_branch ~ else_if_list ~ ("else" ~ statement)?
+    "if" ~ statement_branch ~ else_if_list ~ ("else" ~ statement_body)?
 }
 
 // ------------------------------------------------------
@@ -407,15 +411,15 @@ while_stmt = {
 }
 
 c_for_stmt = {
-    "for" ~ "(" ~ statement ~ statement ~ expr ~ ")" ~ statement
+    "for" ~ "(" ~ statement ~ statement ~ expr ~ ")" ~ statement_body
 }
 
 for_stmt = {
-    "for" ~ "(" ~ mutable? ~ pattern ~ ":" ~ expr ~ ")" ~ statement
+    "for" ~ "(" ~ mutable? ~ pattern ~ ":" ~ expr ~ ")" ~ statement_body
 }
 
 loop_stmt = {
-    "loop" ~ statement
+    "loop" ~ statement_body
 }
 
 // ------------------------------------------------------

crates/parser/src/parser/common/statement.rs

@@ -18,6 +18,24 @@ impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for Block {
     }
 }
 
+impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for StatementBody {
+    type Error = AstError<'a, Self>;
+
+    fn try_from(pair: pest::iterators::Pair<'a, Rule>) -> Result<Self, Self::Error> {
+        let mut inner = pair.clone().into_inner();
+
+        ast_ensure!(pair, Rule::statement_body => {
+            let i = inner.next().unwrap();
+
+            match i.as_rule() {
+                Rule::expr => ast_expr!(StatementBody::Expression(i.try_into())),
+                Rule::statement_wrapper => ast_expr!(StatementBody::Statement(i.try_into().map(Box::new).get_map(Box::new))),
+                _ => AstError::bug_unimplemented(i),
+            }
+        })
+    }
+}
+
 impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for StatementBranch {
     type Error = AstError<'a, Self>;
 
@@ -63,7 +81,7 @@ impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for Statement {
                 ast_expr!(Statement::If {
                     initial: inner.next().unwrap().try_into(),
                     else_if: collect_recovered(inner.next().unwrap().into_inner()),
-                    else_branch: inner.next().map(Expression::try_from).transpose(),
+                    else_branch: inner.next().map(StatementBody::try_from).transpose(),
                 })
             }