Use-elim and loop-uni on fragments

Author: swernli

compiler/qsc_passes/src/lib.rs

@@ -12,6 +12,7 @@ pub mod loop_unification;
 pub mod replace_qubit_allocation;
 pub mod spec_gen;
 
+use loop_unification::LoopUni;
 use miette::Diagnostic;
 use qsc_frontend::{
     compile::{CompileUnit, PackageStore},
@@ -21,7 +22,9 @@ use qsc_hir::{
     assigner::Assigner,
     global::Table,
     hir::{Item, ItemKind},
+    mut_visit::MutVisitor,
 };
+use replace_qubit_allocation::ReplaceQubitAllocation;
 use thiserror::Error;
 
 #[derive(Clone, Debug, Diagnostic, Error)]
@@ -38,8 +41,11 @@ pub fn run_default_passes(core: &Table, unit: &mut CompileUnit) -> Vec<Error> {
     let spec_errors = spec_gen::generate_specs(core, unit);
     let conjugate_errors = conjugate_invert::invert_conjugate_exprs(core, unit);
 
-    loop_unification::loop_unification(core, unit);
-    replace_qubit_allocation::replace_qubit_allocation(unit, store.core());
+    LoopUni {
+        assigner: &mut unit.assigner,
+    }
+    .visit_package(&mut unit.package);
+    ReplaceQubitAllocation::new(&mut unit.assigner, store.core()).visit_package(&mut unit.package);
 
     spec_errors
         .into_iter()
@@ -49,12 +55,16 @@ pub fn run_default_passes(core: &Table, unit: &mut CompileUnit) -> Vec<Error> {
 }
 
 pub fn run_core_passes(unit: &mut CompileUnit) {
-    loop_unification::loop_unification(unit);
+    LoopUni {
+        assigner: &mut unit.assigner,
+    }
+    .visit_package(&mut unit.package);
 }
 
 pub fn run_default_passes_for_fragment(
     core: &Table,
     assigner: &mut Assigner,
+    core_table: &Table,
     fragment: &mut Fragment,
 ) -> Vec<Error> {
     let mut errors = Vec::new();
@@ -66,7 +76,8 @@ pub fn run_default_passes_for_fragment(
                     .into_iter()
                     .map(Error::ConjInvert),
             );
-            loop_unification::loop_unification_for_stmt(assigner, stmt);
+            LoopUni { assigner }.visit_stmt(stmt);
+            ReplaceQubitAllocation::new(assigner, core_table).visit_stmt(stmt);
         }
         Fragment::Item(Item {
             kind: ItemKind::Callable(decl),
@@ -82,7 +93,8 @@ pub fn run_default_passes_for_fragment(
                     .into_iter()
                     .map(Error::ConjInvert),
             );
-            loop_unification::loop_unification_for_callable(assigner, decl);
+            LoopUni { assigner }.visit_callable_decl(decl);
+            ReplaceQubitAllocation::new(assigner, core_table).visit_callable_decl(decl);
         }
         Fragment::Item(_) | Fragment::Error(_) => {}
     }

compiler/qsc_passes/src/loop_unification.rs

@@ -5,13 +5,12 @@ use core::panic;
 use std::{mem::take, rc::Rc};
 
 use qsc_data_structures::span::Span;
-use qsc_frontend::compile::CompileUnit;
 use qsc_hir::{
     assigner::Assigner,
     global::Table,
     hir::{
-        BinOp, Block, CallableDecl, Expr, ExprKind, Lit, Mutability, NodeId, Pat, PrimField,
-        PrimTy, Stmt, StmtKind, Ty, UnOp,
+        BinOp, Block, Expr, ExprKind, Lit, Mutability, NodeId, Pat, PrimField, PrimTy, Stmt,
+        StmtKind, Ty, UnOp,
     },
     mut_visit::{walk_expr, MutVisitor},
 };
@@ -21,29 +20,9 @@ use crate::common::{create_gen_core_ref, IdentTemplate};
 #[cfg(test)]
 mod tests;
 
-pub fn loop_unification(core: &Table, unit: &mut CompileUnit) {
-    let mut pass = LoopUni {
-        core,
-        assigner: &mut unit.assigner,
-    };
-    pass.visit_package(&mut unit.package);
-}
-
-#[allow(clippy::module_name_repetitions)]
-pub fn loop_unification_for_stmt(assigner: &mut Assigner, stmt: &mut Stmt) {
-    let mut pass = LoopUni { assigner };
-    pass.visit_stmt(stmt);
-}
-
-#[allow(clippy::module_name_repetitions)]
-pub fn loop_unification_for_callable(assigner: &mut Assigner, decl: &mut CallableDecl) {
-    let mut pass = LoopUni { assigner };
-    pass.visit_callable_decl(decl);
-}
-
-struct LoopUni<'a> {
-    core: &'a Table,
-    assigner: &'a mut Assigner,
+pub(crate) struct LoopUni<'a> {
+    pub(crate) core: &'a Table,
+    pub(crate) assigner: &'a mut Assigner,
 }
 
 impl LoopUni<'_> {

compiler/qsc_passes/src/loop_unification/tests.rs

@@ -6,15 +6,19 @@
 use expect_test::{expect, Expect};
 use indoc::indoc;
 use qsc_frontend::compile::{self, compile, PackageStore, SourceMap};
+use qsc_hir::mut_visit::MutVisitor;
 
-use crate::loop_unification::loop_unification;
+use crate::loop_unification::LoopUni;
 
 fn check(file: &str, expect: &Expect) {
     let store = PackageStore::new(compile::core());
     let sources = SourceMap::new([("test".into(), file.into())], None);
     let mut unit = compile(&store, &[], sources);
     assert!(unit.errors.is_empty(), "{:?}", unit.errors);
-    loop_unification(store.core(), &mut unit);
+    LoopUni {
+        assigner: &mut unit.assigner,
+    }
+    .visit_package(&mut unit.package);
     expect.assert_eq(&unit.package.to_string());
 }
 

compiler/qsc_passes/src/replace_qubit_allocation.rs

@@ -5,7 +5,6 @@
 mod tests;
 
 use qsc_data_structures::span::Span;
-use qsc_frontend::compile::CompileUnit;
 use qsc_hir::{
     assigner::Assigner,
     global::Table,
@@ -19,31 +18,30 @@ use std::{mem::take, rc::Rc};
 
 use crate::common::{create_gen_core_ref, IdentTemplate};
 
-pub fn replace_qubit_allocation(unit: &mut CompileUnit, core_table: &Table) {
-    let mut pass = ReplaceQubitAllocation {
-        assigner: &mut unit.assigner,
-        core_table,
-        qubits_curr_callable: Vec::new(),
-        qubits_curr_block: Vec::new(),
-        prefix_qubits: Vec::new(),
-    };
-    pass.visit_package(&mut unit.package);
-}
-
 struct QubitIdent {
     id: IdentTemplate,
     is_array: bool,
 }
 
-struct ReplaceQubitAllocation<'a> {
+pub(crate) struct ReplaceQubitAllocation<'a> {
     assigner: &'a mut Assigner,
     core_table: &'a Table,
     qubits_curr_callable: Vec<Vec<QubitIdent>>,
     qubits_curr_block: Vec<QubitIdent>,
     prefix_qubits: Vec<QubitIdent>,
 }
 
-impl ReplaceQubitAllocation<'_> {
+impl<'a> ReplaceQubitAllocation<'a> {
+    pub(crate) fn new(assigner: &'a mut Assigner, core_table: &'a Table) -> Self {
+        Self {
+            assigner,
+            core_table,
+            qubits_curr_callable: Vec::new(),
+            qubits_curr_block: Vec::new(),
+            prefix_qubits: Vec::new(),
+        }
+    }
+
     fn process_qubit_stmt(
         &mut self,
         stmt_span: Span,

compiler/qsc_passes/src/replace_qubit_allocation/tests.rs

@@ -1,17 +1,18 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use crate::replace_qubit_allocation::replace_qubit_allocation;
+use crate::replace_qubit_allocation::ReplaceQubitAllocation;
 use expect_test::{expect, Expect};
 use indoc::indoc;
 use qsc_frontend::compile::{self, compile, PackageStore, SourceMap};
+use qsc_hir::mut_visit::MutVisitor;
 
 fn check(file: &str, expect: &Expect) {
     let store = PackageStore::new(compile::core());
     let sources = SourceMap::new([("test".into(), file.into())], None);
     let mut unit = compile(&store, &[], sources);
     assert!(unit.errors.is_empty(), "{:?}", unit.errors);
-    replace_qubit_allocation(&mut unit, store.core());
+    ReplaceQubitAllocation::new(&mut unit.assigner, store.core()).visit_package(&mut unit.package);
     expect.assert_eq(&unit.package.to_string());
 }