Generics forwarding to wrapped functions and invert

Renmusxd · Sep 1, 2021 · 32044a5 · 32044a5
1 parent 2a91f6a
commit 32044a5
Show file tree

Hide file tree

Showing 4 changed files with 158 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -164,3 +164,24 @@ let r = program!(&mut b, r;
 )?;
 ```
 
+
+Generics can be used by substituting the usual angle brackets for square.
+```rust
+use qip::*;
+
+let mut b = OpBuilder::new();
+let r = b.qubit();
+
+fn rz<T: Into<f64>>(b: &mut dyn UnitaryBuilder, r: Register, theta: T) -> Register {
+    b.rz(r, theta.into())
+}
+
+wrap_fn!(rz_op[T: Into<f64>](theta: T), rz, ra);
+invert_fn!(inv_rz_op[T: Into<f64>](theta: T), rz_op);
+
+let r = program!(&mut b, r;
+    rz_op(3.141_f32) r;
+    inv_rz_op(3.141_f32) r;
+)?;
+```
+
diff --git a/src/lib.rs b/src/lib.rs
@@ -195,6 +195,29 @@
 //! # Ok(())
 //! # }
 //! ```
+//!
+//! Generics can be used by substituting the usual angle brackets for square.
+//! ```rust
+//! use qip::*;
+//! # fn main() -> Result<(), CircuitError> {
+//!
+//! let mut b = OpBuilder::new();
+//! let r = b.qubit();
+//!
+//! fn rz<T: Into<f64>>(b: &mut dyn UnitaryBuilder, r: Register, theta: T) -> Register {
+//!     b.rz(r, theta.into())
+//! }
+//!
+//! wrap_fn!(rz_op[T: Into<f64>](theta: T), rz, ra);
+//! invert_fn!(inv_rz_op[T: Into<f64>](theta: T), rz_op);
+//!
+//! let r = program!(&mut b, r;
+//!     rz_op(3.141_f32) r;
+//!     inv_rz_op(3.141_f32) r;
+//! )?;
+//! # Ok(())
+//! # }
+//! ```
 
 pub use self::builders::*;
 pub use self::common_circuits::*;

diff --git a/src/macros/inverter.rs b/src/macros/inverter.rs
@@ -27,6 +27,22 @@ use crate::{CircuitError, OpBuilder, Register, UnitaryBuilder};
 /// ```
 #[macro_export]
 macro_rules! wrap_and_invert {
+    (pub $newfunc:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), pub $newinvert:ident, $($tail:tt)*) => {
+        wrap_fn!(pub $newfunc[$($typetail)*]($arg: $argtype), $($tail)*);
+        invert_fn!(pub $newinvert[$($typetail)*]($arg: $argtype), $newfunc);
+    };
+    ($newfunc:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), pub $newinvert:ident, $($tail:tt)*) => {
+        wrap_fn!($newfunc[$($typetail)*]($arg: $argtype), $($tail)*);
+        invert_fn!(pub $newinvert[$($typetail)*]($arg: $argtype), $newfunc);
+    };
+    (pub $newfunc:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), $newinvert:ident, $($tail:tt)*) => {
+        wrap_fn!(pub $newfunc[$($typetail)*]($arg: $argtype), $($tail)*);
+        invert_fn!($newinvert[$($typetail)*]($arg: $argtype), $newfunc);
+    };
+    ($newfunc:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), $newinvert:ident, $($tail:tt)*) => {
+        wrap_fn!($newfunc[$($typetail)*]($arg: $argtype), $($tail)*);
+        invert_fn!($newinvert[$($typetail)*]($arg: $argtype), $newfunc);
+    };
     (pub $newfunc:ident($arg:ident: $argtype:ident), pub $newinvert:ident, $($tail:tt)*) => {
         wrap_fn!(pub $newfunc($arg: $argtype), $($tail)*);
         invert_fn!(pub $newinvert($arg: $argtype), $newfunc);
@@ -87,14 +103,33 @@ macro_rules! wrap_and_invert {
 /// ```
 #[macro_export]
 macro_rules! invert_fn {
+    (pub $newinvert:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), $func:expr) => {
+        /// Invert the given function.
+        pub fn $newinvert<$($typetail)*>(
+            b: &mut dyn $crate::UnitaryBuilder,
+            rs: Vec<Register>,
+            $arg: $argtype,
+        ) -> Result<Vec<Register>, $crate::CircuitError> {
+            $crate::inverter_args(b, rs, $func, $arg)
+        }
+    };
+    ($newinvert:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), $func:expr) => {
+        fn $newinvert<$($typetail)*>(
+            b: &mut dyn $crate::UnitaryBuilder,
+            rs: Vec<Register>,
+            $arg: $argtype,
+        ) -> Result<Vec<Register>, $crate::CircuitError> {
+            $crate::inverter_args(b, rs, $func, $arg)
+        }
+    };
     (pub $newinvert:ident($arg:ident: $argtype:ident), $func:expr) => {
         /// Invert the given function.
         pub fn $newinvert(
             b: &mut dyn $crate::UnitaryBuilder,
             rs: Vec<Register>,
             $arg: $argtype,
         ) -> Result<Vec<Register>, $crate::CircuitError> {
-            $crate::inverter(b, rs, |b, rs| $func(b, rs, $arg))
+            $crate::inverter_args(b, rs, $func, $arg)
         }
     };
     ($newinvert:ident($arg:ident: $argtype:ident), $func:expr) => {
@@ -103,7 +138,7 @@ macro_rules! invert_fn {
             rs: Vec<Register>,
             $arg: $argtype,
         ) -> Result<Vec<Register>, $crate::CircuitError> {
-            $crate::inverter(b, rs, |b, rs| $func(b, rs, $arg))
+            $crate::inverter_args(b, rs, $func, $arg)
         }
     };
     (pub $newinvert:ident, $func:expr) => {
@@ -128,11 +163,24 @@ macro_rules! invert_fn {
 /// Invert a circuit applied via the function f.
 pub fn inverter<F>(
     b: &mut dyn UnitaryBuilder,
-    mut rs: Vec<Register>,
+    rs: Vec<Register>,
     f: F,
 ) -> Result<Vec<Register>, CircuitError>
 where
     F: Fn(&mut dyn UnitaryBuilder, Vec<Register>) -> Result<Vec<Register>, CircuitError>,
+{
+    inverter_args(b, rs, |b, rs, _| f(b, rs), ())
+}
+
+/// Invert a circuit applied via the function f.
+pub fn inverter_args<T, F>(
+    b: &mut dyn UnitaryBuilder,
+    mut rs: Vec<Register>,
+    f: F,
+    args: T,
+) -> Result<Vec<Register>, CircuitError>
+where
+    F: Fn(&mut dyn UnitaryBuilder, Vec<Register>, T) -> Result<Vec<Register>, CircuitError>,
 {
     let original_indices: Vec<_> = rs.iter().map(|r| r.indices.clone()).collect();
     let mut inv_builder = OpBuilder::new();
@@ -147,7 +195,7 @@ where
 
     // Call the function and count any qubits allocated inside.
     let before_n = inv_builder.get_qubit_count();
-    let new_rs = f(&mut inv_builder, new_rs)?;
+    let new_rs = f(&mut inv_builder, new_rs, args)?;
     let end_reg = inv_builder.merge(new_rs)?;
     let after_n = inv_builder.get_qubit_count();
 
@@ -375,4 +423,17 @@ mod inverter_test {
         let _rs = inv_rz(&mut b, rs, 1.0)?;
         Ok(())
     }
+
+    #[test]
+    fn test_invert_and_wrap_rz_generic() -> Result<(), CircuitError> {
+        fn rz<T: Into<f64>>(b: &mut dyn UnitaryBuilder, r: Register, theta: T) -> Register {
+            b.rz(r, theta.into())
+        }
+        wrap_and_invert!(rz_op[T: Into<f64>](theta: T), inv_rz, rz, r);
+        let mut b = OpBuilder::new();
+        let r = b.qubit();
+        let rs = rz_op(&mut b, vec![r], 1.0)?;
+        let _rs = inv_rz(&mut b, rs, 1.0)?;
+        Ok(())
+    }
 }
diff --git a/src/macros/program.rs b/src/macros/program.rs
@@ -525,6 +525,28 @@ macro_rules! wrap_fn {
             Ok($rs)
         }
     };
+    (pub $newfunc:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), ($func:expr), $($tail:tt)*) => {
+        /// Wrapped version of function
+        pub fn $newfunc<$($typetail)*>(b: &mut dyn $crate::UnitaryBuilder, mut rs: Vec<$crate::Register>, $arg: $argtype) -> Result<Vec<$crate::Register>, $crate::CircuitError> {
+            wrap_fn!(@result_body(b, $func, rs, $arg) $($tail)*)
+        }
+    };
+    (pub $newfunc:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), $func:expr, $($tail:tt)*) => {
+        /// Wrapped version of function
+        pub fn $newfunc<$($typetail)*>(b: &mut dyn $crate::UnitaryBuilder, mut rs: Vec<$crate::Register>, $arg: $argtype) -> Result<Vec<$crate::Register>, $crate::CircuitError> {
+            wrap_fn!(@raw_body(b, $func, rs, $arg) $($tail)*)
+        }
+    };
+    ($newfunc:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), ($func:expr), $($tail:tt)*) => {
+        fn $newfunc<$($typetail)*>(b: &mut dyn $crate::UnitaryBuilder, mut rs: Vec<$crate::Register>, $arg: $argtype) -> Result<Vec<$crate::Register>, $crate::CircuitError> {
+            wrap_fn!(@result_body(b, $func, rs, $arg) $($tail)*)
+        }
+    };
+    ($newfunc:ident[$($typetail:tt)*]($arg:ident: $argtype:ident), $func:expr, $($tail:tt)*) => {
+        fn $newfunc<$($typetail)*>(b: &mut dyn $crate::UnitaryBuilder, mut rs: Vec<$crate::Register>, $arg: $argtype) -> Result<Vec<$crate::Register>, $crate::CircuitError> {
+            wrap_fn!(@raw_body(b, $func, rs, $arg) $($tail)*)
+        }
+    };
     (pub $newfunc:ident($arg:ident: $argtype:ident), ($func:expr), $($tail:tt)*) => {
         /// Wrapped version of function
         pub fn $newfunc(b: &mut dyn $crate::UnitaryBuilder, mut rs: Vec<$crate::Register>, $arg: $argtype) -> Result<Vec<$crate::Register>, $crate::CircuitError> {
@@ -1328,4 +1350,31 @@ mod common_circuit_tests {
         assert_eq!(macro_circuit, basic_circuit);
         Ok(())
     }
+
+    #[test]
+    fn wrap_unitary_fn_arg_generic() -> Result<(), CircuitError> {
+        fn rz<T: Into<f64>>(b: &mut dyn UnitaryBuilder, r: Register, theta: T) -> Register {
+            b.rz(r, theta.into())
+        }
+        wrap_fn!(wrapped_rz[T: Into<f64>](theta: T), rz, r);
+
+        let mut b = OpBuilder::new();
+        let r = b.qubit();
+
+        let r = program!(&mut b, r;
+            wrapped_rz(1.0) r;
+        )?;
+
+        run_debug(&r)?;
+        // Compare to expected value
+        let macro_circuit = make_circuit_matrix::<f64>(1, &r, true);
+        let mut b = OpBuilder::new();
+        let r = b.qubit();
+
+        let r = b.rz(r, 1.0);
+        run_debug(&r)?;
+        let basic_circuit = make_circuit_matrix::<f64>(1, &r, true);
+        assert_eq!(macro_circuit, basic_circuit);
+        Ok(())
+    }
 }