hugr_core/ops/
handle.rs

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