feat!: Make array.get give back the passed array

hugr-core/src/std_extensions/collections/array.rs

@@ -1,6 +1,7 @@
 //! Fixed-length array type and operations extension.
 
 mod array_clone;
+mod array_conversion;
 mod array_discard;
 mod array_kind;
 mod array_op;
@@ -24,6 +25,7 @@ use crate::types::{CustomCheckFailure, Type, TypeBound, TypeName};
 use crate::Extension;
 
 pub use array_clone::{GenericArrayClone, GenericArrayCloneDef, ARRAY_CLONE_OP_ID};
+pub use array_conversion::{Direction, GenericArrayConvert, GenericArrayConvertDef, FROM, INTO};
 pub use array_discard::{GenericArrayDiscard, GenericArrayDiscardDef, ARRAY_DISCARD_OP_ID};
 pub use array_kind::ArrayKind;
 pub use array_op::{GenericArrayOp, GenericArrayOpDef};

hugr-core/src/std_extensions/collections/array/array_conversion.rs

@@ -0,0 +1,318 @@
+//! Operations for converting between the different array extensions
+
+use std::marker::PhantomData;
+use std::str::FromStr;
+use std::sync::{Arc, Weak};
+
+use crate::extension::simple_op::{
+    HasConcrete, HasDef, MakeExtensionOp, MakeOpDef, MakeRegisteredOp, OpLoadError,
+};
+use crate::extension::{ExtensionId, OpDef, SignatureError, SignatureFunc, TypeDef};
+use crate::ops::{ExtensionOp, NamedOp, OpName};
+use crate::types::type_param::{TypeArg, TypeParam};
+use crate::types::{FuncValueType, PolyFuncTypeRV, Type, TypeBound};
+use crate::Extension;
+
+use super::array_kind::ArrayKind;
+
+/// Array conversion direction.
+///
+/// Either the current array type [INTO] the other one, or the current array type [FROM] the
+/// other one.
+pub type Direction = bool;
+
+/// Array conversion direction to turn the current array type [INTO] the other one.
+pub const INTO: Direction = true;
+
+/// Array conversion direction to obtain the current array type [FROM] the other one.
+pub const FROM: Direction = false;
+
+/// Definition of array conversion operations.
+///
+/// Generic over the concrete array implementation of the extension containing the operation, as
+/// well as over another array implementation that should be converted between. Also generic over
+/// the conversion [Direction].
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+pub struct GenericArrayConvertDef<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind>(
+    PhantomData<AK>,
+    PhantomData<OtherAK>,
+);
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind>
+    GenericArrayConvertDef<AK, DIR, OtherAK>
+{
+    /// Creates a new array conversion definition.
+    pub fn new() -> Self {
+        GenericArrayConvertDef(PhantomData, PhantomData)
+    }
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> Default
+    for GenericArrayConvertDef<AK, DIR, OtherAK>
+{
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> NamedOp
+    for GenericArrayConvertDef<AK, DIR, OtherAK>
+{
+    fn name(&self) -> OpName {
+        match DIR {
+            INTO => format!("to_{}", OtherAK::TYPE_NAME).into(),
+            FROM => format!("from_{}", OtherAK::TYPE_NAME).into(),
+        }
+    }
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> FromStr
+    for GenericArrayConvertDef<AK, DIR, OtherAK>
+{
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let def = GenericArrayConvertDef::new();
+        if s == def.name() {
+            Ok(def)
+        } else {
+            Err(())
+        }
+    }
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind>
+    GenericArrayConvertDef<AK, DIR, OtherAK>
+{
+    /// To avoid recursion when defining the extension, take the type definition as an argument.
+    fn signature_from_def(&self, array_def: &TypeDef) -> SignatureFunc {
+        let params = vec![TypeParam::max_nat(), TypeBound::Any.into()];
+        let size = TypeArg::new_var_use(0, TypeParam::max_nat());
+        let element_ty = Type::new_var_use(1, TypeBound::Any);
+
+        let this_ty = AK::instantiate_ty(array_def, size.clone(), element_ty.clone())
+            .expect("Array type instantiation failed");
+        let other_ty =
+            OtherAK::ty_parametric(size, element_ty).expect("Array type instantiation failed");
+
+        let sig = match DIR {
+            INTO => FuncValueType::new(this_ty, other_ty),
+            FROM => FuncValueType::new(other_ty, this_ty),
+        };
+        PolyFuncTypeRV::new(params, sig).into()
+    }
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> MakeOpDef
+    for GenericArrayConvertDef<AK, DIR, OtherAK>
+{
+    fn from_def(op_def: &OpDef) -> Result<Self, OpLoadError>
+    where
+        Self: Sized,
+    {
+        crate::extension::simple_op::try_from_name(op_def.name(), op_def.extension_id())
+    }
+
+    fn init_signature(&self, _extension_ref: &Weak<Extension>) -> SignatureFunc {
+        self.signature_from_def(AK::type_def())
+    }
+
+    fn extension_ref(&self) -> Weak<Extension> {
+        Arc::downgrade(AK::extension())
+    }
+
+    fn extension(&self) -> ExtensionId {
+        AK::EXTENSION_ID
+    }
+
+    fn description(&self) -> String {
+        match DIR {
+            INTO => format!("Turns `{}` into `{}`", AK::TYPE_NAME, OtherAK::TYPE_NAME),
+            FROM => format!("Turns `{}` into `{}`", OtherAK::TYPE_NAME, AK::TYPE_NAME),
+        }
+    }
+
+    /// Add an operation implemented as a [MakeOpDef], which can provide the data
+    /// required to define an [OpDef], to an extension.
+    //
+    // This method is re-defined here since we need to pass the array type def while
+    // computing the signature, to avoid recursive loops initializing the extension.
+    fn add_to_extension(
+        &self,
+        extension: &mut Extension,
+        extension_ref: &Weak<Extension>,
+    ) -> Result<(), crate::extension::ExtensionBuildError> {
+        let sig = self.signature_from_def(extension.get_type(&AK::TYPE_NAME).unwrap());
+        let def = extension.add_op(self.name(), self.description(), sig, extension_ref)?;
+        self.post_opdef(def);
+        Ok(())
+    }
+}
+
+/// Definition of the array conversion op.
+///
+/// Generic over the concrete array implementation of the extension containing the operation, as
+/// well as over another array implementation that should be converted between. Also generic over
+/// the conversion [Direction].
+#[derive(Clone, Debug, PartialEq)]
+pub struct GenericArrayConvert<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> {
+    /// The element type of the array.
+    pub elem_ty: Type,
+    /// Size of the array.
+    pub size: u64,
+    _kind: PhantomData<AK>,
+    _other_kind: PhantomData<OtherAK>,
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind>
+    GenericArrayConvert<AK, DIR, OtherAK>
+{
+    /// Creates a new array conversion op.
+    pub fn new(elem_ty: Type, size: u64) -> Self {
+        GenericArrayConvert {
+            elem_ty,
+            size,
+            _kind: PhantomData,
+            _other_kind: PhantomData,
+        }
+    }
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> NamedOp
+    for GenericArrayConvert<AK, DIR, OtherAK>
+{
+    fn name(&self) -> OpName {
+        match DIR {
+            INTO => format!("to_{}", OtherAK::TYPE_NAME).into(),
+            FROM => format!("from_{}", OtherAK::TYPE_NAME).into(),
+        }
+    }
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> MakeExtensionOp
+    for GenericArrayConvert<AK, DIR, OtherAK>
+{
+    fn from_extension_op(ext_op: &ExtensionOp) -> Result<Self, OpLoadError>
+    where
+        Self: Sized,
+    {
+        let def = GenericArrayConvertDef::<AK, DIR, OtherAK>::from_def(ext_op.def())?;
+        def.instantiate(ext_op.args())
+    }
+
+    fn type_args(&self) -> Vec<TypeArg> {
+        vec![
+            TypeArg::BoundedNat { n: self.size },
+            self.elem_ty.clone().into(),
+        ]
+    }
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> MakeRegisteredOp
+    for GenericArrayConvert<AK, DIR, OtherAK>
+{
+    fn extension_id(&self) -> ExtensionId {
+        AK::EXTENSION_ID
+    }
+
+    fn extension_ref(&self) -> Weak<Extension> {
+        Arc::downgrade(AK::extension())
+    }
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> HasDef
+    for GenericArrayConvert<AK, DIR, OtherAK>
+{
+    type Def = GenericArrayConvertDef<AK, DIR, OtherAK>;
+}
+
+impl<AK: ArrayKind, const DIR: Direction, OtherAK: ArrayKind> HasConcrete
+    for GenericArrayConvertDef<AK, DIR, OtherAK>
+{
+    type Concrete = GenericArrayConvert<AK, DIR, OtherAK>;
+
+    fn instantiate(&self, type_args: &[TypeArg]) -> Result<Self::Concrete, OpLoadError> {
+        match type_args {
+            [TypeArg::BoundedNat { n }, TypeArg::Type { ty }] => {
+                Ok(GenericArrayConvert::new(ty.clone(), *n))
+            }
+            _ => Err(SignatureError::InvalidTypeArgs.into()),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use crate::extension::prelude::bool_t;
+    use crate::ops::{OpTrait, OpType};
+    use crate::std_extensions::collections::array::Array;
+    use crate::std_extensions::collections::value_array::ValueArray;
+
+    use super::*;
+
+    #[rstest]
+    #[case(ValueArray, Array)]
+    fn test_convert_from_def<AK: ArrayKind, OtherAK: ArrayKind>(
+        #[case] _kind: AK,
+        #[case] _other_kind: OtherAK,
+    ) {
+        let op = GenericArrayConvert::<AK, FROM, OtherAK>::new(bool_t(), 2);
+        let optype: OpType = op.clone().into();
+        let new_op: GenericArrayConvert<AK, FROM, OtherAK> = optype.cast().unwrap();
+        assert_eq!(new_op, op);
+    }
+
+    #[rstest]
+    #[case(ValueArray, Array)]
+    fn test_convert_into_def<AK: ArrayKind, OtherAK: ArrayKind>(
+        #[case] _kind: AK,
+        #[case] _other_kind: OtherAK,
+    ) {
+        let op = GenericArrayConvert::<AK, INTO, OtherAK>::new(bool_t(), 2);
+        let optype: OpType = op.clone().into();
+        let new_op: GenericArrayConvert<AK, INTO, OtherAK> = optype.cast().unwrap();
+        assert_eq!(new_op, op);
+    }
+
+    #[rstest]
+    #[case(ValueArray, Array)]
+    fn test_convert_from<AK: ArrayKind, OtherAK: ArrayKind>(
+        #[case] _kind: AK,
+        #[case] _other_kind: OtherAK,
+    ) {
+        let size = 2;
+        let element_ty = bool_t();
+        let op = GenericArrayConvert::<AK, FROM, OtherAK>::new(element_ty.clone(), size);
+        let optype: OpType = op.into();
+        let sig = optype.dataflow_signature().unwrap();
+        assert_eq!(
+            sig.io(),
+            (
+                &vec![OtherAK::ty(size, element_ty.clone())].into(),
+                &vec![AK::ty(size, element_ty.clone())].into(),
+            )
+        );
+    }
+
+    #[rstest]
+    #[case(ValueArray, Array)]
+    fn test_convert_into<AK: ArrayKind, OtherAK: ArrayKind>(
+        #[case] _kind: AK,
+        #[case] _other_kind: OtherAK,
+    ) {
+        let size = 2;
+        let element_ty = bool_t();
+        let op = GenericArrayConvert::<AK, INTO, OtherAK>::new(element_ty.clone(), size);
+        let optype: OpType = op.into();
+        let sig = optype.dataflow_signature().unwrap();
+        assert_eq!(
+            sig.io(),
+            (
+                &vec![AK::ty(size, element_ty.clone())].into(),
+                &vec![OtherAK::ty(size, element_ty.clone())].into(),
+            )
+        );
+    }
+}

hugr-core/src/std_extensions/collections/array/array_op.rs

@@ -31,7 +31,7 @@ pub enum GenericArrayOpDef<AK: ArrayKind> {
     /// where `SIZE` must be statically known (not a variable)
     new_array,
     /// Copies an element out of the array ([TypeBound::Copyable] elements only):
-    /// `get<size,elemty>: array<size, elemty>, index -> option<elemty>`
+    /// `get<size,elemty>: array<size, elemty>, index -> option<elemty>, array`
     get,
     /// Exchanges an element of the array with an external value:
     /// `set<size, elemty>: array<size, elemty>, index, elemty -> either(elemty, array | elemty, array)`
@@ -148,7 +148,10 @@ impl<AK: ArrayKind> GenericArrayOpDef<AK> {
                     let option_type: Type = option_type(copy_elem_ty).into();
                     PolyFuncTypeRV::new(
                         params,
-                        FuncValueType::new(vec![copy_array_ty, usize_t], option_type),
+                        FuncValueType::new(
+                            vec![copy_array_ty.clone(), usize_t],
+                            vec![option_type, copy_array_ty],
+                        ),
                     )
                 }
                 set => {
@@ -385,7 +388,11 @@ mod tests {
             sig.io(),
             (
                 &vec![AK::ty(size, element_ty.clone()), usize_t()].into(),
-                &vec![option_type(element_ty.clone()).into()].into()
+                &vec![
+                    option_type(element_ty.clone()).into(),
+                    AK::ty(size, element_ty.clone())
+                ]
+                .into()
             )
         );
     }
@@ -500,7 +507,11 @@ mod tests {
             sig.io(),
             (
                 &vec![AK::ty(size, element_ty.clone()), usize_t()].into(),
-                &vec![option_type(element_ty.clone()).into()].into()
+                &vec![
+                    option_type(element_ty.clone()).into(),
+                    AK::ty(size, element_ty.clone())
+                ]
+                .into()
             )
         );
     }