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( &mut self, reg: &ExtensionRegistry, @@ -75,3 +79,30 @@ pub trait UnwrapBuilder: Dataflow { } impl 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]