From 177cccd206c39f1c5450a3b2b77b0fbfff74667b Mon Sep 17 00:00:00 2001
From: Agustin Borgna <agustin.borgna@quantinuum.com>
Date: Wed, 9 Aug 2023 16:47:53 +0100
Subject: [PATCH] Automatically add resource req for custom ops

---
 src/builder/circuit.rs | 15 ++++++++++-----
 src/ops/custom.rs      | 28 ++++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/builder/circuit.rs b/src/builder/circuit.rs
index c9677014b5..f29e3ea31b 100644
--- a/src/builder/circuit.rs
+++ b/src/builder/circuit.rs
@@ -134,6 +134,7 @@ mod test {
     use super::*;
     use cool_asserts::assert_matches;
 
+    use crate::resource::ResourceSet;
     use crate::{
         builder::{
             test::{build_main, BIT, NAT, QB},
@@ -178,12 +179,16 @@ mod test {
             "MyOp",
             "unknown op".to_string(),
             vec![],
-            Some(AbstractSignature::new(vec![QB, NAT], vec![QB], vec![])),
+            Some(
+                AbstractSignature::new(vec![QB, NAT], vec![QB], vec![])
+                    .with_resource_delta(&ResourceSet::singleton(&"MissingRsrc".into())),
+            ),
         )));
-        let build_res = build_main(
-            AbstractSignature::new_df(type_row![QB, QB, NAT], type_row![QB, QB, BIT]).pure(),
-            |mut f_build| {
-                let [q0, q1, angle]: [Wire; 3] = f_build.input_wires_arr();
+        let main_sig = AbstractSignature::new_df(type_row![QB, QB, NAT], type_row![QB, QB, BIT])
+            .with_resource_delta(&ResourceSet::singleton(&"MissingRsrc".into()))
+            .pure();
+        let build_res = build_main(main_sig, |mut f_build| {
+            let [q0, q1, angle]: [Wire; 3] = f_build.input_wires_arr();
 
                 let mut linear = f_build.as_circuit(vec![q0, q1]);
 
diff --git a/src/ops/custom.rs b/src/ops/custom.rs
index a829b1dab2..95454e4c04 100644
--- a/src/ops/custom.rs
+++ b/src/ops/custom.rs
@@ -6,7 +6,7 @@ use std::sync::Arc;
 use thiserror::Error;
 
 use crate::hugr::{HugrMut, HugrView, NodeType};
-use crate::resource::{OpDef, ResourceId, SignatureError};
+use crate::resource::{OpDef, ResourceId, ResourceSet, SignatureError};
 use crate::types::{type_param::TypeArg, AbstractSignature, SignatureDescription};
 use crate::{Hugr, Node, Resource};
 
@@ -104,8 +104,12 @@ pub struct ResourceOp {
 
 impl ResourceOp {
     /// Create a new ResourceOp given the type arguments and specified input resources
+    ///
+    /// Automatically the OpDef's `resource` in the signature.
     pub fn new(def: Arc<OpDef>, args: &[TypeArg]) -> Result<Self, SignatureError> {
-        let signature = def.compute_signature(args)?;
+        let signature = def
+            .compute_signature(args)?
+            .with_resource_delta(&ResourceSet::singleton(def.resource()));
         Ok(Self {
             def,
             args: args.to_vec(),
@@ -165,6 +169,8 @@ fn qualify_name(res_id: &ResourceId, op_name: &SmolStr) -> SmolStr {
 
 impl OpaqueOp {
     /// Creates a new OpaqueOp from all the fields we'd expect to serialize.
+    ///
+    /// Automatically includes `resource` in the signature if `signature` is provided.
     pub fn new(
         resource: ResourceId,
         op_name: impl Into<SmolStr>,
@@ -172,6 +178,8 @@ impl OpaqueOp {
         args: impl Into<Vec<TypeArg>>,
         signature: Option<AbstractSignature>,
     ) -> Self {
+        let signature =
+            signature.map(|s| s.with_resource_delta(&ResourceSet::singleton(&resource)));
         Self {
             resource,
             op_name: op_name.into(),
@@ -273,4 +281,20 @@ mod test {
         assert_eq!(op.description(), "desc");
         assert_eq!(op.args(), &[TypeArg::Type(HashableType::USize.into())]);
     }
+
+    #[test]
+    fn new_opaque_op_with_signature() {
+        let op = OpaqueOp::new(
+            "res".into(),
+            "op",
+            "desc".into(),
+            vec![TypeArg::Type(HashableType::USize.into())],
+            Some(AbstractSignature::new_linear(vec![])),
+        );
+        let op: ExternalOp = op.into();
+        assert_eq!(
+            op.signature().resource_reqs,
+            ResourceSet::singleton(&"res".into())
+        );
+    }
 }