feat: Add hashable Angle type to Quantum extension

specification/hugr.md

@@ -1740,7 +1740,7 @@ Note there are also `measurez: Qubit -> (i1, Qubit)` and on supported
 targets `reset: Qubit -> Qubit` operations to measure or reset a qubit
 without losing a handle to it.
 
-**Dynamic vs static allocation**
+#### Dynamic vs static allocation
 
 With these operations the programmer/front-end can request dynamic qubit
 allocation, and the compiler can add/remove/move these operations to use
@@ -1763,6 +1763,28 @@ allocation for all `Qubit` wires.
 If further the program does not contain any `qalloc` or `qfree`
 operations we can state the program only uses `N` qubits.
 
+#### Angles
+
+The Quantum extension also defines a specialized `angle<N>` type which is used
+to express parameters of rotation gates. The type is parametrized by the
+_log-denominator_, which is an integer $N \in [0, 53]$; angles with
+log-denominator $N$ are multiples of $2 \pi / 2^N$, where the multiplier is an
+unsigned `int<N>` in the range $[0, 2^N]$. The maximum log-denominator $53$
+effectively gives the resolution of a `float64` value; but note that unlike
+`float64` all angle values are equatable and hashable; and two `angle<N>` that
+differ by a multiple of $2 \pi$ are _equal_.
+
+The following operations are defined:
+
+| Name           | Inputs     | Outputs    | Meaning |
+| -------------- | ---------- | ---------- | ------- |
+| `aconst<N, x>` | none       | `angle<N>` | const node producing angle $2 \pi x / 2^N$ (where $0 \leq x \lt 2^N$) |
+| `atrunc<M,N>`  | `angle<M>` | `angle<N>` | round `angle<M>` to `angle<N>`, where $M \geq N$, rounding down in $[0, 2\pi)$ if necessary |
+| `aconvert<M,N>`  | `angle<M>` | `Sum(angle<N>, ErrorType)` | convert `angle<M>` to `angle<N>`, returning an error if $M \gt N$ and exact conversion is impossible |
+| `aadd<M,N>`    | `angle<M>`, `angle<N>` | `angle<max(M,N)>` | add two angles |
+| `asub<M,N>`    | `angle<M>`, `angle<N>` | `angle<max(M,N)>` | subtract the second angle from the first |
+| `aneg<N>`      | `angle<N>` | `angle<N>` | negate an angle |
+
 ### Higher-order (Tierkreis) Extension
 
 In **some** contexts, notably the Tierkreis runtime, higher-order

src/std_extensions/quantum.rs

@@ -1,33 +1,253 @@
 //! Basic HUGR quantum operations
 
+use std::cmp::max;
+use std::f64::consts::TAU;
+use std::num::NonZeroU64;
+
 use smol_str::SmolStr;
 
-use crate::extension::prelude::{BOOL_T, QB_T};
+use crate::extension::prelude::{BOOL_T, ERROR_TYPE, QB_T};
 use crate::extension::{ExtensionId, SignatureError};
 use crate::std_extensions::arithmetic::float_types::FLOAT64_TYPE;
 use crate::type_row;
-use crate::types::type_param::TypeArg;
-use crate::types::FunctionType;
+use crate::types::type_param::{TypeArg, TypeArgError, TypeParam};
+use crate::types::{ConstTypeError, CustomCheckFailure, CustomType, FunctionType, Type, TypeBound};
+use crate::utils::collect_array;
+use crate::values::CustomConst;
 use crate::Extension;
 
 use lazy_static::lazy_static;
 
 /// The extension identifier.
 pub const EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("quantum");
-fn one_qb_func(_: &[TypeArg]) -> Result<FunctionType, SignatureError> {
-    Ok(FunctionType::new(type_row![QB_T], type_row![QB_T]))
+
+/// Identifier for the angle type.
+const ANGLE_TYPE_ID: SmolStr = SmolStr::new_inline("angle");
+
+fn angle_custom_type(log_denom_arg: TypeArg) -> CustomType {
+    CustomType::new(ANGLE_TYPE_ID, [log_denom_arg], EXTENSION_ID, TypeBound::Eq)
 }
 
-fn two_qb_func(_: &[TypeArg]) -> Result<FunctionType, SignatureError> {
+/// Angle type with a given log-denominator (specified by the TypeArg).
+///
+/// This type is capable of representing angles that are multiples of 2π / 2^N where N is the
+/// log-denominator.
+pub(super) fn angle_type(log_denom_arg: TypeArg) -> Type {
+    Type::new_extension(angle_custom_type(log_denom_arg))
+}
+
+/// The largest permitted log-denominator.
+pub const LOG_DENOM_MAX: u8 = 53;
+
+const fn is_valid_log_denom(n: u8) -> bool {
+    n <= LOG_DENOM_MAX
+}
+
+/// Type parameter for the log-denominator of an angle.
+#[allow(clippy::assertions_on_constants)]
+pub const LOG_DENOM_TYPE_PARAM: TypeParam =
+    TypeParam::bounded_nat(NonZeroU64::MIN.saturating_add(LOG_DENOM_MAX as u64));
+
+/// Get the log-denominator of the specified type argument or error if the argument is invalid.
+pub(super) fn get_log_denom(arg: &TypeArg) -> Result<u8, TypeArgError> {
+    match arg {
+        TypeArg::BoundedNat { n } if is_valid_log_denom(*n as u8) => Ok(*n as u8),
+        _ => Err(TypeArgError::TypeMismatch {
+            arg: arg.clone(),
+            param: LOG_DENOM_TYPE_PARAM,
+        }),
+    }
+}
+
+pub(super) const fn type_arg(log_denom: u8) -> TypeArg {
+    TypeArg::BoundedNat {
+        n: log_denom as u64,
+    }
+}
+
+/// An angle
+#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
+pub struct ConstAngle {
+    log_denom: u8,
+    value: u64,
+}
+
+impl ConstAngle {
+    /// Create a new [`ConstAngle`] from a log-denominator and a numerator
+    pub fn new(log_denom: u8, value: u64) -> Result<Self, ConstTypeError> {
+        if !is_valid_log_denom(log_denom) {
+            return Err(ConstTypeError::CustomCheckFail(
+                crate::types::CustomCheckFailure::Message(
+                    "Invalid angle log-denominator.".to_owned(),
+                ),
+            ));
+        }
+        if value >= (1u64 << log_denom) {
+            return Err(ConstTypeError::CustomCheckFail(
+                crate::types::CustomCheckFailure::Message(
+                    "Invalid unsigned integer value.".to_owned(),
+                ),
+            ));
+        }
+        Ok(Self { log_denom, value })
+    }
+
+    /// Create a new [`ConstAngle`] from a log-denominator and a floating-point value in radians,
+    /// rounding to the nearest corresponding value. (Ties round away from zero.)
+    pub fn from_radians_rounding(log_denom: u8, theta: f64) -> Result<Self, ConstTypeError> {
+        if !is_valid_log_denom(log_denom) {
+            return Err(ConstTypeError::CustomCheckFail(
+                crate::types::CustomCheckFailure::Message(
+                    "Invalid angle log-denominator.".to_owned(),
+                ),
+            ));
+        }
+        let a = (((1u64 << log_denom) as f64) * theta / TAU).round() as i64;
+        Ok(Self {
+            log_denom,
+            value: a.rem_euclid(1i64 << log_denom) as u64,
+        })
+    }
+
+    /// Returns the value of the constant
+    pub fn value(&self) -> u64 {
+        self.value
+    }
+
+    /// Returns the log-denominator of the constant
+    pub fn log_denom(&self) -> u8 {
+        self.log_denom
+    }
+}
+
+#[typetag::serde]
+impl CustomConst for ConstAngle {
+    fn name(&self) -> SmolStr {
+        format!("a(2π*{}/2^{})", self.value, self.log_denom).into()
+    }
+    fn check_custom_type(&self, typ: &CustomType) -> Result<(), CustomCheckFailure> {
+        if typ.clone() == angle_custom_type(type_arg(self.log_denom)) {
+            Ok(())
+        } else {
+            Err(CustomCheckFailure::Message(
+                "Angle constant type mismatch.".into(),
+            ))
+        }
+    }
+    fn equal_consts(&self, other: &dyn CustomConst) -> bool {
+        crate::values::downcast_equal_consts(self, other)
+    }
+}
+
+fn atrunc_sig(arg_values: &[TypeArg]) -> Result<FunctionType, SignatureError> {
+    let [arg0, arg1] = collect_array(arg_values);
+    let m: u8 = get_log_denom(arg0)?;
+    let n: u8 = get_log_denom(arg1)?;
+    if m < n {
+        return Err(SignatureError::InvalidTypeArgs);
+    }
+    Ok(FunctionType::new(
+        vec![angle_type(arg0.clone())],
+        vec![angle_type(arg1.clone())],
+    ))
+}
+
+fn aconvert_sig(arg_values: &[TypeArg]) -> Result<FunctionType, SignatureError> {
+    let [arg0, arg1] = collect_array(arg_values);
     Ok(FunctionType::new(
-        type_row![QB_T, QB_T],
-        type_row![QB_T, QB_T],
+        vec![angle_type(arg0.clone())],
+        vec![Type::new_sum(vec![angle_type(arg1.clone()), ERROR_TYPE])],
     ))
 }
 
+fn abinop_sig(arg_values: &[TypeArg]) -> Result<FunctionType, SignatureError> {
+    let [arg0, arg1] = collect_array(arg_values);
+    let m: u8 = get_log_denom(arg0)?;
+    let n: u8 = get_log_denom(arg1)?;
+    let l: u8 = max(m, n);
+    Ok(FunctionType::new(
+        vec![
+            angle_type(TypeArg::BoundedNat { n: m as u64 }),
+            angle_type(TypeArg::BoundedNat { n: n as u64 }),
+        ],
+        vec![angle_type(TypeArg::BoundedNat { n: l as u64 })],
+    ))
+}
+
+fn aunop_sig(arg_values: &[TypeArg]) -> Result<FunctionType, SignatureError> {
+    let [arg] = collect_array(arg_values);
+    Ok(FunctionType::new_linear(vec![angle_type(arg.clone())]))
+}
+
+fn one_qb_func(_: &[TypeArg]) -> Result<FunctionType, SignatureError> {
+    Ok(FunctionType::new_linear(type_row![QB_T]))
+}
+
+fn two_qb_func(_: &[TypeArg]) -> Result<FunctionType, SignatureError> {
+    Ok(FunctionType::new_linear(type_row![QB_T, QB_T]))
+}
+
 fn extension() -> Extension {
     let mut extension = Extension::new(EXTENSION_ID);
 
+    extension
+        .add_type(
+            ANGLE_TYPE_ID,
+            vec![LOG_DENOM_TYPE_PARAM],
+            "angle value with a given log-denominator".to_owned(),
+            TypeBound::Eq.into(),
+        )
+        .unwrap();
+
+    extension
+        .add_op_custom_sig_simple(
+            "atrunc".into(),
+            "truncate an angle to one with a lower log-denominator with the same value, rounding \
+            down in [0, 2π) if necessary"
+                .to_owned(),
+            vec![LOG_DENOM_TYPE_PARAM, LOG_DENOM_TYPE_PARAM],
+            atrunc_sig,
+        )
+        .unwrap();
+
+    extension
+        .add_op_custom_sig_simple(
+            "aconvert".into(),
+            "convert an angle to one with another log-denominator having the same value, if \
+            possible, otherwise return an error"
+                .to_owned(),
+            vec![LOG_DENOM_TYPE_PARAM, LOG_DENOM_TYPE_PARAM],
+            aconvert_sig,
+        )
+        .unwrap();
+
+    extension
+        .add_op_custom_sig_simple(
+            "aadd".into(),
+            "addition of angles".to_owned(),
+            vec![LOG_DENOM_TYPE_PARAM],
+            abinop_sig,
+        )
+        .unwrap();
+
+    extension
+        .add_op_custom_sig_simple(
+            "asub".into(),
+            "subtraction of the second angle from the first".to_owned(),
+            vec![LOG_DENOM_TYPE_PARAM],
+            abinop_sig,
+        )
+        .unwrap();
+
+    extension
+        .add_op_custom_sig_simple(
+            "aneg".into(),
+            "negation of an angle".to_owned(),
+            vec![LOG_DENOM_TYPE_PARAM],
+            aunop_sig,
+        )
+        .unwrap();
+
     extension
         .add_op_custom_sig_simple(
             SmolStr::new_inline("H"),
@@ -78,7 +298,16 @@ lazy_static! {
 
 #[cfg(test)]
 pub(crate) mod test {
-    use crate::{extension::EMPTY_REG, ops::LeafOp};
+    use std::f64::consts::TAU;
+
+    use cool_asserts::assert_matches;
+
+    use crate::{
+        extension::EMPTY_REG,
+        ops::LeafOp,
+        std_extensions::quantum::{get_log_denom, ConstAngle},
+        types::{type_param::TypeArgError, ConstTypeError, TypeArg},
+    };
 
     use super::EXTENSION;
 
@@ -100,4 +329,37 @@ pub(crate) mod test {
     pub(crate) fn measure() -> LeafOp {
         get_gate("Measure")
     }
+
+    #[test]
+    fn test_angle_log_denoms() {
+        let type_arg_53 = TypeArg::BoundedNat { n: 53 };
+        assert_matches!(get_log_denom(&type_arg_53), Ok(53));
+
+        let type_arg_54 = TypeArg::BoundedNat { n: 54 };
+        assert_matches!(
+            get_log_denom(&type_arg_54),
+            Err(TypeArgError::TypeMismatch { .. })
+        );
+    }
+
+    #[test]
+    fn test_angle_consts() {
+        let const_a32_7 = ConstAngle::new(5, 7).unwrap();
+        let const_a33_7 = ConstAngle::new(6, 7).unwrap();
+        let const_a32_8 = ConstAngle::new(6, 8).unwrap();
+        assert_ne!(const_a32_7, const_a33_7);
+        assert_ne!(const_a32_7, const_a32_8);
+        assert_eq!(const_a32_7, ConstAngle::new(5, 7).unwrap());
+        assert_matches!(
+            ConstAngle::new(3, 256),
+            Err(ConstTypeError::CustomCheckFail(_))
+        );
+        assert_matches!(
+            ConstAngle::new(54, 256),
+            Err(ConstTypeError::CustomCheckFail(_))
+        );
+        let const_af1 = ConstAngle::from_radians_rounding(5, 0.21874 * TAU).unwrap();
+        assert_eq!(const_af1.value(), 7);
+        assert_eq!(const_af1.log_denom(), 5);
+    }
 }