From eecdb22c957a2b7d0576644e60aae2e3222cefe8 Mon Sep 17 00:00:00 2001
From: Alan Lawrence <alan.lawrence@cambridgequantum.com>
Date: Wed, 4 Dec 2024 10:30:31 +0000
Subject: [PATCH] interpret_leaf_op is &mut to allow impls that do caching etc.

---
 hugr-passes/src/dataflow.rs         | 4 ++--
 hugr-passes/src/dataflow/datalog.rs | 6 +++---
 hugr-passes/src/dataflow/test.rs    | 3 +--
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/hugr-passes/src/dataflow.rs b/hugr-passes/src/dataflow.rs
index 5faf3d733..bb3023c38 100644
--- a/hugr-passes/src/dataflow.rs
+++ b/hugr-passes/src/dataflow.rs
@@ -20,7 +20,7 @@ use hugr_core::{Hugr, Node};
 /// must implement this trait (including providing an appropriate domain type `V`).
 pub trait DFContext<V>: ConstLoader<V> {
     /// Given lattice values for each input, update lattice values for the (dataflow) outputs.
-    /// For extension ops only, excluding [MakeTuple] and [UnpackTuple].
+    /// For extension ops only, excluding [MakeTuple] and [UnpackTuple] which are handled automatically.
     /// `_outs` is an array with one element per dataflow output, each initialized to [PartialValue::Top]
     /// which is the correct value to leave if nothing can be deduced about that output.
     /// (The default does nothing, i.e. leaves `Top` for all outputs.)
@@ -28,7 +28,7 @@ pub trait DFContext<V>: ConstLoader<V> {
     /// [MakeTuple]: hugr_core::extension::prelude::MakeTuple
     /// [UnpackTuple]: hugr_core::extension::prelude::UnpackTuple
     fn interpret_leaf_op(
-        &self,
+        &mut self,
         _node: Node,
         _e: &ExtensionOp,
         _ins: &[PartialValue<V>],
diff --git a/hugr-passes/src/dataflow/datalog.rs b/hugr-passes/src/dataflow/datalog.rs
index 490a0fb36..172d87c26 100644
--- a/hugr-passes/src/dataflow/datalog.rs
+++ b/hugr-passes/src/dataflow/datalog.rs
@@ -127,7 +127,7 @@ impl<H: HugrView, V: AbstractValue> Machine<H, V> {
 }
 
 pub(super) fn run_datalog<V: AbstractValue, H: HugrView>(
-    ctx: impl DFContext<V>,
+    mut ctx: impl DFContext<V>,
     hugr: H,
     in_wire_value_proto: Vec<(Node, IncomingPort, PV<V>)>,
 ) -> AnalysisResults<V, H> {
@@ -187,7 +187,7 @@ pub(super) fn run_datalog<V: AbstractValue, H: HugrView>(
            if !op_t.is_container(),
            if let Some(sig) = op_t.dataflow_signature(),
            node_in_value_row(n, vs),
-           if let Some(outs) = propagate_leaf_op(&ctx, &hugr, *n, &vs[..], sig.output_count()),
+           if let Some(outs) = propagate_leaf_op(&mut ctx, &hugr, *n, &vs[..], sig.output_count()),
            for (p, v) in (0..).map(OutgoingPort::from).zip(outs);
 
         // DFG --------------------
@@ -329,7 +329,7 @@ pub(super) fn run_datalog<V: AbstractValue, H: HugrView>(
 }
 
 fn propagate_leaf_op<V: AbstractValue>(
-    ctx: &impl DFContext<V>,
+    ctx: &mut impl DFContext<V>,
     hugr: &impl HugrView,
     n: Node,
     ins: &[PV<V>],
diff --git a/hugr-passes/src/dataflow/test.rs b/hugr-passes/src/dataflow/test.rs
index 6df953070..13815d186 100644
--- a/hugr-passes/src/dataflow/test.rs
+++ b/hugr-passes/src/dataflow/test.rs
@@ -387,8 +387,7 @@ fn test_cfg(
     xor_and_cfg: Hugr,
 ) {
     let root = xor_and_cfg.root();
-    let results =
-        Machine::new(&xor_and_cfg).run(TestContext, [(0.into(), inp0), (1.into(), inp1)]);
+    let results = Machine::new(&xor_and_cfg).run(TestContext, [(0.into(), inp0), (1.into(), inp1)]);
 
     assert_eq!(results.read_out_wire(Wire::new(root, 0)).unwrap(), out0);
     assert_eq!(results.read_out_wire(Wire::new(root, 1)).unwrap(), out1);