1
1
use std:: iter;
2
2
3
- use hugr_core :: {
3
+ use crate :: {
4
4
builder:: { BuildError , BuildHandle , Dataflow , DataflowSubContainer , SubContainer } ,
5
5
extension:: {
6
6
prelude:: { ConstError , PANIC_OP_ID , PRELUDE_ID } ,
@@ -12,7 +12,9 @@ use hugr_core::{
12
12
} ;
13
13
use itertools:: { zip_eq, Itertools as _} ;
14
14
15
+ /// Extend dataflow builders with methods for building unwrap operations.
15
16
pub trait UnwrapBuilder : Dataflow {
17
+ /// Add a panic operation to the dataflow with the given error.
16
18
fn add_panic (
17
19
& mut self ,
18
20
reg : & ExtensionRegistry ,
@@ -37,6 +39,8 @@ pub trait UnwrapBuilder: Dataflow {
37
39
self . add_dataflow_op ( op, iter:: once ( err) . chain ( input_wires) )
38
40
}
39
41
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.
40
44
fn build_unwrap_sum < const N : usize > (
41
45
& mut self ,
42
46
reg : & ExtensionRegistry ,
@@ -75,3 +79,30 @@ pub trait UnwrapBuilder: Dataflow {
75
79
}
76
80
77
81
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
+ }
0 commit comments