Use correct IL offset when reporting dataflow warnings for return sta…

…tements (dotnet#105661)

Instead of merging all returned values into a single return value and
reporting warnings based on that, report warnings per return
statement, using the IL offset at that point. This allows us to fix
the message origin to match the IL offset of each individual return
statement, so we can remove a workaround from the pattern store.

---------

Co-authored-by: Michal Strehovský <MichalStrehovsky@users.noreply.github.com>

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/Dataflow/MethodBodyScanner.cs

@@ -52,8 +52,6 @@ internal abstract partial class MethodBodyScanner
 
         protected readonly FlowAnnotations _annotations;
 
-        internal MultiValue ReturnValue { get; private set; }
-
         protected MethodBodyScanner(FlowAnnotations annotations)
         {
             _annotations = annotations;
@@ -356,7 +354,6 @@ protected virtual void Scan(MethodIL methodBody, ref InterproceduralState interp
 
             BasicBlockIterator blockIterator = new BasicBlockIterator(methodBody);
 
-            ReturnValue = default(MultiValue);
             ILReader reader = new ILReader(methodBody.GetILBytes());
             while (reader.HasNext)
             {
@@ -797,10 +794,12 @@ protected virtual void Scan(MethodIL methodBody, ref InterproceduralState interp
                             }
                             if (hasReturnValue)
                             {
-                                StackSlot retValue = PopUnknown(currentStack, 1, methodBody, offset);
+                                StackSlot retStackSlot = PopUnknown(currentStack, 1, methodBody, offset);
                                 // If the return value is a reference, treat it as the value itself for now
                                 // We can handle ref return values better later
-                                ReturnValue = MultiValueLattice.Meet(ReturnValue, DereferenceValue(methodBody, offset, retValue.Value, locals, ref interproceduralState));
+                                MultiValue retValue = DereferenceValue(methodBody, offset, retStackSlot.Value, locals, ref interproceduralState);
+                                var methodReturnValue = GetReturnValue(methodBody);
+                                HandleReturnValue(methodBody, offset, methodReturnValue, retValue);
                                 ValidateNoReferenceToReference(locals, methodBody, offset);
                             }
                             ClearStack(ref currentStack);
@@ -871,6 +870,8 @@ private static void ScanExceptionInformation(Dictionary<int, Stack<StackSlot>> k
 
         protected abstract SingleValue GetMethodParameterValue(ParameterProxy parameter);
 
+        protected abstract MethodReturnValue GetReturnValue(MethodIL method);
+
         private void ScanLdarg(ILOpcode opcode, int parameterIndex, Stack<StackSlot> currentStack, MethodDesc thisMethod)
         {
             ParameterIndex paramNum = (ParameterIndex)parameterIndex;
@@ -1036,7 +1037,7 @@ when GetMethodParameterValue(parameterReference.Parameter) is MethodParameterVal
                         break;
                     case MethodReturnValue methodReturnValue:
                         // Ref returns don't have special ReferenceValue values, so assume if the target here is a MethodReturnValue then it must be a ref return value
-                        HandleStoreMethodReturnValue(method, offset, methodReturnValue, source);
+                        HandleReturnValue(method, offset, methodReturnValue, source);
                         break;
                     case FieldValue fieldValue:
                         HandleStoreField(method, offset, fieldValue, DereferenceValue(method, offset, source, locals, ref ipState));
@@ -1115,7 +1116,7 @@ protected virtual void HandleStoreParameter(MethodIL method, int offset, MethodP
         {
         }
 
-        protected virtual void HandleStoreMethodReturnValue(MethodIL method, int offset, MethodReturnValue thisParameter, MultiValue sourceValue)
+        protected virtual void HandleReturnValue(MethodIL method, int offset, MethodReturnValue thisParameter, MultiValue valueToReturn)
         {
         }
 

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/Dataflow/ReflectionMethodBodyScanner.cs

@@ -114,14 +114,6 @@ protected override void Scan(MethodIL methodBody, ref InterproceduralState inter
         {
             _origin = new MessageOrigin(methodBody.OwningMethod);
             base.Scan(methodBody, ref interproceduralState);
-
-            if (!methodBody.OwningMethod.Signature.ReturnType.IsVoid)
-            {
-                var method = methodBody.OwningMethod;
-                var methodReturnValue = _annotations.GetMethodReturnValue(method, isNewObj: false);
-                if (methodReturnValue.DynamicallyAccessedMemberTypes != 0)
-                    HandleAssignmentPattern(_origin, ReturnValue, methodReturnValue, method.GetDisplayName());
-            }
         }
 
         public static DependencyList ScanAndProcessReturnValue(NodeFactory factory, FlowAnnotations annotations, Logger logger, MethodIL methodBody, out List<INodeWithRuntimeDeterminedDependencies> runtimeDependencies)
@@ -195,6 +187,8 @@ protected override ValueWithDynamicallyAccessedMembers GetMethodParameterValue(P
         private MethodParameterValue GetMethodParameterValue(ParameterProxy parameter, DynamicallyAccessedMemberTypes dynamicallyAccessedMemberTypes)
             => _annotations.GetMethodParameterValue(parameter, dynamicallyAccessedMemberTypes);
 
+        protected override MethodReturnValue GetReturnValue(MethodIL method) => _annotations.GetMethodReturnValue(method.OwningMethod, isNewObj: false);
+
         /// <summary>
         /// HandleGetField is called every time the scanner needs to represent a value of the field
         /// either as a source or target. It is not called when just a reference to field is created,
@@ -219,7 +213,7 @@ private void HandleStoreValueWithDynamicallyAccessedMembers(MethodIL methodBody,
             if (targetValue.DynamicallyAccessedMemberTypes != 0)
             {
                 _origin = _origin.WithInstructionOffset(methodBody, offset);
-                HandleAssignmentPattern(_origin, sourceValue, targetValue, reason);
+                TrimAnalysisPatterns.Add(new TrimAnalysisAssignmentPattern(sourceValue, targetValue, _origin, reason));
             }
         }
 
@@ -229,7 +223,7 @@ protected override void HandleStoreField(MethodIL methodBody, int offset, FieldV
         protected override void HandleStoreParameter(MethodIL methodBody, int offset, MethodParameterValue parameter, MultiValue valueToStore)
             => HandleStoreValueWithDynamicallyAccessedMembers(methodBody, offset, parameter, valueToStore, parameter.Parameter.Method.GetDisplayName());
 
-        protected override void HandleStoreMethodReturnValue(MethodIL methodBody, int offset, MethodReturnValue returnValue, MultiValue valueToStore)
+        protected override void HandleReturnValue(MethodIL methodBody, int offset, MethodReturnValue returnValue, MultiValue valueToStore)
             => HandleStoreValueWithDynamicallyAccessedMembers(methodBody, offset, returnValue, valueToStore, returnValue.Method.GetDisplayName());
 
         protected override void HandleTypeTokenAccess(MethodIL methodBody, int offset, TypeDesc accessedType)
@@ -414,15 +408,6 @@ private static bool IsComInterop(MarshalAsDescriptor? marshalInfoProvider, TypeD
             return false;
         }
 
-        private void HandleAssignmentPattern(
-            in MessageOrigin origin,
-            in MultiValue value,
-            ValueWithDynamicallyAccessedMembers targetValue,
-            string reason)
-        {
-            TrimAnalysisPatterns.Add(new TrimAnalysisAssignmentPattern(value, targetValue, origin, reason));
-        }
-
         private void ProcessGenericArgumentDataFlow(MethodDesc method)
         {
             // We only need to validate static methods and then all generic methods

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/Dataflow/TrimAnalysisPatternStore.cs

@@ -13,7 +13,7 @@ namespace ILCompiler.Dataflow
 {
     public readonly struct TrimAnalysisPatternStore
     {
-        private readonly Dictionary<(MessageOrigin, bool), TrimAnalysisAssignmentPattern> AssignmentPatterns;
+        private readonly Dictionary<MessageOrigin, TrimAnalysisAssignmentPattern> AssignmentPatterns;
         private readonly Dictionary<MessageOrigin, TrimAnalysisMethodCallPattern> MethodCallPatterns;
         private readonly Dictionary<(MessageOrigin, TypeSystemEntity), TrimAnalysisTokenAccessPattern> TokenAccessPatterns;
         private readonly Dictionary<(MessageOrigin, TypeSystemEntity), TrimAnalysisGenericInstantiationAccessPattern> GenericInstantiations;
@@ -23,7 +23,7 @@ public readonly struct TrimAnalysisPatternStore
 
         public TrimAnalysisPatternStore(ValueSetLattice<SingleValue> lattice, Logger logger)
         {
-            AssignmentPatterns = new Dictionary<(MessageOrigin, bool), TrimAnalysisAssignmentPattern>();
+            AssignmentPatterns = new Dictionary<MessageOrigin, TrimAnalysisAssignmentPattern>();
             MethodCallPatterns = new Dictionary<MessageOrigin, TrimAnalysisMethodCallPattern>();
             TokenAccessPatterns = new Dictionary<(MessageOrigin, TypeSystemEntity), TrimAnalysisTokenAccessPattern>();
             GenericInstantiations = new Dictionary<(MessageOrigin, TypeSystemEntity), TrimAnalysisGenericInstantiationAccessPattern>();
@@ -34,19 +34,13 @@ public TrimAnalysisPatternStore(ValueSetLattice<SingleValue> lattice, Logger log
 
         public void Add(TrimAnalysisAssignmentPattern pattern)
         {
-            // While trimming, each pattern should have a unique origin (which has ILOffset)
-            // but we don't track the correct ILOffset for return instructions.
-            // https://github.com/dotnet/linker/issues/2778
-            // For now, work around it with a separate bit.
-            bool isReturnValue = pattern.Target.AsSingleValue() is MethodReturnValue;
-
-            if (!AssignmentPatterns.TryGetValue((pattern.Origin, isReturnValue), out var existingPattern))
+            if (!AssignmentPatterns.TryGetValue(pattern.Origin, out var existingPattern))
             {
-                AssignmentPatterns.Add((pattern.Origin, isReturnValue), pattern);
+                AssignmentPatterns.Add(pattern.Origin, pattern);
                 return;
             }
 
-            AssignmentPatterns[(pattern.Origin, isReturnValue)] = pattern.Merge(Lattice, existingPattern);
+            AssignmentPatterns[pattern.Origin] = pattern.Merge(Lattice, existingPattern);
         }
 
         public void Add(TrimAnalysisMethodCallPattern pattern)

src/tools/illink/src/ILLink.RoslynAnalyzer/TrimAnalysis/TrimAnalysisPatternStore.cs

@@ -12,7 +12,7 @@ namespace ILLink.RoslynAnalyzer.TrimAnalysis
 {
 	internal readonly struct TrimAnalysisPatternStore
 	{
-		readonly Dictionary<(IOperation, bool), TrimAnalysisAssignmentPattern> AssignmentPatterns;
+		readonly Dictionary<IOperation, TrimAnalysisAssignmentPattern> AssignmentPatterns;
 		readonly Dictionary<IOperation, TrimAnalysisFieldAccessPattern> FieldAccessPatterns;
 		readonly Dictionary<IOperation, TrimAnalysisGenericInstantiationPattern> GenericInstantiationPatterns;
 		readonly Dictionary<IOperation, TrimAnalysisMethodCallPattern> MethodCallPatterns;
@@ -25,7 +25,7 @@ public TrimAnalysisPatternStore (
 			ValueSetLattice<SingleValue> lattice,
 			FeatureContextLattice featureContextLattice)
 		{
-			AssignmentPatterns = new Dictionary<(IOperation, bool), TrimAnalysisAssignmentPattern> ();
+			AssignmentPatterns = new Dictionary<IOperation, TrimAnalysisAssignmentPattern> ();
 			FieldAccessPatterns = new Dictionary<IOperation, TrimAnalysisFieldAccessPattern> ();
 			GenericInstantiationPatterns = new Dictionary<IOperation, TrimAnalysisGenericInstantiationPattern> ();
 			MethodCallPatterns = new Dictionary<IOperation, TrimAnalysisMethodCallPattern> ();
@@ -35,7 +35,7 @@ public TrimAnalysisPatternStore (
 			FeatureContextLattice = featureContextLattice;
 		}
 
-		public void Add (TrimAnalysisAssignmentPattern trimAnalysisPattern, bool isReturnValue)
+		public void Add (TrimAnalysisAssignmentPattern trimAnalysisPattern)
 		{
 			// Finally blocks will be analyzed multiple times, once for normal control flow and once
 			// for exceptional control flow, and these separate analyses could produce different
@@ -45,12 +45,12 @@ public void Add (TrimAnalysisAssignmentPattern trimAnalysisPattern, bool isRetur
 			// of the normal control-flow state.
 			// We still add patterns to the operation, rather than replacing, to make this resilient to
 			// changes in the analysis algorithm.
-			if (!AssignmentPatterns.TryGetValue ((trimAnalysisPattern.Operation, isReturnValue), out var existingPattern)) {
-				AssignmentPatterns.Add ((trimAnalysisPattern.Operation, isReturnValue), trimAnalysisPattern);
+			if (!AssignmentPatterns.TryGetValue (trimAnalysisPattern.Operation, out var existingPattern)) {
+				AssignmentPatterns.Add (trimAnalysisPattern.Operation, trimAnalysisPattern);
 				return;
 			}
 
-			AssignmentPatterns[(trimAnalysisPattern.Operation, isReturnValue)] = trimAnalysisPattern.Merge (Lattice, FeatureContextLattice, existingPattern);
+			AssignmentPatterns[trimAnalysisPattern.Operation] = trimAnalysisPattern.Merge (Lattice, FeatureContextLattice, existingPattern);
 		}
 
 		public void Add (TrimAnalysisFieldAccessPattern pattern)

src/tools/illink/src/ILLink.RoslynAnalyzer/TrimAnalysis/TrimAnalysisVisitor.cs

@@ -238,9 +238,7 @@ public override void HandleAssignment (MultiValue source, MultiValue target, IOp
 			// annotated with DAMT.
 			TrimAnalysisPatterns.Add (
 				// This will copy the values if necessary
-				new TrimAnalysisAssignmentPattern (source, target, operation, OwningSymbol, featureContext),
-				isReturnValue: false
-			);
+				new TrimAnalysisAssignmentPattern (source, target, operation, OwningSymbol, featureContext));
 		}
 
 		public override MultiValue HandleArrayElementRead (MultiValue arrayValue, MultiValue indexValue, IOperation operation)
@@ -350,9 +348,7 @@ public override void HandleReturnValue (MultiValue returnValue, IOperation opera
 				var returnParameter = new MethodReturnValue (method, isNewObj: false);
 
 				TrimAnalysisPatterns.Add (
-					new TrimAnalysisAssignmentPattern (returnValue, returnParameter, operation, OwningSymbol, featureContext),
-					isReturnValue: true
-				);
+					new TrimAnalysisAssignmentPattern (returnValue, returnParameter, operation, OwningSymbol, featureContext));
 			}
 		}
 

src/tools/illink/src/linker/Linker.Dataflow/MethodBodyScanner.cs

@@ -54,8 +54,6 @@ protected MethodBodyScanner (LinkContext context)
 			this.InterproceduralStateLattice = new InterproceduralStateLattice (default, default, context);
 		}
 
-		internal MultiValue ReturnValue { get; private set; }
-
 		protected virtual void WarnAboutInvalidILInMethod (MethodBody method, int ilOffset)
 		{
 		}
@@ -289,7 +287,6 @@ protected virtual void Scan (MethodIL methodIL, ref InterproceduralState interpr
 
 			BasicBlockIterator blockIterator = new BasicBlockIterator (methodIL);
 
-			ReturnValue = default;
 			foreach (Instruction operation in methodIL.Instructions) {
 				int curBasicBlock = blockIterator.MoveNext (operation);
 
@@ -655,10 +652,12 @@ protected virtual void Scan (MethodIL methodIL, ref InterproceduralState interpr
 							WarnAboutInvalidILInMethod (methodBody, operation.Offset);
 						}
 						if (hasReturnValue) {
-							StackSlot retValue = PopUnknown (currentStack, 1, methodBody, operation.Offset);
+							StackSlot retStackSlot = PopUnknown (currentStack, 1, methodBody, operation.Offset);
 							// If the return value is a reference, treat it as the value itself for now
 							// We can handle ref return values better later
-							ReturnValue = MultiValueLattice.Meet (ReturnValue, DereferenceValue (retValue.Value, locals, ref interproceduralState));
+							MultiValue retValue = DereferenceValue (retStackSlot.Value, locals, ref interproceduralState);
+							var methodReturnValue = GetReturnValue (methodBody.Method);
+							HandleReturnValue (thisMethod, methodReturnValue, operation, retValue);
 							ValidateNoReferenceToReference (locals, methodBody.Method, operation.Offset);
 						}
 						ClearStack (ref currentStack);
@@ -719,6 +718,8 @@ private static void ScanExceptionInformation (Dictionary<int, Stack<StackSlot>>
 
 		protected abstract SingleValue GetMethodParameterValue (ParameterProxy parameter);
 
+		protected abstract MethodReturnValue GetReturnValue (MethodDefinition method);
+
 		private void ScanLdarg (Instruction operation, Stack<StackSlot> currentStack, MethodDefinition thisMethod)
 		{
 			Code code = operation.OpCode.Code;
@@ -873,7 +874,7 @@ when GetMethodParameterValue (parameterReference.Parameter) is MethodParameterVa
 					break;
 				case MethodReturnValue methodReturnValue:
 					// Ref returns don't have special ReferenceValue values, so assume if the target here is a MethodReturnValue then it must be a ref return value
-					HandleStoreMethodReturnValue (method, methodReturnValue, operation, source);
+					HandleReturnValue (method, methodReturnValue, operation, source);
 					break;
 				case FieldValue fieldValue:
 					HandleStoreField (method, fieldValue, operation, DereferenceValue (source, locals, ref ipState));
@@ -934,7 +935,7 @@ protected virtual void HandleStoreParameter (MethodDefinition method, MethodPara
 		{
 		}
 
-		protected virtual void HandleStoreMethodReturnValue (MethodDefinition method, MethodReturnValue thisParameter, Instruction operation, MultiValue sourceValue)
+		protected virtual void HandleReturnValue (MethodDefinition method, MethodReturnValue thisParameter, Instruction operation, MultiValue valueToReturn)
 		{
 		}
 

src/tools/illink/src/linker/Linker.Dataflow/ReflectionMethodBodyScanner.cs

@@ -72,13 +72,6 @@ protected override void Scan (MethodIL methodIL, ref InterproceduralState interp
 		{
 			_origin = new MessageOrigin (methodIL.Method);
 			base.Scan (methodIL, ref interproceduralState);
-
-			if (!methodIL.Method.ReturnsVoid ()) {
-				var method = methodIL.Method;
-				var methodReturnValue = _annotations.GetMethodReturnValue (method, isNewObj: false);
-				if (methodReturnValue.DynamicallyAccessedMemberTypes != 0)
-					HandleAssignmentPattern (_origin, ReturnValue, methodReturnValue);
-			}
 		}
 
 		protected override void WarnAboutInvalidILInMethod (MethodBody method, int ilOffset)
@@ -100,11 +93,13 @@ MethodParameterValue GetMethodParameterValue (ParameterProxy parameter, Dynamica
 
 		protected override MultiValue GetFieldValue (FieldDefinition field) => _annotations.GetFieldValue (field);
 
+		protected override MethodReturnValue GetReturnValue (MethodDefinition method) => _annotations.GetMethodReturnValue (method, isNewObj: false);
+
 		private void HandleStoreValueWithDynamicallyAccessedMembers (ValueWithDynamicallyAccessedMembers targetValue, Instruction operation, MultiValue sourceValue)
 		{
 			if (targetValue.DynamicallyAccessedMemberTypes != 0) {
 				_origin = _origin.WithInstructionOffset (operation.Offset);
-				HandleAssignmentPattern (_origin, sourceValue, targetValue);
+				TrimAnalysisPatterns.Add (new TrimAnalysisAssignmentPattern (sourceValue, targetValue, _origin));
 			}
 		}
 
@@ -114,7 +109,7 @@ protected override void HandleStoreField (MethodDefinition method, FieldValue fi
 		protected override void HandleStoreParameter (MethodDefinition method, MethodParameterValue parameter, Instruction operation, MultiValue valueToStore)
 			=> HandleStoreValueWithDynamicallyAccessedMembers (parameter, operation, valueToStore);
 
-		protected override void HandleStoreMethodReturnValue (MethodDefinition method, MethodReturnValue returnValue, Instruction operation, MultiValue valueToStore)
+		protected override void HandleReturnValue (MethodDefinition method, MethodReturnValue returnValue, Instruction operation, MultiValue valueToStore)
 			=> HandleStoreValueWithDynamicallyAccessedMembers (returnValue, operation, valueToStore);
 
 		public override MultiValue HandleCall (MethodBody callingMethodBody, MethodReference calledMethod, Instruction operation, ValueNodeList methodParams)
@@ -248,14 +243,6 @@ static bool IsComInterop (IMarshalInfoProvider marshalInfoProvider, TypeReferenc
 			return false;
 		}
 
-		void HandleAssignmentPattern (
-			in MessageOrigin origin,
-			in MultiValue value,
-			ValueWithDynamicallyAccessedMembers targetValue)
-		{
-			TrimAnalysisPatterns.Add (new TrimAnalysisAssignmentPattern (value, targetValue, origin));
-		}
-
 		internal static bool IsPInvokeDangerous (MethodDefinition methodDefinition, LinkContext context, out bool comDangerousMethod)
 		{
 			// The method in ILLink only detects one condition - COM Dangerous, but it's structured like this