From 4edd8a42d6c0c553658a9ffd7f6e98adb70ac52b Mon Sep 17 00:00:00 2001
From: Rawi01 <Rawi01@users.noreply.github.com>
Date: Sun, 16 Apr 2023 14:14:26 +0200
Subject: [PATCH] [fixes #3402] Delay replacement of extension methods in
 arguments

---
 .../javac/handlers/HandleExtensionMethod.java |  5 ++-
 .../ExtensionMethodInLambda.java              | 34 ++++++++++++++++++
 .../after-ecj/ExtensionMethodInLambda.java    | 33 +++++++++++++++++
 .../before/ExtensionMethodInLambda.java       | 35 +++++++++++++++++++
 .../ExtensionMethodInLambda.java.messages     |  1 +
 .../ExtensionMethodInLambda.java.messages     |  1 +
 .../ExtensionMethodInLambda.java.messages     |  1 +
 7 files changed, 109 insertions(+), 1 deletion(-)
 create mode 100644 test/transform/resource/after-delombok/ExtensionMethodInLambda.java
 create mode 100644 test/transform/resource/after-ecj/ExtensionMethodInLambda.java
 create mode 100644 test/transform/resource/before/ExtensionMethodInLambda.java
 create mode 100644 test/transform/resource/messages-delombok/ExtensionMethodInLambda.java.messages
 create mode 100644 test/transform/resource/messages-ecj/ExtensionMethodInLambda.java.messages
 create mode 100644 test/transform/resource/messages-idempotent/ExtensionMethodInLambda.java.messages

diff --git a/src/core/lombok/javac/handlers/HandleExtensionMethod.java b/src/core/lombok/javac/handlers/HandleExtensionMethod.java
index af03d000..cb797f9e 100644
--- a/src/core/lombok/javac/handlers/HandleExtensionMethod.java
+++ b/src/core/lombok/javac/handlers/HandleExtensionMethod.java
@@ -152,8 +152,10 @@ public class HandleExtensionMethod extends JavacAnnotationHandler<ExtensionMetho
 		
 		@Override
 		public Void visitMethodInvocation(final MethodInvocationTree tree, final Void p) {
-			super.visitMethodInvocation(tree, p);
+			scan(tree.getTypeArguments(), p);
+			scan(tree.getMethodSelect(), p);
 			handleMethodCall((JCMethodInvocation) tree);
+			scan(tree.getArguments(), p);
 			return null;
 		}
 		
@@ -181,6 +183,7 @@ public class HandleExtensionMethod extends JavacAnnotationHandler<ExtensionMetho
 			JCTree resolvedReceiver = resolution.get(receiver);
 			if (resolvedReceiver == null || resolvedReceiver.type == null) return;
 			Type receiverType = resolvedReceiver.type;
+			if (receiverType.isErroneous()) return;
 			
 			// Skip static method access
 			Symbol sym = null;
diff --git a/test/transform/resource/after-delombok/ExtensionMethodInLambda.java b/test/transform/resource/after-delombok/ExtensionMethodInLambda.java
new file mode 100644
index 00000000..56f374df
--- /dev/null
+++ b/test/transform/resource/after-delombok/ExtensionMethodInLambda.java
@@ -0,0 +1,34 @@
+// version 8:
+import java.util.function.Function;
+
+public class ExtensionMethodInLambda {
+	public void testSimple() {
+		String test = "test";
+		test = ExtensionMethodInLambda.Extensions.map(test, s -> ExtensionMethodInLambda.Extensions.reverse(s));
+	}
+
+	public void testSameName() {
+		String test = "test";
+		test = ExtensionMethodInLambda.Extensions.map(test, s -> s.trim());
+	}
+
+	public void testArgumentOfInvalidMethod() {
+		String test = "test";
+		test.invalid(s -> s.reverse());
+	}
+
+
+	static class Extensions {
+		public static <T, R> R map(T value, Function<T, R> mapper) {
+			return mapper.apply(value);
+		}
+
+		public static String reverse(String string) {
+			return new StringBuilder(string).reverse().toString();
+		}
+
+		public static String trim(Integer integer) {
+			return "0";
+		}
+	}
+}
diff --git a/test/transform/resource/after-ecj/ExtensionMethodInLambda.java b/test/transform/resource/after-ecj/ExtensionMethodInLambda.java
new file mode 100644
index 00000000..ef2b2cac
--- /dev/null
+++ b/test/transform/resource/after-ecj/ExtensionMethodInLambda.java
@@ -0,0 +1,33 @@
+import java.util.function.Function;
+import lombok.experimental.ExtensionMethod;
+public @ExtensionMethod(value = ExtensionMethodInLambda.Extensions.class) class ExtensionMethodInLambda {
+  static class Extensions {
+    Extensions() {
+      super();
+    }
+    public static <T, R>R map(T value, Function<T, R> mapper) {
+      return mapper.apply(value);
+    }
+    public static String reverse(String string) {
+      return new StringBuilder(string).reverse().toString();
+    }
+    public static String trim(Integer integer) {
+      return "0";
+    }
+  }
+  public ExtensionMethodInLambda() {
+    super();
+  }
+  public void testSimple() {
+    String test = "test";
+    test = ExtensionMethodInLambda.Extensions.map(test, (<no type> s) -> ExtensionMethodInLambda.Extensions.reverse(s));
+  }
+  public void testSameName() {
+    String test = "test";
+    test = ExtensionMethodInLambda.Extensions.map(test, (<no type> s) -> s.trim());
+  }
+  public void testArgumentOfInvalidMethod() {
+    String test = "test";
+    test.invalid((<no type> s) -> s.reverse());
+  }
+}
diff --git a/test/transform/resource/before/ExtensionMethodInLambda.java b/test/transform/resource/before/ExtensionMethodInLambda.java
new file mode 100644
index 00000000..c441a924
--- /dev/null
+++ b/test/transform/resource/before/ExtensionMethodInLambda.java
@@ -0,0 +1,35 @@
+// version 8:
+import java.util.function.Function;
+import lombok.experimental.ExtensionMethod;
+
+@ExtensionMethod(value = ExtensionMethodInLambda.Extensions.class)
+public class ExtensionMethodInLambda {
+	public void testSimple() {
+		String test = "test";
+		test = test.map(s -> s.reverse());
+	}
+	
+	public void testSameName() {
+		String test = "test";
+		test = test.map(s -> s.trim());
+	}
+	
+	public void testArgumentOfInvalidMethod() {
+		String test = "test";
+		test.invalid(s -> s.reverse());
+	}
+	
+	static class Extensions {
+		public static <T, R> R map(T value, Function<T, R> mapper) {
+			return mapper.apply(value);
+		}
+		
+		public static String reverse(String string) {
+			return new StringBuilder(string).reverse().toString();
+		}
+		
+		public static String trim(Integer integer) {
+			return "0";
+		}
+	}
+}
diff --git a/test/transform/resource/messages-delombok/ExtensionMethodInLambda.java.messages b/test/transform/resource/messages-delombok/ExtensionMethodInLambda.java.messages
new file mode 100644
index 00000000..7605def9
--- /dev/null
+++ b/test/transform/resource/messages-delombok/ExtensionMethodInLambda.java.messages
@@ -0,0 +1 @@
+19 cannot find symbol symbol: method invalid((s)->s.reverse()) location: variable test of type java.lang.String
\ No newline at end of file
diff --git a/test/transform/resource/messages-ecj/ExtensionMethodInLambda.java.messages b/test/transform/resource/messages-ecj/ExtensionMethodInLambda.java.messages
new file mode 100644
index 00000000..56888c34
--- /dev/null
+++ b/test/transform/resource/messages-ecj/ExtensionMethodInLambda.java.messages
@@ -0,0 +1 @@
+19 The method invalid((<no type> s) -> {}) is undefined for the type String
diff --git a/test/transform/resource/messages-idempotent/ExtensionMethodInLambda.java.messages b/test/transform/resource/messages-idempotent/ExtensionMethodInLambda.java.messages
new file mode 100644
index 00000000..15d70875
--- /dev/null
+++ b/test/transform/resource/messages-idempotent/ExtensionMethodInLambda.java.messages
@@ -0,0 +1 @@
+17 cannot find symbol symbol: method invalid((s)->s.reverse()) location: variable test of type java.lang.String
\ No newline at end of file