Handle extension methods which create delegates

GrahamTheCoder · GrahamTheCoder · commit f4759c061534 · 2020-06-14T16:37:28.000+01:00
The first parameter appears as the "instance" for a reduced extension method
diff --git a/NullabilityInference.Tests/InferenceTests.cs b/NullabilityInference.Tests/InferenceTests.cs
@@ -591,6 +591,41 @@ public class SomeContext
 ");
         }
 
+        /// <summary>
+        /// FUTURE: NotNullWhenNotNullAttribute on IfNull would allow Array to be marked non-null
+        /// </summary>
+        [Fact]
+        public void ExtensionMethodDelegate()
+        {
+            AssertNullabilityInference(@"
+using System.Linq;
+
+public static class SomeStatic
+{
+    public static void Main()
+    {
+        object?[] x = new[] { new object() }.Select(7.IfNull).ToArray();
+    }
+    public static object? IfNull(this object? a, object? b) => b ?? a;
+}");
+        }
+
+        [Fact]
+        public void InstanceMethodDelegate()
+        {
+            AssertNullabilityInference(@"
+using System.Linq;
+
+public class SomeInstance
+{
+    public void Caller()
+    {
+        object[] x = new[] { new object() }.Select(IfNull).ToArray();
+    }
+    public object IfNull(object? b) => b ?? this;
+}");
+        }
+
         [Fact]
         public void CompareNullableValueTypeToNullDoesNotThrow()
         {
diff --git a/NullabilityInference/EdgeBuildingOperationVisitor.cs b/NullabilityInference/EdgeBuildingOperationVisitor.cs
@@ -957,7 +957,7 @@ private TypeWithNode[] ExtendMethodTypeArguments(IMethodSymbol targetMethod, Typ
             return methodTypeArgNodes;
         }
 
-        private void HandleMethodGroup(IMethodReferenceOperation operation, TypeWithNode delegateReturnType, TypeWithNode[] delegateParameters)
+        private void HandleMethodGroup(IMethodReferenceOperation operation, TypeWithNode delegateReturnType, IReadOnlyCollection<TypeWithNode> delegateParameters)
         {
             var receiverType = GetReceiverType(operation);
             var classTypeArgNodes = ClassTypeArgumentsForMemberAccess(receiverType, operation.Method);
@@ -970,7 +970,7 @@ private void HandleMethodGroup(IMethodReferenceOperation operation, TypeWithNode
             var substitution = new TypeSubstitution(classTypeArgNodes, methodTypeArgNodes);
             EdgeLabel label = new EdgeLabel($"MethodGroup", operation);
 
-            Debug.Assert(operation.Method.Parameters.Length == delegateParameters.Length);
+            Debug.Assert(operation.Method.Parameters.Length == delegateParameters.Count);
             foreach (var (methodParam, delegateParam) in operation.Method.Parameters.Zip(delegateParameters)) {
                 var methodParamType = typeSystem.GetSymbolType(methodParam.OriginalDefinition);
                 methodParamType = methodParamType.WithSubstitution(methodParam.Type, substitution, tsBuilder);
@@ -1608,7 +1608,7 @@ public override TypeWithNode VisitDelegateCreation(IDelegateCreationOperation op
             var delegateReturnType = typeSystem.GetSymbolType(delegateType.DelegateInvokeMethod.OriginalDefinition);
             delegateReturnType = delegateReturnType.WithSubstitution(delegateType.DelegateInvokeMethod.ReturnType, substitution, tsBuilder);
             var delegateParameters = delegateType.DelegateInvokeMethod.Parameters
-                .Select(p => typeSystem.GetSymbolType(p.OriginalDefinition).WithSubstitution(p.Type, substitution, tsBuilder)).ToArray();
+                .Select(p => typeSystem.GetSymbolType(p.OriginalDefinition).WithSubstitution(p.Type, substitution, tsBuilder)).ToList();
             switch (operation.Target) {
                 case IAnonymousFunctionOperation lambda:
                     // Create edges for lambda parameters
@@ -1629,7 +1629,7 @@ public override TypeWithNode VisitDelegateCreation(IDelegateCreationOperation op
                         };
                     }
                     if (parameterList != null) {
-                        Debug.Assert(parameterList.Count == delegateParameters.Length);
+                        Debug.Assert(parameterList.Count == delegateParameters.Count);
                         foreach (var (lambdaParamSyntax, invokeParam) in parameterList.Zip(delegateParameters)) {
                             if (lambdaParamSyntax.Type != null) {
                                 var paramType = lambdaParamSyntax.Type.Accept(syntaxVisitor);
@@ -1665,6 +1665,7 @@ public override TypeWithNode VisitDelegateCreation(IDelegateCreationOperation op
                     }
                     break;
                 case IMethodReferenceOperation methodReference:
+                    if (methodReference.Instance != null && methodReference.Method.IsExtensionMethod) delegateParameters.Insert(0, methodReference.Instance.Accept(this, EdgeBuildingContext.Normal));
                     HandleMethodGroup(methodReference, delegateReturnType, delegateParameters);
                     break;
                 default:
