fix(jdk-codemodel): resolve receiver type for unqualified method calls (#13)

Author: deer

jdk-codemodel/src/main/java/build/codemodel/jdk/JdkExpressionConverter.java

@@ -119,6 +119,7 @@ public class JdkExpressionConverter
     private Trees trees;
     private CompilationUnitTree compilationUnit;
     private Function<TypeMirror, TypeUsage> typeResolver;
+    private TypeUsage enclosingType;
 
     /**
      * Creates a {@link JdkExpressionConverter}.
@@ -155,6 +156,16 @@ public void setTypeContext(final Trees trees,
         this.typeResolver = typeResolver;
     }
 
+    /**
+     * Sets the type of the class currently being processed, used as the implicit receiver type for
+     * unqualified method calls (e.g. {@code foo()} where the receiver is {@code this}).
+     *
+     * @param enclosingType the {@link TypeUsage} of the enclosing class
+     */
+    public void setEnclosingType(final TypeUsage enclosingType) {
+        this.enclosingType = enclosingType;
+    }
+
     /**
      * Converts the given {@link ExpressionTree} to a model {@link Expression}.
      * Returns {@link NullLiteral} for a {@code null} tree.
@@ -223,6 +234,7 @@ public Expression visitMethodInvocation(final MethodInvocationTree t, final Void
         } else {
             target = null;
             methodName = t.getMethodSelect().toString();
+            receiverType = Optional.ofNullable(enclosingType);
         }
         final var args = t.getArguments().stream()
             .map(this::convert)

jdk-codemodel/src/main/java/build/codemodel/jdk/JdkInitializer.java

@@ -214,6 +214,7 @@ public Void visitClass(final ClassTree classTree, final Void unused) {
                         final var typeElement = (TypeElement) trees.getElement(classPath);
                         if (typeElement != null
                                 && !typeElement.getQualifiedName().toString().isEmpty()) {
+                            exprConverter.setEnclosingType(resolveTypeUsage(typeElement.asType(), null));
                             processTypeElement(typeElement, classTree, cut);
                         }
                         return super.visitClass(classTree, unused);

jdk-codemodel/src/test/java/build/codemodel/jdk/MethodInvocationReceiverTypeTests.java

@@ -91,7 +91,7 @@ public void run(StringBuilder sb) {
     }
 
     @Test
-    void shouldReturnEmptyReceiverTypeForUnqualifiedCall() {
+    void shouldResolveReceiverTypeForUnqualifiedCallToEnclosingClass() {
         final var source = JavaFileObjects.forSourceString(
             "build.codemodel.jdk.example.Caller", """
                 package build.codemodel.jdk.example;
@@ -119,6 +119,46 @@ public void run() {
         final var invocation = (MethodInvocation) stmt.expression();
 
         assertThat(invocation.methodName()).isEqualTo("bar");
-        assertThat(invocation.receiverType()).isEmpty();
+        assertThat(invocation.receiverType()).isPresent();
+        assertThat(invocation.receiverType().get()).isInstanceOf(NamedTypeUsage.class);
+        final var receiverTypeName = ((NamedTypeUsage) invocation.receiverType().get()).typeName();
+        assertThat(receiverTypeName.name().toString()).isEqualTo("Caller");
+    }
+
+    @Test
+    void shouldResolveReceiverTypeForUnqualifiedCallInInnerClass() {
+        final var source = JavaFileObjects.forSourceString(
+            "build.codemodel.jdk.example.Outer", """
+                package build.codemodel.jdk.example;
+                public class Outer {
+                    public class Inner {
+                        public void helper() {}
+                        public void run() {
+                            helper();
+                        }
+                    }
+                }
+                """);
+
+        final var codeModel = JdkInitializerTests.runInternal(
+            new JdkInitializer(List.of(), List.of(), List.of(source)));
+
+        final var innerTypeName = codeModel.getNameProvider()
+            .getTypeName(Optional.empty(), "build.codemodel.jdk.example.Outer.Inner");
+        final var innerDescriptor = codeModel.getTypeDescriptor(innerTypeName).orElseThrow();
+
+        final var run = innerDescriptor.traits(MethodDescriptor.class)
+            .filter(m -> m.methodName().name().toString().equals("run"))
+            .findFirst().orElseThrow();
+
+        final var body = run.getTrait(MethodBodyDescriptor.class).orElseThrow().body();
+        final var stmt = (ExpressionStatement) body.statements().findFirst().orElseThrow();
+        final var invocation = (MethodInvocation) stmt.expression();
+
+        assertThat(invocation.methodName()).isEqualTo("helper");
+        assertThat(invocation.receiverType()).isPresent();
+        assertThat(invocation.receiverType().get()).isInstanceOf(NamedTypeUsage.class);
+        final var receiverTypeName = ((NamedTypeUsage) invocation.receiverType().get()).typeName();
+        assertThat(receiverTypeName.name().toString()).isEqualTo("Inner");
     }
 }