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

hugr_core/
ops.rs
