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    ///
305    /// Returns `None` if the operation is not of the requested type.
306    pub fn cast<T: MakeExtensionOp>(&self) -> Option<T> {
307        self.as_extension_op().and_then(ExtensionOp::cast)
308    }
309
310    /// Returns the extension where the operation is defined, if any.
311    pub fn extension_id(&self) -> Option<&ExtensionId> {
312        match self {
313            OpType::OpaqueOp(opaque) => Some(opaque.extension()),
314            OpType::ExtensionOp(e) => Some(e.def().extension_id()),
315            _ => None,
316        }
317    }
318
319    /// Returns a registry with all the extensions required by the operation.
320    ///
321    /// This includes the operation extension in [`OpType::extension_id`], and any
322    /// extension required by the operation's signature types.
323    pub fn used_extensions(&self) -> Result<ExtensionRegistry, ExtensionCollectionError> {
324        // Collect extensions on the types.
325        let mut reg = collect_op_types_extensions(None, self)?;
326        // And on the operation definition itself.
327        if let Some(ext) = collect_op_extension(None, self)? {
328            reg.register_updated(ext);
329        }
330        Ok(reg)
331    }
332}
333
334/// Macro used by operations that want their
335/// name to be the same as their type name
336macro_rules! impl_op_name {
337    ($i: ident) => {
338        impl $crate::ops::NamedOp for $i {
339            fn name(&self) -> $crate::ops::OpName {
340                stringify!($i).into()
341            }
342        }
343    };
344}
345
346use impl_op_name;
347
348/// A unique identifier for a operation.
349pub type OpName = SmolStr;
350
351/// Slice of a [`OpName`] operation identifier.
352pub type OpNameRef = str;
353
354#[enum_dispatch]
355/// Trait for setting name of OpType variants.
356// Separate to OpTrait to allow simple definition via impl_op_name
357pub trait NamedOp {
358    /// The name of the operation.
359    fn name(&self) -> OpName;
360}
361
362/// Trait statically querying the tag of an operation.
363///
364/// This is implemented by all OpType variants, and always contains the dynamic
365/// tag returned by `OpType::tag(&self)`.
366pub trait StaticTag {
367    /// The name of the operation.
368    const TAG: OpTag;
369}
370
371#[enum_dispatch]
372/// Trait implemented by all OpType variants.
373pub trait OpTrait: Sized + Clone {
374    /// A human-readable description of the operation.
375    fn description(&self) -> &str;
376
377    /// Tag identifying the operation.
378    fn tag(&self) -> OpTag;
379
380    /// The signature of the operation.
381    ///
382    /// Only dataflow operations have a signature, otherwise returns None.
383    fn dataflow_signature(&self) -> Option<Cow<'_, Signature>> {
384        None
385    }
386
387    /// The delta between the input extensions specified for a node,
388    /// and the output extensions calculated for that node
389    fn extension_delta(&self) -> ExtensionSet {
390        ExtensionSet::new()
391    }
392
393    /// The edge kind for the non-dataflow inputs of the operation,
394    /// not described by the signature.
395    ///
396    /// If not None, a single extra input port of that kind will be
397    /// present.
398    fn other_input(&self) -> Option<EdgeKind> {
399        None
400    }
401
402    /// The edge kind for the non-dataflow outputs of the operation, not
403    /// described by the signature.
404    ///
405    /// If not None, a single extra output port of that kind will be
406    /// present.
407    fn other_output(&self) -> Option<EdgeKind> {
408        None
409    }
410
411    /// The edge kind for a single constant input of the operation, not
412    /// described by the dataflow signature.
413    ///
414    /// If not None, an extra input port of that kind will be present after the
415    /// dataflow input ports and before any [`OpTrait::other_input`] ports.
416    fn static_input(&self) -> Option<EdgeKind> {
417        None
418    }
419
420    /// The edge kind for a single constant output of the operation, not
421    /// described by the dataflow signature.
422    ///
423    /// If not None, an extra output port of that kind will be present after the
424    /// dataflow input ports and before any [`OpTrait::other_output`] ports.
425    fn static_output(&self) -> Option<EdgeKind> {
426        None
427    }
428
429    /// Get the number of non-dataflow multiports.
430    fn non_df_port_count(&self, dir: Direction) -> usize {
431        match dir {
432            Direction::Incoming => self.other_input(),
433            Direction::Outgoing => self.other_output(),
434        }
435        .is_some() as usize
436    }
437
438    /// Apply a type-level substitution to this OpType, i.e. replace
439    /// [type variables](crate::types::TypeArg::new_var_use) with new types.
440    fn substitute(&self, _subst: &Substitution) -> Self {
441        self.clone()
442    }
443}
444
445/// Properties of child graphs of ops, if the op has children.
446#[enum_dispatch]
447pub trait OpParent {
448    /// The inner function type of the operation, if it has a child dataflow
449    /// sibling graph.
450    ///
451    /// Non-container ops like `FuncDecl` return `None` even though they represent a function.
452    fn inner_function_type(&self) -> Option<Cow<'_, Signature>> {
453        None
454    }
455}
456
457impl<T: DataflowParent> OpParent for T {
458    fn inner_function_type(&self) -> Option<Cow<'_, Signature>> {
459        Some(DataflowParent::inner_signature(self))
460    }
461}
462
463impl OpParent for Module {}
464impl OpParent for AliasDecl {}
465impl OpParent for AliasDefn {}
466impl OpParent for Const {}
467impl OpParent for Input {}
468impl OpParent for Output {}
469impl OpParent for Call {}
470impl OpParent for CallIndirect {}
471impl OpParent for LoadConstant {}
472impl OpParent for LoadFunction {}
473impl OpParent for ExtensionOp {}
474impl OpParent for OpaqueOp {}
475impl OpParent for Tag {}
476impl OpParent for CFG {}
477impl OpParent for Conditional {}
478impl OpParent for FuncDecl {}
479impl OpParent for ExitBlock {}
480
481#[enum_dispatch]
482/// Methods for Ops to validate themselves and children
483pub trait ValidateOp {
484    /// Returns a set of flags describing the validity predicates for this operation.
485    #[inline]
486    fn validity_flags(&self) -> validate::OpValidityFlags {
487        Default::default()
488    }
489
490    /// Validate the ordered list of children.
491    #[inline]
492    fn validate_op_children<'a>(
493        &self,
494        _children: impl DoubleEndedIterator<Item = (NodeIndex, &'a OpType)>,
495    ) -> Result<(), validate::ChildrenValidationError> {
496        Ok(())
497    }
498}
499
500/// Macro used for default implementation of ValidateOp
501macro_rules! impl_validate_op {
502    ($i: ident) => {
503        impl $crate::ops::ValidateOp for $i {}
504    };
505}
506
507use impl_validate_op;
hugr_core/ops.rs

hugr_core/
ops.rs
