Adjust arguments for local state tracing. Add tests for CodeCoverageInstrumenter

Author: Julien Couvreur

src/Compilers/CSharp/Portable/FlowAnalysis/NullableWalker.cs

@@ -6660,14 +6660,15 @@ static ImmutableArray<ParameterSymbol> getParameters(ImmutableArray<ParameterSym
             }
         }
 
-        private static ImmutableArray<RefKind> GetArgumentRefKinds(ImmutableArray<RefKind> argumentRefKindsOpt, bool isNewExtension,
+        internal static ImmutableArray<RefKind> GetArgumentRefKinds(ImmutableArray<RefKind> argumentRefKindsOpt, bool adjustForNewExtension,
             MethodSymbol method, int argumentCount)
         {
-            if (!isNewExtension)
+            if (!adjustForNewExtension)
             {
                 return argumentRefKindsOpt;
             }
 
+            Debug.Assert(method.GetIsNewExtensionMember());
             RefKind receiverRefKind = GetExtensionReceiverRefKind(method);
 
             if (argumentRefKindsOpt.IsDefault)

src/Compilers/CSharp/Portable/Lowering/ExtensionMethodBodyRewriter.cs

@@ -66,7 +66,7 @@ public ExtensionMethodBodyRewriter(MethodSymbol sourceMethod, SourceExtensionImp
 
                 foreach (var parameter in symbol.Parameters)
                 {
-                    builder.Add(parameter, rewrittenParameters[parameter.Ordinal]); // TODO2
+                    builder.Add(parameter, rewrittenParameters[parameter.Ordinal]);
                 }
 
                 _symbolMap = builder.ToImmutable();
@@ -215,10 +215,7 @@ public override BoundNode VisitUnaryOperator(BoundUnaryOperator node)
 
         public override BoundNode? VisitMethodDefIndex(BoundMethodDefIndex node)
         {
-            MethodSymbol method = node.Method;
-            Debug.Assert(method.IsDefinition);
-
-            method = ExtensionMethodReferenceRewriter.VisitMethodSymbolWithExtensionRewrite(this, method);
+            MethodSymbol method = ExtensionMethodReferenceRewriter.VisitMethodSymbolWithExtensionRewrite(this, node.Method);
             TypeSymbol? type = this.VisitType(node.Type);
             return node.Update(method, type);
         }

src/Compilers/CSharp/Portable/Lowering/ExtensionMethodReferenceRewriter.cs

@@ -194,10 +194,7 @@ method.OriginalDefinition is ErrorMethodSymbol ||
 
         public override BoundNode? VisitMethodDefIndex(BoundMethodDefIndex node)
         {
-            MethodSymbol method = node.Method;
-            Debug.Assert(method.IsDefinition);
-            // TODO2 Is it possible to have a MethodDef to extension method, from outside the method?
-
+            MethodSymbol method = VisitMethodSymbolWithExtensionRewrite(this, node.Method);
             TypeSymbol? type = this.VisitType(node.Type);
             return node.Update(method, type);
         }

src/Compilers/CSharp/Portable/Lowering/Instrumentation/CodeCoverageInstrumenter.cs

@@ -263,7 +263,7 @@ public override void InstrumentBlock(BoundBlock original, LocalRewriter rewriter
                     _methodBodyFactory.Local(_methodPayload),
                     _methodBodyFactory.ArrayAccess(
                         _methodBodyFactory.InstrumentationPayloadRoot(analysisKind, modulePayloadType),
-                        ImmutableArray.Create(_methodBodyFactory.MethodDefIndex(_method)))); // TODO2 test this scenario
+                        ImmutableArray.Create(_methodBodyFactory.MethodDefIndex(_method))));
 
             BoundExpression mvid = _methodBodyFactory.ModuleVersionId();
             BoundExpression methodToken = _methodBodyFactory.MethodDefIndex(_method);

src/Compilers/CSharp/Portable/Lowering/Instrumentation/LocalStateTracingInstrumenter.cs

@@ -309,7 +309,6 @@ public override void InstrumentBlock(BoundBlock original, LocalRewriter rewriter
                     continue;
                 }
 
-                // TODO2
                 var parameterLogger = GetLocalOrParameterStoreLogger(parameter.Type, parameter, refAssignmentSourceIsLocal: null, _factory.Syntax);
                 if (parameterLogger != null)
                 {
@@ -372,7 +371,7 @@ public override BoundExpression InstrumentUserDefinedLocalAssignment(BoundAssign
                 else if (original.Right is BoundParameter rightParameter)
                 {
                     refAssignmentSourceIsLocal = false;
-                    refAssignmentSourceIndex = _factory.ParameterId(rightParameter.ParameterSymbol); // TODO2
+                    refAssignmentSourceIndex = _factory.ParameterId(rightParameter.ParameterSymbol);
                 }
                 else
                 {
@@ -423,7 +422,7 @@ private bool TryGetLocalOrParameterInfo(BoundNode node, [NotNullWhen(true)] out
 
                 symbol = parameterSymbol;
                 type = parameterSymbol.Type;
-                indexExpression = _factory.ParameterId(parameterSymbol); // TODO2
+                indexExpression = _factory.ParameterId(parameterSymbol);
                 return true;
             }
 
@@ -543,7 +542,20 @@ public override void InstrumentCatchBlock(
         }
 
         public override BoundExpression InstrumentCall(BoundCall original, BoundExpression rewritten)
-            => InstrumentCall(base.InstrumentCall(original, rewritten), original.Arguments, original.ArgumentRefKindsOpt); // TODO2 adjust arguments?
+        {
+            ImmutableArray<BoundExpression> arguments = original.Arguments;
+            MethodSymbol method = original.Method;
+            bool adjustForNewExtension = method.GetIsNewExtensionMember() && !method.IsStatic;
+            ImmutableArray<RefKind> argumentRefKindsOpt = NullableWalker.GetArgumentRefKinds(original.ArgumentRefKindsOpt, adjustForNewExtension, method, arguments.Length);
+
+            if (adjustForNewExtension)
+            {
+                Debug.Assert(original.ReceiverOpt is not null);
+                arguments = [original.ReceiverOpt, .. arguments];
+            }
+
+            return InstrumentCall(base.InstrumentCall(original, rewritten), arguments, argumentRefKindsOpt);
+        }
 
         public override BoundExpression InstrumentObjectCreationExpression(BoundObjectCreationExpression original, BoundExpression rewritten)
             => InstrumentCall(base.InstrumentObjectCreationExpression(original, rewritten), original.Arguments, original.ArgumentRefKindsOpt);
@@ -553,7 +565,6 @@ public override BoundExpression InstrumentFunctionPointerInvocation(BoundFunctio
 
         private BoundExpression InstrumentCall(BoundExpression invocation, ImmutableArray<BoundExpression> arguments, ImmutableArray<RefKind> refKinds)
         {
-            // TODO2
             Debug.Assert(refKinds.IsDefault || arguments.Length == refKinds.Length);
             Debug.Assert(invocation.Type is not null);
 
@@ -575,6 +586,7 @@ private BoundExpression InstrumentCall(BoundExpression invocation, ImmutableArra
                 builder.Add(invocation);
             }
 
+            // Record outbound assignments
             for (int i = 0; i < arguments.Length; i++)
             {
                 if (refKinds[i] is not (RefKind.Ref or RefKind.Out))

src/Compilers/CSharp/Portable/Symbols/Extensions/SourceExtensionImplementationMethodSymbol.cs

@@ -125,9 +125,6 @@ protected override ImmutableArray<ParameterSymbol> MakeParameters()
 
             if (!_originalMethod.IsStatic)
             {
-                // TODO2
-                // Tracked by https://github.com/dotnet/roslyn/issues/78962 : Need to confirm if this rewrite going to break LocalStateTracingInstrumenter
-                //            Specifically BoundParameterId, etc.   
                 parameters.Add(new ExtensionMetadataMethodParameterSymbol(this, ((SourceNamedTypeSymbol)_originalMethod.ContainingType).ExtensionParameter!));
             }