hugr_core/std_extensions/collections/array/
array_discard.rs

1//! Definition of the array discard operation.
2
3use std::marker::PhantomData;
4use std::str::FromStr;
5use std::sync::{Arc, Weak};
6
7use crate::extension::simple_op::{
8    HasConcrete, HasDef, MakeExtensionOp, MakeOpDef, MakeRegisteredOp, OpLoadError,
9};
10use crate::extension::{ExtensionId, OpDef, SignatureError, SignatureFunc, TypeDef};
11use crate::ops::{ExtensionOp, OpName};
12use crate::types::type_param::{TypeArg, TypeParam};
13use crate::types::{FuncValueType, PolyFuncTypeRV, Type, TypeBound};
14use crate::{Extension, type_row};
15
16use super::array_kind::ArrayKind;
17
18/// Name of the operation to discard an array
19pub const ARRAY_DISCARD_OP_ID: OpName = OpName::new_inline("discard");
20
21/// Definition of the array discard op. Generic over the concrete array implementation.
22#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
23pub struct GenericArrayDiscardDef<AK: ArrayKind>(PhantomData<AK>);
24
25impl<AK: ArrayKind> GenericArrayDiscardDef<AK> {
26    /// Creates a new array discard operation definition.
27    #[must_use]
28    pub fn new() -> Self {
29        GenericArrayDiscardDef(PhantomData)
30    }
31}
32
33impl<AK: ArrayKind> Default for GenericArrayDiscardDef<AK> {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl<AK: ArrayKind> FromStr for GenericArrayDiscardDef<AK> {
40    type Err = ();
41
42    fn from_str(s: &str) -> Result<Self, Self::Err> {
43        if s == ARRAY_DISCARD_OP_ID {
44            Ok(GenericArrayDiscardDef::new())
45        } else {
46            Err(())
47        }
48    }
49}
50
51impl<AK: ArrayKind> GenericArrayDiscardDef<AK> {
52    /// To avoid recursion when defining the extension, take the type definition as an argument.
53    fn signature_from_def(&self, array_def: &TypeDef) -> SignatureFunc {
54        let params = vec![TypeParam::max_nat_type(), TypeBound::Copyable.into()];
55        let size = TypeArg::new_var_use(0, TypeParam::max_nat_type());
56        let element_ty = Type::new_var_use(1, TypeBound::Copyable);
57        let array_ty = AK::instantiate_ty(array_def, size, element_ty)
58            .expect("Array type instantiation failed");
59        PolyFuncTypeRV::new(params, FuncValueType::new(array_ty, type_row![])).into()
60    }
61}
62
63impl<AK: ArrayKind> MakeOpDef for GenericArrayDiscardDef<AK> {
64    fn opdef_id(&self) -> OpName {
65        ARRAY_DISCARD_OP_ID
66    }
67
68    fn from_def(op_def: &OpDef) -> Result<Self, OpLoadError>
69    where
70        Self: Sized,
71    {
72        crate::extension::simple_op::try_from_name(op_def.name(), op_def.extension_id())
73    }
74
75    fn init_signature(&self, _extension_ref: &Weak<Extension>) -> SignatureFunc {
76        self.signature_from_def(AK::type_def())
77    }
78
79    fn extension_ref(&self) -> Weak<Extension> {
80        Arc::downgrade(AK::extension())
81    }
82
83    fn extension(&self) -> ExtensionId {
84        AK::EXTENSION_ID
85    }
86
87    fn description(&self) -> String {
88        "Discards an array with copyable elements".into()
89    }
90
91    /// Add an operation implemented as a [`MakeOpDef`], which can provide the data
92    /// required to define an [`OpDef`], to an extension.
93    //
94    // This method is re-defined here since we need to pass the array type def while
95    // computing the signature, to avoid recursive loops initializing the extension.
96    fn add_to_extension(
97        &self,
98        extension: &mut Extension,
99        extension_ref: &Weak<Extension>,
100    ) -> Result<(), crate::extension::ExtensionBuildError> {
101        let sig = self.signature_from_def(extension.get_type(&AK::TYPE_NAME).unwrap());
102        let def = extension.add_op(self.opdef_id(), self.description(), sig, extension_ref)?;
103        self.post_opdef(def);
104        Ok(())
105    }
106}
107
108/// Definition of the array discard op. Generic over the concrete array implementation.
109#[derive(Clone, Debug, PartialEq)]
110pub struct GenericArrayDiscard<AK: ArrayKind> {
111    /// The element type of the array.
112    pub elem_ty: Type,
113    /// Size of the array.
114    pub size: u64,
115    _kind: PhantomData<AK>,
116}
117
118impl<AK: ArrayKind> GenericArrayDiscard<AK> {
119    /// Creates a new array discard op.
120    #[must_use]
121    pub fn new(elem_ty: Type, size: u64) -> Option<Self> {
122        elem_ty.copyable().then_some(GenericArrayDiscard {
123            elem_ty,
124            size,
125            _kind: PhantomData,
126        })
127    }
128}
129
130impl<AK: ArrayKind> MakeExtensionOp for GenericArrayDiscard<AK> {
131    fn op_id(&self) -> OpName {
132        GenericArrayDiscardDef::<AK>::default().opdef_id()
133    }
134
135    fn from_extension_op(ext_op: &ExtensionOp) -> Result<Self, OpLoadError>
136    where
137        Self: Sized,
138    {
139        let def = GenericArrayDiscardDef::<AK>::from_def(ext_op.def())?;
140        def.instantiate(ext_op.args())
141    }
142
143    fn type_args(&self) -> Vec<TypeArg> {
144        vec![self.size.into(), self.elem_ty.clone().into()]
145    }
146}
147
148impl<AK: ArrayKind> MakeRegisteredOp for GenericArrayDiscard<AK> {
149    fn extension_id(&self) -> ExtensionId {
150        AK::EXTENSION_ID
151    }
152
153    fn extension_ref(&self) -> Weak<Extension> {
154        Arc::downgrade(AK::extension())
155    }
156}
157
158impl<AK: ArrayKind> HasDef for GenericArrayDiscard<AK> {
159    type Def = GenericArrayDiscardDef<AK>;
160}
161
162impl<AK: ArrayKind> HasConcrete for GenericArrayDiscardDef<AK> {
163    type Concrete = GenericArrayDiscard<AK>;
164
165    fn instantiate(&self, type_args: &[TypeArg]) -> Result<Self::Concrete, OpLoadError> {
166        match type_args {
167            [TypeArg::BoundedNat(n), TypeArg::Runtime(ty)] if ty.copyable() => {
168                Ok(GenericArrayDiscard::new(ty.clone(), *n).unwrap())
169            }
170            _ => Err(SignatureError::InvalidTypeArgs.into()),
171        }
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use rstest::rstest;
178
179    use crate::extension::prelude::bool_t;
180    use crate::std_extensions::collections::array::Array;
181    use crate::std_extensions::collections::borrow_array::BorrowArray;
182    use crate::{
183        extension::prelude::qb_t,
184        ops::{OpTrait, OpType},
185    };
186
187    use super::*;
188
189    #[rstest]
190    #[case(Array)]
191    #[case(BorrowArray)]
192    fn test_discard_def<AK: ArrayKind>(#[case] _kind: AK) {
193        let op = GenericArrayDiscard::<AK>::new(bool_t(), 2).unwrap();
194        let optype: OpType = op.clone().into();
195        let new_op: GenericArrayDiscard<AK> = optype.cast().unwrap();
196        assert_eq!(new_op, op);
197
198        assert_eq!(GenericArrayDiscard::<AK>::new(qb_t(), 2), None);
199    }
200
201    #[rstest]
202    #[case(Array)]
203    #[case(BorrowArray)]
204    fn test_discard<AK: ArrayKind>(#[case] _kind: AK) {
205        let size = 2;
206        let element_ty = bool_t();
207        let op = GenericArrayDiscard::<AK>::new(element_ty.clone(), size).unwrap();
208        let optype: OpType = op.into();
209        let sig = optype.dataflow_signature().unwrap();
210        assert_eq!(
211            sig.io(),
212            (
213                &vec![AK::ty(size, element_ty.clone())].into(),
214                &vec![].into(),
215            )
216        );
217    }
218}
hugr_core/std_extensions/collections/array/array_discard.rs

hugr_core/std_extensions/collections/array/
array_discard.rs
