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,24 @@ 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 = call_argument_closure ( 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
+ /// Call the argument closure, assigning its output to `args`.
153
+ pub fn call_argument_closure ( types : & [ DataType ] ) -> TokenStream {
174
154
match types. len ( ) {
175
- // Don't bother with arguments if there are none .
176
- 0 => quote ! { __usdt_private_args_lambda #maybe_clone ( ) ; } ,
155
+ // Don't bother with any closure if there are no arguments .
156
+ 0 => quote ! { } ,
177
157
// Wrap a single argument in a tuple.
178
- 1 => quote ! { let args = ( __usdt_private_args_lambda #maybe_clone ( ) , ) ; } ,
158
+ 1 => quote ! { let args = ( ( $args_lambda ) ( ) , ) ; } ,
179
159
// General case.
180
- _ => quote ! { let args = __usdt_private_args_lambda #maybe_clone ( ) ; } ,
160
+ _ => quote ! { let args = ( $args_lambda ) ( ) ; } ,
181
161
}
182
162
}
183
163
@@ -218,17 +198,18 @@ fn asm_type_convert(typ: &DataType, input: TokenStream) -> (TokenStream, TokenSt
218
198
}
219
199
}
220
200
201
+ /// Create the top-level probe macro.
202
+ ///
203
+ /// This takes the implementation block constructed elsewhere, and builds out
204
+ /// the actual macro users call in their code to fire the probe.
221
205
pub ( crate ) fn build_probe_macro (
222
206
config : & crate :: CompileProvidersConfig ,
223
- provider : & Provider ,
224
207
probe_name : & str ,
225
208
types : & [ DataType ] ,
226
209
impl_block : TokenStream ,
227
210
) -> TokenStream {
228
211
let module = config. module_ident ( ) ;
229
212
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
213
let no_args_match = if types. is_empty ( ) {
233
214
quote ! { ( ) => { crate :: #module:: #macro_name!( || ( ) ) } ; }
234
215
} else {
@@ -243,7 +224,6 @@ pub(crate) fn build_probe_macro(
243
224
} ;
244
225
( $args_lambda: expr) => {
245
226
{
246
- #type_check_block
247
227
#impl_block
248
228
}
249
229
} ;
@@ -263,20 +243,16 @@ mod tests {
263
243
use dtrace_parser:: Sign ;
264
244
265
245
#[ test]
266
- fn test_generate_type_check_empty ( ) {
267
- let types = & [ ] ;
246
+ fn test_construct_type_check_empty ( ) {
268
247
let expected = quote ! {
269
- let __usdt_private_args_lambda = $args_lambda;
270
- let _ = || {
271
- let _: ( ) = __usdt_private_args_lambda( ) ;
272
- } ;
248
+ let _ : ( ) = ( $args_lambda) ( ) ;
273
249
} ;
274
- let block = generate_type_check ( "" , & [ ] , "" , types ) ;
250
+ let block = construct_type_check ( "" , "" , & [ ] , & [ ] ) ;
275
251
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
276
252
}
277
253
278
254
#[ test]
279
- fn test_generate_type_check_native ( ) {
255
+ fn test_construct_type_check_native ( ) {
280
256
let provider = "provider" ;
281
257
let probe = "probe" ;
282
258
let types = & [
@@ -290,80 +266,72 @@ mod tests {
290
266
} ) ) ,
291
267
] ;
292
268
let expected = quote ! {
293
- let __usdt_private_args_lambda = $args_lambda;
294
269
#[ allow( unused_imports) ]
295
270
#[ allow( non_snake_case) ]
296
271
fn __usdt_private_provider_probe_type_check(
297
272
_: impl :: std:: borrow:: Borrow <u8 >,
298
273
_: impl :: std:: borrow:: Borrow <i64 >
299
274
) { }
300
275
let _ = || {
301
- let args = __usdt_private_args_lambda. clone( ) ( ) ;
302
276
__usdt_private_provider_probe_type_check( args. 0 , args. 1 ) ;
303
277
} ;
304
278
} ;
305
- let block = generate_type_check ( provider, & [ ] , probe , types) ;
279
+ let block = construct_type_check ( provider, probe , & [ ] , types) ;
306
280
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
307
281
}
308
282
309
283
#[ test]
310
- fn test_generate_type_check_with_string ( ) {
284
+ fn test_construct_type_check_with_string ( ) {
311
285
let provider = "provider" ;
312
286
let probe = "probe" ;
313
287
let types = & [ DataType :: Native ( dtrace_parser:: DataType :: String ) ] ;
314
288
let use_statements = vec ! [ ] ;
315
289
let expected = quote ! {
316
- let __usdt_private_args_lambda = $args_lambda;
317
290
#[ allow( unused_imports) ]
318
291
#[ allow( non_snake_case) ]
319
292
fn __usdt_private_provider_probe_type_check( _: impl AsRef <str >) { }
320
293
let _ = || {
321
- let args = ( __usdt_private_args_lambda. clone( ) ( ) , ) ;
322
294
__usdt_private_provider_probe_type_check( args. 0 ) ;
323
295
} ;
324
296
} ;
325
- let block = generate_type_check ( provider, & use_statements , probe , types) ;
297
+ let block = construct_type_check ( provider, probe , & use_statements , types) ;
326
298
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
327
299
}
328
300
329
301
#[ test]
330
- fn test_generate_type_check_with_shared_slice ( ) {
302
+ fn test_construct_type_check_with_shared_slice ( ) {
331
303
let provider = "provider" ;
332
304
let probe = "probe" ;
333
305
let types = & [ DataType :: Serializable ( syn:: parse_str ( "&[u8]" ) . unwrap ( ) ) ] ;
334
306
let use_statements = vec ! [ ] ;
335
307
let expected = quote ! {
336
- let __usdt_private_args_lambda = $args_lambda;
337
308
#[ allow( unused_imports) ]
338
309
#[ allow( non_snake_case) ]
339
310
fn __usdt_private_provider_probe_type_check( _: impl AsRef <[ u8 ] >) { }
340
311
let _ = || {
341
- let args = ( __usdt_private_args_lambda. clone( ) ( ) , ) ;
342
312
__usdt_private_provider_probe_type_check( args. 0 ) ;
343
313
} ;
344
314
} ;
345
- let block = generate_type_check ( provider, & use_statements , probe , types) ;
315
+ let block = construct_type_check ( provider, probe , & use_statements , types) ;
346
316
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
347
317
}
348
318
349
319
#[ test]
350
- fn test_generate_type_check_with_custom_type ( ) {
320
+ fn test_construct_type_check_with_custom_type ( ) {
351
321
let provider = "provider" ;
352
322
let probe = "probe" ;
353
323
let types = & [ DataType :: Serializable ( syn:: parse_str ( "MyType" ) . unwrap ( ) ) ] ;
354
324
let use_statements = vec ! [ syn:: parse2( quote! { use my_module:: MyType ; } ) . unwrap( ) ] ;
355
325
let expected = quote ! {
356
- let __usdt_private_args_lambda = $args_lambda;
357
326
#[ allow( unused_imports) ]
358
327
use my_module:: MyType ;
359
328
#[ allow( non_snake_case) ]
360
329
fn __usdt_private_provider_probe_type_check( _: impl :: std:: borrow:: Borrow <MyType >) { }
361
330
let _ = || {
362
- let args = ( __usdt_private_args_lambda. clone( ) ( ) , ) ;
363
331
__usdt_private_provider_probe_type_check( args. 0 ) ;
364
332
} ;
365
333
} ;
366
- let block = generate_type_check ( provider, & use_statements , probe , types) ;
334
+ let block = construct_type_check ( provider, probe , & use_statements , types) ;
367
335
assert_eq ! ( block. to_string( ) , expected. to_string( ) ) ;
368
336
}
369
337
@@ -382,7 +350,7 @@ mod tests {
382
350
let registers = [ "x0" , "x1" ] ;
383
351
let ( args, regs) = construct_probe_args ( types) ;
384
352
let expected = quote ! {
385
- let args = __usdt_private_args_lambda ( ) ;
353
+ let args = ( $args_lambda ) ( ) ;
386
354
let arg_0 = ( * <_ as :: std:: borrow:: Borrow <* const u8 >>:: borrow( & args. 0 ) as usize ) ;
387
355
let arg_1 = [ ( args. 1 . as_ref( ) as & str ) . as_bytes( ) , & [ 0_u8 ] ] . concat( ) ;
388
356
} ;
0 commit comments