Improve implementation of `asMemberOf`

PiperOrigin-RevId: 487077153

6 files changed, 145 insertions, 67 deletions

diff --git a/java/com/google/turbine/processing/TurbineElements.java b/java/com/google/turbine/processing/TurbineElements.java
index b5fd7f4..9b3ea26 100644
--- a/java/com/google/turbine/processing/TurbineElements.java
+++ b/java/com/google/turbine/processing/TurbineElements.java
@@ -384,7 +384,7 @@ public class TurbineElements implements Elements {
       return false;
     }
     TypeMirror a = overrider.asType();
-    TypeMirror b = types.asMemberOf((DeclaredType) type.asType(), overridden);
+    TypeMirror b = types.asMemberOfInternal((DeclaredType) type.asType(), overridden);
     if (b == null) {
       return false;
     }
diff --git a/java/com/google/turbine/processing/TurbineTypes.java b/java/com/google/turbine/processing/TurbineTypes.java
index d2068dd..467059c 100644
--- a/java/com/google/turbine/processing/TurbineTypes.java
+++ b/java/com/google/turbine/processing/TurbineTypes.java
@@ -26,18 +26,11 @@ import com.google.common.collect.ImmutableMap;
 import com.google.turbine.binder.bound.TypeBoundClass;
 import com.google.turbine.binder.bound.TypeBoundClass.TyVarInfo;
 import com.google.turbine.binder.sym.ClassSymbol;
-import com.google.turbine.binder.sym.FieldSymbol;
-import com.google.turbine.binder.sym.MethodSymbol;
-import com.google.turbine.binder.sym.ParamSymbol;
-import com.google.turbine.binder.sym.RecordComponentSymbol;
 import com.google.turbine.binder.sym.Symbol;
 import com.google.turbine.binder.sym.TyVarSymbol;
 import com.google.turbine.model.TurbineConstantTypeKind;
 import com.google.turbine.model.TurbineTyKind;
-import com.google.turbine.processing.TurbineElement.TurbineExecutableElement;
-import com.google.turbine.processing.TurbineElement.TurbineFieldElement;
 import com.google.turbine.processing.TurbineElement.TurbineTypeElement;
-import com.google.turbine.processing.TurbineElement.TurbineTypeParameterElement;
 import com.google.turbine.processing.TurbineTypeMirror.TurbineDeclaredType;
 import com.google.turbine.processing.TurbineTypeMirror.TurbineErrorType;
 import com.google.turbine.processing.TurbineTypeMirror.TurbineTypeVariable;
@@ -1127,41 +1120,6 @@ public class TurbineTypes implements Types {
                     .build()));
   }
 
-  private static ClassSymbol enclosingClass(Symbol symbol) {
-    switch (symbol.symKind()) {
-      case CLASS:
-        return (ClassSymbol) symbol;
-      case TY_PARAM:
-        return enclosingClass(((TyVarSymbol) symbol).owner());
-      case METHOD:
-        return ((MethodSymbol) symbol).owner();
-      case FIELD:
-        return ((FieldSymbol) symbol).owner();
-      case PARAMETER:
-        return ((ParamSymbol) symbol).owner().owner();
-      case RECORD_COMPONENT:
-        return ((RecordComponentSymbol) symbol).owner();
-      case MODULE:
-      case PACKAGE:
-        throw new IllegalArgumentException(symbol.symKind().toString());
-    }
-    throw new AssertionError(symbol.symKind());
-  }
-
-  private static Type type(Element element) {
-    switch (element.getKind()) {
-      case TYPE_PARAMETER:
-        return TyVar.create(((TurbineTypeParameterElement) element).sym(), ImmutableList.of());
-      case FIELD:
-        return ((TurbineFieldElement) element).info().type();
-      case METHOD:
-      case CONSTRUCTOR:
-        return ((TurbineExecutableElement) element).info().asType();
-      default:
-        throw new UnsupportedOperationException(element.toString());
-    }
-  }
-
   /**
    * Returns the {@link TypeMirror} of the given {@code element} as a member of {@code containing},
    * or else {@code null} if it is not a member.
@@ -1171,13 +1129,24 @@ public class TurbineTypes implements Types {
    */
   @Override
   public TypeMirror asMemberOf(DeclaredType containing, Element element) {
+    TypeMirror result = asMemberOfInternal(containing, element);
+    if (result == null) {
+      throw new IllegalArgumentException(String.format("asMemberOf(%s, %s)", containing, element));
+    }
+    return result;
+  }
+
+  public @Nullable TypeMirror asMemberOfInternal(DeclaredType containing, Element element) {
     ClassTy c = ((TurbineDeclaredType) containing).asTurbineType();
-    ClassSymbol symbol = enclosingClass(((TurbineElement) element).sym());
-    ImmutableList<ClassTy> path = factory.cha().search(c, enclosingClass(symbol));
+    Symbol enclosing = ((TurbineElement) element.getEnclosingElement()).sym();
+    if (!enclosing.symKind().equals(Symbol.Kind.CLASS)) {
+      return null;
+    }
+    ImmutableList<ClassTy> path = factory.cha().search(c, (ClassSymbol) enclosing);
     if (path.isEmpty()) {
       return null;
     }
-    Type type = type(element);
+    Type type = asTurbineType(element.asType());
     for (ClassTy ty : path) {
       ImmutableMap<TyVarSymbol, Type> mapping = getMapping(ty);
       if (mapping == null) {
diff --git a/javatests/com/google/turbine/processing/AbstractTurbineTypesBiFunctionTest.java b/javatests/com/google/turbine/processing/AbstractTurbineTypesBiFunctionTest.java
new file mode 100644
index 0000000..e00673d
--- /dev/null
+++ b/javatests/com/google/turbine/processing/AbstractTurbineTypesBiFunctionTest.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright 2022 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.turbine.processing;
+
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import javax.lang.model.element.Element;
+import javax.lang.model.type.DeclaredType;
+import javax.lang.model.util.Types;
+
+/**
+ * A combo test for {@link TurbineTypes} that compares the behaviour of bifunctions like {@link
+ * Types#asMemberOf(DeclaredType, Element)} with javac's implementation.
+ */
+abstract class AbstractTurbineTypesBiFunctionTest<T> extends AbstractTurbineTypesTest {
+
+  final String testDescription;
+  final TypesBiFunctionInput javacInput;
+  final TypesBiFunctionInput turbineInput;
+
+  public AbstractTurbineTypesBiFunctionTest(
+      String testDescription, TypesBiFunctionInput javacInput, TypesBiFunctionInput turbineInput) {
+    this.testDescription = testDescription;
+    this.javacInput = javacInput;
+    this.turbineInput = turbineInput;
+  }
+
+  protected void test(String symbol, TypeBiFunction<T> predicate) {
+    assertWithMessage("%s = %s", javacInput.format(symbol), turbineInput.format(symbol))
+        .that(turbineInput.apply(predicate))
+        .isEqualTo(javacInput.apply(predicate));
+  }
+}
diff --git a/javatests/com/google/turbine/processing/AbstractTurbineTypesBiPredicateTest.java b/javatests/com/google/turbine/processing/AbstractTurbineTypesBiPredicateTest.java
index 6ea6e72..08891eb 100644
--- a/javatests/com/google/turbine/processing/AbstractTurbineTypesBiPredicateTest.java
+++ b/javatests/com/google/turbine/processing/AbstractTurbineTypesBiPredicateTest.java
@@ -16,8 +16,6 @@
 
 package com.google.turbine.processing;
 
-import static com.google.common.truth.Truth.assertWithMessage;
-
 import javax.lang.model.type.TypeMirror;
 import javax.lang.model.util.Types;
 
@@ -25,22 +23,10 @@ import javax.lang.model.util.Types;
  * A combo test for {@link TurbineTypes} that compares the behaviour of bipredicates like {@link
  * Types#isSubtype(TypeMirror, TypeMirror)} with javac's implementation.
  */
-abstract class AbstractTurbineTypesBiPredicateTest extends AbstractTurbineTypesTest {
-
-  final String testDescription;
-  final TypesBiFunctionInput javacInput;
-  final TypesBiFunctionInput turbineInput;
-
+abstract class AbstractTurbineTypesBiPredicateTest
+    extends AbstractTurbineTypesBiFunctionTest<Boolean> {
   public AbstractTurbineTypesBiPredicateTest(
       String testDescription, TypesBiFunctionInput javacInput, TypesBiFunctionInput turbineInput) {
-    this.testDescription = testDescription;
-    this.javacInput = javacInput;
-    this.turbineInput = turbineInput;
-  }
-
-  protected void test(String symbol, TypeBiPredicate predicate) {
-    assertWithMessage("%s = %s", javacInput.format(symbol), turbineInput.format(symbol))
-        .that(turbineInput.apply(predicate))
-        .isEqualTo(javacInput.apply(predicate));
+    super(testDescription, javacInput, turbineInput);
   }
 }
diff --git a/javatests/com/google/turbine/processing/AbstractTurbineTypesTest.java b/javatests/com/google/turbine/processing/AbstractTurbineTypesTest.java
index 2ad041c..02df1ec 100644
--- a/javatests/com/google/turbine/processing/AbstractTurbineTypesTest.java
+++ b/javatests/com/google/turbine/processing/AbstractTurbineTypesTest.java
@@ -94,8 +94,8 @@ class AbstractTurbineTypesTest {
     }
   }
 
-  protected interface TypeBiPredicate {
-    boolean apply(Types types, TypeMirror a, TypeMirror b);
+  protected interface TypeBiFunction<T> {
+    T apply(Types types, TypeMirror a, TypeMirror b);
   }
 
   static class TypesBiFunctionInput {
@@ -109,8 +109,8 @@ class AbstractTurbineTypesTest {
       this.rhs = rhs;
     }
 
-    boolean apply(TypeBiPredicate predicate) {
-      return predicate.apply(types, lhs, rhs);
+    <T> T apply(TypeBiFunction<T> function) {
+      return function.apply(types, lhs, rhs);
     }
 
     String format(String symbol) {
diff --git a/javatests/com/google/turbine/processing/TurbineTypesAsMemberOfTest.java b/javatests/com/google/turbine/processing/TurbineTypesAsMemberOfTest.java
new file mode 100644
index 0000000..1a368c9
--- /dev/null
+++ b/javatests/com/google/turbine/processing/TurbineTypesAsMemberOfTest.java
@@ -0,0 +1,76 @@
+/*
+ * Copyright 2022 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.turbine.processing;
+
+import static com.google.common.truth.TruthJUnit.assume;
+import static org.junit.Assert.assertThrows;
+
+import javax.lang.model.element.Element;
+import javax.lang.model.type.DeclaredType;
+import javax.lang.model.type.TypeKind;
+import javax.lang.model.type.TypeMirror;
+import javax.lang.model.type.TypeVariable;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class TurbineTypesAsMemberOfTest extends AbstractTurbineTypesBiFunctionTest<String> {
+
+  @Parameters(name = "{index}: {0}")
+  public static Iterable<Object[]> parameters() throws Exception {
+    return binaryParameters();
+  }
+
+  public TurbineTypesAsMemberOfTest(
+      String testDescription, TypesBiFunctionInput javacInput, TypesBiFunctionInput turbineInput) {
+    super(testDescription, javacInput, turbineInput);
+  }
+
+  @Test
+  public void asMemberOf() {
+    assume().that(javacInput.lhs.getKind()).isEqualTo(TypeKind.DECLARED);
+    assume().that(javacInput.rhs.getKind()).isAnyOf(TypeKind.TYPEVAR, TypeKind.DECLARED);
+
+    TypeBiFunction<String> predicate =
+        (types, lhs, rhs) -> types.asMemberOf((DeclaredType) lhs, element(rhs)).toString();
+
+    try {
+      String unused = javacInput.apply(predicate);
+    } catch (IllegalArgumentException e) {
+      assertThrows(
+          turbineInput.format("asMemberOf"),
+          IllegalArgumentException.class,
+          () -> turbineInput.apply(predicate));
+      return;
+    }
+
+    test("asMemberOf", predicate);
+  }
+
+  private static Element element(TypeMirror rhs) {
+    switch (rhs.getKind()) {
+      case TYPEVAR:
+        return ((TypeVariable) rhs).asElement();
+      case DECLARED:
+        return ((DeclaredType) rhs).asElement();
+      default:
+        throw new AssertionError(rhs.getKind());
+    }
+  }
+}