hugr_passes/
dead_code.rs

1//! Pass for removing dead code, i.e. that computes values that are then discarded
2
3use hugr_core::hugr::internal::HugrInternals;
4use hugr_core::{HugrView, Node, hugr::hugrmut::HugrMut, ops::OpType};
5use std::collections::{HashMap, HashSet, VecDeque};
6use std::fmt::{Debug, Display, Formatter};
7use std::sync::Arc;
8
9use crate::composable::WithScope;
10use crate::{ComposablePass, PassScope};
11
12/// Configuration for Dead Code Elimination pass
13#[derive(Clone)]
14#[deprecated(
15    note = "`hugr-passes` is deprecated. Use tket::passes instead",
16    since = "0.26.2"
17)]
18pub struct DeadCodeElimPass<H: HugrView> {
19    /// Nodes that are definitely needed - e.g. `FuncDefns`, but could be anything.
20    /// Hugr Root is assumed to be an entry point even if not mentioned here.
21    entry_points: Vec<H::Node>,
22    /// If None, use entrypoint-subtree (even if module root)
23    scope: Option<PassScope>,
24    /// Callback identifying nodes that must be preserved even if their
25    /// results are not used. Defaults to [`PreserveNode::default_for`].
26    preserve_callback: Arc<PreserveCallback<H>>,
27}
28
29impl<H: HugrView + 'static> Default for DeadCodeElimPass<H> {
30    fn default() -> Self {
31        Self {
32            entry_points: Default::default(),
33            // Preserve pre-PassScope behaviour of affecting entrypoint subtree only:
34            scope: None,
35            preserve_callback: Arc::new(PreserveNode::default_for),
36        }
37    }
38}
39
40impl<H: HugrView> Debug for DeadCodeElimPass<H> {
41    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
42        // Use "derive Debug" by defining an identical struct without the unprintable fields
43
44        #[expect(unused)] // Rust ignores the derive-Debug in figuring out what's used
45        #[derive(Debug)]
46        struct DCEDebug<'a, N> {
47            entry_points: &'a Vec<N>,
48            scope: &'a Option<PassScope>,
49        }
50
51        Debug::fmt(
52            &DCEDebug {
53                entry_points: &self.entry_points,
54                scope: &self.scope,
55            },
56            f,
57        )
58    }
59}
60
61/// Callback that identifies nodes that must be preserved even if their
62/// results are not used. For example, (the default) [`PreserveNode::default_for`].
63#[deprecated(
64    note = "`hugr-passes` is deprecated. Use tket::passes instead",
65    since = "0.26.2"
66)]
67pub type PreserveCallback<H> = dyn Fn(&H, <H as HugrInternals>::Node) -> PreserveNode;
68
69#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
70/// Signal that a node must be preserved even when its result is not used
71#[deprecated(
72    note = "`hugr-passes` is deprecated. Use tket::passes instead",
73    since = "0.26.2"
74)]
75pub enum PreserveNode {
76    /// The node must be kept (nodes inside it may be removed)
77    MustKeep,
78    /// The node can be removed, even if nodes inside it must be kept
79    /// - this will remove the descendants too, so use with care.
80    CanRemoveIgnoringChildren,
81    /// The node may be removed if-and-only-if all of its children can
82    /// (must be kept iff any of its children must be kept).
83    DeferToChildren,
84}
85
86impl PreserveNode {
87    /// A conservative default for a given node. Just examines the node's [`OpType`]:
88    /// * Assumes all Calls must be preserved. (One could scan the called `FuncDefn`, but would
89    ///   also need to check for cycles in the [`ModuleGraph`](hugr_core::module_graph::ModuleGraph).)
90    /// * Assumes all CFGs must be preserved. (One could, for example, allow acyclic
91    ///   CFGs to be removed.)
92    /// * Assumes all `TailLoops` must be preserved. (One could, for example, use dataflow
93    ///   analysis to allow removal of `TailLoops` that never [Continue](hugr_core::ops::TailLoop::CONTINUE_TAG).)
94    pub fn default_for<H: HugrView>(h: &H, n: H::Node) -> PreserveNode {
95        match h.get_optype(n) {
96            OpType::CFG(_) | OpType::TailLoop(_) | OpType::Call(_) => PreserveNode::MustKeep,
97            _ => Self::DeferToChildren,
98        }
99    }
100}
101
102/// Errors from [DeadCodeElimPass]
103#[derive(Clone, Debug, thiserror::Error, PartialEq)]
104#[non_exhaustive]
105#[deprecated(
106    note = "`hugr-passes` is deprecated. Use tket::passes instead",
107    since = "0.26.2"
108)]
109pub enum DeadCodeElimError<N: Display = Node> {
110    /// A node specified to [DeadCodeElimPass::with_entry_points] was not found
111    #[error("Node {_0} does not exist in the Hugr")]
112    NodeNotFound(N),
113}
114
115impl<H: HugrView> DeadCodeElimPass<H> {
116    /// Allows setting a callback that determines whether a node must be preserved
117    /// (even when its result is not used)
118    pub fn set_preserve_callback(mut self, cb: Arc<PreserveCallback<H>>) -> Self {
119        self.preserve_callback = cb;
120        self
121    }
122
123    /// Mark some nodes as starting points for analysis, i.e. so we cannot eliminate any code
124    /// used to evaluate these nodes. (E.g. nodes at which we may start executing the Hugr.)
125    ///
126    /// Other starting points are added according to the [PassScope].
127    // TODO should we deprecate this? i.e. require use of PreserveCallback / Hugr edges?
128    pub fn with_entry_points(mut self, entry_points: impl IntoIterator<Item = H::Node>) -> Self {
129        self.entry_points.extend(entry_points);
130        self
131    }
132
133    fn find_needed_nodes(&self, h: &H) -> Result<HashSet<H::Node>, DeadCodeElimError<H::Node>> {
134        let mut must_preserve = HashMap::new();
135        let mut needed = HashSet::new();
136        let mut q = VecDeque::from_iter(self.entry_points.iter().copied());
137
138        match &self.scope {
139            None => q.push_back(h.entrypoint()),
140            Some(scope) => q.extend(scope.preserve_interface(h)),
141        };
142        while let Some(n) = q.pop_front() {
143            if !h.contains_node(n) {
144                return Err(DeadCodeElimError::NodeNotFound(n));
145            }
146            if !needed.insert(n) {
147                continue;
148            }
149            // Ensure no orphans, e.g. when preserving an entrypoint deep within a Hugr
150            // being globally optimized. We could remove more from parent, but would require transforming
151            // (e.g. removing individual Output ports) not just deleting, so don't.
152            q.extend(h.get_parent(n));
153            for (i, ch) in h.children(n).enumerate() {
154                if self.must_preserve(h, &mut must_preserve, ch)
155                    || match h.get_optype(ch) {
156                        OpType::Case(_)  // Include all Cases in Conditionals
157                        | OpType::ExitBlock(_)
158                        | OpType::AliasDecl(_) // and all Aliases (we do not track their uses in types)
159                        | OpType::AliasDefn(_)
160                        | OpType::Input(_) // Also Dataflow input/output, these are necessary for legality
161                        | OpType::Output(_) => true,
162                        // Assumes entry block is always the first child of a CFG.
163                        OpType::DataflowBlock(_) => h.get_optype(n).is_cfg() && i == 0,
164                        // Do not include FuncDecl / FuncDefn / Const,
165                        // unless reachable by static edges (from Call/LoadConst/LoadFunction)
166                        _ => false,
167                    }
168                {
169                    q.push_back(ch);
170                }
171            }
172            if matches!(
173                h.get_optype(n),
174                OpType::DataflowBlock(_) | OpType::ExitBlock(_)
175            ) {
176                // Follow control flow forwards to find reachable basic blocks besides entry and exit.
177                q.extend(h.output_neighbours(n))
178            } else {
179                // Follow dataflow demand (including e.g edges from Call to FuncDefn) backwards.
180                q.extend(h.input_neighbours(n));
181            }
182            // Also keep consumers of any linear outputs
183            if let Some(sig) = h.signature(n) {
184                for op in sig.output_ports() {
185                    if !sig.out_port_type(op).unwrap().copyable() {
186                        q.extend(h.linked_inputs(n, op).map(|(n, _inp)| n))
187                    }
188                }
189            }
190        }
191        Ok(needed)
192    }
193
194    fn must_preserve(&self, h: &H, cache: &mut HashMap<H::Node, bool>, n: H::Node) -> bool {
195        if let Some(res) = cache.get(&n) {
196            return *res;
197        }
198        let res = match self.preserve_callback.as_ref()(h, n) {
199            PreserveNode::MustKeep => true,
200            PreserveNode::CanRemoveIgnoringChildren => false,
201            PreserveNode::DeferToChildren => {
202                h.children(n).any(|ch| self.must_preserve(h, cache, ch))
203            }
204        };
205        cache.insert(n, res);
206        res
207    }
208}
209
210impl<H: HugrMut> ComposablePass<H> for DeadCodeElimPass<H> {
211    type Error = DeadCodeElimError<H::Node>;
212    type Result = ();
213
214    fn run(&self, hugr: &mut H) -> Result<(), Self::Error> {
215        let root = match &self.scope {
216            None => hugr.entrypoint(),
217            Some(scope) => match scope.root(hugr) {
218                Some(root) => root,
219                None => return Ok(()),
220            },
221        };
222        let needed = self.find_needed_nodes(&*hugr)?;
223        let remove = hugr
224            .descendants(root)
225            .filter(|n| !needed.contains(n))
226            .collect::<Vec<_>>();
227        for n in remove {
228            hugr.remove_node(n);
229        }
230        Ok(())
231    }
232}
233
234impl<H: HugrMut> WithScope for DeadCodeElimPass<H> {
235    fn with_scope(mut self, scope: impl Into<PassScope>) -> Self {
236        self.scope = Some(scope.into());
237        self
238    }
239}
240
241#[cfg(test)]
242mod test {
243    use std::sync::Arc;
244
245    use hugr_core::builder::{
246        CFGBuilder, Container, DFGBuilder, Dataflow, DataflowHugr, DataflowSubContainer,
247        HugrBuilder, endo_sig, inout_sig,
248    };
249    use hugr_core::extension::prelude::{ConstUsize, bool_t, qb_t, usize_t};
250    use hugr_core::extension::{ExtensionId, Version};
251    use hugr_core::ops::{ExtensionOp, OpType};
252    use hugr_core::ops::{OpTag, OpTrait, handle::NodeHandle};
253    use hugr_core::types::Signature;
254    use hugr_core::{Extension, Hugr};
255    use hugr_core::{HugrView, ops::Value, type_row};
256    use itertools::Itertools;
257
258    use crate::ComposablePass;
259
260    use super::{DeadCodeElimPass, PreserveNode};
261
262    #[test]
263    fn test_cfg_callback() {
264        let mut cb = CFGBuilder::new(Signature::new_endo(type_row![])).unwrap();
265        let cst_unused = cb.add_constant(Value::from(ConstUsize::new(3)));
266        let cst_used_in_dfg = cb.add_constant(Value::from(ConstUsize::new(5)));
267        let cst_used = cb.add_constant(Value::unary_unit_sum());
268        let mut block = cb.entry_builder([type_row![]], type_row![]).unwrap();
269        let mut dfg_unused = block
270            .dfg_builder(Signature::new(type_row![], [usize_t()]), [])
271            .unwrap();
272        let lc_unused = dfg_unused.load_const(&cst_unused);
273        let lc1 = dfg_unused.load_const(&cst_used_in_dfg);
274        let dfg_unused = dfg_unused.finish_with_outputs([lc1]).unwrap().node();
275        let pred = block.load_const(&cst_used);
276        let block = block.finish_with_outputs(pred, []).unwrap();
277        let exit = cb.exit_block();
278        cb.branch(&block, 0, &exit).unwrap();
279        let orig = cb.finish_hugr().unwrap();
280
281        // Callbacks that allow removing the DFG (and cst_unused)
282        for dce in [
283            DeadCodeElimPass::<Hugr>::default(),
284            // keep the node inside the DFG, but remove the DFG without checking its children:
285            DeadCodeElimPass::default().set_preserve_callback(Arc::new(move |h, n| {
286                if n == dfg_unused || h.get_optype(n).is_const() {
287                    PreserveNode::CanRemoveIgnoringChildren
288                } else {
289                    PreserveNode::MustKeep
290                }
291            })),
292        ] {
293            let mut h = orig.clone();
294            dce.run(&mut h).unwrap();
295            assert_eq!(
296                h.children(h.entrypoint()).collect_vec(),
297                [block.node(), exit.node(), cst_used.node()]
298            );
299            assert_eq!(
300                h.children(block.node())
301                    .map(|n| h.get_optype(n).tag())
302                    .collect_vec(),
303                [OpTag::Input, OpTag::Output, OpTag::LoadConst]
304            );
305        }
306
307        // Callbacks that prevent removing any node...
308        fn keep_if(b: bool) -> PreserveNode {
309            if b {
310                PreserveNode::MustKeep
311            } else {
312                PreserveNode::DeferToChildren
313            }
314        }
315        for dce in [
316            DeadCodeElimPass::<Hugr>::default()
317                .set_preserve_callback(Arc::new(|_, _| PreserveNode::MustKeep)),
318            // keeping the unused node in the DFG, means keeping the DFG (which uses its other children)
319            DeadCodeElimPass::default()
320                .set_preserve_callback(Arc::new(move |_, n| keep_if(n == lc_unused.node()))),
321        ] {
322            let mut h = orig.clone();
323            dce.run(&mut h).unwrap();
324            assert_eq!(orig, h);
325        }
326
327        // Callbacks that keep the DFG but allow removing the unused constant
328        for dce in [
329            DeadCodeElimPass::<Hugr>::default()
330                .set_preserve_callback(Arc::new(move |_, n| keep_if(n == dfg_unused))),
331            DeadCodeElimPass::default()
332                .set_preserve_callback(Arc::new(move |_, n| keep_if(n == lc1.node()))),
333        ] {
334            let mut h = orig.clone();
335            dce.run(&mut h).unwrap();
336            assert_eq!(
337                h.children(h.entrypoint()).collect_vec(),
338                [
339                    block.node(),
340                    exit.node(),
341                    cst_used_in_dfg.node(),
342                    cst_used.node()
343                ]
344            );
345            assert_eq!(
346                h.children(block.node()).skip(2).collect_vec(),
347                [dfg_unused, pred.node()]
348            );
349            assert_eq!(
350                h.children(dfg_unused.node())
351                    .map(|n| h.get_optype(n).tag())
352                    .collect_vec(),
353                [OpTag::Input, OpTag::Output, OpTag::LoadConst]
354            );
355        }
356
357        // Callback that allows removing the DFG but require keeping cst_unused
358        {
359            let cst_unused = cst_unused.node();
360            let mut h = orig.clone();
361            DeadCodeElimPass::<Hugr>::default()
362                .set_preserve_callback(Arc::new(move |_, n| keep_if(n == cst_unused)))
363                .run(&mut h)
364                .unwrap();
365            assert_eq!(
366                h.children(h.entrypoint()).collect_vec(),
367                [block.node(), exit.node(), cst_unused, cst_used.node()]
368            );
369            assert_eq!(
370                h.children(block.node())
371                    .map(|n| h.get_optype(n).tag())
372                    .collect_vec(),
373                [OpTag::Input, OpTag::Output, OpTag::LoadConst]
374            );
375        }
376    }
377
378    #[test]
379    fn preserve_linear() {
380        // A simple linear alloc/measure. Note we do *not* model ordering among allocations for this test.
381        let test_ext = Extension::new_arc(
382            ExtensionId::new_unchecked("test_qext"),
383            Version::new(0, 0, 0),
384            |e, w| {
385                e.add_op("new".into(), "".into(), inout_sig(vec![], [qb_t()]), w)
386                    .unwrap();
387                e.add_op("gate".into(), "".into(), endo_sig([qb_t()]), w)
388                    .unwrap();
389                e.add_op(
390                    "measure".into(),
391                    "".into(),
392                    inout_sig([qb_t()], [bool_t()]),
393                    w,
394                )
395                .unwrap();
396                e.add_op("not".into(), "".into(), endo_sig([bool_t()]), w)
397                    .unwrap();
398            },
399        );
400        let [new, gate, measure, not] = ["new", "gate", "measure", "not"]
401            .map(|n| ExtensionOp::new(test_ext.get_op(n).unwrap().clone(), []).unwrap());
402        let mut dfb = DFGBuilder::new(endo_sig([qb_t()])).unwrap();
403        // Unused new...measure, can be removed
404        let qn = dfb.add_dataflow_op(new.clone(), []).unwrap().outputs();
405        let [_] = dfb
406            .add_dataflow_op(measure.clone(), qn)
407            .unwrap()
408            .outputs_arr();
409
410        // Free (measure) the input, so not connected to the output
411        let [q_in] = dfb.input_wires_arr();
412        let [h_in] = dfb
413            .add_dataflow_op(gate.clone(), [q_in])
414            .unwrap()
415            .outputs_arr();
416        let [b] = dfb.add_dataflow_op(measure, [h_in]).unwrap().outputs_arr();
417        // Operate on the bool only, can be removed as not linear:
418        dfb.add_dataflow_op(not, [b]).unwrap();
419
420        // Alloc a new qubit and output that
421        let q = dfb.add_dataflow_op(new, []).unwrap().outputs();
422        let outs = dfb.add_dataflow_op(gate, q).unwrap().outputs();
423        let mut h = dfb.finish_hugr_with_outputs(outs).unwrap();
424        DeadCodeElimPass::default().run(&mut h).unwrap();
425        // This was failing before https://github.com/CQCL/hugr/pull/2560:
426        h.validate().unwrap();
427
428        // Remove one new and measure, and a "not"; keep both gates
429        // (cannot remove the other gate or measure even tho results not needed).
430        // Removing the gate because the measure-result is not used is beyond (current) DeadCodeElim.
431        let ext_ops = h
432            .nodes()
433            .filter_map(|n| h.get_optype(n).as_extension_op())
434            .map(ExtensionOp::unqualified_id);
435        assert_eq!(
436            ext_ops.sorted().collect_vec(),
437            ["gate", "gate", "measure", "new"]
438        );
439    }
440
441    #[test]
442    fn remove_unreachable_bb() {
443        let mut cb = CFGBuilder::new(Signature::new_endo(type_row![])).unwrap();
444
445        let cst_unused = cb.add_constant(Value::from(ConstUsize::new(3)));
446        let b1_pred = cb.add_constant(Value::unary_unit_sum());
447        let b2_pred = cb.add_constant(Value::unit_sum(0, 2).expect("0 < 2"));
448
449        // Entry block
450        let mut entry = cb.entry_builder([type_row![]], type_row![]).unwrap();
451        let pred1 = entry.load_const(&b1_pred);
452        let entry = entry.finish_with_outputs(pred1, []).unwrap();
453
454        // Reachable block
455        let mut block_reachable = cb
456            .simple_block_builder(Signature::new(type_row![], type_row![]), 1)
457            .unwrap();
458        let pred2 = block_reachable.load_const(&b1_pred);
459        let block_reachable = block_reachable.finish_with_outputs(pred2, []).unwrap();
460
461        // Unreachable block
462        let mut block_unreachable = cb
463            .simple_block_builder(Signature::new(type_row![], type_row![]), 2)
464            .unwrap();
465        let _ = block_unreachable.load_const(&cst_unused);
466        let pred3 = block_unreachable.load_const(&b2_pred);
467        let block_unreachable = block_unreachable.finish_with_outputs(pred3, []).unwrap();
468
469        // Exit block
470        let exit = cb.exit_block();
471
472        // Construct CFG
473        cb.branch(&entry, 0, &block_reachable).unwrap();
474        cb.branch(&block_reachable, 0, &exit).unwrap();
475        cb.branch(&block_unreachable, 0, &exit).unwrap();
476        // Addtionally add a loop to check it works with a cycle
477        cb.branch(&block_unreachable, 1, &block_unreachable)
478            .unwrap();
479        let mut h = cb.finish_hugr().unwrap();
480        h.validate().unwrap();
481        let num_nodes_before = h.nodes().count();
482        let cfg_node = h.entrypoint();
483        let num_cfg_children_before: usize = h
484            .children(cfg_node)
485            .filter(|child| matches!(h.get_optype(*child), OpType::DataflowBlock(_)))
486            .count();
487
488        // Run pass and check that unreachable block is removed
489        DeadCodeElimPass::default().run(&mut h).unwrap();
490        h.validate().unwrap();
491
492        // Check we removed the expected number of nodes.
493        // 7 nodes removed:
494        // - 1 block (block_unreachable)
495        // - 2 constants (cst_unused, b2_pred)
496        // - 4 ops in the unreachable block (2 LoadConst, the block's Input and Output)
497        let num_nodes_after = h.nodes().count();
498        assert_eq!(num_nodes_before - num_nodes_after, 7);
499
500        // Check that `block_unreachable` is no longer a valid node.
501        assert!(!h.contains_node(block_unreachable.node()));
502
503        // CFG checks: should still be a CFG and have one less dataflow block child.
504        assert!(h.get_optype(cfg_node).is_cfg());
505        let num_cfg_children_after: usize = h
506            .children(cfg_node)
507            .filter(|child| matches!(h.get_optype(*child), OpType::DataflowBlock(_)))
508            .count();
509        assert_eq!(num_cfg_children_after, num_cfg_children_before - 1);
510
511        // Also the exit block should only have one predecessor now.
512        let exit_preds = h.input_neighbours(exit.node()).collect_vec();
513        assert_eq!(exit_preds.len(), 1);
514    }
515}
hugr_passes/dead_code.rs

hugr_passes/
dead_code.rs
