hugr-persistent 0.6.1

Persistent IR structure for Quantinuum's HUGR
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
//! Incremental traversal and construction of [`PersistentHugr`]s from
//! [`CommitStateSpace`]s.
//!
//! This module provides the [`Walker`] type, which enables exploration of a
//! [`CommitStateSpace`] by traversing wires and gradually pinning (selecting)
//! nodes that match a desired pattern. Unlike direct manipulation of a
//! [`CommitStateSpace`], which may contain many alternative and conflicting
//! versions of the graph data, the Walker presents a familiar nodes-and-wires
//! interface that feels similar to walking a standard (HUGR) graph.
//!
//! The key concept is that of "pinning" nodes - marking specific nodes as fixed
//! points in the exploration. As more nodes are pinned, the space of possible
//! HUGRs that are considered narrows; a Walker instance where all nodes are
//! pinned then corresponds to a [`PersistentHugr`], which is equivalent to a
//! HUGR graph. By choosing which nodes are pinned, the user in effect specifies
//! and extracts the [`PersistentHugr`]s of interest from within the
//! [`CommitStateSpace`].
//!
//! To establish which nodes can be pinned and explore the possible
//! alternatives, walkers are "expanded" along wires. Given a walker and a wire,
//! along with, optionally, a direction of expansion, [`Walker::expand`] returns
//! an iterator over all possible walkers that can be obtained by pinning
//! exactly one additional node from the wire. Repeated expansions result in
//! walkers with an increasing number of pinned nodes, narrowing down the commit
//! state space of interest at each step.
//!
//! ## Typical usage
//!
//! Such incremental exploration is particularly suited for pattern matching
//! and the creation of replacements (rewrite factories).
//!
//! A typical usage flow involves:
//! 1. Create a [`Walker`] over a [`CommitStateSpace`].
//! 2. Pin (at least one) initial node of interest using
//!    [`Walker::try_pin_node`] or [`Walker::from_pinned_node`].
//! 3. Traverse wires starting from a pinned node to explore the neighbourhood
//!    in the graph using [`Walker::expand`]. Each new walker thus obtained will
//!    correspond to a different selction of commits, which will result in
//!    several alternative HUGRs that must be considered in parallel. Thus, the
//!    exploration typically branches out into several parallel walkers.
//! 4. As walkers are expanded, more nodes get pinned, narrowing down the space
//!    of possible HUGRs.
//! 5. Once exploration is complete (e.g. a pattern was fully matched), the
//!    walker can be converted into a [`PersistentHugr`] instance using
//!    [`Walker::into_persistent_hugr`]. The matched nodes and ports can then be
//!    used to create a
//!    [`SiblingSubgraph`](hugr_core::hugr::views::SiblingSubgraph) object,
//!    which can then be used to create new
//!    [`SimpleReplacement`](hugr_core::SimpleReplacement) instances---and
//!    possibly in turn be added to the commit state space and the exploration
//!    of the state space restarted!
//!
//! This approach allows efficiently finding patterns across many potential
//! versions of the graph simultaneously, without having to materialize
//! each version separately.

use std::collections::BTreeSet;

use hugr_core::Node;
use hugr_core::hugr::patch::simple_replace::BoundaryMode;
use hugr_core::ops::handle::DataflowParentID;
use itertools::{Either, Itertools};
use thiserror::Error;

use hugr_core::{Direction, Hugr, HugrView, Port, PortIndex, hugr::views::RootChecked};

use crate::{Commit, PersistentReplacement, PinnedSubgraph};

use crate::PersistentWire;

use super::{CommitStateSpace, InvalidCommit, PatchNode, PersistentHugr, state_space::CommitId};

/// A walker over a [`CommitStateSpace`].
///
/// A walker is given by a set of selected commits, along with a set of pinned
/// nodes that belong to those selected commits. The selected commits determine
/// which replacements in the state space are applied. Meanwhile, the pinned
/// nodes within these commits determine which nodes are "frozen" in the
/// exploration: no further commit can be selected that would invalidate any
/// pinned node.
///
/// The set of selected commits of a walker defines a [`PersistentHugr`]
/// instance that can be retrieved by calling [`Walker::into_persistent_hugr`].
/// As the walker is expanded and more of the state space is explored, more
/// commits are selected, and the [`PersistentHugr`] will change accordingly.
/// Pinned nodes (and pinned ports, i.e. ports at pinned nodes) are guaranteed
/// to be valid in all [`PersistentHugr`] instances obtained as a result of
/// expansions of the current walker.
/// current walker.
#[derive(Debug, Clone)]
pub struct Walker<'a> {
    /// The state space being traversed.
    state_space: &'a CommitStateSpace,
    /// The subset of compatible commits in `state_space` that are currently
    /// selected.
    // Note that we could store this as a set of `CommitId`s, but it is very
    // convenient to have access to all the methods of PersistentHugr (on top
    // of guaranteeing the compatibility invariant). The tradeoff is more
    // memory consumption.
    selected_commits: PersistentHugr,
    /// The set of nodes that have been traversed by the walker and can no
    /// longer be rewritten.
    pinned_nodes: BTreeSet<PatchNode>,
}

impl<'a> Walker<'a> {
    /// Create a new [`Walker`] over the given state space.
    ///
    /// No nodes are pinned initially. The [`Walker`] starts with only the base
    /// Hugr `state_space.base_hugr()` selected.
    ///
    /// # Panics
    /// Panics if the commit state space is empty.
    pub fn new(state_space: &'a CommitStateSpace) -> Self {
        let base = state_space.base_commit().expect("non-empty state space");
        let selected_commits: PersistentHugr = PersistentHugr::from_commit(base);
        Self {
            state_space,
            selected_commits,
            pinned_nodes: BTreeSet::new(),
        }
    }

    /// Create a new [`Walker`] with a single pinned node.
    pub fn from_pinned_node(node: PatchNode, state_space: &'a CommitStateSpace) -> Self {
        let mut walker = Self::new(state_space);
        walker
            .try_pin_node(node)
            .expect("node is valid and not deleted");
        walker
    }

    /// Pin a node in the [`Walker`].
    ///
    /// This method allows pinning a specific node in the Walker, which
    /// restricts the possible paths that can be explored in the Walker.
    ///
    /// Return true if `node` was not pinned already, and false otherwise.
    ///
    /// If the node belongs to a commit that isn't currently selected, the
    /// commit is added to the Walker, or an error is returned if this is
    /// not possible. The node is then added to the set of pinned nodes. If
    /// the node is already pinned, this method is a no-op.
    pub fn try_pin_node(&mut self, node: PatchNode) -> Result<bool, PinNodeError> {
        let commit_id = node.0;
        if self.selected_commits.contains_id(commit_id) {
            if !self.selected_commits.contains_node(node) {
                return Err(PinNodeError::AlreadyDeleted(node));
            }
        } else {
            let commit = self
                .state_space
                .try_upgrade(commit_id)
                .ok_or(PinNodeError::UnknownCommitId(commit_id))?;
            self.try_select_commit(commit)?;
        }
        Ok(self.pinned_nodes.insert(node))
    }

    /// Add a commit to the selected commits of the Walker.
    ///
    /// Return the ID of the added commit if it was added successfully, or the
    /// existing ID of the commit if it was already selected.
    ///
    /// Return an error if the commit is not compatible with the current set of
    /// selected commits, or if the commit deletes an already pinned node.
    pub fn try_select_commit(&mut self, commit: Commit) -> Result<CommitId, PinNodeError> {
        // TODO: we should be able to check for an AlreadyPinned error at
        // the same time that we check the ancestors are compatible in
        // `PersistentHugr`, with e.g. a callback, instead of storing a backup
        let backup = self.selected_commits.clone();
        let commit_id = self.selected_commits.try_add_commit(commit)?;
        if let Some(&pinned_node) = self
            .pinned_nodes
            .iter()
            .find(|&&n| !self.selected_commits.contains_node(n))
        {
            self.selected_commits = backup;
            return Err(PinNodeError::AlreadyPinned(pinned_node));
        }
        Ok(commit_id)
    }

    /// Expand the Walker by pinning a node connected to the given wire.
    ///
    /// To understand how Walkers are expanded, it is useful to understand how
    /// in a walker, the HUGR graph is partitioned into two parts:
    ///  - a subgraph made of pinned nodes: this part of the HUGR is frozen: it
    ///    cannot be modified by further expansions the Walker.
    ///  - the complement subgraph: the unpinned part of the HUGR has not been
    ///    explored yet. Multiple alternative HUGRs can be obtained depending on
    ///    which commits are selected.
    ///
    /// To every walker thus corresponds a space of possible HUGRs that can be
    /// obtained, depending on which commits are selected and which further
    /// nodes are pinned. The expansion of a walker returns a set of
    /// walkers, which together cover the same space of possible HUGRs, each
    /// having a different additional node pinned.
    ///
    /// If the wire is not complete yet, return an iterator over all possible
    /// [`Walker`]s that can be created by pinning exactly one additional
    /// node (or one additional commit with an empty wire) connected to
    /// `wire`. Each returned [`Walker`] represents a different alternative
    /// Hugr in the exploration space.
    ///
    /// If the wire is already complete, return an iterator containing one
    /// walker: the current walker unchanged.
    ///
    /// Optionally, the expansion can be restricted to only ports with the given
    /// direction (incoming or outgoing).
    ///
    /// Repeatedly calling `expand` on a wire will progressively pin all its
    /// endpoints. The multiple alternative [`Walker`]s returned then form
    /// a branching search space to be explored by the user. When a wire is
    /// fully pinned in the specified direction, i.e. `wire.is_complete(dir)` is
    /// true, then an empty iterator is returned.
    pub fn expand<'b>(
        &'b self,
        wire: &'b PersistentWire,
        dir: impl Into<Option<Direction>>,
    ) -> impl Iterator<Item = Walker<'a>> + 'b {
        let dir = dir.into();

        if self.is_complete(wire, dir) {
            return Either::Left(std::iter::once(self.clone()));
        }

        // Find unpinned ports on the wire (satisfying the direction constraint)
        let unpinned_ports = self.wire_unpinned_ports(wire, dir);

        // Obtain set of pinnable nodes by considering all ports (in descendant
        // commits) equivalent to currently unpinned ports.
        let pinnable_nodes = unpinned_ports
            .flat_map(|(node, port)| self.equivalent_descendant_ports(node, port))
            .map(|(n, _, commits)| (n, commits))
            .unique();

        let new_walkers = pinnable_nodes.filter_map(|(pinnable_node, new_commits)| {
            let contains_new_commit = || {
                new_commits
                    .iter()
                    .any(|&cm| !self.selected_commits.contains_id(cm))
            };
            debug_assert!(
                !self.is_pinned(pinnable_node) || contains_new_commit(),
                "trying to pin already pinned node and no new commit is selected"
            );

            // Upgrade the commit IDs
            let new_commits = new_commits
                .iter()
                .map(|&id| self.state_space.try_upgrade(id))
                .collect::<Option<Vec<_>>>()?;

            // Update the selected commits to include the new commits.
            let new_selected_commits = {
                let mut phugr = self.selected_commits.clone();
                phugr.try_add_commits(new_commits).ok()?;
                phugr
            };
            // .try_create(self.selected_commits.all_commit_ids().chain(new_commits))
            // .ok()?;

            // Make sure that the pinned nodes are still valid after including the new
            // selected commits.
            if self
                .pinned_nodes
                .iter()
                .any(|&pnode| !new_selected_commits.contains_node(pnode))
            {
                return None;
            }

            // Construct a new walker and pin `pinnable_node`.
            let mut new_walker = Walker {
                state_space: self.state_space,
                selected_commits: new_selected_commits,
                pinned_nodes: self.pinned_nodes.clone(),
            };
            new_walker.try_pin_node(pinnable_node).ok()?;
            Some(new_walker)
        });

        Either::Right(new_walkers)
    }

    /// Create a new commit from a set of complete pinned wires and a
    /// replacement.
    ///
    /// The subgraph of the commit is the subgraph given by the set of edges
    /// in `wires`. `map_boundary` must provide a map from the boundary ports
    /// of the subgraph to the inputs/output ports in `repl`. The returned port
    /// must be of the opposite direction as the port passed as argument:
    ///  - an incoming subgraph port must be mapped to an outgoing port of the
    ///    input node of `repl`
    /// - an outgoing subgraph port must be mapped to an incoming port of the
    ///   output node of `repl`
    ///
    /// ## Panics
    ///
    /// This will panic if repl is not a DFG graph.
    pub fn try_create_commit(
        &self,
        subgraph: impl Into<PinnedSubgraph>,
        mut repl: RootChecked<Hugr, DataflowParentID>,
        map_boundary: impl Fn(PatchNode, Port) -> Port,
    ) -> Result<Commit<'a>, InvalidCommit> {
        let pinned_subgraph = subgraph.into();
        let subgraph = pinned_subgraph.to_sibling_subgraph(self.as_hugr_view())?;
        let selected_commits = pinned_subgraph
            .selected_commits()
            .map(|id| self.selected_commits.get_commit(id).clone());

        let repl = {
            let new_inputs = subgraph
                .incoming_ports()
                .iter()
                .flatten() // because of singleton-vec wrapping above
                .map(|&(n, p)| {
                    map_boundary(n, p.into())
                        .as_outgoing()
                        .expect("unexpected port direction returned by map_boundary")
                        .index()
                })
                .collect_vec();
            let new_outputs = subgraph
                .outgoing_ports()
                .iter()
                .map(|&(n, p)| {
                    map_boundary(n, p.into())
                        .as_incoming()
                        .expect("unexpected port direction returned by map_boundary")
                        .index()
                })
                .collect_vec();
            repl.map_function_type(&new_inputs, &new_outputs)?;
            PersistentReplacement::try_new(subgraph, self.as_hugr_view(), repl.into_hugr())?
        };

        Commit::try_new(repl, selected_commits, self.state_space)
    }

    /// Get the wire connected to a specified port of a pinned node.
    ///
    /// # Panics
    /// Panics if `node` is not already pinned in this Walker.
    pub fn get_wire(&self, node: PatchNode, port: impl Into<Port>) -> PersistentWire {
        assert!(self.is_pinned(node), "node must be pinned");
        self.selected_commits.get_wire(node, port)
    }

    /// Materialise the [`PersistentHugr`] containing all the compatible commits
    /// that have been selected during exploration.
    pub fn into_persistent_hugr(self) -> PersistentHugr {
        self.selected_commits
    }

    /// View the [`PersistentHugr`] containing all the compatible commits that
    /// have been selected so far during exploration.
    ///
    /// Of the space of all possible HUGRs that can be obtained from future
    /// expansions of the walker, this is the HUGR corresponding to selecting
    /// as few commits as possible (i.e. all the commits that have been selected
    /// so far and no more).
    pub fn as_hugr_view(&self) -> &PersistentHugr {
        &self.selected_commits
    }

    /// Check if a node is pinned in the [`Walker`].
    pub fn is_pinned(&self, node: PatchNode) -> bool {
        self.pinned_nodes.contains(&node)
    }

    /// Iterate over all pinned nodes in the [`Walker`].
    pub fn pinned_nodes(&self) -> impl Iterator<Item = PatchNode> + '_ {
        self.pinned_nodes.iter().copied()
    }

    /// Get all equivalent ports among the commits that are descendants of the
    /// current commit.
    ///
    /// The ports in the returned iterator will be in the same direction as
    /// `port`. For each equivalent port, also return the set of empty commits
    /// that were visited to find it.
    fn equivalent_descendant_ports(
        &self,
        node: PatchNode,
        port: Port,
    ) -> Vec<(PatchNode, Port, BTreeSet<CommitId>)> {
        // Now, perform a BFS to find all equivalent ports
        let mut all_ports = vec![(node, port, BTreeSet::new())];
        let mut index = 0;
        while index < all_ports.len() {
            let (node, port, empty_commits) = all_ports[index].clone();
            let Some(commit) = self.state_space.try_upgrade(node.owner()) else {
                continue;
            };
            index += 1;

            for (child, (opp_node, opp_port)) in
                commit.children_at_boundary_port(node.1, port, self.state_space)
            {
                for (node, port) in
                    commit.linked_child_ports(opp_node, opp_port, &child, BoundaryMode::SnapToHost)
                {
                    let mut empty_commits = empty_commits.clone();
                    if node.owner() != child.id() {
                        empty_commits.insert(child.id());
                    }
                    all_ports.push((node, port, empty_commits));
                }
            }
        }
        all_ports
    }
}

#[cfg(test)]
impl Walker<'_> {
    // Check walker equality using pointer equality component-wise. For testing
    // purposes.
    fn component_wise_ptr_eq(&self, other: &Self) -> bool {
        self.state_space == other.state_space
            && self.pinned_nodes == other.pinned_nodes
            && BTreeSet::from_iter(self.selected_commits.all_commit_ids())
                == BTreeSet::from_iter(other.selected_commits.all_commit_ids())
    }

    /// Check if the Walker cannot be expanded further, i.e. expanding it
    /// returns the same Walker.
    fn no_more_expansion(&self, wire: &PersistentWire, dir: impl Into<Option<Direction>>) -> bool {
        let Some([new_walker]) = self.expand(wire, dir).collect_array() else {
            return false;
        };
        new_walker.component_wise_ptr_eq(self)
    }
}

impl<'a> Commit<'a> {
    /// Given a node and port in `self`, return all child commits of `self`
    /// in the state space that delete `node` but keep at least one port linked
    /// to `(node, port)`. In other words, (node, port) is a boundary port
    /// of the subgraph of the child replacement.
    ///
    /// Return all tuples of children and linked port of (node, port) in `self`
    /// that is outside of the subgraph of the child. The returned ports are
    /// opposite to the direction of `port`.
    fn children_at_boundary_port(
        &self,
        node: Node,
        port: Port,
        state_space: &'a CommitStateSpace,
    ) -> impl Iterator<Item = (Commit<'a>, (Node, Port))> + '_ {
        let linked_ports = self.commit_hugr().linked_ports(node, port).collect_vec();

        self.children(state_space).flat_map(move |child| {
            let deleted_nodes: BTreeSet<_> = child.deleted_parent_nodes().collect();
            if !deleted_nodes.contains(&self.to_patch_node(node)) {
                vec![]
            } else {
                linked_ports
                    .iter()
                    .filter_map(move |&(linked_node, linked_port)| {
                        (!deleted_nodes.contains(&self.to_patch_node(linked_node)))
                            .then_some((child.clone(), (linked_node, linked_port)))
                    })
                    .collect_vec()
            }
        })
    }
}

/// An error that occurs when trying to pin a node in a [`Walker`].
#[derive(Debug, Clone, Error)]
#[non_exhaustive]
pub enum PinNodeError {
    /// Commit is not compatible with the current set of selected commits.
    #[error("cannot add commit to pin node: {0}")]
    InvalidNewCommit(InvalidCommit),
    /// Node to pin is already deleted in the current set of selected commits.
    #[error("cannot pin deleted node: {0}")]
    AlreadyDeleted(PatchNode),
    /// A node that would have to be deleted to pin the given node is already
    /// pinned.
    #[error("cannot delete already pinned node: {0}")]
    AlreadyPinned(PatchNode),
    /// The commit ID is not in the state space or has been deleted.
    #[error("unknown commit ID: {0:?}")]
    UnknownCommitId(CommitId),
}

impl From<InvalidCommit> for PinNodeError {
    fn from(value: InvalidCommit) -> Self {
        PinNodeError::InvalidNewCommit(value)
    }
}

impl<'w> hugr_core::hugr::views::NodesIter for Walker<'w> {
    type Node = PatchNode;

    fn nodes(&self) -> impl Iterator<Item = Self::Node> + '_ {
        <PersistentHugr as HugrView>::nodes(self.as_hugr_view())
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeSet;

    use hugr_core::{
        Direction, HugrView, IncomingPort, OutgoingPort,
        builder::{DFGBuilder, Dataflow, DataflowHugr, endo_sig},
        extension::prelude::bool_t,
        std_extensions::logic::LogicOp,
    };
    use itertools::Itertools;
    use rstest::rstest;

    use super::*;
    use crate::{
        PersistentHugr, Walker,
        state_space::CommitId,
        tests::{TestStateSpace, persistent_hugr_empty_child, test_state_space},
    };

    #[rstest]
    fn test_walker_base_or_child_expansion(test_state_space: TestStateSpace) {
        let [commit1, _commit2, _commit3, _commit4] = test_state_space.commits();
        let state_space = commit1.state_space();
        let base_commit = commit1.base_commit();

        // Get an initial node to pin
        let base_and_node = {
            let base_hugr = base_commit.commit_hugr();
            let and_node = base_hugr
                .nodes()
                .find(|&n| base_hugr.get_optype(n) == &LogicOp::And.into())
                .unwrap();
            base_commit.to_patch_node(and_node)
        };
        let walker = Walker::from_pinned_node(base_and_node, &state_space);
        assert!(walker.is_pinned(base_and_node));

        let in0 = walker.get_wire(base_and_node, IncomingPort::from(0));

        // a single incoming port (already pinned) => no more expansion
        assert!(walker.no_more_expansion(&in0, Direction::Incoming));

        // commit 2 cannot be applied, because AND is pinned
        // => only base commit, or commit1
        let out_walkers = walker.expand(&in0, Direction::Outgoing).collect_vec();
        assert_eq!(out_walkers.len(), 2);
        for new_walker in out_walkers {
            // new wire is complete (and thus cannot be expanded)
            let in0 = new_walker.get_wire(base_and_node, IncomingPort::from(0));
            assert!(new_walker.is_complete(&in0, None));
            assert!(new_walker.no_more_expansion(&in0, None));

            // all nodes on wire are pinned
            let (not_node, _) = in0.single_outgoing_port(new_walker.as_hugr_view()).unwrap();
            assert!(new_walker.is_pinned(base_and_node));
            assert!(new_walker.is_pinned(not_node));

            // not node is either in commit1 or the base
            assert!([commit1.id(), base_commit.id()].contains(&not_node.0));

            // not node is a NOT gate
            assert_eq!(
                new_walker.as_hugr_view().get_optype(not_node),
                &LogicOp::Not.into()
            );

            let persistent_hugr = new_walker.into_persistent_hugr();
            let hugr = persistent_hugr.get_commit(not_node.owner()).commit_hugr();
            assert_eq!(hugr.get_optype(not_node.1), &LogicOp::Not.into());
        }
    }

    /// Test that a walker can be expanded in various ways. Starting from
    /// commit4, it is possible to expand the pinned nodes in three ways:
    ///  - base
    ///  - commit1 without commit2
    ///  - commit1 and commit2
    ///
    /// commit3 is not an option as it conflicts with commit4.
    ///
    /// ```
    ///                         ------- base -------
    ///                        /         |          \
    ///                       /          |           \
    ///                      /           |            \
    ///                  commit1      c̶o̶m̶m̶i̶t̶3̶        commit4 <--- start here
    ///                     |     (incompatible)
    ///                     |
    ///                     |
    ///                  commit2
    /// ```
    #[rstest]
    fn test_walker_disjoint_nephew_expansion(test_state_space: TestStateSpace) {
        let [commit1, commit2, commit3, commit4] = test_state_space.commits();
        let base_commit = commit1.base_commit();
        let state_space = commit4.state_space();

        // Initially, pin the second NOT of commit4
        let not4_node = {
            let repl4 = commit4.replacement().unwrap();
            let hugr4 = commit4.commit_hugr();
            let [_, output] = repl4.get_replacement_io();
            let (second_not_node, _) = hugr4.single_linked_output(output, 0).unwrap();
            commit4.to_patch_node(second_not_node)
        };
        let walker = Walker::from_pinned_node(not4_node, &state_space);
        assert!(walker.is_pinned(not4_node));

        let not4_out = walker.get_wire(not4_node, OutgoingPort::from(0));
        // a single outgoing port (already pinned) => no more expansion
        assert!(walker.no_more_expansion(&not4_out, Direction::Outgoing));

        // Three options:
        // - AND gate from base
        // - XOR gate from commit3
        // - XOR gate from commit2 (which implies commit1)
        let mut exp_options = BTreeSet::from_iter([
            BTreeSet::from_iter([base_commit.id(), commit4.id()]),
            BTreeSet::from_iter([base_commit.id(), commit3.id(), commit4.id()]),
            BTreeSet::from_iter([base_commit.id(), commit1.id(), commit2.id(), commit4.id()]),
        ]);
        for new_walker in walker.expand(&not4_out, None) {
            // selected commits must be one of the valid options
            let commit_ids = new_walker
                .as_hugr_view()
                .all_commit_ids()
                .collect::<BTreeSet<_>>();
            assert!(
                exp_options.remove(&commit_ids),
                "{commit_ids:?} not an expected set of commit IDs (or duplicate)"
            );

            // new wire is complete (and thus cannot be expanded)
            let not4_out = new_walker.get_wire(not4_node, OutgoingPort::from(0));
            assert!(new_walker.is_complete(&not4_out, None));
            assert!(new_walker.no_more_expansion(&not4_out, None));

            // all nodes on wire are pinned
            let (next_node, _) = not4_out
                .all_incoming_ports(new_walker.as_hugr_view())
                .exactly_one()
                .ok()
                .unwrap();
            assert!(new_walker.is_pinned(not4_node));
            assert!(new_walker.is_pinned(next_node));

            let persistent_hugr = new_walker.into_persistent_hugr();

            // next_node is either in base (AND gate), in commit3 (XOR gate), or
            // in commit2 (XOR gate)
            let expected_optype = match next_node.0 {
                commit_id if commit_id == base_commit.id() => LogicOp::And,
                commit_id if [commit2.id(), commit3.id()].contains(&commit_id) => LogicOp::Xor,
                _ => panic!("neighbour of not4 must be in base, commit2 or commit3"),
            };
            assert_eq!(
                persistent_hugr.get_optype(next_node),
                &expected_optype.into()
            );
        }

        assert!(
            exp_options.is_empty(),
            "missing expected options: {exp_options:?}"
        );
    }

    #[rstest]
    fn test_get_wire_endpoints(test_state_space: TestStateSpace) {
        let [commit1, commit2, _commit3, commit4] = test_state_space.commits();
        let base_commit = commit1.base_commit();

        let base_and_node = {
            let base_hugr = base_commit.commit_hugr();
            let and_node = base_hugr
                .nodes()
                .find(|&n| base_hugr.get_optype(n) == &LogicOp::And.into())
                .unwrap();
            base_commit.to_patch_node(and_node)
        };

        let hugr = PersistentHugr::try_new([commit4.clone()]).unwrap();
        let (second_not_node, out_port) =
            hugr.single_outgoing_port(base_and_node, IncomingPort::from(1));
        assert_eq!(second_not_node.0, commit4.id());
        assert_eq!(out_port, OutgoingPort::from(0));

        let hugr =
            PersistentHugr::try_new([commit1.clone(), commit2.clone(), commit4.clone()]).unwrap();
        let (new_and_node, in_port) = hugr
            .all_incoming_ports(second_not_node, out_port)
            .exactly_one()
            .ok()
            .unwrap();
        assert_eq!(new_and_node.0, commit2.id());
        assert_eq!(in_port, 1.into());
    }

    /// Test that the walker handles empty replacements correctly.
    ///
    /// The base hugr is a sequence of 3 NOT gates, with a single input/output
    /// boolean. A single replacement exists in the state space, which replaces
    /// the middle NOT gate with nothing.
    #[rstest]
    fn test_walk_over_empty_repls(
        persistent_hugr_empty_child: (PersistentHugr, [CommitId; 2], [PatchNode; 3]),
    ) {
        let (hugr, [base_commit, empty_commit], [not0, not1, not2]) = persistent_hugr_empty_child;
        let state_space = hugr.state_space();
        let walker = Walker::from_pinned_node(not0, state_space);

        let not0_outwire = walker.get_wire(not0, OutgoingPort::from(0));
        let expanded_wires = walker
            .expand(&not0_outwire, Direction::Incoming)
            .collect_vec();

        assert_eq!(expanded_wires.len(), 2);

        let connected_inports: BTreeSet<_> = expanded_wires
            .iter()
            .map(|new_walker| {
                let wire = new_walker.get_wire(not0, OutgoingPort::from(0));
                wire.all_incoming_ports(new_walker.as_hugr_view())
                    .exactly_one()
                    .ok()
                    .unwrap()
            })
            .collect();

        assert_eq!(
            connected_inports,
            BTreeSet::from_iter([(not1, IncomingPort::from(0)), (not2, IncomingPort::from(0))])
        );

        let traversed_commits: BTreeSet<BTreeSet<_>> = expanded_wires
            .iter()
            .map(|new_walker| {
                let wire = new_walker.get_wire(not0, OutgoingPort::from(0));
                wire.owners().collect()
            })
            .collect();

        assert_eq!(
            traversed_commits,
            BTreeSet::from_iter([
                BTreeSet::from_iter([base_commit]),
                BTreeSet::from_iter([base_commit, empty_commit])
            ])
        );
    }

    #[rstest]
    fn test_create_commit_over_empty(
        persistent_hugr_empty_child: (PersistentHugr, [CommitId; 2], [PatchNode; 3]),
    ) {
        let (mut hugr, [base_commit, empty_commit], [not0, _not1, not2]) =
            persistent_hugr_empty_child;
        let state_space = hugr.state_space().clone();
        let mut walker = Walker {
            state_space: &state_space,
            selected_commits: hugr.clone(),
            pinned_nodes: BTreeSet::from_iter([not0]),
        };

        // wire: Not0 -> Not2 (bridging over Not1)
        let wire = walker.get_wire(not0, OutgoingPort::from(0));
        walker = walker.expand(&wire, None).exactly_one().ok().unwrap();
        let wire = walker.get_wire(not0, OutgoingPort::from(0));
        assert!(walker.is_complete(&wire, None));

        let empty_hugr = {
            let dfg_builder = DFGBuilder::new(endo_sig([bool_t()])).unwrap();
            let inputs = dfg_builder.input_wires();
            dfg_builder.finish_hugr_with_outputs(inputs).unwrap()
        };
        let commit = walker
            .try_create_commit(
                PinnedSubgraph::try_from_pinned(std::iter::empty(), [wire], &walker).unwrap(),
                RootChecked::try_new(empty_hugr).expect("Root should be DFG."),
                |node, port| {
                    assert_eq!(port.index(), 0);
                    assert!([not0, not2].contains(&node));
                    match port.direction() {
                        Direction::Incoming => OutgoingPort::from(0).into(),
                        Direction::Outgoing => IncomingPort::from(0).into(),
                    }
                },
            )
            .unwrap();

        let commit_id = hugr.try_add_commit(commit.clone()).unwrap();
        assert_eq!(
            hugr.parent_commits(commit_id).collect::<BTreeSet<_>>(),
            BTreeSet::from_iter([base_commit, empty_commit])
        );

        let res_hugr: PersistentHugr = PersistentHugr::from_commit(commit);
        assert!(res_hugr.validate().is_ok());

        // should be an empty DFG hugr
        // module root + function def + func I/O nodes + DFG entrypoint + I/O nodes
        assert_eq!(res_hugr.num_nodes(), 1 + 1 + 2 + 1 + 2);
    }

    /// Test that the walker handles empty replacements correctly.
    ///
    /// The base hugr is a sequence of 3 NOT gates, with a single input/output
    /// boolean. A single replacement exists in the state space, which replaces
    /// the middle NOT gate with nothing.
    ///
    /// In this test, we pin both the first and third NOT and see if the walker
    /// suggests to possible wires as outgoing from the first NOT. This tests
    /// the edge case in which a new wire already has all its ports pinned.
    #[rstest]
    fn test_walk_over_two_pinned_nodes(
        persistent_hugr_empty_child: (PersistentHugr, [CommitId; 2], [PatchNode; 3]),
    ) {
        let (hugr, [base_commit, empty_commit], [not0, _not1, not2]) = persistent_hugr_empty_child;
        let mut walker = Walker::from_pinned_node(not0, hugr.state_space());
        assert!(walker.try_pin_node(not2).unwrap());

        let not0_outwire = walker.get_wire(not0, OutgoingPort::from(0));
        let expanded_walkers = walker.expand(&not0_outwire, Direction::Incoming);

        let expanded_wires: BTreeSet<BTreeSet<_>> = expanded_walkers
            .map(|new_walker| {
                new_walker
                    .get_wire(not0, OutgoingPort::from(0))
                    .owners()
                    .collect()
            })
            .collect();

        assert_eq!(
            expanded_wires,
            BTreeSet::from_iter([
                BTreeSet::from_iter([base_commit]),
                BTreeSet::from_iter([base_commit, empty_commit])
            ])
        );
    }
}