From a4db7b9ed6026398dab479b8eb2db93999944af7 Mon Sep 17 00:00:00 2001
From: Seyon Sivarajah <seyon.sivarajah@quantinuum.com>
Date: Fri, 22 Nov 2024 15:41:26 +0000
Subject: [PATCH] feat: move unwrap builder to hugr core (#1674)

Not marking as breaking as this is not breaking since the last release.
---
 hugr-core/src/extension/prelude.rs            |  4 +++
 .../src/extension/prelude}/unwrap_builder.rs  | 33 ++++++++++++++++++-
 hugr-llvm/src/extension/prelude/array.rs      |  8 ++---
 hugr-llvm/src/utils.rs                        |  2 --
 hugr-llvm/src/utils/array_op_builder.rs       |  6 ++--
 5 files changed, 43 insertions(+), 10 deletions(-)
 rename {hugr-llvm/src/utils => hugr-core/src/extension/prelude}/unwrap_builder.rs (73%)

diff --git a/hugr-core/src/extension/prelude.rs b/hugr-core/src/extension/prelude.rs
index 8e0eb0f4b..d48ad2815 100644
--- a/hugr-core/src/extension/prelude.rs
+++ b/hugr-core/src/extension/prelude.rs
@@ -25,6 +25,10 @@ use strum_macros::{EnumIter, EnumString, IntoStaticStr};
 
 use super::ExtensionRegistry;
 
+mod unwrap_builder;
+
+pub use unwrap_builder::UnwrapBuilder;
+
 /// Array type and operations.
 pub mod array;
 pub use array::{array_type, new_array_op, ArrayOp, ArrayOpDef, ARRAY_TYPE_NAME, NEW_ARRAY_OP_ID};
diff --git a/hugr-llvm/src/utils/unwrap_builder.rs b/hugr-core/src/extension/prelude/unwrap_builder.rs
similarity index 73%
rename from hugr-llvm/src/utils/unwrap_builder.rs
rename to hugr-core/src/extension/prelude/unwrap_builder.rs
index eef5d8527..7d7f3f5b3 100644
--- a/hugr-llvm/src/utils/unwrap_builder.rs
+++ b/hugr-core/src/extension/prelude/unwrap_builder.rs
@@ -1,6 +1,6 @@
 use std::iter;
 
-use hugr_core::{
+use crate::{
     builder::{BuildError, BuildHandle, Dataflow, DataflowSubContainer, SubContainer},
     extension::{
         prelude::{ConstError, PANIC_OP_ID, PRELUDE_ID},
@@ -12,7 +12,9 @@ use hugr_core::{
 };
 use itertools::{zip_eq, Itertools as _};
 
+/// Extend dataflow builders with methods for building unwrap operations.
 pub trait UnwrapBuilder: Dataflow {
+    /// Add a panic operation to the dataflow with the given error.
     fn add_panic(
         &mut self,
         reg: &ExtensionRegistry,
@@ -37,6 +39,8 @@ pub trait UnwrapBuilder: Dataflow {
         self.add_dataflow_op(op, iter::once(err).chain(input_wires))
     }
 
+    /// Build an unwrap operation for a sum type to extract the variant at the given tag
+    /// or panic if the tag is not the expected value.
     fn build_unwrap_sum<const N: usize>(
         &mut self,
         reg: &ExtensionRegistry,
@@ -75,3 +79,30 @@ pub trait UnwrapBuilder: Dataflow {
 }
 
 impl<D: Dataflow> UnwrapBuilder for D {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        builder::{DFGBuilder, DataflowHugr},
+        extension::{
+            prelude::{option_type, BOOL_T},
+            PRELUDE_REGISTRY,
+        },
+        types::Signature,
+    };
+
+    #[test]
+    fn test_build_unwrap() {
+        let mut builder =
+            DFGBuilder::new(Signature::new(Type::from(option_type(BOOL_T)), BOOL_T).with_prelude())
+                .unwrap();
+
+        let [opt] = builder.input_wires_arr();
+
+        let [res] = builder
+            .build_unwrap_sum(&PRELUDE_REGISTRY, 1, option_type(BOOL_T), opt)
+            .unwrap();
+        builder.finish_prelude_hugr_with_outputs([res]).unwrap();
+    }
+}
diff --git a/hugr-llvm/src/extension/prelude/array.rs b/hugr-llvm/src/extension/prelude/array.rs
index 8aff7eb60..96f5d8c24 100644
--- a/hugr-llvm/src/extension/prelude/array.rs
+++ b/hugr-llvm/src/extension/prelude/array.rs
@@ -389,7 +389,9 @@ mod test {
     use hugr_core::{
         builder::{Dataflow, DataflowSubContainer, SubContainer},
         extension::{
-            prelude::{self, array_type, option_type, ConstUsize, BOOL_T, USIZE_T},
+            prelude::{
+                self, array_type, option_type, ConstUsize, UnwrapBuilder as _, BOOL_T, USIZE_T,
+            },
             ExtensionRegistry,
         },
         ops::Value,
@@ -410,9 +412,7 @@ mod test {
         custom::CodegenExtsBuilder,
         emit::test::SimpleHugrConfig,
         test::{exec_ctx, llvm_ctx, TestContext},
-        utils::{
-            array_op_builder, ArrayOpBuilder, IntOpBuilder, LogicOpBuilder, UnwrapBuilder as _,
-        },
+        utils::{array_op_builder, ArrayOpBuilder, IntOpBuilder, LogicOpBuilder},
     };
 
     #[rstest]
diff --git a/hugr-llvm/src/utils.rs b/hugr-llvm/src/utils.rs
index 07d683877..7285aa141 100644
--- a/hugr-llvm/src/utils.rs
+++ b/hugr-llvm/src/utils.rs
@@ -6,10 +6,8 @@ pub mod inline_constant_functions;
 pub mod int_op_builder;
 pub mod logic_op_builder;
 pub mod type_map;
-pub mod unwrap_builder;
 
 pub use array_op_builder::ArrayOpBuilder;
 pub use inline_constant_functions::inline_constant_functions;
 pub use int_op_builder::IntOpBuilder;
 pub use logic_op_builder::LogicOpBuilder;
-pub use unwrap_builder::UnwrapBuilder;
diff --git a/hugr-llvm/src/utils/array_op_builder.rs b/hugr-llvm/src/utils/array_op_builder.rs
index 3fd8ba782..eb4f542c6 100644
--- a/hugr-llvm/src/utils/array_op_builder.rs
+++ b/hugr-llvm/src/utils/array_op_builder.rs
@@ -118,7 +118,9 @@ pub mod test {
     use hugr_core::{
         builder::{DFGBuilder, HugrBuilder},
         extension::{
-            prelude::{array_type, either_type, option_type, ConstUsize, USIZE_T},
+            prelude::{
+                array_type, either_type, option_type, ConstUsize, UnwrapBuilder as _, USIZE_T,
+            },
             PRELUDE_REGISTRY,
         },
         types::Signature,
@@ -126,8 +128,6 @@ pub mod test {
     };
     use rstest::rstest;
 
-    use crate::utils::UnwrapBuilder as _;
-
     use super::*;
 
     #[rstest::fixture]