@@ -245,6 +245,10 @@ private static void EmitNoArgsMethodCallPopReturn(
245245 private TypeBuilder runnableBuilder ;
246246 private ConsumableTypeInfo consumableInfo ;
247247 private ConsumeEmitter consumeEmitter ;
248+ private ConsumableTypeInfo globalSetupReturnInfo ;
249+ private ConsumableTypeInfo globalCleanupReturnInfo ;
250+ private ConsumableTypeInfo iterationSetupReturnInfo ;
251+ private ConsumableTypeInfo iterationCleanupReturnInfo ;
248252
249253 private FieldBuilder awaitHelperField ;
250254 private FieldBuilder globalSetupActionField ;
@@ -358,13 +362,22 @@ private void InitForEmitRunnable(BenchmarkBuildInfo newBenchmark)
358362
359363 consumableInfo = new ConsumableTypeInfo ( benchmark . BenchmarkCase . Descriptor . WorkloadMethod . ReturnType ) ;
360364 consumeEmitter = ConsumeEmitter . GetConsumeEmitter ( consumableInfo ) ;
365+ globalSetupReturnInfo = GetConsumableTypeInfo ( benchmark . BenchmarkCase . Descriptor . GlobalSetupMethod ? . ReturnType ) ;
366+ globalCleanupReturnInfo = GetConsumableTypeInfo ( benchmark . BenchmarkCase . Descriptor . GlobalCleanupMethod ? . ReturnType ) ;
367+ iterationSetupReturnInfo = GetConsumableTypeInfo ( benchmark . BenchmarkCase . Descriptor . IterationSetupMethod ? . ReturnType ) ;
368+ iterationCleanupReturnInfo = GetConsumableTypeInfo ( benchmark . BenchmarkCase . Descriptor . IterationCleanupMethod ? . ReturnType ) ;
361369
362370 // Init types
363371 runnableBuilder = DefineRunnableTypeBuilder ( benchmark , moduleBuilder ) ;
364372 overheadDelegateType = EmitOverheadDelegateType ( ) ;
365373 workloadDelegateType = EmitWorkloadDelegateType ( ) ;
366374 }
367375
376+ private static ConsumableTypeInfo GetConsumableTypeInfo ( Type methodReturnType )
377+ {
378+ return methodReturnType == null ? null : new ConsumableTypeInfo ( methodReturnType ) ;
379+ }
380+
368381 private Type EmitOverheadDelegateType ( )
369382 {
370383 // .class public auto ansi sealed BenchmarkDotNet.Autogenerated.Runnable_0OverheadDelegate
@@ -908,34 +921,89 @@ private void EmitSetupCleanupMethods()
908921 {
909922 // Emit Setup/Cleanup methods
910923 // We emit empty method instead of EmptyAction = "() => { }"
911- globalSetupMethod = EmitWrapperMethod (
912- GlobalSetupMethodName ,
913- Descriptor . GlobalSetupMethod ) ;
914- globalCleanupMethod = EmitWrapperMethod (
915- GlobalCleanupMethodName ,
916- Descriptor . GlobalCleanupMethod ) ;
917- iterationSetupMethod = EmitWrapperMethod (
918- IterationSetupMethodName ,
919- Descriptor . IterationSetupMethod ) ;
920- iterationCleanupMethod = EmitWrapperMethod (
921- IterationCleanupMethodName ,
922- Descriptor . IterationCleanupMethod ) ;
924+ globalSetupMethod = EmitWrapperMethod ( GlobalSetupMethodName , Descriptor . GlobalSetupMethod , globalSetupReturnInfo ) ;
925+ globalCleanupMethod = EmitWrapperMethod ( GlobalCleanupMethodName , Descriptor . GlobalCleanupMethod , globalCleanupReturnInfo ) ;
926+ iterationSetupMethod = EmitWrapperMethod ( IterationSetupMethodName , Descriptor . IterationSetupMethod , iterationSetupReturnInfo ) ;
927+ iterationCleanupMethod = EmitWrapperMethod ( IterationCleanupMethodName , Descriptor . IterationCleanupMethod , iterationCleanupReturnInfo ) ;
923928 }
924929
925- private MethodBuilder EmitWrapperMethod ( string methodName , MethodInfo optionalTargetMethod )
930+ private MethodBuilder EmitWrapperMethod ( string methodName , MethodInfo optionalTargetMethod , ConsumableTypeInfo returnTypeInfo )
926931 {
927932 var methodBuilder = runnableBuilder . DefinePrivateVoidInstanceMethod ( methodName ) ;
928933
929934 var ilBuilder = methodBuilder . GetILGenerator ( ) ;
930935
931936 if ( optionalTargetMethod != null )
932- EmitNoArgsMethodCallPopReturn ( methodBuilder , optionalTargetMethod , ilBuilder , forceDirectCall : true ) ;
937+ {
938+ if ( returnTypeInfo ? . IsAwaitable == true )
939+ {
940+ EmitAwaitableSetupTeardown ( methodBuilder , optionalTargetMethod , ilBuilder , returnTypeInfo ) ;
941+ }
942+ else
943+ {
944+ EmitNoArgsMethodCallPopReturn ( methodBuilder , optionalTargetMethod , ilBuilder , forceDirectCall : true ) ;
945+ }
946+ }
933947
934948 ilBuilder . EmitVoidReturn ( methodBuilder ) ;
935949
936950 return methodBuilder ;
937951 }
938952
953+ private void EmitAwaitableSetupTeardown (
954+ MethodBuilder methodBuilder ,
955+ MethodInfo targetMethod ,
956+ ILGenerator ilBuilder ,
957+ ConsumableTypeInfo returnTypeInfo )
958+ {
959+ if ( targetMethod == null )
960+ throw new ArgumentNullException ( nameof ( targetMethod ) ) ;
961+
962+ if ( returnTypeInfo . WorkloadMethodReturnType == typeof ( void ) )
963+ {
964+ ilBuilder . Emit ( OpCodes . Ldarg_0 ) ;
965+ }
966+ /*
967+ IL_0000: ldarg.0
968+ IL_0001: ldfld class BenchmarkDotNet.Helpers.AwaitHelper BenchmarkDotNet.Helpers.Runnable_0::awaitHelper
969+ */
970+ ilBuilder . Emit ( OpCodes . Ldarg_0 ) ;
971+ ilBuilder . Emit ( OpCodes . Ldfld , awaitHelperField ) ;
972+ /*
973+ // call for instance
974+ // GlobalSetup();
975+ IL_0006: ldarg.0
976+ IL_0007: call instance void [BenchmarkDotNet]BenchmarkDotNet.Samples.SampleBenchmark::GlobalSetup()
977+ */
978+ /*
979+ // call for static
980+ // GlobalSetup();
981+ IL_0006: call string [BenchmarkDotNet]BenchmarkDotNet.Samples.SampleBenchmark::GlobalCleanup()
982+ */
983+ if ( targetMethod . IsStatic )
984+ {
985+ ilBuilder . Emit ( OpCodes . Call , targetMethod ) ;
986+
987+ }
988+ else if ( methodBuilder . IsStatic )
989+ {
990+ throw new InvalidOperationException (
991+ $ "[BUG] Static method { methodBuilder . Name } tries to call instance member { targetMethod . Name } ") ;
992+ }
993+ else
994+ {
995+ ilBuilder . Emit ( OpCodes . Ldarg_0 ) ;
996+ ilBuilder . Emit ( OpCodes . Call , targetMethod ) ;
997+ }
998+
999+ /*
1000+ // awaitHelper.GetResult(...);
1001+ IL_000e: callvirt instance void BenchmarkDotNet.Helpers.AwaitHelper::GetResult(class [System.Private.CoreLib]System.Threading.Tasks.Task)
1002+ */
1003+ ilBuilder . Emit ( OpCodes . Callvirt , returnTypeInfo . GetResultMethod ) ;
1004+ ilBuilder . Emit ( OpCodes . Pop ) ;
1005+ }
1006+
9391007 private void EmitCtorBody ( )
9401008 {
9411009 var ilBuilder = ctorMethod . GetILGenerator ( ) ;
0 commit comments