microsoft/qdk
Publicmirrored from https://github.com/microsoft/qdkAvailable
compiler/qsc_passes/src/spec_gen/ctl_gen.rs
79lines · modecode
| 1 | // Copyright (c) Microsoft Corporation. |
| 2 | // Licensed under the MIT License. |
| 3 | |
| 4 | use miette::Diagnostic; |
| 5 | use qsc_data_structures::span::Span; |
| 6 | use qsc_hir::{ |
| 7 | hir::{CallableKind, Expr, ExprKind, Functor, NodeId, Res, UnOp}, |
| 8 | mut_visit::{walk_expr, MutVisitor}, |
| 9 | ty::{Arrow, Prim, Ty}, |
| 10 | }; |
| 11 | use thiserror::Error; |
| 12 | |
| 13 | #[derive(Clone, Debug, Diagnostic, Error)] |
| 14 | pub enum Error { |
| 15 | #[error("operation does not support the controlled functor")] |
| 16 | #[diagnostic(help("each operation called inside an operation with compiler-generated controlled specializations must support the controlled functor"))] |
| 17 | #[diagnostic(code("Qsc.CtlGen.MissingCtlFunctor"))] |
| 18 | MissingCtlFunctor(#[label] Span), |
| 19 | } |
| 20 | |
| 21 | pub(super) struct CtlDistrib { |
| 22 | pub(super) ctls: Res, |
| 23 | pub(super) errors: Vec<Error>, |
| 24 | } |
| 25 | |
| 26 | impl MutVisitor for CtlDistrib { |
| 27 | fn visit_expr(&mut self, expr: &mut Expr) { |
| 28 | match &mut expr.kind { |
| 29 | ExprKind::Call(op, args) => { |
| 30 | match &op.ty { |
| 31 | Ty::Arrow(arrow) if arrow.kind == CallableKind::Operation => { |
| 32 | let functors = arrow |
| 33 | .functors |
| 34 | .expect_value("arrow type should have concrete functors"); |
| 35 | |
| 36 | if functors.contains(&Functor::Ctl) { |
| 37 | op.kind = ExprKind::UnOp(UnOp::Functor(Functor::Ctl), op.clone()); |
| 38 | op.id = NodeId::default(); |
| 39 | op.ty = Ty::Arrow(Box::new(Arrow { |
| 40 | kind: CallableKind::Operation, |
| 41 | input: Box::new(Ty::Tuple(vec![ |
| 42 | Ty::Array(Box::new(Ty::Prim(Prim::Qubit))), |
| 43 | Ty::clone(&arrow.input), |
| 44 | ])), |
| 45 | output: arrow.output.clone(), |
| 46 | functors: arrow.functors, |
| 47 | })); |
| 48 | |
| 49 | args.kind = ExprKind::Tuple(vec![ |
| 50 | Expr { |
| 51 | id: NodeId::default(), |
| 52 | span: args.span, |
| 53 | ty: Ty::Array(Box::new(Ty::Prim(Prim::Qubit))), |
| 54 | kind: ExprKind::Var(self.ctls, Vec::new()), |
| 55 | }, |
| 56 | Expr::clone(args), |
| 57 | ]); |
| 58 | args.ty = Ty::Tuple(vec![ |
| 59 | Ty::Array(Box::new(Ty::Prim(Prim::Qubit))), |
| 60 | Ty::clone(&args.ty), |
| 61 | ]); |
| 62 | args.id = NodeId::default(); |
| 63 | } else { |
| 64 | self.errors.push(Error::MissingCtlFunctor(op.span)); |
| 65 | } |
| 66 | } |
| 67 | _ => {} |
| 68 | } |
| 69 | |
| 70 | walk_expr(self, expr); |
| 71 | } |
| 72 | ExprKind::Conjugate(_, apply) => { |
| 73 | // Only transform the apply block, the within block can remain as-is. |
| 74 | self.visit_block(apply); |
| 75 | } |
| 76 | _ => walk_expr(self, expr), |
| 77 | } |
| 78 | } |
| 79 | } |
| 80 | |