feat: move unwrap builder to hugr core (#1674)

Not marking as breaking as this is not breaking since the last release.
CQCL · Nov 22, 2024 · a4db7b9 · a4db7b9
1 parent 78435b5
commit a4db7b9
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 10 deletions.
diff --git 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 → ...e/src/extension/prelude/unwrap_builder.rs b/hugr-llvm/src/utils/unwrap_builder.rs → ...e/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
@@ -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
@@ -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
@@ -118,16 +118,16 @@ 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,
         Hugr,
     };
     use rstest::rstest;
 
-    use crate::utils::UnwrapBuilder as _;
-
     use super::*;
 
     #[rstest::fixture]