hugr_core/
ops.rs

1//! The operation types for the HUGR.
2
3pub mod constant;
4pub mod controlflow;
5pub mod custom;
6pub mod dataflow;
7pub mod handle;
8pub mod module;
9pub mod sum;
10pub mod tag;
11pub mod validate;
12use crate::extension::resolution::{
13    collect_op_extension, collect_op_types_extensions, ExtensionCollectionError,
14};
15use std::borrow::Cow;
16
17use crate::extension::simple_op::MakeExtensionOp;
18use crate::extension::{ExtensionId, ExtensionRegistry, ExtensionSet};
19use crate::types::{EdgeKind, Signature, Substitution};
20use crate::{Direction, OutgoingPort, Port};
21use crate::{IncomingPort, PortIndex};
22use derive_more::Display;
23use paste::paste;
24use portgraph::NodeIndex;
25
26use enum_dispatch::enum_dispatch;
27
28pub use constant::{Const, Value};
29pub use controlflow::{BasicBlock, Case, Conditional, DataflowBlock, ExitBlock, TailLoop, CFG};
30pub use custom::{ExtensionOp, OpaqueOp};
31pub use dataflow::{
32    Call, CallIndirect, DataflowOpTrait, DataflowParent, Input, LoadConstant, LoadFunction, Output,
33    DFG,
34};
35pub use module::{AliasDecl, AliasDefn, FuncDecl, FuncDefn, Module};
36use smol_str::SmolStr;
37pub use sum::Tag;
38pub use tag::OpTag;
39
40#[enum_dispatch(OpTrait, NamedOp, ValidateOp, OpParent)]
41#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
42#[cfg_attr(test, derive(proptest_derive::Arbitrary))]
43/// The concrete operation types for a node in the HUGR.
44// TODO: Link the NodeHandles to the OpType.
45#[non_exhaustive]
46#[allow(missing_docs)]
47#[serde(tag = "op")]
48pub enum OpType {
49    Module,
50    FuncDefn,
51    FuncDecl,
52    AliasDecl,
53    AliasDefn,
54    Const,
55    Input,
56    Output,
57    Call,
58    CallIndirect,
59    LoadConstant,
60    LoadFunction,
61    DFG,
62    #[serde(skip_deserializing, rename = "Extension")]
63    ExtensionOp,
64    #[serde(rename = "Extension")]
65    OpaqueOp,
66    Tag,
67    DataflowBlock,
68    ExitBlock,
69    TailLoop,
70    CFG,
71    Conditional,
72    Case,
73}
74
75macro_rules! impl_op_ref_try_into {
76    ($Op: tt, $sname:ident) => {
77        paste! {
78            impl OpType {
79                #[doc = "If is an instance of `" $Op "` return a reference to it."]
80                pub fn [<as_ $sname:snake>](&self) -> Option<&$Op> {
81                    TryInto::<&$Op>::try_into(self).ok()
82                }
83
84                #[doc = "Returns `true` if the operation is an instance of `" $Op "`."]
85                pub fn [<is_ $sname:snake>](&self) -> bool {
86                    self.[<as_ $sname:snake>]().is_some()
87                }
88            }
89
90            impl<'a> TryFrom<&'a OpType> for &'a $Op {
91                type Error = ();
92                fn try_from(optype: &'a OpType) -> Result<Self, Self::Error> {
93                    if let OpType::$Op(l) = optype {
94                        Ok(l)
95                    } else {
96                        Err(())
97                    }
98                }
99            }
100        }
101    };
102    ($Op:tt) => {
103        impl_op_ref_try_into!($Op, $Op);
104    };
105}
106
107impl_op_ref_try_into!(Module);
108impl_op_ref_try_into!(FuncDefn);
109impl_op_ref_try_into!(FuncDecl);
110impl_op_ref_try_into!(AliasDecl);
111impl_op_ref_try_into!(AliasDefn);
112impl_op_ref_try_into!(Const);
113impl_op_ref_try_into!(Input);
114impl_op_ref_try_into!(Output);
115impl_op_ref_try_into!(Call);
116impl_op_ref_try_into!(CallIndirect);
117impl_op_ref_try_into!(LoadConstant);
118impl_op_ref_try_into!(LoadFunction);
119impl_op_ref_try_into!(DFG, dfg);
120impl_op_ref_try_into!(ExtensionOp);
121impl_op_ref_try_into!(Tag);
122impl_op_ref_try_into!(DataflowBlock);
123impl_op_ref_try_into!(ExitBlock);
124impl_op_ref_try_into!(TailLoop);
125impl_op_ref_try_into!(CFG, cfg);
126impl_op_ref_try_into!(Conditional);
127impl_op_ref_try_into!(Case);
128
129/// The default OpType (as returned by [Default::default])
130pub const DEFAULT_OPTYPE: OpType = OpType::Module(Module::new());
131
132impl Default for OpType {
133    fn default() -> Self {
134        DEFAULT_OPTYPE
135    }
136}
137
138impl Display for OpType {
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        write!(f, "{}", self.name())
141    }
142}
143
144impl OpType {
145    /// The edge kind for the non-dataflow ports of the operation, not described
146    /// by the signature.
147    ///
148    /// If not None, a single extra port of that kind will be present on
149    /// the given direction after any dataflow or constant ports.
150    #[inline]
151    pub fn other_port_kind(&self, dir: Direction) -> Option<EdgeKind> {
152        match dir {
153            Direction::Incoming => self.other_input(),
154            Direction::Outgoing => self.other_output(),
155        }
156    }
157
158    /// The edge kind for the static ports of the operation, not described by
159    /// the dataflow signature.
160    ///
161    /// If not None, an extra input port of that kind will be present on the
162    /// given direction after any dataflow ports and before any
163    /// [`OpType::other_port_kind`] ports.
164    #[inline]
165    pub fn static_port_kind(&self, dir: Direction) -> Option<EdgeKind> {
166        match dir {
167            Direction::Incoming => self.static_input(),
168            Direction::Outgoing => self.static_output(),
169        }
170    }
171
172    /// Returns the edge kind for the given port.
173    ///
174    /// The result may be a value port, a static port, or a non-dataflow port.
175    /// See [`OpType::dataflow_signature`], [`OpType::static_port_kind`], and
176    /// [`OpType::other_port_kind`].
177    pub fn port_kind(&self, port: impl Into<Port>) -> Option<EdgeKind> {
178        let signature = self.dataflow_signature().unwrap_or_default();
179        let port: Port = port.into();
180        let dir = port.direction();
181        let port_count = signature.port_count(dir);
182
183        // Dataflow ports
184        if port.index() < port_count {
185            return signature.port_type(port).cloned().map(EdgeKind::Value);
186        }
187
188        // Constant port
189        let static_kind = self.static_port_kind(dir);
190        if port.index() == port_count {
191            if let Some(kind) = static_kind {
192                return Some(kind);
193            }
194        }
195
196        // Non-dataflow ports
197        self.other_port_kind(dir)
198    }
199
200    /// The non-dataflow port for the operation, not described by the signature.
201    /// See `[OpType::other_port_kind]`.
202    ///
203    /// Returns None if there is no such port, or if the operation defines multiple non-dataflow ports.
204    pub fn other_port(&self, dir: Direction) -> Option<Port> {
205        let df_count = self.value_port_count(dir);
206        let non_df_count = self.non_df_port_count(dir);
207        // if there is a static input it comes before the non_df_ports
208        let static_input =
209            (dir == Direction::Incoming && OpTag::StaticInput.is_superset(self.tag())) as usize;
210        if self.other_port_kind(dir).is_some() && non_df_count >= 1 {
211            Some(Port::new(dir, df_count + static_input))
212        } else {
213            None
214        }
215    }
216
217    /// The non-dataflow input port for the operation, not described by the signature.
218    /// See `[OpType::other_port]`.
219    #[inline]
220    pub fn other_input_port(&self) -> Option<IncomingPort> {
221        self.other_port(Direction::Incoming)
222            .map(|p| p.as_incoming().unwrap())
223    }
224
225    /// The non-dataflow output port for the operation, not described by the signature.
226    /// See `[OpType::other_port]`.
227    #[inline]
228    pub fn other_output_port(&self) -> Option<OutgoingPort> {
229        self.other_port(Direction::Outgoing)
230            .map(|p| p.as_outgoing().unwrap())
231    }
232
233    /// If the op has a static port, the port of that input.
234    ///
235    /// See [`OpType::static_input_port`] and [`OpType::static_output_port`].
236    #[inline]
237    pub fn static_port(&self, dir: Direction) -> Option<Port> {
238        self.static_port_kind(dir)?;
239        Some(Port::new(dir, self.value_port_count(dir)))
240    }
241
242    /// If the op has a static input ([`Call`], [`LoadConstant`], and [`LoadFunction`]), the port of
243    /// that input.
244    #[inline]
245    pub fn static_input_port(&self) -> Option<IncomingPort> {
246        self.static_port(Direction::Incoming)
247            .map(|p| p.as_incoming().unwrap())
248    }
249
250    /// If the op has a static output ([`Const`], [`FuncDefn`], [`FuncDecl`]), the port of that output.
251    #[inline]
252    pub fn static_output_port(&self) -> Option<OutgoingPort> {
253        self.static_port(Direction::Outgoing)
254            .map(|p| p.as_outgoing().unwrap())
255    }
256
257    /// The number of Value ports in given direction.
258    #[inline]
259    pub fn value_port_count(&self, dir: portgraph::Direction) -> usize {
260        self.dataflow_signature()
261            .map(|sig| sig.port_count(dir))
262            .unwrap_or(0)
263    }
264
265    /// The number of Value input ports.
266    #[inline]
267    pub fn value_input_count(&self) -> usize {
268        self.value_port_count(Direction::Incoming)
269    }
270
271    /// The number of Value output ports.
272    #[inline]
273    pub fn value_output_count(&self) -> usize {
274        self.value_port_count(Direction::Outgoing)
275    }
276
277    /// Returns the number of ports for the given direction.
278    #[inline]
279    pub fn port_count(&self, dir: Direction) -> usize {
280        let has_static_port = self.static_port_kind(dir).is_some();
281        let non_df_count = self.non_df_port_count(dir);
282        self.value_port_count(dir) + has_static_port as usize + non_df_count
283    }
284
285    /// Returns the number of inputs ports for the operation.
286    #[inline]
287    pub fn input_count(&self) -> usize {
288        self.port_count(Direction::Incoming)
289    }
290
291    /// Returns the number of outputs ports for the operation.
292    #[inline]
293    pub fn output_count(&self) -> usize {
294        self.port_count(Direction::Outgoing)
295    }
296
297    /// Checks whether the operation can contain children nodes.
298    #[inline]
299    pub fn is_container(&self) -> bool {
300        self.validity_flags().allowed_children != OpTag::None
301    }
302
303    /// Cast to an extension operation.
304    pub fn cast<T: MakeExtensionOp>(&self) -> Option<T> {
305        self.as_extension_op()
306            .and_then(|o| T::from_extension_op(o).ok())
307    }
308
309    /// Returns the extension where the operation is defined, if any.
310    pub fn extension_id(&self) -> Option<&ExtensionId> {
311        match self {
312            OpType::OpaqueOp(opaque) => Some(opaque.extension()),
313            OpType::ExtensionOp(e) => Some(e.def().extension_id()),
314            _ => None,
315        }
316    }
317
318    /// Returns a registry with all the extensions required by the operation.
319    ///
320    /// This includes the operation extension in [`OpType::extension_id`], and any
321    /// extension required by the operation's signature types.
322    pub fn used_extensions(&self) -> Result<ExtensionRegistry, ExtensionCollectionError> {
323        // Collect extensions on the types.
324        let mut reg = collect_op_types_extensions(None, self)?;
325        // And on the operation definition itself.
326        if let Some(ext) = collect_op_extension(None, self)? {
327            reg.register_updated(ext);
328        }
329        Ok(reg)
330    }
331}
332
333/// Macro used by operations that want their
334/// name to be the same as their type name
335macro_rules! impl_op_name {
336    ($i: ident) => {
337        impl $crate::ops::NamedOp for $i {
338            fn name(&self) -> $crate::ops::OpName {
339                stringify!($i).into()
340            }
341        }
342    };
343}
344
345use impl_op_name;
346
347/// A unique identifier for a operation.
348pub type OpName = SmolStr;
349
350/// Slice of a [`OpName`] operation identifier.
351pub type OpNameRef = str;
352
353#[enum_dispatch]
354/// Trait for setting name of OpType variants.
355// Separate to OpTrait to allow simple definition via impl_op_name
356pub trait NamedOp {
357    /// The name of the operation.
358    fn name(&self) -> OpName;
359}
360
361/// Trait statically querying the tag of an operation.
362///
363/// This is implemented by all OpType variants, and always contains the dynamic
364/// tag returned by `OpType::tag(&self)`.
365pub trait StaticTag {
366    /// The name of the operation.
367    const TAG: OpTag;
368}
369
370#[enum_dispatch]
371/// Trait implemented by all OpType variants.
372pub trait OpTrait: Sized + Clone {
373    /// A human-readable description of the operation.
374    fn description(&self) -> &str;
375
376    /// Tag identifying the operation.
377    fn tag(&self) -> OpTag;
378
379    /// The signature of the operation.
380    ///
381    /// Only dataflow operations have a signature, otherwise returns None.
382    fn dataflow_signature(&self) -> Option<Cow<'_, Signature>> {
383        None
384    }
385
386    /// The delta between the input extensions specified for a node,
387    /// and the output extensions calculated for that node
388    fn extension_delta(&self) -> ExtensionSet {
389        ExtensionSet::new()
390    }
391
392    /// The edge kind for the non-dataflow inputs of the operation,
393    /// not described by the signature.
394    ///
395    /// If not None, a single extra input port of that kind will be
396    /// present.
397    fn other_input(&self) -> Option<EdgeKind> {
398        None
399    }
400
401    /// The edge kind for the non-dataflow outputs of the operation, not
402    /// described by the signature.
403    ///
404    /// If not None, a single extra output port of that kind will be
405    /// present.
406    fn other_output(&self) -> Option<EdgeKind> {
407        None
408    }
409
410    /// The edge kind for a single constant input of the operation, not
411    /// described by the dataflow signature.
412    ///
413    /// If not None, an extra input port of that kind will be present after the
414    /// dataflow input ports and before any [`OpTrait::other_input`] ports.
415    fn static_input(&self) -> Option<EdgeKind> {
416        None
417    }
418
419    /// The edge kind for a single constant output of the operation, not
420    /// described by the dataflow signature.
421    ///
422    /// If not None, an extra output port of that kind will be present after the
423    /// dataflow input ports and before any [`OpTrait::other_output`] ports.
424    fn static_output(&self) -> Option<EdgeKind> {
425        None
426    }
427
428    /// Get the number of non-dataflow multiports.
429    fn non_df_port_count(&self, dir: Direction) -> usize {
430        match dir {
431            Direction::Incoming => self.other_input(),
432            Direction::Outgoing => self.other_output(),
433        }
434        .is_some() as usize
435    }
436
437    /// Apply a type-level substitution to this OpType, i.e. replace
438    /// [type variables](crate::types::TypeArg::new_var_use) with new types.
439    fn substitute(&self, _subst: &Substitution) -> Self {
440        self.clone()
441    }
442}
443
444/// Properties of child graphs of ops, if the op has children.
445#[enum_dispatch]
446pub trait OpParent {
447    /// The inner function type of the operation, if it has a child dataflow
448    /// sibling graph.
449    ///
450    /// Non-container ops like `FuncDecl` return `None` even though they represent a function.
451    fn inner_function_type(&self) -> Option<Cow<'_, Signature>> {
452        None
453    }
454}
455
456impl<T: DataflowParent> OpParent for T {
457    fn inner_function_type(&self) -> Option<Cow<'_, Signature>> {
458        Some(DataflowParent::inner_signature(self))
459    }
460}
461
462impl OpParent for Module {}
463impl OpParent for AliasDecl {}
464impl OpParent for AliasDefn {}
465impl OpParent for Const {}
466impl OpParent for Input {}
467impl OpParent for Output {}
468impl OpParent for Call {}
469impl OpParent for CallIndirect {}
470impl OpParent for LoadConstant {}
471impl OpParent for LoadFunction {}
472impl OpParent for ExtensionOp {}
473impl OpParent for OpaqueOp {}
474impl OpParent for Tag {}
475impl OpParent for CFG {}
476impl OpParent for Conditional {}
477impl OpParent for FuncDecl {}
478impl OpParent for ExitBlock {}
479
480#[enum_dispatch]
481/// Methods for Ops to validate themselves and children
482pub trait ValidateOp {
483    /// Returns a set of flags describing the validity predicates for this operation.
484    #[inline]
485    fn validity_flags(&self) -> validate::OpValidityFlags {
486        Default::default()
487    }
488
489    /// Validate the ordered list of children.
490    #[inline]
491    fn validate_op_children<'a>(
492        &self,
493        _children: impl DoubleEndedIterator<Item = (NodeIndex, &'a OpType)>,
494    ) -> Result<(), validate::ChildrenValidationError> {
495        Ok(())
496    }
497}
498
499/// Macro used for default implementation of ValidateOp
500macro_rules! impl_validate_op {
501    ($i: ident) => {
502        impl $crate::ops::ValidateOp for $i {}
503    };
504}
505
506use impl_validate_op;
hugr_core/ops.rs

hugr_core/
ops.rs
