Handle complex capture assignment

aromaa · Jan 17, 2024 · 9e456f7 · 9e456f7
1 parent 3796893
commit 9e456f7
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 6 deletions.
diff --git a/src/FunctionalInterfaces.SourceGenerator/FunctionalInterfacesGenerator.cs b/src/FunctionalInterfaces.SourceGenerator/FunctionalInterfacesGenerator.cs
@@ -133,12 +133,19 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 							}
 							else if (modified is IdentifierNameSyntax { Parent: not MemberAccessExpressionSyntax { Expression: IdentifierNameSyntax { Identifier.Text: "__functionalInterface" } } } identifier)
 							{
-								bool captured = dataFlowAnalysis.Captured.Any(i => i.Name == identifier.Identifier.Text);
-								if (captured)
+								ISymbol? captured = dataFlowAnalysis.Captured.SingleOrDefault(i => i is ILocalSymbol && i.Name == identifier.Identifier.Text);
+								if (captured is not null)
 								{
-									return SyntaxFactory.MemberAccessExpression(
-										SyntaxKind.SimpleMemberAccessExpression,
-										SyntaxFactory.IdentifierName("__functionalInterface"), identifier);
+									foreach (SyntaxReference reference in captured.DeclaringSyntaxReferences)
+									{
+										SyntaxNode declaration = reference.GetSyntax();
+										if (declaration is VariableDeclaratorSyntax)
+										{
+											return SyntaxFactory.MemberAccessExpression(
+												SyntaxKind.SimpleMemberAccessExpression,
+												SyntaxFactory.IdentifierName("__functionalInterface"), identifier);
+										}
+									}
 								}
 							}
 

diff --git a/src/FunctionalInterfaces.TestAssembly/ActionTests.cs b/src/FunctionalInterfaces.TestAssembly/ActionTests.cs
@@ -31,7 +31,7 @@ public static void CallActionWithCapturedIntReferenceOutside()
 	{
 		int param = 50;
 
-		_ = param;
+		Assert.Equal(50, param);
 
 		ActionTests.Invoke(() =>
 		{
@@ -54,6 +54,26 @@ public static void CallActionWithCapturedIntTwoTimes()
 		});
 	}
 
+	public static void CallActionWithComplexCaptureAssignment()
+	{
+		DataHolder param = new(50);
+		if (param.Data is not int intParam)
+		{
+			Assert.Fail("Unreachable");
+
+			return;
+		}
+
+		Assert.Equal(50, intParam);
+
+		ActionTests.Invoke(() =>
+		{
+			Assert.Equal(50, intParam);
+		});
+	}
+
+	private sealed record DataHolder(object Data);
+
 	public static void CallVirtualActionWithCapturedInt()
 	{
 		int param = 50;