1
1
//! Shared code used in both the linker and no-linker implementations of this crate.
2
2
3
- // Copyright 2021 Oxide Computer Company
3
+ // Copyright 2024 Oxide Computer Company
4
4
//
5
5
// Licensed under the Apache License, Version 2.0 (the "License");
6
6
// you may not use this file except in compliance with the License.
14
14
// See the License for the specific language governing permissions and
15
15
// limitations under the License.
16
16
17
- use crate :: { DataType , Provider } ;
17
+ use crate :: DataType ;
18
18
use proc_macro2:: TokenStream ;
19
19
use quote:: { format_ident, quote} ;
20
20
21
- // Construct function call that is used internally in the UDST-generated macros, to allow
22
- // compile-time type checking of the lambda arguments.
23
- pub fn generate_type_check (
21
+ /// Construct a function to type-check the argument closure.
22
+ ///
23
+ /// This constructs a function that is never called, but is used to ensure that
24
+ /// the closure provided to each probe macro returns arguments of the right
25
+ /// type.
26
+ pub fn construct_type_check (
24
27
provider_name : & str ,
25
- use_statements : & [ syn:: ItemUse ] ,
26
28
probe_name : & str ,
29
+ use_statements : & [ syn:: ItemUse ] ,
27
30
types : & [ DataType ] ,
28
31
) -> TokenStream {
29
- // If the probe has zero arguments, verify that the result of calling the closure is `()`
30
- // Note that there's no need to clone the closure here, since () is Copy .
32
+ // If there are zero arguments, we need to make sure we can assign the
33
+ // result of the closure to () .
31
34
if types. is_empty ( ) {
32
35
return quote ! {
33
- let __usdt_private_args_lambda = $args_lambda;
34
- let _ = || {
35
- let _: ( ) = __usdt_private_args_lambda( ) ;
36
- } ;
36
+ let _: ( ) = ( $args_lambda) ( ) ;
37
37
} ;
38
38
}
39
-
40
- // For one or more arguments, verify that we can unpack the closure into a tuple of type
41
- // `(arg0, arg1, ...)`. We verify that we can pass those arguments to a type-check function
42
- // that is _similar to_, but not exactly the probe function signature. In particular, we try to
43
- // support passing things by value or reference, and take some form of reference to that thing.
44
- // The mapping is generally:
45
- //
46
- // T or &T -> Borrow<T>
47
- // Strings -> AsRef<str>
48
- // [T; N] or &[T] -> AsRef<[T]>
49
- let type_check_args = types
39
+ let type_check_params = types
50
40
. iter ( )
51
41
. map ( |typ| match typ {
52
42
DataType :: Serializable ( ty) => {
@@ -84,28 +74,22 @@ pub fn generate_type_check(
84
74
} )
85
75
. collect :: < Vec < _ > > ( ) ;
86
76
87
- // Unpack the tuple from the closure to `args.0, args.1, ...`.
88
- let expanded_lambda_args = ( 0 ..types. len ( ) )
77
+ // Create a list of arguments `arg.0`, `arg.1`, ... to pass to the check
78
+ // function.
79
+ let type_check_args = ( 0 ..types. len ( ) )
89
80
. map ( |i| {
90
81
let index = syn:: Index :: from ( i) ;
91
82
quote ! { args. #index }
92
83
} )
93
84
. collect :: < Vec < _ > > ( ) ;
94
85
95
- let preamble = unpack_argument_lambda ( types, /* clone = */ true ) ;
96
-
97
- let type_check_function =
98
- format_ident ! ( "__usdt_private_{}_{}_type_check" , provider_name, probe_name) ;
86
+ let type_check_fn = format_ident ! ( "__usdt_private_{}_{}_type_check" , provider_name, probe_name) ;
99
87
quote ! {
100
- let __usdt_private_args_lambda = $args_lambda;
101
88
#[ allow( unused_imports) ]
102
89
#( #use_statements) *
103
90
#[ allow( non_snake_case) ]
104
- fn #type_check_function ( #( #type_check_args) , * ) { }
105
- let _ = || {
106
- #preamble
107
- #type_check_function( #( #expanded_lambda_args) , * ) ;
108
- } ;
91
+ fn #type_check_fn( #( #type_check_params) , * ) { }
92
+ let _ = || { #type_check_fn( #( #type_check_args) , * ) ; } ;
109
93
}
110
94
}
111
95
@@ -156,28 +140,23 @@ pub fn construct_probe_args(types: &[DataType]) -> (TokenStream, TokenStream) {
156
140
( destructured_arg, register_arg)
157
141
} )
158
142
. unzip ( ) ;
159
- let preamble = unpack_argument_lambda ( types, /* clone = */ false ) ;
143
+ let arg_lambda = unpack_argument_lambda ( types) ;
160
144
let unpacked_args = quote ! {
161
- #preamble
145
+ #arg_lambda
162
146
#( #unpacked_args) *
163
147
} ;
164
148
let in_regs = quote ! { #( #in_regs, ) * } ;
165
149
( unpacked_args, in_regs)
166
150
}
167
151
168
- fn unpack_argument_lambda ( types : & [ DataType ] , clone : bool ) -> TokenStream {
169
- let maybe_clone = if clone {
170
- quote ! { . clone( ) }
171
- } else {
172
- quote ! { }
173
- } ;
152
+ fn unpack_argument_lambda ( types : & [ DataType ] ) -> TokenStream {
174
153
match types. len ( ) {
175
- // Don't bother with arguments if there are none .
176
- 0 => quote ! { __usdt_private_args_lambda #maybe_clone ( ) ; } ,
154
+ // Don't bother with any closure if there are no arguments .
155
+ 0 => quote ! { } ,
177
156
// Wrap a single argument in a tuple.
178
- 1 => quote ! { let args = ( __usdt_private_args_lambda #maybe_clone ( ) , ) ; } ,
157
+ 1 => quote ! { let args = ( ( $args_lambda ) ( ) , ) ; } ,
179
158
// General case.
180
- _ => quote ! { let args = __usdt_private_args_lambda #maybe_clone ( ) ; } ,
159
+ _ => quote ! { let args = ( $args_lambda ) ( ) ; } ,
181
160
}
182
161
}
183
162
@@ -218,17 +197,18 @@ fn asm_type_convert(typ: &DataType, input: TokenStream) -> (TokenStream, TokenSt
218
197
}
219
198
}
220
199
200
+ /// Create the top-level probe macro.
201
+ ///
202
+ /// This takes the implementation block constructed elsewhere, and builds out
203
+ /// the actual macro users call in their code to fire the probe.
221
204
pub ( crate ) fn build_probe_macro (
222
205
config : & crate :: CompileProvidersConfig ,
223
- provider : & Provider ,
224
206
probe_name : & str ,
225
207
types : & [ DataType ] ,
226
208
impl_block : TokenStream ,
227
209
) -> TokenStream {
228
210
let module = config. module_ident ( ) ;
229
211
let macro_name = config. probe_ident ( probe_name) ;
230
- let type_check_block =
231
- generate_type_check ( & provider. name , & provider. use_statements , probe_name, types) ;
232
212
let no_args_match = if types. is_empty ( ) {
233
213
quote ! { ( ) => { crate :: #module:: #macro_name!( || ( ) ) } ; }
234
214
} else {
@@ -243,7 +223,6 @@ pub(crate) fn build_probe_macro(
243
223
} ;
244
224
( $args_lambda: expr) => {
245
225
{
246
- #type_check_block
247
226
#impl_block
248
227
}
249
228
} ;
@@ -263,20 +242,16 @@ mod tests {
263
242
use dtrace_parser:: Sign ;
264
243
265
244
#[ test]
266
- fn test_generate_type_check_empty ( ) {
267
- let types = & [ ] ;
245
+ fn test_construct_type_check_empty ( ) {
268
246
let expected = quote ! {
269
- let __usdt_private_args_lambda = $args_lambda;
270
- let _ = || {
271
- let _: ( ) = __usdt_private_args_lambda( ) ;
272
- } ;
247
+ let _ : ( ) = ( $args_lambda) ( ) ;
273
248
} ;
274
- let block = generate_type_check ( "" , & [ ] , "" , types ) ;
249
+ let block = construct_type_check ( "" , "" , & [ ] , & [ ] ) ;
275
250
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
276
251
}
277
252
278
253
#[ test]
279
- fn test_generate_type_check_native ( ) {
254
+ fn test_construct_type_check_native ( ) {
280
255
let provider = "provider" ;
281
256
let probe = "probe" ;
282
257
let types = & [
@@ -290,80 +265,72 @@ mod tests {
290
265
} ) ) ,
291
266
] ;
292
267
let expected = quote ! {
293
- let __usdt_private_args_lambda = $args_lambda;
294
268
#[ allow( unused_imports) ]
295
269
#[ allow( non_snake_case) ]
296
270
fn __usdt_private_provider_probe_type_check(
297
271
_: impl :: std:: borrow:: Borrow <u8 >,
298
272
_: impl :: std:: borrow:: Borrow <i64 >
299
273
) { }
300
274
let _ = || {
301
- let args = __usdt_private_args_lambda. clone( ) ( ) ;
302
275
__usdt_private_provider_probe_type_check( args. 0 , args. 1 ) ;
303
276
} ;
304
277
} ;
305
- let block = generate_type_check ( provider, & [ ] , probe , types) ;
278
+ let block = construct_type_check ( provider, probe , & [ ] , types) ;
306
279
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
307
280
}
308
281
309
282
#[ test]
310
- fn test_generate_type_check_with_string ( ) {
283
+ fn test_construct_type_check_with_string ( ) {
311
284
let provider = "provider" ;
312
285
let probe = "probe" ;
313
286
let types = & [ DataType :: Native ( dtrace_parser:: DataType :: String ) ] ;
314
287
let use_statements = vec ! [ ] ;
315
288
let expected = quote ! {
316
- let __usdt_private_args_lambda = $args_lambda;
317
289
#[ allow( unused_imports) ]
318
290
#[ allow( non_snake_case) ]
319
291
fn __usdt_private_provider_probe_type_check( _: impl AsRef <str >) { }
320
292
let _ = || {
321
- let args = ( __usdt_private_args_lambda. clone( ) ( ) , ) ;
322
293
__usdt_private_provider_probe_type_check( args. 0 ) ;
323
294
} ;
324
295
} ;
325
- let block = generate_type_check ( provider, & use_statements , probe , types) ;
296
+ let block = construct_type_check ( provider, probe , & use_statements , types) ;
326
297
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
327
298
}
328
299
329
300
#[ test]
330
- fn test_generate_type_check_with_shared_slice ( ) {
301
+ fn test_construct_type_check_with_shared_slice ( ) {
331
302
let provider = "provider" ;
332
303
let probe = "probe" ;
333
304
let types = & [ DataType :: Serializable ( syn:: parse_str ( "&[u8]" ) . unwrap ( ) ) ] ;
334
305
let use_statements = vec ! [ ] ;
335
306
let expected = quote ! {
336
- let __usdt_private_args_lambda = $args_lambda;
337
307
#[ allow( unused_imports) ]
338
308
#[ allow( non_snake_case) ]
339
309
fn __usdt_private_provider_probe_type_check( _: impl AsRef <[ u8 ] >) { }
340
310
let _ = || {
341
- let args = ( __usdt_private_args_lambda. clone( ) ( ) , ) ;
342
311
__usdt_private_provider_probe_type_check( args. 0 ) ;
343
312
} ;
344
313
} ;
345
- let block = generate_type_check ( provider, & use_statements , probe , types) ;
314
+ let block = construct_type_check ( provider, probe , & use_statements , types) ;
346
315
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
347
316
}
348
317
349
318
#[ test]
350
- fn test_generate_type_check_with_custom_type ( ) {
319
+ fn test_construct_type_check_with_custom_type ( ) {
351
320
let provider = "provider" ;
352
321
let probe = "probe" ;
353
322
let types = & [ DataType :: Serializable ( syn:: parse_str ( "MyType" ) . unwrap ( ) ) ] ;
354
323
let use_statements = vec ! [ syn:: parse2( quote! { use my_module:: MyType ; } ) . unwrap( ) ] ;
355
324
let expected = quote ! {
356
- let __usdt_private_args_lambda = $args_lambda;
357
325
#[ allow( unused_imports) ]
358
326
use my_module:: MyType ;
359
327
#[ allow( non_snake_case) ]
360
328
fn __usdt_private_provider_probe_type_check( _: impl :: std:: borrow:: Borrow <MyType >) { }
361
329
let _ = || {
362
- let args = ( __usdt_private_args_lambda. clone( ) ( ) , ) ;
363
330
__usdt_private_provider_probe_type_check( args. 0 ) ;
364
331
} ;
365
332
} ;
366
- let block = generate_type_check ( provider, & use_statements , probe , types) ;
333
+ let block = construct_type_check ( provider, probe , & use_statements , types) ;
367
334
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
368
335
}
369
336
@@ -382,7 +349,7 @@ mod tests {
382
349
let registers = [ "x0" , "x1" ] ;
383
350
let ( args, regs) = construct_probe_args ( types) ;
384
351
let expected = quote ! {
385
- let args = __usdt_private_args_lambda ( ) ;
352
+ let args = ( $args_lambda ) ( ) ;
386
353
let arg_0 = ( * <_ as :: std:: borrow:: Borrow <* const u8 >>:: borrow( & args. 0 ) as usize ) ;
387
354
let arg_1 = [ ( args. 1 . as_ref( ) as & str ) . as_bytes( ) , & [ 0_u8 ] ] . concat( ) ;
388
355
} ;
0 commit comments