hugr_core/ops/
handle.rs

1//! Handles to nodes in HUGR.
2use crate::types::{Type, TypeBound};
3use crate::Node;
4
5use derive_more::From as DerFrom;
6use smol_str::SmolStr;
7
8use super::{AliasDecl, OpTag};
9
10/// Common trait for handles to a node.
11/// Typically wrappers around [`Node`].
12pub trait NodeHandle: Clone {
13    /// The most specific operation tag associated with the handle.
14    const TAG: OpTag;
15
16    /// Index of underlying node.
17    fn node(&self) -> Node;
18
19    /// Operation tag for the handle.
20    #[inline]
21    fn tag(&self) -> OpTag {
22        Self::TAG
23    }
24
25    /// Cast the handle to a different more general tag.
26    fn try_cast<T: NodeHandle + From<Node>>(&self) -> Option<T> {
27        T::TAG.is_superset(Self::TAG).then(|| self.node().into())
28    }
29
30    /// Checks whether the handle can hold an operation with the given tag.
31    fn can_hold(tag: OpTag) -> bool {
32        Self::TAG.is_superset(tag)
33    }
34}
35
36/// Trait for handles that contain children.
37///
38/// The allowed children handles are defined by the associated type.
39pub trait ContainerHandle: NodeHandle {
40    /// Handle type for the children of this node.
41    type ChildrenHandle: NodeHandle;
42}
43
44#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
45/// Handle to a [DataflowOp](crate::ops::dataflow).
46pub struct DataflowOpID(Node);
47
48#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
49/// Handle to a [DFG](crate::ops::DFG) node.
50pub struct DfgID(Node);
51
52#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
53/// Handle to a [CFG](crate::ops::CFG) node.
54pub struct CfgID(Node);
55
56#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
57/// Handle to a module [Module](crate::ops::Module) node.
58pub struct ModuleRootID(Node);
59
60#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
61/// Handle to a [module op](crate::ops::module) node.
62pub struct ModuleID(Node);
63
64#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
65/// Handle to a [def](crate::ops::OpType::FuncDefn)
66/// or [declare](crate::ops::OpType::FuncDecl) node.
67///
68/// The `DEF` const generic is used to indicate whether the function is
69/// defined or just declared.
70pub struct FuncID<const DEF: bool>(Node);
71
72#[derive(Debug, Clone, PartialEq, Eq)]
73/// Handle to an [AliasDefn](crate::ops::OpType::AliasDefn)
74/// or [AliasDecl](crate::ops::OpType::AliasDecl) node.
75///
76/// The `DEF` const generic is used to indicate whether the function is
77/// defined or just declared.
78pub struct AliasID<const DEF: bool> {
79    node: Node,
80    name: SmolStr,
81    bound: TypeBound,
82}
83
84impl<const DEF: bool> AliasID<DEF> {
85    /// Construct new AliasID
86    pub fn new(node: Node, name: SmolStr, bound: TypeBound) -> Self {
87        Self { node, name, bound }
88    }
89
90    /// Construct new AliasID
91    pub fn get_alias_type(&self) -> Type {
92        Type::new_alias(AliasDecl::new(self.name.clone(), self.bound))
93    }
94    /// Retrieve the underlying core type
95    pub fn get_name(&self) -> &SmolStr {
96        &self.name
97    }
98}
99
100#[derive(DerFrom, Debug, Clone, PartialEq, Eq)]
101/// Handle to a [Const](crate::ops::OpType::Const) node.
102pub struct ConstID(Node);
103
104#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
105/// Handle to a [DataflowBlock](crate::ops::DataflowBlock) or [Exit](crate::ops::ExitBlock) node.
106pub struct BasicBlockID(Node);
107
108#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
109/// Handle to a [Case](crate::ops::Case) node.
110pub struct CaseID(Node);
111
112#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
113/// Handle to a [TailLoop](crate::ops::TailLoop) node.
114pub struct TailLoopID(Node);
115
116#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
117/// Handle to a [Conditional](crate::ops::Conditional) node.
118pub struct ConditionalID(Node);
119
120#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, DerFrom, Debug)]
121/// Handle to a dataflow container node.
122pub struct DataflowParentID(Node);
123
124/// Implements the `NodeHandle` trait for a tuple struct that contains just a
125/// NodeIndex. Takes the name of the struct, and the corresponding OpTag.
126///
127/// Optionally, the name of the field containing the NodeIndex can be specified
128/// as a third argument. Otherwise, it is assumed to be a tuple struct 0th item.
129macro_rules! impl_nodehandle {
130    ($name:ident, $tag:expr) => {
131        impl_nodehandle!($name, $tag, 0);
132    };
133    ($name:ident, $tag:expr, $node_attr:tt) => {
134        impl NodeHandle for $name {
135            const TAG: OpTag = $tag;
136
137            #[inline]
138            fn node(&self) -> Node {
139                self.$node_attr
140            }
141        }
142    };
143}
144
145impl_nodehandle!(DataflowParentID, OpTag::DataflowParent);
146impl_nodehandle!(DataflowOpID, OpTag::DataflowChild);
147impl_nodehandle!(ConditionalID, OpTag::Conditional);
148impl_nodehandle!(CaseID, OpTag::Case);
149impl_nodehandle!(DfgID, OpTag::Dfg);
150impl_nodehandle!(TailLoopID, OpTag::TailLoop);
151impl_nodehandle!(CfgID, OpTag::Cfg);
152
153impl_nodehandle!(ModuleRootID, OpTag::ModuleRoot);
154impl_nodehandle!(ModuleID, OpTag::ModuleOp);
155impl_nodehandle!(ConstID, OpTag::Const);
156
157impl_nodehandle!(BasicBlockID, OpTag::DataflowBlock);
158
159impl<const DEF: bool> NodeHandle for FuncID<DEF> {
160    const TAG: OpTag = OpTag::Function;
161    #[inline]
162    fn node(&self) -> Node {
163        self.0
164    }
165}
166
167impl<const DEF: bool> NodeHandle for AliasID<DEF> {
168    const TAG: OpTag = OpTag::Alias;
169    #[inline]
170    fn node(&self) -> Node {
171        self.node
172    }
173}
174
175impl NodeHandle for Node {
176    const TAG: OpTag = OpTag::Any;
177    #[inline]
178    fn node(&self) -> Node {
179        *self
180    }
181}
182
183/// Implements the `ContainerHandle` trait, with the given child handle type.
184macro_rules! impl_containerHandle {
185    ($name:path, $children:ident) => {
186        impl ContainerHandle for $name {
187            type ChildrenHandle = $children;
188        }
189    };
190}
191
192impl_containerHandle!(DataflowParentID, DataflowOpID);
193impl_containerHandle!(DfgID, DataflowOpID);
194impl_containerHandle!(TailLoopID, DataflowOpID);
195impl_containerHandle!(ConditionalID, CaseID);
196impl_containerHandle!(CaseID, DataflowOpID);
197impl_containerHandle!(ModuleRootID, ModuleID);
198impl_containerHandle!(CfgID, BasicBlockID);
199impl_containerHandle!(BasicBlockID, DataflowOpID);
200impl_containerHandle!(FuncID<true>, DataflowOpID);
201impl_containerHandle!(AliasID<true>, DataflowOpID);