@@ -35,9 +35,9 @@ pub fn seamock(_args: TokenStream, input: TokenStream) -> TokenStream {
35
35
generate_attr_names ( method, & [ "times" ] )
36
36
} ) ;
37
37
38
- let mut returning_attrs = vec ! { } ;
39
- let mut with_attrs = vec ! { } ;
40
- let mut with_methods = vec ! { } ;
38
+ let mut returning_attrs = vec ! [ ] ;
39
+ let mut with_attrs = vec ! [ ] ;
40
+ let mut with_methods = vec ! [ ] ;
41
41
42
42
let ret = trait_methods. clone ( ) . flat_map ( |method| {
43
43
let method_output = & method. sig . output ;
@@ -134,8 +134,7 @@ pub fn seamock(_args: TokenStream, input: TokenStream) -> TokenStream {
134
134
let method_name = & method. sig . ident ;
135
135
let method_output = & method. sig . output ;
136
136
let method_inputs = & method. sig . inputs ;
137
- let mut params = vec ! { } ;
138
- // For each argument, create WithVal<T> where T is the argument type
137
+ let mut params = vec ! [ ] ;
139
138
for arg in method_inputs. iter ( ) {
140
139
if let syn:: FnArg :: Typed ( pat_type) = arg {
141
140
let arg_name = match & * pat_type. pat {
@@ -149,15 +148,49 @@ pub fn seamock(_args: TokenStream, input: TokenStream) -> TokenStream {
149
148
let times_attr = Ident :: new ( & format ! ( "times_{}" , & method. sig. ident) , method. sig . ident . span ( ) ) ;
150
149
let max_times_attr = Ident :: new ( & format ! ( "max_times_{}" , & method. sig. ident) , method. sig . ident . span ( ) ) ;
151
150
let ret_func = Ident :: new ( & format ! ( "val_returning_{}" , & method. sig. ident) , method. sig . ident . span ( ) ) ;
152
- let error = format ! ( "Hit times limit for {}" , & method. sig. ident) ;
151
+ let max_attr_error = format ! ( "Hit times limit for {}" , & method. sig. ident) ;
152
+ let val_attr = Ident :: new ( & format ! ( "val_with_{}" , & method. sig. ident) , method. sig . ident . span ( ) ) ;
153
+ let val_error = format ! ( "Called {} with incorrect parameters" , & method. sig. ident) ;
154
+
155
+ let expected_val_logic = params. iter ( ) . enumerate ( ) . map ( |( i, _) | {
156
+ let idx = syn:: Index :: from ( i) ;
157
+ Some ( quote ! {
158
+ let with_val = & tuple. #idx;
159
+ let input = & params. #idx;
160
+ val_match = val_match && match with_val {
161
+ WithVal :: Gt ( val) => input > val,
162
+ WithVal :: Gte ( val) => input >= val,
163
+ WithVal :: Lt ( val) => input < val,
164
+ WithVal :: Lte ( val) => input <= val,
165
+ WithVal :: Eq ( val) => input == val,
166
+ } ;
167
+ } )
168
+ } ) ;
169
+
170
+ let with_matching = if params. len ( ) > 0 {
171
+ quote ! {
172
+ if let Some ( tuple) = & self . #val_attr {
173
+ let params = ( #( #params) * ) ;
174
+ let mut val_match = true ;
175
+ #(
176
+ #expected_val_logic
177
+ ) *
178
+ if !val_match {
179
+ sea:: sea_printf!( #val_error, & self . #val_attr) ;
180
+ verifier:: vassert!( false ) ;
181
+ }
182
+ }
183
+ }
184
+ } else { quote ! { } } ;
153
185
154
186
Some ( quote ! {
155
187
fn #method_name( #method_inputs) #method_output {
156
188
self . #times_attr. replace_with( |& mut old| old + 1 ) ;
157
- if ( * self . #times_attr. borrow( ) > self . #max_times_attr) {
158
- sea:: sea_printf!( #error , self . #max_times_attr) ;
189
+ if * self . #times_attr. borrow( ) > self . #max_times_attr {
190
+ sea:: sea_printf!( #max_attr_error , self . #max_times_attr) ;
159
191
verifier:: vassert!( false ) ;
160
192
}
193
+ #with_matching
161
194
( self . #ret_func) ( #( #params) * )
162
195
}
163
196
} )
@@ -220,6 +253,8 @@ pub fn seamock(_args: TokenStream, input: TokenStream) -> TokenStream {
220
253
// Combine the generated tokens
221
254
let expanded = quote ! {
222
255
use core:: cell:: RefCell ;
256
+ #[ macro_use]
257
+ extern crate alloc;
223
258
enum WithVal <T > {
224
259
Gt ( T ) ,
225
260
Gte ( T ) ,
0 commit comments