Skip to content

Commit cc44a11

Browse files
authored
feat: move unwrap builder to hugr core (#1674)
Not marking as breaking as this is not breaking since the last release.
1 parent 72efe47 commit cc44a11

File tree

5 files changed

+43
-10
lines changed

5 files changed

+43
-10
lines changed

hugr-core/src/extension/prelude.rs

+4
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@ use strum_macros::{EnumIter, EnumString, IntoStaticStr};
2525

2626
use super::ExtensionRegistry;
2727

28+
mod unwrap_builder;
29+
30+
pub use unwrap_builder::UnwrapBuilder;
31+
2832
/// Array type and operations.
2933
pub mod array;
3034
pub use array::{array_type, new_array_op, ArrayOp, ArrayOpDef, ARRAY_TYPE_NAME, NEW_ARRAY_OP_ID};

hugr-llvm/src/utils/unwrap_builder.rs renamed to hugr-core/src/extension/prelude/unwrap_builder.rs

+32-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use std::iter;
22

3-
use hugr_core::{
3+
use crate::{
44
builder::{BuildError, BuildHandle, Dataflow, DataflowSubContainer, SubContainer},
55
extension::{
66
prelude::{ConstError, PANIC_OP_ID, PRELUDE_ID},
@@ -12,7 +12,9 @@ use hugr_core::{
1212
};
1313
use itertools::{zip_eq, Itertools as _};
1414

15+
/// Extend dataflow builders with methods for building unwrap operations.
1516
pub trait UnwrapBuilder: Dataflow {
17+
/// Add a panic operation to the dataflow with the given error.
1618
fn add_panic(
1719
&mut self,
1820
reg: &ExtensionRegistry,
@@ -37,6 +39,8 @@ pub trait UnwrapBuilder: Dataflow {
3739
self.add_dataflow_op(op, iter::once(err).chain(input_wires))
3840
}
3941

42+
/// Build an unwrap operation for a sum type to extract the variant at the given tag
43+
/// or panic if the tag is not the expected value.
4044
fn build_unwrap_sum<const N: usize>(
4145
&mut self,
4246
reg: &ExtensionRegistry,
@@ -75,3 +79,30 @@ pub trait UnwrapBuilder: Dataflow {
7579
}
7680

7781
impl<D: Dataflow> UnwrapBuilder for D {}
82+
83+
#[cfg(test)]
84+
mod tests {
85+
use super::*;
86+
use crate::{
87+
builder::{DFGBuilder, DataflowHugr},
88+
extension::{
89+
prelude::{option_type, BOOL_T},
90+
PRELUDE_REGISTRY,
91+
},
92+
types::Signature,
93+
};
94+
95+
#[test]
96+
fn test_build_unwrap() {
97+
let mut builder =
98+
DFGBuilder::new(Signature::new(Type::from(option_type(BOOL_T)), BOOL_T).with_prelude())
99+
.unwrap();
100+
101+
let [opt] = builder.input_wires_arr();
102+
103+
let [res] = builder
104+
.build_unwrap_sum(&PRELUDE_REGISTRY, 1, option_type(BOOL_T), opt)
105+
.unwrap();
106+
builder.finish_prelude_hugr_with_outputs([res]).unwrap();
107+
}
108+
}

hugr-llvm/src/extension/prelude/array.rs

+4-4
Original file line numberDiff line numberDiff line change
@@ -389,7 +389,9 @@ mod test {
389389
use hugr_core::{
390390
builder::{Dataflow, DataflowSubContainer, SubContainer},
391391
extension::{
392-
prelude::{self, array_type, option_type, ConstUsize, BOOL_T, USIZE_T},
392+
prelude::{
393+
self, array_type, option_type, ConstUsize, UnwrapBuilder as _, BOOL_T, USIZE_T,
394+
},
393395
ExtensionRegistry,
394396
},
395397
ops::Value,
@@ -410,9 +412,7 @@ mod test {
410412
custom::CodegenExtsBuilder,
411413
emit::test::SimpleHugrConfig,
412414
test::{exec_ctx, llvm_ctx, TestContext},
413-
utils::{
414-
array_op_builder, ArrayOpBuilder, IntOpBuilder, LogicOpBuilder, UnwrapBuilder as _,
415-
},
415+
utils::{array_op_builder, ArrayOpBuilder, IntOpBuilder, LogicOpBuilder},
416416
};
417417

418418
#[rstest]

hugr-llvm/src/utils.rs

-2
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,8 @@ pub mod inline_constant_functions;
66
pub mod int_op_builder;
77
pub mod logic_op_builder;
88
pub mod type_map;
9-
pub mod unwrap_builder;
109

1110
pub use array_op_builder::ArrayOpBuilder;
1211
pub use inline_constant_functions::inline_constant_functions;
1312
pub use int_op_builder::IntOpBuilder;
1413
pub use logic_op_builder::LogicOpBuilder;
15-
pub use unwrap_builder::UnwrapBuilder;

hugr-llvm/src/utils/array_op_builder.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -118,16 +118,16 @@ pub mod test {
118118
use hugr_core::{
119119
builder::{DFGBuilder, HugrBuilder},
120120
extension::{
121-
prelude::{array_type, either_type, option_type, ConstUsize, USIZE_T},
121+
prelude::{
122+
array_type, either_type, option_type, ConstUsize, UnwrapBuilder as _, USIZE_T,
123+
},
122124
PRELUDE_REGISTRY,
123125
},
124126
types::Signature,
125127
Hugr,
126128
};
127129
use rstest::rstest;
128130

129-
use crate::utils::UnwrapBuilder as _;
130-
131131
use super::*;
132132

133133
#[rstest::fixture]

0 commit comments

Comments
 (0)