@@ -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