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 derive_more::Display;
24use handle::NodeHandle;
25use paste::paste;
26
27use enum_dispatch::enum_dispatch;
28
29pub use constant::{Const, Value};
30pub use controlflow::{BasicBlock, CFG, Case, Conditional, DataflowBlock, ExitBlock, TailLoop};
31pub use custom::{ExtensionOp, OpaqueOp};
32pub use dataflow::{
33    Call, CallIndirect, DFG, DataflowOpTrait, DataflowParent, Input, LoadConstant, LoadFunction,
34    Output,
35};
36pub use module::{AliasDecl, AliasDefn, FuncDecl, FuncDefn, Module};
37use smol_str::SmolStr;
38pub use sum::Tag;
39pub use tag::OpTag;
40
41#[enum_dispatch(OpTrait, NamedOp, ValidateOp, OpParent)]
42#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
43#[cfg_attr(test, derive(proptest_derive::Arbitrary))]
44/// The concrete operation types for a node in the HUGR.
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                #[must_use] 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                #[must_use] 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    #[must_use]
152    pub fn other_port_kind(&self, dir: Direction) -> Option<EdgeKind> {
153        match dir {
154            Direction::Incoming => self.other_input(),
155            Direction::Outgoing => self.other_output(),
156        }
157    }
158
159    /// The edge kind for the static ports of the operation, not described by
160    /// the dataflow signature.
161    ///
162    /// If not None, an extra input port of that kind will be present on the
163    /// given direction after any dataflow ports and before any
164    /// [`OpType::other_port_kind`] ports.
165    #[inline]
166    #[must_use]
167    pub fn static_port_kind(&self, dir: Direction) -> Option<EdgeKind> {
168        match dir {
169            Direction::Incoming => self.static_input(),
170            Direction::Outgoing => self.static_output(),
171        }
172    }
173
174    /// Returns the edge kind for the given port.
175    ///
176    /// The result may be a value port, a static port, or a non-dataflow port.
177    /// See [`OpType::dataflow_signature`], [`OpType::static_port_kind`], and
178    /// [`OpType::other_port_kind`].
179    pub fn port_kind(&self, port: impl Into<Port>) -> Option<EdgeKind> {
180        let signature = self.dataflow_signature().unwrap_or_default();
181        let port: Port = port.into();
182        let dir = port.direction();
183        let port_count = signature.port_count(dir);
184
185        // Dataflow ports
186        if port.index() < port_count {
187            return signature.port_type(port).cloned().map(EdgeKind::Value);
188        }
189
190        // Constant port
191        let static_kind = self.static_port_kind(dir);
192        if port.index() == port_count {
193            if let Some(kind) = static_kind {
194                return Some(kind);
195            }
196        }
197
198        // Non-dataflow ports
199        self.other_port_kind(dir)
200    }
201
202    /// The non-dataflow port for the operation, not described by the signature.
203    /// See `[OpType::other_port_kind]`.
204    ///
205    /// Returns None if there is no such port, or if the operation defines multiple non-dataflow ports.
206    #[must_use]
207    pub fn other_port(&self, dir: Direction) -> Option<Port> {
208        let df_count = self.value_port_count(dir);
209        let non_df_count = self.non_df_port_count(dir);
210        // if there is a static input it comes before the non_df_ports
211        let static_input =
212            usize::from(dir == Direction::Incoming && OpTag::StaticInput.is_superset(self.tag()));
213        if self.other_port_kind(dir).is_some() && non_df_count >= 1 {
214            Some(Port::new(dir, df_count + static_input))
215        } else {
216            None
217        }
218    }
219
220    /// The non-dataflow input port for the operation, not described by the signature.
221    /// See `[OpType::other_port]`.
222    #[inline]
223    #[must_use]
224    pub fn other_input_port(&self) -> Option<IncomingPort> {
225        self.other_port(Direction::Incoming)
226            .map(|p| p.as_incoming().unwrap())
227    }
228
229    /// The non-dataflow output port for the operation, not described by the signature.
230    /// See `[OpType::other_port]`.
231    #[inline]
232    #[must_use]
233    pub fn other_output_port(&self) -> Option<OutgoingPort> {
234        self.other_port(Direction::Outgoing)
235            .map(|p| p.as_outgoing().unwrap())
236    }
237
238    /// If the op has a static port, the port of that input.
239    ///
240    /// See [`OpType::static_input_port`] and [`OpType::static_output_port`].
241    #[inline]
242    #[must_use]
243    pub fn static_port(&self, dir: Direction) -> Option<Port> {
244        self.static_port_kind(dir)?;
245        Some(Port::new(dir, self.value_port_count(dir)))
246    }
247
248    /// If the op has a static input ([`Call`], [`LoadConstant`], and [`LoadFunction`]), the port of
249    /// that input.
250    #[inline]
251    #[must_use]
252    pub fn static_input_port(&self) -> Option<IncomingPort> {
253        self.static_port(Direction::Incoming)
254            .map(|p| p.as_incoming().unwrap())
255    }
256
257    /// If the op has a static output ([`Const`], [`FuncDefn`], [`FuncDecl`]), the port of that output.
258    #[inline]
259    #[must_use]
260    pub fn static_output_port(&self) -> Option<OutgoingPort> {
261        self.static_port(Direction::Outgoing)
262            .map(|p| p.as_outgoing().unwrap())
263    }
264
265    /// The number of Value ports in given direction.
266    #[inline]
267    #[must_use]
268    pub fn value_port_count(&self, dir: portgraph::Direction) -> usize {
269        self.dataflow_signature()
270            .map_or(0, |sig| sig.port_count(dir))
271    }
272
273    /// The number of Value input ports.
274    #[inline]
275    #[must_use]
276    pub fn value_input_count(&self) -> usize {
277        self.value_port_count(Direction::Incoming)
278    }
279
280    /// The number of Value output ports.
281    #[inline]
282    #[must_use]
283    pub fn value_output_count(&self) -> usize {
284        self.value_port_count(Direction::Outgoing)
285    }
286
287    /// Returns the number of ports for the given direction.
288    #[inline]
289    #[must_use]
290    pub fn port_count(&self, dir: Direction) -> usize {
291        let has_static_port = self.static_port_kind(dir).is_some();
292        let non_df_count = self.non_df_port_count(dir);
293        self.value_port_count(dir) + usize::from(has_static_port) + non_df_count
294    }
295
296    /// Returns the number of inputs ports for the operation.
297    #[inline]
298    #[must_use]
299    pub fn input_count(&self) -> usize {
300        self.port_count(Direction::Incoming)
301    }
302
303    /// Returns the number of outputs ports for the operation.
304    #[inline]
305    #[must_use]
306    pub fn output_count(&self) -> usize {
307        self.port_count(Direction::Outgoing)
308    }
309
310    /// Checks whether the operation can contain children nodes.
311    #[inline]
312    #[must_use]
313    pub fn is_container(&self) -> bool {
314        self.validity_flags::<Node>().allowed_children != OpTag::None
315    }
316
317    /// Cast to an extension operation.
318    ///
319    /// Returns `None` if the operation is not of the requested type.
320    pub fn cast<T: MakeExtensionOp>(&self) -> Option<T> {
321        self.as_extension_op().and_then(ExtensionOp::cast)
322    }
323
324    /// Returns the extension where the operation is defined, if any.
325    #[must_use]
326    pub fn extension_id(&self) -> Option<&ExtensionId> {
327        match self {
328            OpType::OpaqueOp(opaque) => Some(opaque.extension()),
329            OpType::ExtensionOp(e) => Some(e.def().extension_id()),
330            _ => None,
331        }
332    }
333
334    /// Returns a registry with all the extensions required by the operation.
335    ///
336    /// This includes the operation extension in [`OpType::extension_id`], and any
337    /// extension required by the operation's signature types.
338    pub fn used_extensions(&self) -> Result<ExtensionRegistry, ExtensionCollectionError> {
339        // Collect extensions on the types.
340        let mut reg = collect_op_types_extensions(None, self)?;
341        // And on the operation definition itself.
342        if let Some(ext) = collect_op_extension(None, self)? {
343            reg.register_updated(ext);
344        }
345        Ok(reg)
346    }
347}
348
349/// Macro used by operations that want their
350/// name to be the same as their type name
351macro_rules! impl_op_name {
352    ($i: ident) => {
353        impl $crate::ops::NamedOp for $i {
354            fn name(&self) -> $crate::ops::OpName {
355                stringify!($i).into()
356            }
357        }
358    };
359}
360
361use impl_op_name;
362
363/// A unique identifier for a operation.
364pub type OpName = SmolStr;
365
366/// Slice of a [`OpName`] operation identifier.
367pub type OpNameRef = str;
368
369#[enum_dispatch]
370/// Trait for setting name of `OpType` variants.
371// Separate to OpTrait to allow simple definition via impl_op_name
372pub(crate) trait NamedOp {
373    /// The name of the operation.
374    fn name(&self) -> OpName;
375}
376
377/// Trait statically querying the tag of an operation.
378///
379/// This is implemented by all `OpType` variants, and always contains the dynamic
380/// tag returned by `OpType::tag(&self)`.
381pub trait StaticTag {
382    /// The name of the operation.
383    const TAG: OpTag;
384}
385
386#[enum_dispatch]
387/// Trait implemented by all `OpType` variants.
388pub trait OpTrait: Sized + Clone {
389    /// A human-readable description of the operation.
390    fn description(&self) -> &str;
391
392    /// Tag identifying the operation.
393    fn tag(&self) -> OpTag;
394
395    /// Tries to create a specific [`NodeHandle`] for a node with this operation
396    /// type.
397    ///
398    /// Fails if the operation's [`OpTrait::tag`] does not match the
399    /// [`NodeHandle::TAG`] of the requested handle.
400    fn try_node_handle<N, H>(&self, node: N) -> Option<H>
401    where
402        N: HugrNode,
403        H: NodeHandle<N> + From<N>,
404    {
405        H::TAG.is_superset(self.tag()).then(|| node.into())
406    }
407
408    /// The signature of the operation.
409    ///
410    /// Only dataflow operations have a signature, otherwise returns None.
411    fn dataflow_signature(&self) -> Option<Cow<'_, Signature>> {
412        None
413    }
414
415    /// The edge kind for the non-dataflow inputs of the operation,
416    /// not described by the signature.
417    ///
418    /// If not None, a single extra input port of that kind will be
419    /// present.
420    fn other_input(&self) -> Option<EdgeKind> {
421        None
422    }
423
424    /// The edge kind for the non-dataflow outputs of the operation, not
425    /// described by the signature.
426    ///
427    /// If not None, a single extra output port of that kind will be
428    /// present.
429    fn other_output(&self) -> Option<EdgeKind> {
430        None
431    }
432
433    /// The edge kind for a single constant input of the operation, not
434    /// described by the dataflow signature.
435    ///
436    /// If not None, an extra input port of that kind will be present after the
437    /// dataflow input ports and before any [`OpTrait::other_input`] ports.
438    fn static_input(&self) -> Option<EdgeKind> {
439        None
440    }
441
442    /// The edge kind for a single constant output of the operation, not
443    /// described by the dataflow signature.
444    ///
445    /// If not None, an extra output port of that kind will be present after the
446    /// dataflow input ports and before any [`OpTrait::other_output`] ports.
447    fn static_output(&self) -> Option<EdgeKind> {
448        None
449    }
450
451    /// Get the number of non-dataflow multiports.
452    fn non_df_port_count(&self, dir: Direction) -> usize {
453        usize::from(
454            match dir {
455                Direction::Incoming => self.other_input(),
456                Direction::Outgoing => self.other_output(),
457            }
458            .is_some(),
459        )
460    }
461
462    /// Apply a type-level substitution to this `OpType`, i.e. replace
463    /// [type variables](crate::types::TypeArg::new_var_use) with new types.
464    fn substitute(&self, _subst: &Substitution) -> Self {
465        self.clone()
466    }
467}
468
469/// Properties of child graphs of ops, if the op has children.
470#[enum_dispatch]
471pub trait OpParent {
472    /// The inner function type of the operation, if it has a child dataflow
473    /// sibling graph.
474    ///
475    /// Non-container ops like `FuncDecl` return `None` even though they represent a function.
476    fn inner_function_type(&self) -> Option<Cow<'_, Signature>> {
477        None
478    }
479}
480
481impl<T: DataflowParent> OpParent for T {
482    fn inner_function_type(&self) -> Option<Cow<'_, Signature>> {
483        Some(DataflowParent::inner_signature(self))
484    }
485}
486
487impl OpParent for Module {}
488impl OpParent for AliasDecl {}
489impl OpParent for AliasDefn {}
490impl OpParent for Const {}
491impl OpParent for Input {}
492impl OpParent for Output {}
493impl OpParent for Call {}
494impl OpParent for CallIndirect {}
495impl OpParent for LoadConstant {}
496impl OpParent for LoadFunction {}
497impl OpParent for ExtensionOp {}
498impl OpParent for OpaqueOp {}
499impl OpParent for Tag {}
500impl OpParent for CFG {}
501impl OpParent for Conditional {}
502impl OpParent for FuncDecl {}
503impl OpParent for ExitBlock {}
504
505#[enum_dispatch]
506/// Methods for Ops to validate themselves and children
507pub trait ValidateOp {
508    /// Returns a set of flags describing the validity predicates for this operation.
509    #[inline]
510    fn validity_flags<N: HugrNode>(&self) -> validate::OpValidityFlags<N> {
511        Default::default()
512    }
513
514    /// Validate the ordered list of children.
515    #[inline]
516    fn validate_op_children<'a, N: HugrNode>(
517        &self,
518        _children: impl DoubleEndedIterator<Item = (N, &'a OpType)>,
519    ) -> Result<(), validate::ChildrenValidationError<N>> {
520        Ok(())
521    }
522}
523
524/// Macro used for default implementation of `ValidateOp`
525macro_rules! impl_validate_op {
526    ($i: ident) => {
527        impl $crate::ops::ValidateOp for $i {}
528    };
529}
530
531use impl_validate_op;
hugr_core/ops.rs

hugr_core/
ops.rs
