Refactor BeginEndStatements into a reusable struct, then use for functions

- this lets us disacard the former (unfortunate, bespoke) multi statement parsing because we can just use `parse_statement_list`
- however, `parse_statement_list` also needed a small change to allow subsequent statements to come after the final `END`

Author: aharpervc

src/ast/ddl.rs

@@ -2277,11 +2277,9 @@ impl fmt::Display for CreateFunction {
         if let Some(CreateFunctionBody::AsAfterOptions(function_body)) = &self.function_body {
             write!(f, " AS {function_body}")?;
         }
-        if let Some(CreateFunctionBody::MultiStatement(statements)) = &self.function_body {
+        if let Some(CreateFunctionBody::AsBeginEnd(bes)) = &self.function_body {
             write!(f, " AS")?;
-            write!(f, " BEGIN")?;
-            write!(f, " {}", display_separated(statements, "; "))?;
-            write!(f, " END")?;
+            write!(f, " {}", bes)?;
         }
         Ok(())
     }

src/ast/mod.rs

@@ -2292,18 +2292,14 @@ pub enum ConditionalStatements {
     /// SELECT 1; SELECT 2; SELECT 3; ...
     Sequence { statements: Vec<Statement> },
     /// BEGIN SELECT 1; SELECT 2; SELECT 3; ... END
-    BeginEnd {
-        begin_token: AttachedToken,
-        statements: Vec<Statement>,
-        end_token: AttachedToken,
-    },
+    BeginEnd(BeginEndStatements),
 }
 
 impl ConditionalStatements {
     pub fn statements(&self) -> &Vec<Statement> {
         match self {
             ConditionalStatements::Sequence { statements } => statements,
-            ConditionalStatements::BeginEnd { statements, .. } => statements,
+            ConditionalStatements::BeginEnd(bes) => &bes.statements,
         }
     }
 }
@@ -2317,12 +2313,34 @@ impl fmt::Display for ConditionalStatements {
                 }
                 Ok(())
             }
-            ConditionalStatements::BeginEnd { statements, .. } => {
-                write!(f, "BEGIN ")?;
-                format_statement_list(f, statements)?;
-                write!(f, " END")
-            }
+            ConditionalStatements::BeginEnd(bes) => write!(f, "{}", bes),
+        }
+    }
+}
+
+/// A shared representation of `BEGIN`, multiple statements, and `END` tokens.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct BeginEndStatements {
+    pub begin_token: AttachedToken,
+    pub statements: Vec<Statement>,
+    pub end_token: AttachedToken,
+}
+
+impl fmt::Display for BeginEndStatements {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let BeginEndStatements {
+            begin_token: AttachedToken(begin_token),
+            statements,
+            end_token: AttachedToken(end_token),
+        } = self;
+
+        write!(f, "{begin_token} ")?;
+        if !statements.is_empty() {
+            format_statement_list(f, statements)?;
         }
+        write!(f, " {end_token}")
     }
 }
 
@@ -8399,7 +8417,7 @@ pub enum CreateFunctionBody {
     /// ```
     ///
     /// [MsSql]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql
-    MultiStatement(Vec<Statement>),
+    AsBeginEnd(BeginEndStatements),
     /// Function body expression using the 'RETURN' keyword.
     ///
     /// Example:

src/ast/spans.rs

@@ -778,11 +778,9 @@ impl Spanned for ConditionalStatements {
             ConditionalStatements::Sequence { statements } => {
                 union_spans(statements.iter().map(|s| s.span()))
             }
-            ConditionalStatements::BeginEnd {
-                begin_token: AttachedToken(start),
-                statements: _,
-                end_token: AttachedToken(end),
-            } => union_spans([start.span, end.span].into_iter()),
+            ConditionalStatements::BeginEnd(bes) => {
+                union_spans([bes.begin_token.0.span, bes.end_token.0.span].into_iter())
+            }
         }
     }
 }

src/dialect/mssql.rs

@@ -16,7 +16,9 @@
 // under the License.
 
 use crate::ast::helpers::attached_token::AttachedToken;
-use crate::ast::{ConditionalStatementBlock, ConditionalStatements, IfStatement, Statement};
+use crate::ast::{
+    BeginEndStatements, ConditionalStatementBlock, ConditionalStatements, IfStatement, Statement,
+};
 use crate::dialect::Dialect;
 use crate::keywords::{self, Keyword};
 use crate::parser::{Parser, ParserError};
@@ -149,11 +151,11 @@ impl MsSqlDialect {
                 start_token: AttachedToken(if_token),
                 condition: Some(condition),
                 then_token: None,
-                conditional_statements: ConditionalStatements::BeginEnd {
+                conditional_statements: ConditionalStatements::BeginEnd(BeginEndStatements {
                     begin_token: AttachedToken(begin_token),
                     statements,
                     end_token: AttachedToken(end_token),
-                },
+                }),
             }
         } else {
             let stmt = parser.parse_statement()?;
@@ -182,11 +184,11 @@ impl MsSqlDialect {
                     start_token: AttachedToken(else_token),
                     condition: None,
                     then_token: None,
-                    conditional_statements: ConditionalStatements::BeginEnd {
+                    conditional_statements: ConditionalStatements::BeginEnd(BeginEndStatements {
                         begin_token: AttachedToken(begin_token),
                         statements,
                         end_token: AttachedToken(end_token),
-                    },
+                    }),
                 });
             } else {
                 let stmt = parser.parse_statement()?;

src/parser/mod.rs

@@ -485,11 +485,11 @@ impl<'a> Parser<'a> {
 
             let statement = self.parse_statement()?;
             expecting_statement_delimiter = match &statement {
-                Statement::If(s) => match s.if_block.conditional_statements {
-                    ConditionalStatements::BeginEnd { .. } => false,
-                   _ => true,
-                },
-                _ => true
+                Statement::If(s) => !matches!(
+                    s.if_block.conditional_statements,
+                    ConditionalStatements::BeginEnd { .. }
+                ),
+                _ => true,
             };
             stmts.push(statement);
         }
@@ -4460,9 +4460,17 @@ impl<'a> Parser<'a> {
                     break;
                 }
             }
-
             values.push(self.parse_statement()?);
-            self.expect_token(&Token::SemiColon)?;
+
+            let semi_colon_expected = match values.last() {
+                Some(Statement::If(if_statement)) => if_statement.end_token.is_some(),
+                Some(_) => true,
+                None => false,
+            };
+
+            if semi_colon_expected {
+                self.expect_token(&Token::SemiColon)?;
+            }
         }
         Ok(values)
     }
@@ -5175,20 +5183,16 @@ impl<'a> Parser<'a> {
         };
 
         self.expect_keyword_is(Keyword::AS)?;
-        self.expect_keyword_is(Keyword::BEGIN)?;
-        let mut result = self.parse_statements()?;
-        // note: `parse_statements` will consume the `END` token & produce a Commit statement...
-        if let Some(Statement::Commit {
-            chain,
-            end,
-            modifier,
-        }) = result.last()
-        {
-            if *chain == false && *end == true && *modifier == None {
-                result = result[..result.len() - 1].to_vec();
-            }
-        }
-        let function_body = Some(CreateFunctionBody::MultiStatement(result));
+
+        let begin_token = self.expect_keyword(Keyword::BEGIN)?;
+        let statements = self.parse_statement_list(&[Keyword::END])?;
+        let end_token = self.expect_keyword(Keyword::END)?;
+
+        let function_body = Some(CreateFunctionBody::AsBeginEnd(BeginEndStatements {
+            begin_token: AttachedToken(begin_token),
+            statements,
+            end_token: AttachedToken(end_token),
+        }));
 
         Ok(Statement::CreateFunction(CreateFunction {
             or_alter,