From 3ef5bd9380142b62426070a318b6dc3f300d615d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Agust=C3=ADn=20Borgna?=
 <121866228+aborgna-q@users.noreply.github.com>
Date: Mon, 16 Dec 2024 15:03:26 +0000
Subject: [PATCH] feat: Automatically add the custom op's extension to its
 'runtime_reqs' set (#1787)

Ignores custom operations that are missing their own extension as a
runtime req when validating.

drive-by: Update doc that still refered to `extension_requirements`
---
 hugr-core/src/extension/op_def.rs         | 3 +++
 hugr-core/src/extension/resolution.rs     | 2 +-
 hugr-core/src/extension/resolution/ops.rs | 2 +-
 hugr-core/src/hugr.rs                     | 2 +-
 hugr-py/src/hugr/ext.py                   | 7 +++++++
 5 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/hugr-core/src/extension/op_def.rs b/hugr-core/src/extension/op_def.rs
index 31c576391..d5c9a5b5d 100644
--- a/hugr-core/src/extension/op_def.rs
+++ b/hugr-core/src/extension/op_def.rs
@@ -245,6 +245,9 @@ impl SignatureFunc {
             SignatureFunc::MissingValidateFunc(ts) => (ts, args),
         };
         let mut res = pf.instantiate(args)?;
+
+        // Automatically add the extensions where the operation is defined to
+        // the runtime requirements of the op.
         res.runtime_reqs.insert(def.extension.clone());
 
         // If there are any row variables left, this will fail with an error:
diff --git a/hugr-core/src/extension/resolution.rs b/hugr-core/src/extension/resolution.rs
index 0a6b85a52..90eae9422 100644
--- a/hugr-core/src/extension/resolution.rs
+++ b/hugr-core/src/extension/resolution.rs
@@ -10,7 +10,7 @@
 //! required across the HUGR.
 //!
 //! This is distinct from _runtime_ extension requirements, which are defined
-//! more granularly in each function signature by the `required_extensions`
+//! more granularly in each function signature by the `runtime_reqs`
 //! field. See the `extension_inference` feature and related modules for that.
 //!
 //! Note: These procedures are only temporary until `hugr-model` is stabilized.
diff --git a/hugr-core/src/extension/resolution/ops.rs b/hugr-core/src/extension/resolution/ops.rs
index 6af838385..3ab1491b4 100644
--- a/hugr-core/src/extension/resolution/ops.rs
+++ b/hugr-core/src/extension/resolution/ops.rs
@@ -93,7 +93,7 @@ pub(crate) fn resolve_op_extensions<'e>(
             cause: e,
         })?;
 
-    if opaque.signature() != ext_op.signature() {
+    if opaque.signature().io() != ext_op.signature().io() {
         return Err(OpaqueOpError::SignatureMismatch {
             node,
             extension: opaque.extension().clone(),
diff --git a/hugr-core/src/hugr.rs b/hugr-core/src/hugr.rs
index b97435610..57acf9ddc 100644
--- a/hugr-core/src/hugr.rs
+++ b/hugr-core/src/hugr.rs
@@ -197,7 +197,7 @@ impl Hugr {
     ///
     /// This is distinct from _runtime_ extension requirements computed in
     /// [`Hugr::infer_extensions`], which are computed more granularly in each
-    /// function signature by the `required_extensions` field and define the set
+    /// function signature by the `runtime_reqs` field and define the set
     /// of capabilities required by the runtime to execute each function.
     ///
     /// Updates the internal extension registry with the extensions used in the
diff --git a/hugr-py/src/hugr/ext.py b/hugr-py/src/hugr/ext.py
index c22ef32df..494ea3c69 100644
--- a/hugr-py/src/hugr/ext.py
+++ b/hugr-py/src/hugr/ext.py
@@ -236,6 +236,13 @@ def instantiate(
             concrete_signature: Concrete function type of the operation, only required
             if the operation is polymorphic.
         """
+        # Add the extension where the operation is defined as a runtime requirement.
+        # We don't store this in the json definition as it is redundant information.
+        if concrete_signature is not None:
+            concrete_signature = concrete_signature.with_runtime_reqs(
+                [self.get_extension().name]
+            )
+
         return ops.ExtOp(self, concrete_signature, list(args or []))