Wrap tokens attached to AST nodes in AttachedToken.

Author: Roman Borschel

src/ast/mod.rs

@@ -38,7 +38,7 @@ use serde::{Deserialize, Serialize};
 use sqlparser_derive::{Visit, VisitMut};
 
 use crate::keywords::Keyword;
-use crate::tokenizer::{Span, Token, TokenWithSpan};
+use crate::tokenizer::{Span, Token};
 
 pub use self::data_type::{
     ArrayElemTypeDef, BinaryLength, CharLengthUnits, CharacterLength, DataType, EnumMember,
@@ -2120,12 +2120,12 @@ pub enum Password {
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct CaseStatement {
     /// The `CASE` token that starts the statement.
-    pub case_token: TokenWithSpan,
+    pub case_token: AttachedToken,
     pub match_expr: Option<Expr>,
     pub when_blocks: Vec<ConditionalStatements>,
     pub else_block: Option<ConditionalStatements>,
     /// The last token of the statement (`END` or `CASE`).
-    pub end_case_token: TokenWithSpan,
+    pub end_case_token: AttachedToken,
 }
 
 impl fmt::Display for CaseStatement {
@@ -2135,7 +2135,7 @@ impl fmt::Display for CaseStatement {
             match_expr,
             when_blocks,
             else_block,
-            end_case_token,
+            end_case_token: AttachedToken(end),
         } = self;
 
         write!(f, "CASE")?;
@@ -2154,7 +2154,7 @@ impl fmt::Display for CaseStatement {
 
         write!(f, " END")?;
 
-        if let Token::Word(w) = &end_case_token.token {
+        if let Token::Word(w) = &end.token {
             if w.keyword == Keyword::CASE {
                 write!(f, " CASE")?;
             }
@@ -2187,12 +2187,12 @@ pub enum IfStatement {
     /// [Snowflake](https://docs.snowflake.com/en/sql-reference/snowflake-scripting/if)
     IfThenElseEnd {
         /// The `IF` token that starts the statement.
-        if_token: TokenWithSpan,
+        if_token: AttachedToken,
         if_block: ConditionalStatements,
         elseif_blocks: Vec<ConditionalStatements>,
         else_block: Option<ConditionalStatements>,
         /// The `IF` token that ends the statement.
-        end_if_token: TokenWithSpan,
+        end_if_token: AttachedToken,
     },
     /// An MSSQL `IF ... ELSE ...` statement.
     ///
@@ -2203,7 +2203,7 @@ pub enum IfStatement {
     ///
     /// [MSSQL](https://learn.microsoft.com/en-us/sql/t-sql/language-elements/if-else-transact-sql?view=sql-server-ver16)
     MsSqlIfElse {
-        if_token: TokenWithSpan,
+        if_token: AttachedToken,
         condition: Expr,
         if_statements: MsSqlIfStatements,
         else_statements: Option<MsSqlIfStatements>,
@@ -2270,9 +2270,9 @@ pub enum MsSqlIfStatements {
     /// END
     /// ```
     Block {
-        begin_token: TokenWithSpan,
+        begin_token: AttachedToken,
         statements: Vec<Statement>,
-        end_token: TokenWithSpan,
+        end_token: AttachedToken,
     },
 }
 
@@ -2310,7 +2310,7 @@ impl fmt::Display for MsSqlIfStatements {
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct ConditionalStatements {
     /// The start token of the conditional (`WHEN`, `IF`, `ELSEIF` or `ELSE`).
-    pub start_token: TokenWithSpan,
+    pub start_token: AttachedToken,
     /// The condition expression. `None` for `ELSE` statements.
     pub condition: Option<Expr>,
     /// Statement list of the `THEN` clause.
@@ -2320,12 +2320,12 @@ pub struct ConditionalStatements {
 impl fmt::Display for ConditionalStatements {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let ConditionalStatements {
-            start_token,
+            start_token: AttachedToken(start),
             condition,
             statements,
         } = self;
 
-        let keyword = &start_token.token;
+        let keyword = &start.token;
 
         if let Some(expr) = condition {
             write!(f, "{keyword} {expr} THEN")?;

src/ast/spans.rs

@@ -21,7 +21,7 @@ use core::iter;
 use crate::tokenizer::Span;
 
 use super::{
-    dcl::SecondaryRoles, value::ValueWithSpan, AccessExpr, AlterColumnOperation,
+    dcl::SecondaryRoles, value::ValueWithSpan, AccessExpr, AlterColumnOperation, AttachedToken,
     AlterIndexOperation, AlterTableOperation, Array, Assignment, AssignmentTarget, CaseStatement,
     CloseCursor, ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef, ConditionalStatements,
     ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable,
@@ -740,30 +740,30 @@ impl Spanned for CreateIndex {
 impl Spanned for CaseStatement {
     fn span(&self) -> Span {
         let CaseStatement {
-            case_token,
-            end_case_token,
+            case_token: AttachedToken(start),
+            end_case_token: AttachedToken(end),
             ..
         } = self;
 
-        union_spans([case_token.span, end_case_token.span].into_iter())
+        union_spans([start.span, end.span].into_iter())
     }
 }
 
 impl Spanned for IfStatement {
     fn span(&self) -> Span {
         match self {
             IfStatement::IfThenElseEnd {
-                if_token,
-                end_if_token,
+                if_token: AttachedToken(start),
+                end_if_token: AttachedToken(end),
                 ..
-            } => union_spans([if_token.span, end_if_token.span].into_iter()),
+            } => union_spans([start.span, end.span].into_iter()),
             IfStatement::MsSqlIfElse {
-                if_token,
+                if_token: AttachedToken(start),
                 if_statements,
                 else_statements,
                 ..
             } => union_spans(
-                [if_token.span, if_statements.span()]
+                [start.span, if_statements.span()]
                     .into_iter()
                     .chain(else_statements.as_ref().into_iter().map(|s| s.span())),
             ),
@@ -776,24 +776,24 @@ impl Spanned for MsSqlIfStatements {
         match self {
             MsSqlIfStatements::Single(s) => s.span(),
             MsSqlIfStatements::Block {
-                begin_token,
-                end_token,
+                begin_token: AttachedToken(start),
+                end_token: AttachedToken(end),
                 ..
-            } => union_spans([begin_token.span, end_token.span].into_iter()),
+            } => union_spans([start.span, end.span].into_iter()),
         }
     }
 }
 
 impl Spanned for ConditionalStatements {
     fn span(&self) -> Span {
         let ConditionalStatements {
-            start_token,
+            start_token: AttachedToken(start),
             condition,
             statements,
         } = self;
 
         union_spans(
-            iter::once(start_token.span)
+            iter::once(start.span)
                 .chain(condition.as_ref().map(|c| c.span()).into_iter())
                 .chain(statements.iter().map(|s| s.span())),
         )

src/dialect/mssql.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::ast::{IfStatement, MsSqlIfStatements, Statement};
+use crate::ast::helpers::attached_token::AttachedToken;
 use crate::dialect::Dialect;
 use crate::keywords::{self, Keyword};
 use crate::parser::{Parser, ParserError};
@@ -144,9 +145,9 @@ impl MsSqlDialect {
             let statements = self.parse_statement_list(parser, Some(Keyword::END))?;
             let end_token = parser.expect_keyword(Keyword::END)?;
             if_statements = MsSqlIfStatements::Block {
-                begin_token,
+                begin_token: AttachedToken(begin_token),
                 statements,
-                end_token,
+                end_token: AttachedToken(end_token),
             };
         } else {
             let stmt = parser.parse_statement()?;
@@ -160,9 +161,9 @@ impl MsSqlDialect {
                 let statements = self.parse_statement_list(parser, Some(Keyword::END))?;
                 let end_token = parser.expect_keyword(Keyword::END)?;
                 else_statements = Some(MsSqlIfStatements::Block {
-                    begin_token,
+                    begin_token: AttachedToken(begin_token),
                     statements,
-                    end_token,
+                    end_token: AttachedToken(end_token),
                 });
             } else {
                 let stmt = parser.parse_statement()?;
@@ -171,7 +172,7 @@ impl MsSqlDialect {
         }
 
         Ok(Statement::If(IfStatement::MsSqlIfElse {
-            if_token,
+            if_token: AttachedToken(if_token),
             condition,
             if_statements,
             else_statements,

src/parser/mod.rs

@@ -656,11 +656,11 @@ impl<'a> Parser<'a> {
         }
 
         Ok(Statement::Case(CaseStatement {
-            case_token,
+            case_token: AttachedToken(case_token),
             match_expr,
             when_blocks,
             else_block,
-            end_case_token,
+            end_case_token: AttachedToken(end_case_token),
         }))
     }
 
@@ -690,11 +690,11 @@ impl<'a> Parser<'a> {
         let end_if_token = self.expect_keyword(Keyword::IF)?;
 
         Ok(Statement::If(IfStatement::IfThenElseEnd {
-            if_token,
+            if_token: AttachedToken(if_token),
             if_block,
             elseif_blocks,
             else_block,
-            end_if_token,
+            end_if_token: AttachedToken(end_if_token),
         }))
     }
 
@@ -723,7 +723,7 @@ impl<'a> Parser<'a> {
         let statements = self.parse_statement_list(terminal_keywords)?;
 
         Ok(ConditionalStatements {
-            start_token,
+            start_token: AttachedToken(start_token),
             condition,
             statements,
         })