feat: Add ResourceScope

Author: lmondada

Cargo.lock

@@ -87,6 +87,8 @@ typetag = "0.2.20"
 cool_asserts = "2.0.3"
 zstd = "0.13.3"
 anyhow = "1.0.98"
+num-rational = "0.4"
+uuid = { version = "1.0", features = ["v4"] }
 
 [profile.release.package.tket-py]
 # Some configurations to reduce the size of tket wheels

Cargo.toml

@@ -76,6 +76,8 @@ pest = { workspace = true }
 pest_derive = { workspace = true }
 zstd = { workspace = true, optional = true }
 anyhow = { workspace = true, optional = true }
+num-rational = { workspace = true }
+uuid = { workspace = true }
 
 
 [dev-dependencies]
@@ -86,6 +88,7 @@ cool_asserts = { workspace = true }
 # Defined here so it can be overridden by the codspeed CI job
 # using `cargo add`.
 criterion = { version = "0.7.0", features = ["html_reports"] }
+insta = "1.43.1"
 
 [[bench]]
 name = "bench_main"

tket/Cargo.toml

@@ -49,6 +49,7 @@ pub mod extension;
 pub(crate) mod ops;
 pub mod optimiser;
 pub mod passes;
+pub mod resource;
 pub mod rewrite;
 pub mod serialize;
 

tket/src/lib.rs

@@ -0,0 +1,160 @@
+//! Resource tracking within HUGR DFG subgraphs.
+//!
+//! This module implements the resource tracking system. It provides
+//! facilities for tracking linear resources (such as qubits) through quantum
+//! circuits represented as HUGR subgraphs.
+//!
+//! # Overview
+//!
+//! HUGR has a notion of "Value": the data that corresponds to a wire within a
+//! dataflow graph. It further has a notion of "linear value" a.k.a non-copyable
+//! value: a value that cannot be copied or discarded (implicitly).
+//!
+//! As far as HUGR is concerned, a linear value (or any value, for that matter)
+//! is born at an op's output and dies at the next op's input. TKET introduces
+//! the notion of "Resource" to extend the lifetime of a linear value over
+//! multiple ops.
+//!
+//! If a linear value appears both in an op's input and output, we say that it
+//! is "resource-preserving". Using [`ResourceFlow`], we can track resources
+//! as they "flow" through multiple operations. The chains of
+//! resource-preserving ops acting on a same resource form a so-called resource
+//! path.
+//!
+//! # Resources and Copyable Values
+//!
+//! Resource tracking distinguishes between two types of values:
+//!
+//! - **Linear resources**: Non-copyable values that form resource paths through
+//!   the circuit. Each resource has a unique [`ResourceId`] and operations on
+//!   the same resource are ordered by a [`Position`].
+//! - **Copyable values**: Regular values that can be copied and discarded
+//!   freely. Each is identified by a unique [`CopyableValueId`].
+//!
+//! # Resource Scope
+//!
+//! Tracking resources is not free: there is a one-off linear cost to compute
+//! the resource paths, plus a linear memory cost to store them.
+//!
+//! Use a [`SiblingSubgraph`] to define a region of a `HUGR`, within which
+//! resources should be tracked. You can then construct a resource-tracked scope
+//! using [`ResourceScope::new`].
+
+// Public API exports
+pub use flow::{DefaultResourceFlow, ResourceFlow, UnsupportedOp};
+pub use scope::{ResourceScope, ResourceScopeConfig};
+pub use types::{CopyableValueId, OpValue, Position, ResourceAllocator, ResourceId};
+
+// Internal modules
+mod flow;
+mod scope;
+mod types;
+
+#[cfg(test)]
+mod tests {
+    use hugr::{
+        builder::{DFGBuilder, Dataflow, DataflowHugr},
+        extension::prelude::qb_t,
+        hugr::views::SiblingSubgraph,
+        ops::handle::DataflowParentID,
+        types::Signature,
+        CircuitUnit, Hugr,
+    };
+
+    use itertools::Itertools;
+    use rstest::rstest;
+
+    use crate::{
+        extension::rotation::{rotation_type, ConstRotation},
+        resource::scope::tests::ResourceScopeReport,
+        TketOp,
+    };
+
+    use super::ResourceScope;
+
+    // Gate being commuted has a non-linear input
+    fn circ(n_qubits: usize, add_rz: bool, add_const_rz: bool) -> Hugr {
+        let build = || {
+            let out_qb_row = vec![qb_t(); n_qubits];
+            let mut inp_qb_row = out_qb_row.clone();
+            if add_rz {
+                inp_qb_row.push(rotation_type());
+            };
+            let mut dfg = DFGBuilder::new(Signature::new(inp_qb_row, out_qb_row))?;
+
+            let (qubits, f) = if add_rz {
+                let mut inputs = dfg.input_wires().collect_vec();
+                let f = inputs.pop().unwrap();
+                (inputs, Some(f))
+            } else {
+                (dfg.input_wires().collect_vec(), None)
+            };
+
+            let mut circ = dfg.as_circuit(qubits);
+
+            for i in 0..n_qubits {
+                circ.append(TketOp::H, [i])?;
+            }
+            for i in (0..n_qubits).step_by(2) {
+                if i + 1 < n_qubits {
+                    circ.append(TketOp::CX, [i, i + 1])?;
+                }
+            }
+            if let Some(f) = f {
+                for i in 0..n_qubits {
+                    circ.append_and_consume(
+                        TketOp::Rz,
+                        [CircuitUnit::Linear(i), CircuitUnit::Wire(f)],
+                    )?;
+                }
+            }
+            if add_const_rz {
+                let const_angle = circ.add_constant(ConstRotation::PI_2);
+                for i in 0..n_qubits {
+                    circ.append_and_consume(
+                        TketOp::Rz,
+                        [CircuitUnit::Linear(i), CircuitUnit::Wire(const_angle)],
+                    )?;
+                }
+            }
+            let qbs = circ.finish();
+            dfg.finish_hugr_with_outputs(qbs)
+        };
+        build().unwrap()
+    }
+
+    #[rstest]
+    #[case(2, false, false)]
+    #[case(2, true, false)]
+    #[case(2, false, true)]
+    #[case(2, true, true)]
+    #[case(4, false, false)]
+    #[case(4, true, false)]
+    #[case(4, false, true)]
+    #[case(4, true, true)]
+    fn test_resource_scope_creation(
+        #[case] n_qubits: usize,
+        #[case] add_rz: bool,
+        #[case] add_const_rz: bool,
+    ) {
+        let circ = circ(n_qubits, add_rz, add_const_rz);
+        let subgraph =
+            SiblingSubgraph::try_new_dataflow_subgraph::<_, DataflowParentID>(&circ).unwrap();
+        let scope = ResourceScope::new(&circ, subgraph);
+        let info = ResourceScopeReport::from(&scope);
+
+        let mut name = format!("{n_qubits}_qubits");
+        if add_rz {
+            name.push('_');
+            name.push_str("add_rz");
+        }
+        if add_const_rz {
+            name.push('_');
+            name.push_str("add_const_rz");
+        }
+
+        assert_eq!(info.resource_paths.len(), n_qubits);
+        assert_eq!(info.n_copyable, add_const_rz as usize + add_rz as usize);
+        insta::assert_snapshot!(name, info);
+    }
+}

tket/src/resource.rs

@@ -0,0 +1,136 @@
+//! Resource flow logic for tracking how resources move through operations.
+//!
+//! This module defines the [`ResourceFlow`] trait that specifies how resources
+//! flow through operations, along with a default implementation.
+
+use crate::resource::types::ResourceId;
+use derive_more::derive::{Display, Error};
+use hugr::ops::{OpTrait, OpType};
+use hugr::types::Type;
+use itertools::{EitherOrBoth, Itertools};
+
+/// Error type for unsupported operations in ResourceFlow implementations.
+#[derive(Debug, Display, Clone, PartialEq, Error)]
+#[display("Unsupported operation '{op_type:?}'")]
+pub struct UnsupportedOp {
+    /// The operation type that is unsupported.
+    pub op_type: OpType,
+}
+
+/// Trait for specifying how resources flow through operations.
+///
+/// This trait allows different implementations to define how linear resources
+/// are mapped from inputs to outputs through various operation types.
+pub trait ResourceFlow {
+    /// Map resource IDs from operation inputs to outputs.
+    ///
+    /// Takes an operation type and the resource IDs of the operation's inputs.
+    /// The i-th entry is Some(resource_id) if the i-th port is a linear type,
+    /// None otherwise. Returns the resource IDs of the operation's outputs
+    /// in port order. Output resource IDs should be one of the input resource
+    /// IDs for resource-preserving operations, or None for new resources or
+    /// non-linear types.
+    ///
+    /// # Arguments
+    /// * `op` - The operation type
+    /// * `inputs` - Resource IDs for each input port (None for non-linear
+    ///   types)
+    ///
+    /// # Returns
+    /// Resource IDs for each output port, or UnsupportedOp if the operation
+    /// cannot be handled by this implementation.
+    fn map_resources(
+        &self,
+        op: &OpType,
+        inputs: &[Option<ResourceId>],
+    ) -> Result<Vec<Option<ResourceId>>, UnsupportedOp>;
+}
+
+/// Default implementation of ResourceFlow.
+///
+/// This implementation considers that an operation is resource-preserving if
+/// whenever the i-th input or output is linear, then the i-th input type
+/// matches the i-th output. The i-th input is then mapped to the i-th output.
+///
+/// Otherwise, all input resources are discarded and all outputs will be given
+/// fresh resource IDs.
+#[derive(Debug, Clone, Default)]
+pub struct DefaultResourceFlow;
+
+impl DefaultResourceFlow {
+    /// Create a new DefaultResourceFlow instance.
+    pub fn new() -> Self {
+        Self
+    }
+
+    /// Check if a type is linear (non-copyable).
+    fn is_linear_type(ty: &Type) -> bool {
+        !ty.copyable()
+    }
+
+    /// Determine if an operation is resource-preserving based on input/output
+    /// types.
+    fn is_resource_preserving(input_types: &[Type], output_types: &[Type]) -> bool {
+        // An operation is resource-preserving if for each i, if input[i] or
+        // output[i] is linear, then type(input[i]) == type(output[i])
+
+        for io_ty in input_types.iter().zip_longest(output_types.iter()) {
+            let (input_ty, output_ty) = match io_ty {
+                EitherOrBoth::Both(input_ty, output_ty) => (input_ty, output_ty),
+                EitherOrBoth::Left(ty) | EitherOrBoth::Right(ty) => {
+                    if Self::is_linear_type(ty) {
+                        // linear type on one side, nothing on the other
+                        return false;
+                    }
+                    continue;
+                }
+            };
+
+            if Self::is_linear_type(input_ty) || Self::is_linear_type(output_ty) {
+                // If input/output is linear, both must be the same type
+                if input_ty != output_ty {
+                    return false;
+                }
+            }
+        }
+
+        true
+    }
+}
+
+impl ResourceFlow for DefaultResourceFlow {
+    fn map_resources(
+        &self,
+        op: &OpType,
+        inputs: &[Option<ResourceId>],
+    ) -> Result<Vec<Option<ResourceId>>, UnsupportedOp> {
+        let signature = op.dataflow_signature().expect("dataflow op");
+        let input_types = signature.input_types();
+        let output_types = signature.output_types();
+
+        debug_assert_eq!(
+            inputs.len(),
+            input_types.len(),
+            "Input resource array length must match operation input count"
+        );
+
+        if Self::is_resource_preserving(input_types, output_types) {
+            Ok(retain_linear_types(inputs.to_vec(), output_types))
+        } else {
+            // Resource-producing/consuming: all linear outputs are new resources (None)
+            Ok(vec![None; output_types.len()])
+        }
+    }
+}
+
+fn retain_linear_types(
+    mut resources: Vec<Option<ResourceId>>,
+    types: &[Type],
+) -> Vec<Option<ResourceId>> {
+    for (ty, resource) in types.iter().zip(resources.iter_mut()) {
+        if ty.copyable() {
+            *resource = None;
+        }
+    }
+    resources
+}