QASM3 - Parse switch statements

compiler/qsc_qasm3/src/ast.rs

@@ -21,14 +21,19 @@ fn set_indentation<'a, 'b>(
         0 => indent.with_str(""),
         1 => indent.with_str("    "),
         2 => indent.with_str("        "),
+        3 => indent.with_str("            "),
         _ => unimplemented!("indentation level not supported"),
     }
 }
 
 // TODO: profile this with iai-callgrind in a large OpenQASM3
 // sample to verify that is actually faster than using Vec<T>.
 /// An alternative to `Vec<T>` that uses less stack space.
-type List<T> = Box<[Box<T>]>;
+pub(crate) type List<T> = Box<[Box<T>]>;
+
+pub(crate) fn list_from_iter<T>(vals: impl IntoIterator<Item = T>) -> List<T> {
+    vals.into_iter().map(Box::new).collect()
+}
 
 #[derive(Clone, Debug)]
 pub struct Program {
@@ -1551,20 +1556,59 @@ impl Display for EnumerableSet {
 #[derive(Clone, Debug)]
 pub struct SwitchStmt {
     pub span: Span,
-    pub target: ExprStmt,
-    pub cases: List<(List<ExprStmt>, CompoundStmt)>,
+    pub target: Expr,
+    pub cases: List<(List<Expr>, Block)>,
     /// Note that `None` is quite different to `[]` in this case; the latter is
     /// an explicitly empty body, whereas the absence of a default might mean
     /// that the switch is inexhaustive, and a linter might want to complain.
-    pub default: Option<CompoundStmt>,
+    pub default: Option<Block>,
 }
 
 impl Display for SwitchStmt {
-    fn fmt(&self, _f: &mut Formatter<'_>) -> fmt::Result {
-        todo!("SwitchStmt display");
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "SwitchStmt {}:", self.span)?;
+        let mut indent = set_indentation(indented(f), 1);
+        write!(indent, "\nTarget: {}", self.target)?;
+        if self.cases.is_empty() {
+            write!(indent, "\n<no cases>")?;
+        } else {
+            write!(indent, "\nCases:")?;
+            for elt in &self.cases {
+                let (labels, block) = &**elt;
+                indent = display_switch_case(indent, labels, block)?;
+            }
+        }
+        if let Some(default) = &self.default {
+            write!(indent, "\nDefault Case:")?;
+            indent = set_indentation(indent, 2);
+            write!(indent, "\n{default}")?;
+        } else {
+            write!(indent, "\n<no default>")?;
+        }
+        Ok(())
     }
 }
 
+fn display_switch_case<'a, 'b>(
+    mut indent: Indented<'a, Formatter<'b>>,
+    labels: &List<Expr>,
+    block: &Block,
+) -> Result<Indented<'a, Formatter<'b>>, core::fmt::Error> {
+    indent = set_indentation(indent, 2);
+    if labels.is_empty() {
+        write!(indent, "\n<no labels>")?;
+    } else {
+        write!(indent, "\nLabels:")?;
+        indent = set_indentation(indent, 3);
+        for label in labels {
+            write!(indent, "\n{label}")?;
+        }
+    }
+    indent = set_indentation(indent, 2);
+    write!(indent, "\n{block}")?;
+    Ok(indent)
+}
+
 #[derive(Clone, Debug)]
 pub struct ClassicalAssignment {
     pub span: Span,

compiler/qsc_qasm3/src/lex/cooked.rs

@@ -509,11 +509,15 @@ impl<'a> Lexer<'a> {
             raw::TokenKind::Ident => {
                 let ident = &self.input[(token.offset as usize)..(self.offset() as usize)];
                 let cooked_ident = Self::ident(ident);
-                if matches!(cooked_ident, TokenKind::Keyword(Keyword::Pragma)) {
-                    self.eat_to_end_of_line();
-                    Ok(Some(TokenKind::Pragma))
-                } else {
-                    Ok(Some(cooked_ident))
+                match cooked_ident {
+                    // A `dim` token without a `#` in front should be
+                    // treated as an identifier and not as a keyword.
+                    TokenKind::Dim => Ok(Some(TokenKind::Identifier)),
+                    TokenKind::Keyword(Keyword::Pragma) => {
+                        self.eat_to_end_of_line();
+                        Ok(Some(TokenKind::Pragma))
+                    }
+                    _ => Ok(Some(cooked_ident)),
                 }
             }
             raw::TokenKind::HardwareQubit => Ok(Some(TokenKind::HardwareQubit)),
@@ -553,23 +557,26 @@ impl<'a> Lexer<'a> {
                 }
             }
             raw::TokenKind::Single(Single::Sharp) => {
-                let complete = TokenKind::Pragma;
-                self.expect(raw::TokenKind::Ident, complete)?;
+                self.expect(raw::TokenKind::Ident, TokenKind::Identifier)?;
                 let ident = &self.input[(token.offset as usize + 1)..(self.offset() as usize)];
-                if matches!(Self::ident(ident), TokenKind::Keyword(Keyword::Pragma)) {
-                    self.eat_to_end_of_line();
-                    Ok(Some(complete))
-                } else {
-                    let span = Span {
-                        lo: token.offset,
-                        hi: self.offset(),
-                    };
-                    Err(Error::Incomplete(
-                        raw::TokenKind::Ident,
-                        complete,
-                        raw::TokenKind::Ident,
-                        span,
-                    ))
+                match Self::ident(ident) {
+                    TokenKind::Dim => Ok(Some(TokenKind::Dim)),
+                    TokenKind::Keyword(Keyword::Pragma) => {
+                        self.eat_to_end_of_line();
+                        Ok(Some(TokenKind::Pragma))
+                    }
+                    _ => {
+                        let span = Span {
+                            lo: token.offset,
+                            hi: self.offset(),
+                        };
+                        Err(Error::Incomplete(
+                            raw::TokenKind::Ident,
+                            TokenKind::Pragma,
+                            raw::TokenKind::Ident,
+                            span,
+                        ))
+                    }
                 }
             }
             raw::TokenKind::Single(single) => self.single(single).map(Some),

compiler/qsc_qasm3/src/lex/cooked/tests.rs

@@ -1003,7 +1003,7 @@ fn unknown() {
                 Err(
                     Incomplete(
                         Ident,
-                        Pragma,
+                        Identifier,
                         Single(
                             Sharp,
                         ),
@@ -1016,7 +1016,7 @@ fn unknown() {
                 Err(
                     IncompleteEof(
                         Ident,
-                        Pragma,
+                        Identifier,
                         Span {
                             lo: 2,
                             hi: 2,
@@ -1187,3 +1187,43 @@ fn sharp_pragma_ident() {
         "#]],
     );
 }
+
+#[test]
+fn dim() {
+    check(
+        "dim",
+        &expect![[r#"
+        [
+            Ok(
+                Token {
+                    kind: Identifier,
+                    span: Span {
+                        lo: 0,
+                        hi: 3,
+                    },
+                },
+            ),
+        ]
+    "#]],
+    );
+}
+
+#[test]
+fn sharp_dim() {
+    check(
+        "#dim",
+        &expect![[r#"
+        [
+            Ok(
+                Token {
+                    kind: Dim,
+                    span: Span {
+                        lo: 0,
+                        hi: 4,
+                    },
+                },
+            ),
+        ]
+    "#]],
+    );
+}

compiler/qsc_qasm3/src/parser/completion/tests.rs

@@ -36,7 +36,7 @@ fn begin_document() {
         "|OPENQASM 3;",
         &expect![[r#"
             WordKinds(
-                Annotation | Include | Input | OpenQASM | Output | Pragma | Qubit,
+                Annotation | Include | Input | OpenQASM | Output | Pragma | Qubit | Switch,
             )
         "#]],
     );
@@ -48,7 +48,7 @@ fn end_of_version() {
         "OPENQASM 3;|",
         &expect![[r#"
             WordKinds(
-                Annotation | Include | Input | Output | Pragma | Qubit,
+                Annotation | Include | Input | Output | Pragma | Qubit | Switch,
             )
         "#]],
     );

compiler/qsc_qasm3/src/parser/error.rs

@@ -113,6 +113,12 @@ pub enum ErrorKind {
     #[error("missing entry in sequence")]
     #[diagnostic(code("Qasm3.Parse.MissingSeqEntry"))]
     MissingSeqEntry(#[label] Span),
+    #[error("missing switch statement cases")]
+    #[diagnostic(code("Qasm3.Parse.MissingSwitchCases"))]
+    MissingSwitchCases(#[label] Span),
+    #[error("missing switch statement case labels")]
+    #[diagnostic(code("Qasm3.Parse.MissingSwitchCaseLabels"))]
+    MissingSwitchCaseLabels(#[label] Span),
     #[error("expected an item or closing brace, found {0}")]
     #[diagnostic(code("Qasm3.Parse.ExpectedItem"))]
     ExpectedItem(TokenKind, #[label] Span),
@@ -131,6 +137,8 @@ impl ErrorKind {
             Self::MissingParens(span) => Self::MissingParens(span + offset),
             Self::FloatingAnnotation(span) => Self::FloatingAnnotation(span + offset),
             Self::MissingSeqEntry(span) => Self::MissingSeqEntry(span + offset),
+            Self::MissingSwitchCases(span) => Self::MissingSwitchCases(span + offset),
+            Self::MissingSwitchCaseLabels(span) => Self::MissingSwitchCaseLabels(span + offset),
             Self::ExpectedItem(token, span) => Self::ExpectedItem(token, span + offset),
         }
     }

compiler/qsc_qasm3/src/parser/expr.rs

@@ -409,7 +409,7 @@ fn lit_bigint(lexeme: &str, radix: u32) -> Option<BigInt> {
     }
 }
 
-fn paren_expr(s: &mut ParserContext, lo: u32) -> Result<Expr> {
+pub(crate) fn paren_expr(s: &mut ParserContext, lo: u32) -> Result<Expr> {
     let (mut exprs, final_sep) = seq(s, expr)?;
     token(s, TokenKind::Close(Delim::Paren))?;
 

compiler/qsc_qasm3/src/parser/stmt.rs

@@ -17,11 +17,13 @@ use super::{
 };
 use crate::{
     ast::{
-        AngleType, Annotation, ArrayBaseTypeKind, ArrayType, BitType, Block,
-        ClassicalDeclarationStmt, ComplexType, ConstantDeclaration, ExprStmt, FloatType,
-        IODeclaration, IOKeyword, IncludeStmt, IntType, LiteralKind, Pragma, QubitDeclaration,
-        ScalarType, ScalarTypeKind, Stmt, StmtKind, TypeDef, UIntType,
+        list_from_iter, AngleType, Annotation, ArrayBaseTypeKind, ArrayType, BitType, Block,
+        ClassicalDeclarationStmt, ComplexType, ConstantDeclaration, Expr, ExprStmt, FloatType,
+        IODeclaration, IOKeyword, IncludeStmt, IntType, List, LiteralKind, Pragma,
+        QubitDeclaration, ScalarType, ScalarTypeKind, Stmt, StmtKind, SwitchStmt, TypeDef,
+        UIntType,
     },
+    keyword::Keyword,
     lex::{
         cooked::{Literal, Type},
         Delim, TokenKind,
@@ -52,6 +54,8 @@ pub(super) fn parse(s: &mut ParserContext) -> Result<Box<Stmt>> {
         Box::new(v)
     } else if let Some(decl) = opt(s, parse_local)? {
         Box::new(decl)
+    } else if let Some(switch) = opt(s, parse_switch_stmt)? {
+        Box::new(StmtKind::Switch(switch))
     } else {
         return Err(Error::new(ErrorKind::Rule(
             "statement",
@@ -598,3 +602,56 @@ pub(super) fn complex_subtype(s: &mut ParserContext) -> Result<FloatType> {
     token(s, TokenKind::Close(Delim::Bracket))?;
     Ok(ty)
 }
+
+/// The Language Spec and the grammar for switch statements disagree.
+/// We followed the Spec when writing the parser
+/// <https://openqasm.com/language/classical.html#the-switch-statement>.
+pub fn parse_switch_stmt(s: &mut ParserContext) -> Result<SwitchStmt> {
+    let lo = s.peek().span.lo;
+    token(s, TokenKind::Keyword(Keyword::Switch))?;
+
+    // Controlling expression.
+    token(s, TokenKind::Open(Delim::Paren))?;
+    let controlling_expr = expr::paren_expr(s, lo)?;
+
+    // Open cases bracket.
+    token(s, TokenKind::Open(Delim::Brace))?;
+
+    // Cases.
+    let lo_cases = s.peek().span.lo;
+    let cases = list_from_iter(many(s, case_stmt)?);
+    if cases.is_empty() {
+        s.push_error(Error::new(ErrorKind::MissingSwitchCases(s.span(lo_cases))));
+    }
+
+    // Default case.
+    let default = opt(s, default_case_stmt)?;
+
+    // Close cases bracket.
+    recovering_token(s, TokenKind::Close(Delim::Brace));
+
+    Ok(SwitchStmt {
+        span: s.span(lo),
+        target: controlling_expr,
+        cases,
+        default,
+    })
+}
+
+fn case_stmt(s: &mut ParserContext) -> Result<(List<Expr>, Block)> {
+    let lo = s.peek().span.lo;
+    token(s, TokenKind::Keyword(Keyword::Case))?;
+
+    let controlling_label = expr::expr_list(s)?;
+    if controlling_label.is_empty() {
+        s.push_error(Error::new(ErrorKind::MissingSwitchCaseLabels(s.span(lo))));
+    }
+
+    let block = parse_block(s).map(|block| *block)?;
+    Ok((list_from_iter(controlling_label), block))
+}
+
+fn default_case_stmt(s: &mut ParserContext) -> Result<Block> {
+    token(s, TokenKind::Keyword(Keyword::Default))?;
+    parse_block(s).map(|block| *block)
+}

compiler/qsc_qasm3/src/parser/stmt/tests.rs

@@ -6,3 +6,4 @@ mod classical_decl;
 mod io_decl;
 mod pragma;
 mod quantum_decl;
+mod switch_stmt;

compiler/qsc_qasm3/src/parser/stmt/tests/pragma.rs

@@ -89,7 +89,7 @@ fn legacy_pragma_ws_after_hash() {
                     Lex(
                         Incomplete(
                             Ident,
-                            Pragma,
+                            Identifier,
                             Whitespace,
                             Span {
                                 lo: 1,