Ensure interpreter supports global qubits

Author: swernli

compiler/qsc/src/interpret/stateful/tests.rs

@@ -222,6 +222,23 @@ mod given_interpreter {
             );
             is_unit_with_output(&result, &output, "before\nafter");
         }
+
+        #[test]
+        fn global_qubits() {
+            let mut interpreter = get_interpreter();
+            let (result, output) = line(&mut interpreter, "open Microsoft.Quantum.Diagnostics;");
+            is_only_value(&result, &output, &Value::unit());
+            let (result, output) = line(&mut interpreter, "DumpMachine()");
+            is_unit_with_output(&result, &output, "STATE:\n|0⟩: 1+0i");
+            let (result, output) = line(&mut interpreter, "use (q0, qs) = (Qubit(), Qubit[3]);");
+            is_only_value(&result, &output, &Value::unit());
+            let (result, output) = line(&mut interpreter, "DumpMachine()");
+            is_unit_with_output(&result, &output, "STATE:\n|0000⟩: 1+0i");
+            let (result, output) = line(&mut interpreter, "X(q0); X(qs[1]);");
+            is_only_value(&result, &output, &Value::unit());
+            let (result, output) = line(&mut interpreter, "DumpMachine()");
+            is_unit_with_output(&result, &output, "STATE:\n|0101⟩: 1+0i");
+        }
     }
 
     #[cfg(test)]

compiler/qsc_eval/src/lib.rs

@@ -663,6 +663,7 @@ impl<'a, G: GlobalLookup<'a>> State<'a, G> {
     fn handle_stmt(&mut self, stmt: &'a Stmt) {
         match &stmt.kind {
             StmtKind::Expr(expr) => self.push_expr(expr),
+            StmtKind::Err => panic!("err stmt should not be present during evaluation"),
             StmtKind::Item(..) => self.push_val(Value::unit()),
             StmtKind::Local(mutability, pat, expr) => {
                 self.push_cont(Cont::Bind(pat, *mutability));

compiler/qsc_hir/src/hir.rs

@@ -541,10 +541,13 @@ impl Display for Stmt {
 }
 
 /// A statement kind.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, Default, PartialEq)]
 pub enum StmtKind {
     /// An expression without a trailing semicolon.
     Expr(Expr),
+    /// Default, empty statement case.
+    #[default]
+    Err,
     /// An item.
     Item(LocalItemId),
     /// A let or mutable binding: `let a = b;` or `mutable x = b;`.
@@ -560,6 +563,7 @@ impl Display for StmtKind {
         let mut indent = set_indentation(indented(f), 0);
         match self {
             StmtKind::Expr(e) => write!(indent, "Expr: {e}")?,
+            StmtKind::Err => panic!("statement kind Err should not be displayed"),
             StmtKind::Item(item) => write!(indent, "Item: {item}")?,
             StmtKind::Local(m, lhs, rhs) => {
                 write!(indent, "Local ({m:?}):")?;

compiler/qsc_hir/src/mut_visit.rs

@@ -116,7 +116,7 @@ pub fn walk_stmt(vis: &mut impl MutVisitor, stmt: &mut Stmt) {
     vis.visit_span(&mut stmt.span);
 
     match &mut stmt.kind {
-        StmtKind::Item(_) => {}
+        StmtKind::Item(_) | StmtKind::Err => {}
         StmtKind::Expr(expr) | StmtKind::Semi(expr) => vis.visit_expr(expr),
         StmtKind::Local(_, pat, value) => {
             vis.visit_pat(pat);

compiler/qsc_hir/src/visit.rs

@@ -100,7 +100,7 @@ pub fn walk_block<'a>(vis: &mut impl Visitor<'a>, block: &'a Block) {
 
 pub fn walk_stmt<'a>(vis: &mut impl Visitor<'a>, stmt: &'a Stmt) {
     match &stmt.kind {
-        StmtKind::Item(_) => {}
+        StmtKind::Item(_) | StmtKind::Err => {}
         StmtKind::Expr(expr) | StmtKind::Semi(expr) => vis.visit_expr(expr),
         StmtKind::Local(_, pat, value) => {
             vis.visit_pat(pat);

compiler/qsc_passes/src/lib.rs

@@ -59,6 +59,7 @@ pub fn run_core_passes(core: &mut CompileUnit) {
         assigner: &mut core.assigner,
     }
     .visit_package(&mut core.package);
+    ReplaceQubitAllocation::new(&table, &mut core.assigner).visit_package(&mut core.package);
 }
 
 pub fn run_default_passes_for_fragment(

compiler/qsc_passes/src/logic_sep.rs

@@ -69,7 +69,7 @@ impl<'a> Visitor<'a> for SepCheck {
         let mut has_inner_op_call = false;
         for stmt in &block.stmts {
             let has_op_call = match &stmt.kind {
-                StmtKind::Item(_) => false,
+                StmtKind::Item(_) | StmtKind::Err => false,
 
                 StmtKind::Local(..) | StmtKind::Qubit(_, _, _, None) => {
                     self.op_call_allowed = false;

compiler/qsc_passes/src/replace_qubit_allocation.rs

@@ -336,6 +336,63 @@ impl MutVisitor for ReplaceQubitAllocation<'_> {
             _ => walk_expr(self, expr),
         }
     }
+
+    fn visit_stmt(&mut self, stmt: &mut Stmt) {
+        // This function is not called by visit_block above, so the only time it will be used is for
+        // top-level statement fragments. Given that, the qubits allocated will always be live for
+        // the entirety of a global scope, so only qubit allocations need to be generated.
+        match take(&mut stmt.kind) {
+            StmtKind::Qubit(_, pat, qubit_init, None) => {
+                stmt.kind = create_qubit_global_alloc(self.core, pat, qubit_init);
+            }
+            StmtKind::Qubit(_, pat, qubit_init, block) => {
+                let stmts = self.process_qubit_stmt(stmt.span, pat, qubit_init, block);
+                *stmt = stmts[0].clone();
+            }
+            kind => {
+                stmt.kind = kind;
+                walk_stmt(self, stmt);
+            }
+        }
+    }
+}
+
+fn create_qubit_global_alloc(core: &Table, pat: Pat, qubit_init: QubitInit) -> StmtKind {
+    fn qubit_alloc_expr(core: &Table, qubit_init: QubitInit) -> Expr {
+        match qubit_init.kind {
+            QubitInitKind::Array(mut expr) => creat_qubit_alloc_call_expr(
+                qubit_init.span,
+                create_gen_core_ref(core, "QIR.Runtime", "AllocateQubitArray", qubit_init.span),
+                Some(take(&mut expr)),
+            ),
+            QubitInitKind::Single => creat_qubit_alloc_call_expr(
+                qubit_init.span,
+                create_gen_core_ref(
+                    core,
+                    "QIR.Runtime",
+                    "__quantum__rt__qubit_allocate",
+                    qubit_init.span,
+                ),
+                None,
+            ),
+            QubitInitKind::Tuple(tup) => Expr {
+                id: NodeId::default(),
+                span: qubit_init.span,
+                ty: qubit_init.ty,
+                kind: ExprKind::Tuple(
+                    tup.into_iter()
+                        .map(|init| qubit_alloc_expr(core, init))
+                        .collect(),
+                ),
+            },
+        }
+    }
+
+    StmtKind::Local(
+        Mutability::Immutable,
+        pat,
+        qubit_alloc_expr(core, qubit_init),
+    )
 }
 
 fn create_general_alloc_stmt(
@@ -345,23 +402,27 @@ fn create_general_alloc_stmt(
 ) -> Stmt {
     ident.gen_id_init(
         Mutability::Immutable,
-        Expr {
-            id: NodeId::default(),
-            span: ident.span,
-            ty: Ty::Prim(PrimTy::Qubit),
-            kind: ExprKind::Call(
-                Box::new(call_expr),
-                Box::new(array_size.unwrap_or(Expr {
-                    id: NodeId::default(),
-                    span: ident.span,
-                    ty: Ty::UNIT,
-                    kind: ExprKind::Tuple(vec![]),
-                })),
-            ),
-        },
+        creat_qubit_alloc_call_expr(ident.span, call_expr, array_size),
     )
 }
 
+fn creat_qubit_alloc_call_expr(span: Span, call_expr: Expr, array_size: Option<Expr>) -> Expr {
+    Expr {
+        id: NodeId::default(),
+        span,
+        ty: Ty::Prim(PrimTy::Qubit),
+        kind: ExprKind::Call(
+            Box::new(call_expr),
+            Box::new(array_size.unwrap_or(Expr {
+                id: NodeId::default(),
+                span,
+                ty: Ty::UNIT,
+                kind: ExprKind::Tuple(vec![]),
+            })),
+        ),
+    }
+}
+
 fn create_general_dealloc_stmt(call_expr: Expr, ident: &IdentTemplate) -> Stmt {
     Stmt {
         id: NodeId::default(),