Support assignment to multiple refs in trim analyzer (#90287)

Fixes dotnet/linker#3046 by adding
analyzer support for assignments to multiple refs, for example:

```csharp
(b ? ref f1 : ref f2) = v;
```

This also includes support for assignments by reference to array
elements:

```csharp
(b ? ref a1[0] : ref a2[0]) = v;
```

ILLink and ILCompiler produce a different warning in this
case (`stdin.ref` results in the array values being replaced with
`UnknownValue`). I attempted to quirk this in the analyzer, but
this caused more problems in existing tests because it was
difficult to make the quirk specific enough.

Also fixes assignment to multiple arrays on the LHS (this appears
to be an analysis hole in ILLink and ILCompiler):

```csharp
(b ? a1 : a2)[0] = v;
```

Author: sbomer

src/tools/illink/src/ILLink.RoslynAnalyzer/DataFlow/CapturedReferenceValue.cs

@@ -10,7 +10,7 @@ namespace ILLink.RoslynAnalyzer.DataFlow
 {
 	public readonly struct CapturedReferenceValue : IEquatable<CapturedReferenceValue>
 	{
-		public readonly IOperation? Reference;
+		public readonly IOperation Reference;
 
 		public CapturedReferenceValue (IOperation operation)
 		{
@@ -48,23 +48,4 @@ public override bool Equals (object obj)
 		public override int GetHashCode ()
 			=> Reference?.GetHashCode () ?? 0;
 	}
-
-
-	public struct CapturedReferenceLattice : ILattice<CapturedReferenceValue>
-	{
-		public CapturedReferenceValue Top => default;
-
-		public CapturedReferenceValue Meet (CapturedReferenceValue left, CapturedReferenceValue right)
-		{
-			if (left.Equals (right))
-				return left;
-			if (left.Reference == null)
-				return right;
-			if (right.Reference == null)
-				return left;
-			// Both non-null and different shouldn't happen.
-			// We assume that a flow capture can capture only a single property.
-			throw new InvalidOperationException ();
-		}
-	}
 }

src/tools/illink/src/ILLink.RoslynAnalyzer/DataFlow/LValueFlowCaptureProvider.cs

@@ -1,7 +1,7 @@
 // Copyright (c) .NET Foundation and contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
-#nullable disable
+#nullable enable
 
 using System.Collections.Generic;
 using System.Collections.Immutable;
@@ -15,7 +15,7 @@
 
 namespace ILLink.RoslynAnalyzer.DataFlow
 {
-	// Copied from https://github.com/dotnet/roslyn/blob/c8ebc8682889b395fcb84c85bf4ff54577377d26/src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/FlowAnalysis/LValueFlowCaptureProvider.cs
+	// Adapted from https://github.com/dotnet/roslyn/blob/c8ebc8682889b395fcb84c85bf4ff54577377d26/src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/FlowAnalysis/LValueFlowCaptureProvider.cs
 	/// <summary>
 	/// Helper class to detect <see cref="IFlowCaptureOperation"/>s that are l-value captures.
 	/// L-value captures are essentially captures of a symbol's location/address.
@@ -38,6 +38,22 @@ namespace ILLink.RoslynAnalyzer.DataFlow
 	/// </remarks>
 	internal static class LValueFlowCapturesProvider
 	{
+		static bool IsLValueFlowCapture (IFlowCaptureReferenceOperation flowCaptureReference, out IAssignmentOperation? assignment)
+		{
+			assignment = flowCaptureReference.Parent as IAssignmentOperation;
+			if (assignment?.Target == flowCaptureReference)
+				return true;
+
+			if (flowCaptureReference.Parent is IArrayElementReferenceOperation arrayAlementRef) {
+				assignment = arrayAlementRef.Parent as IAssignmentOperation;
+				if (assignment?.Target == arrayAlementRef)
+					return true;
+			}
+
+			assignment = null;
+			return flowCaptureReference.IsInLeftOfDeconstructionAssignment (out _);
+		}
+
 		public static ImmutableDictionary<CaptureId, FlowCaptureKind> CreateLValueFlowCaptures (ControlFlowGraph cfg)
 		{
 			// This method identifies flow capture reference operations that are target of an assignment
@@ -47,15 +63,13 @@ public static ImmutableDictionary<CaptureId, FlowCaptureKind> CreateLValueFlowCa
 			// the flow graph. Specifically, for an ICoalesceOperation a flow capture acts
 			// as both an r-value and l-value flow capture.
 
-			ImmutableDictionary<CaptureId, FlowCaptureKind>.Builder lvalueFlowCaptureIdBuilder = null;
+			ImmutableDictionary<CaptureId, FlowCaptureKind>.Builder? lvalueFlowCaptureIdBuilder = null;
 			var rvalueFlowCaptureIds = new HashSet<CaptureId> ();
 
 			foreach (var flowCaptureReference in cfg.DescendantOperations<IFlowCaptureReferenceOperation> (OperationKind.FlowCaptureReference)) {
-				if (flowCaptureReference.Parent is IAssignmentOperation assignment &&
-					assignment.Target == flowCaptureReference ||
-					flowCaptureReference.IsInLeftOfDeconstructionAssignment (out _)) {
+				if (IsLValueFlowCapture (flowCaptureReference, out IAssignmentOperation? assignment)) {
 					lvalueFlowCaptureIdBuilder ??= ImmutableDictionary.CreateBuilder<CaptureId, FlowCaptureKind> ();
-					var captureKind = flowCaptureReference.Parent.IsAnyCompoundAssignment () || rvalueFlowCaptureIds.Contains (flowCaptureReference.Id)
+					var captureKind = assignment?.IsAnyCompoundAssignment () == true || rvalueFlowCaptureIds.Contains (flowCaptureReference.Id)
 						? FlowCaptureKind.LValueAndRValueCapture
 						: FlowCaptureKind.LValueCapture;
 					lvalueFlowCaptureIdBuilder.Add (flowCaptureReference.Id, captureKind);

src/tools/illink/src/ILLink.RoslynAnalyzer/DataFlow/LocalDataFlowVisitor.cs

@@ -96,7 +96,7 @@ public void Transfer (BlockProxy block, LocalDataFlowState<TValue, TValueLattice
 
 		public abstract TValue HandleArrayElementRead (TValue arrayValue, TValue indexValue, IOperation operation);
 
-		public abstract void HandleArrayElementWrite (TValue arrayValue, TValue indexValue, TValue valueToWrite, IOperation operation);
+		public abstract void HandleArrayElementWrite (TValue arrayValue, TValue indexValue, TValue valueToWrite, IOperation operation, bool merge);
 
 		// This takes an IOperation rather than an IReturnOperation because the return value
 		// may (must?) come from BranchValue of an operation whose FallThroughSuccessor is the exit block.
@@ -139,7 +139,7 @@ TValue GetLocal (ILocalReferenceOperation operation, LocalDataFlowState<TValue,
 			return state.Get (local);
 		}
 
-		void SetLocal (ILocalReferenceOperation operation, TValue value, LocalDataFlowState<TValue, TValueLattice> state)
+		void SetLocal (ILocalReferenceOperation operation, TValue value, LocalDataFlowState<TValue, TValueLattice> state, bool merge = false)
 		{
 			var local = new LocalKey (operation.Local);
 			if (IsReferenceToCapturedVariable (operation))
@@ -149,27 +149,14 @@ void SetLocal (ILocalReferenceOperation operation, TValue value, LocalDataFlowSt
 			if (InterproceduralState.TrySetHoistedLocal (local, value))
 				return;
 
-			state.Set (local, value);
+			var newValue = merge
+				? state.Lattice.Lattice.ValueLattice.Meet (state.Get (local), value)
+				: value;
+			state.Set (local, newValue);
 		}
 
-		public override TValue VisitSimpleAssignment (ISimpleAssignmentOperation operation, LocalDataFlowState<TValue, TValueLattice> state)
+		TValue ProcessSingleTargetAssignment (IOperation targetOperation, ISimpleAssignmentOperation operation, LocalDataFlowState<TValue, TValueLattice> state, bool merge)
 		{
-			var targetOperation = operation.Target;
-			if (targetOperation is IFlowCaptureReferenceOperation flowCaptureReference) {
-				Debug.Assert (IsLValueFlowCapture (flowCaptureReference.Id));
-				Debug.Assert (!flowCaptureReference.GetValueUsageInfo (Method).HasFlag (ValueUsageInfo.Read));
-				var capturedReference = state.Current.CapturedReferences.Get (flowCaptureReference.Id).Reference;
-				targetOperation = capturedReference;
-				if (targetOperation == null)
-					throw new InvalidOperationException ();
-
-				// Note: technically we should avoid visiting the target operation below when assigning to a flow capture reference,
-				// because this should be done when the capture is created. For example, a flow capture used as both an LValue and a RValue
-				// should only evaluate the expression that computes the object instance of a property reference once.
-				// However, we just visit the instance again below for simplicity. This could be generalized if we encounter a dataflow
-				// behavior where this makes a difference.
-			}
-
 			switch (targetOperation) {
 			case IFieldReferenceOperation:
 			case IParameterReferenceOperation: {
@@ -237,17 +224,52 @@ public override TValue VisitSimpleAssignment (ISimpleAssignmentOperation operati
 			// TODO: when setting a property in an attribute, target is an IPropertyReference.
 			case ILocalReferenceOperation localRef: {
 					TValue value = Visit (operation.Value, state);
-					SetLocal (localRef, value, state);
+					SetLocal (localRef, value, state, merge);
 					return value;
 				}
 			case IArrayElementReferenceOperation arrayElementRef: {
 					if (arrayElementRef.Indices.Length != 1)
 						break;
 
-					TValue arrayRef = Visit (arrayElementRef.ArrayReference, state);
-					TValue index = Visit (arrayElementRef.Indices[0], state);
-					TValue value = Visit (operation.Value, state);
-					HandleArrayElementWrite (arrayRef, index, value, operation);
+					// Similarly to VisitSimpleAssignment, this needs to handle cases where the array reference
+					// is a captured variable, even if the target of the assignment (the array element reference) is not.
+
+					TValue arrayRef;
+					TValue index;
+					TValue value;
+					if (arrayElementRef.ArrayReference is not IFlowCaptureReferenceOperation captureReference) {
+						arrayRef = Visit (arrayElementRef.ArrayReference, state);
+						index = Visit (arrayElementRef.Indices[0], state);
+						value = Visit (operation.Value, state);
+						HandleArrayElementWrite (arrayRef, index, value, operation, merge: merge);
+						return value;
+					}
+
+					index = Visit (arrayElementRef.Indices[0], state);
+					value = Visit (operation.Value, state);
+
+					var capturedReferences = state.Current.CapturedReferences.Get (captureReference.Id);
+					if (!capturedReferences.HasMultipleValues) {
+						// Single captured reference. Treat this as an overwriting assignment,
+						// unless the caller already told us to merge values because this is an
+						// assignment to one of multiple captured array element references.
+						var enumerator = capturedReferences.GetEnumerator ();
+						enumerator.MoveNext ();
+						var capture = enumerator.Current;
+						arrayRef = Visit (capture.Reference, state);
+						HandleArrayElementWrite (arrayRef, index, value, operation, merge: merge);
+						return value;
+					}
+
+					// The capture id may have captured multiple references, as in:
+					// (b ? arr1 : arr2)[0] = value;
+					// We treat this as possible write to each of the captured references,
+					// which requires merging with the previous values of each.
+
+					foreach (var capture in state.Current.CapturedReferences.Get (captureReference.Id)) {
+						arrayRef = Visit (capture.Reference, state);
+						HandleArrayElementWrite (arrayRef, index, value, operation, merge: true);
+					}
 					return value;
 				}
 			case IDiscardOperation:
@@ -293,8 +315,50 @@ public override TValue VisitSimpleAssignment (ISimpleAssignmentOperation operati
 			return Visit (operation.Value, state);
 		}
 
-		// Similar to VisitLocalReference
-		public override TValue VisitFlowCaptureReference (IFlowCaptureReferenceOperation operation, LocalDataFlowState<TValue, TValueLattice> state)
+		public override TValue VisitSimpleAssignment (ISimpleAssignmentOperation operation, LocalDataFlowState<TValue, TValueLattice> state)
+		{
+			var targetOperation = operation.Target;
+			if (targetOperation is not IFlowCaptureReferenceOperation flowCaptureReference)
+				return ProcessSingleTargetAssignment (targetOperation, operation, state, merge: false);
+
+			// Note: technically we should avoid visiting the target operation in ProcessNonCapturedAssignment when assigning
+			// to a flow capture reference, because this should be done when the capture is created.
+			// For example, a flow capture used as both an LValue and a RValue should only evaluate the expression that
+			// computes the object instance of a property reference once. However, we just visit the instance again below
+			// for simplicity. This could be generalized if we encounter a dataflow behavior where this makes a difference.
+
+			Debug.Assert (IsLValueFlowCapture (flowCaptureReference.Id));
+			Debug.Assert (!flowCaptureReference.GetValueUsageInfo (Method).HasFlag (ValueUsageInfo.Read));
+			var capturedReferences = state.Current.CapturedReferences.Get (flowCaptureReference.Id);
+			if (!capturedReferences.HasMultipleValues) {
+				// Single captured reference. Treat this as an overwriting assignment.
+				var enumerator = capturedReferences.GetEnumerator ();
+				enumerator.MoveNext ();
+				targetOperation = enumerator.Current.Reference;
+				return ProcessSingleTargetAssignment (targetOperation, operation, state, merge: false);
+			}
+
+			// The capture id may have captured multiple references, as in:
+			// (b ? ref v1 : ref v2) = value;
+			// We treat this as a possible write to each of the captured references,
+			// which requires merging with the previous values of each.
+
+			// Note: technically this should only visit the RHS of the assignment once.
+			// For now we visit the RHS in ProcessSingleTargetAssignment for simplicity, and
+			// rely on the warning deduplication to prevent this from producing multiple warnings
+			// if the RHS has dataflow warnings.
+
+			TValue value = TopValue;
+			foreach (var capturedReference in capturedReferences) {
+				targetOperation = capturedReference.Reference;
+				var singleValue = ProcessSingleTargetAssignment (targetOperation, operation, state, merge: true);
+				value = LocalStateLattice.Lattice.ValueLattice.Meet (value, singleValue);
+			}
+
+			return value;
+		}
+
+		TValue GetFlowCaptureValue (IFlowCaptureReferenceOperation operation, LocalDataFlowState<TValue, TValueLattice> state)
 		{
 			if (!operation.GetValueUsageInfo (Method).HasFlag (ValueUsageInfo.Read)) {
 				// There are known cases where this assert doesn't hold, because LValueFlowCaptureProvider
@@ -304,10 +368,21 @@ public override TValue VisitFlowCaptureReference (IFlowCaptureReferenceOperation
 				return TopValue;
 			}
 
-			Debug.Assert (IsRValueFlowCapture (operation.Id));
+			// This assert is incorrect for cases like (b ? arr1 : arr2)[0] = v;
+			// Here the ValueUsageInfo shows that the value usage is for reading (this is probably wrong!)
+			// but the value is actually an LValueFlowCapture.
+			// Let's just disable the assert for now.
+			// Debug.Assert (IsRValueFlowCapture (operation.Id));
+
 			return state.Get (new LocalKey (operation.Id));
 		}
 
+		// Similar to VisitLocalReference
+		public override TValue VisitFlowCaptureReference (IFlowCaptureReferenceOperation operation, LocalDataFlowState<TValue, TValueLattice> state)
+		{
+			return GetFlowCaptureValue (operation, state);
+		}
+
 		// Similar to VisitSimpleAssignment when assigning to a local, but for values which are captured without a
 		// corresponding local variable. The "flow capture" is like a local assignment, and the "flow capture reference"
 		// is like a local reference.

src/tools/illink/src/ILLink.RoslynAnalyzer/DataFlow/LocalStateLattice.cs

@@ -43,22 +43,22 @@ public struct LocalState<TValue> : IEquatable<LocalState<TValue>>
 		// Stores any operations which are captured by reference in a FlowCaptureOperation.
 		// Only stores captures which are assigned through. Captures of the values of operations
 		// are tracked as part of the dictionary of values, keyed by LocalKey.
-		public DefaultValueDictionary<CaptureId, CapturedReferenceValue> CapturedReferences;
+		public DefaultValueDictionary<CaptureId, ValueSet<CapturedReferenceValue>> CapturedReferences;
 
 		public LocalState (TValue defaultValue)
 			: this (new DefaultValueDictionary<LocalKey, TValue> (defaultValue),
-				new DefaultValueDictionary<CaptureId, CapturedReferenceValue> (default (CapturedReferenceValue)))
+				new DefaultValueDictionary<CaptureId, ValueSet<CapturedReferenceValue>> (default (ValueSet<CapturedReferenceValue>)))
 		{
 		}
 
-		public LocalState (DefaultValueDictionary<LocalKey, TValue> dictionary, DefaultValueDictionary<CaptureId, CapturedReferenceValue> capturedReferences)
+		public LocalState (DefaultValueDictionary<LocalKey, TValue> dictionary, DefaultValueDictionary<CaptureId, ValueSet<CapturedReferenceValue>> capturedReferences)
 		{
 			Dictionary = dictionary;
 			CapturedReferences = capturedReferences;
 		}
 
 		public LocalState (DefaultValueDictionary<LocalKey, TValue> dictionary)
-			: this (dictionary, new DefaultValueDictionary<CaptureId, CapturedReferenceValue> (default (CapturedReferenceValue)))
+			: this (dictionary, new DefaultValueDictionary<CaptureId, ValueSet<CapturedReferenceValue>> (default (ValueSet<CapturedReferenceValue>)))
 		{
 		}
 
@@ -83,12 +83,12 @@ public override int GetHashCode ()
 		where TValueLattice : ILattice<TValue>
 	{
 		public readonly DictionaryLattice<LocalKey, TValue, TValueLattice> Lattice;
-		public readonly DictionaryLattice<CaptureId, CapturedReferenceValue, CapturedReferenceLattice> CapturedReferenceLattice;
+		public readonly DictionaryLattice<CaptureId, ValueSet<CapturedReferenceValue>, ValueSetLattice<CapturedReferenceValue>> CapturedReferenceLattice;
 
 		public LocalStateLattice (TValueLattice valueLattice)
 		{
 			Lattice = new DictionaryLattice<LocalKey, TValue, TValueLattice> (valueLattice);
-			CapturedReferenceLattice = new DictionaryLattice<CaptureId, CapturedReferenceValue, CapturedReferenceLattice> (default (CapturedReferenceLattice));
+			CapturedReferenceLattice = new DictionaryLattice<CaptureId, ValueSet<CapturedReferenceValue>, ValueSetLattice<CapturedReferenceValue>> (default (ValueSetLattice<CapturedReferenceValue>));
 			Top = new (Lattice.Top);
 		}
 

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

@@ -219,20 +219,19 @@ public override MultiValue HandleArrayElementRead (MultiValue arrayValue, MultiV
 			return result.Equals (TopValue) ? UnknownValue.Instance : result;
 		}
 
-		public override void HandleArrayElementWrite (MultiValue arrayValue, MultiValue indexValue, MultiValue valueToWrite, IOperation operation)
+		public override void HandleArrayElementWrite (MultiValue arrayValue, MultiValue indexValue, MultiValue valueToWrite, IOperation operation, bool merge)
 		{
 			int? index = indexValue.AsConstInt ();
 			foreach (var arraySingleValue in arrayValue) {
 				if (arraySingleValue is ArrayValue arr) {
 					if (index == null) {
 						// Reset the array to all unknowns - since we don't know which index is being assigned
 						arr.IndexValues.Clear ();
-					} else {
-						if (arr.IndexValues.TryGetValue (index.Value, out _)) {
-							arr.IndexValues[index.Value] = ArrayValue.SanitizeArrayElementValue(valueToWrite);
-						} else if (arr.IndexValues.Count < MaxTrackedArrayValues) {
-							arr.IndexValues[index.Value] = ArrayValue.SanitizeArrayElementValue(valueToWrite);
-						}
+					} else if (arr.IndexValues.TryGetValue (index.Value, out _) || arr.IndexValues.Count < MaxTrackedArrayValues) {
+						var sanitizedValue = ArrayValue.SanitizeArrayElementValue(valueToWrite);
+						arr.IndexValues[index.Value] = merge
+							? _multiValueLattice.Meet (arr.IndexValues[index.Value], sanitizedValue)
+							: sanitizedValue;
 					}
 				}
 			}

src/tools/illink/src/ILLink.Shared/DataFlow/ValueSet.cs

@@ -132,6 +132,8 @@ public void Reset ()
 
 		public static implicit operator ValueSet<TValue> (TValue value) => new (value);
 
+		public bool HasMultipleValues => _values is EnumerableValues;
+
 		public override bool Equals (object? obj) => obj is ValueSet<TValue> other && Equals (other);
 
 		public bool Equals (ValueSet<TValue> other)

src/tools/illink/test/Mono.Linker.Tests.Cases/Attributes/OnlyKeepUsed/MethodWithUnmanagedConstraint.cs

@@ -4,7 +4,6 @@
 namespace Mono.Linker.Tests.Cases.Attributes.OnlyKeepUsed
 {
 	[SetupCSharpCompilerToUse ("csc")]
-	[SetupCompileArgument ("/langversion:7.3")]
 	[SetupLinkerArgument ("--used-attrs-only", "true")]
 	public class MethodWithUnmanagedConstraint
 	{