From fe32414f3828c8f36f9fa73c2aa5ebb4e151f578 Mon Sep 17 00:00:00 2001
From: Suzanne Millstein <smillst@cs.washington.edu>
Date: Thu, 17 Jun 2021 09:05:06 -0700
Subject: [PATCH] Clean up getTypeFactoryOfSubchecker. (#4729)

---
 .../checker/nullness/NullnessTransfer.java    |  3 +-
 .../common/basetype/BaseTypeChecker.java      | 19 ++++++-------
 .../type/GenericAnnotatedTypeFactory.java     | 28 +++++++++++++------
 3 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/checker/src/main/java/org/checkerframework/checker/nullness/NullnessTransfer.java b/checker/src/main/java/org/checkerframework/checker/nullness/NullnessTransfer.java
index f0d4c21e813..8a505b85a17 100644
--- a/checker/src/main/java/org/checkerframework/checker/nullness/NullnessTransfer.java
+++ b/checker/src/main/java/org/checkerframework/checker/nullness/NullnessTransfer.java
@@ -100,7 +100,8 @@ public NullnessTransfer(NullnessAnalysis analysis) {
     } else {
       // It is error-prone to put a type factory in a field.  It is OK here because
       // keyForTypeFactory is used only to call methods isMapGet() and isKeyForMap().
-      this.keyForTypeFactory = checker.getTypeFactoryOfSubchecker(KeyForSubchecker.class);
+      this.keyForTypeFactory =
+          nullnessTypeFactory.getTypeFactoryOfSubchecker(KeyForSubchecker.class);
     }
 
     NONNULL = AnnotationBuilder.fromClass(elements, NonNull.class);
diff --git a/framework/src/main/java/org/checkerframework/common/basetype/BaseTypeChecker.java b/framework/src/main/java/org/checkerframework/common/basetype/BaseTypeChecker.java
index 475beeb8e87..a224e7fe400 100644
--- a/framework/src/main/java/org/checkerframework/common/basetype/BaseTypeChecker.java
+++ b/framework/src/main/java/org/checkerframework/common/basetype/BaseTypeChecker.java
@@ -402,18 +402,17 @@ public <T extends BaseTypeChecker> T getSubchecker(Class<T> checkerClass) {
    * Returns the type factory used by a subchecker. Returns null if no matching subchecker was found
    * or if the type factory is null. The caller must know the exact checker class to request.
    *
-   * @param checkerClass the class of the subchecker
+   * <p>Because the visitor state is copied, call this method each time a subfactory is needed
+   * rather than store the returned subfactory in a field.
+   *
+   * @param subCheckerClass the class of the subchecker
+   * @param <T> the type of {@code subCheckerClass}'s {@link AnnotatedTypeFactory}
    * @return the type factory of the requested subchecker or null if not found
    */
-  @SuppressWarnings({"unchecked", "TypeParameterUnusedInFormals"}) // Intentional abuse
-  public <T extends GenericAnnotatedTypeFactory<?, ?, ?, ?>, U extends BaseTypeChecker>
-      T getTypeFactoryOfSubchecker(Class<U> checkerClass) {
-    BaseTypeChecker checker = getSubchecker(checkerClass);
-    if (checker != null) {
-      return (T) checker.getTypeFactory();
-    }
-
-    return null;
+  @SuppressWarnings("TypeParameterUnusedInFormals") // Intentional abuse
+  public <T extends GenericAnnotatedTypeFactory<?, ?, ?, ?>> @Nullable T getTypeFactoryOfSubchecker(
+      Class<? extends BaseTypeChecker> subCheckerClass) {
+    return getTypeFactory().getTypeFactoryOfSubchecker(subCheckerClass);
   }
 
   /*
diff --git a/framework/src/main/java/org/checkerframework/framework/type/GenericAnnotatedTypeFactory.java b/framework/src/main/java/org/checkerframework/framework/type/GenericAnnotatedTypeFactory.java
index f304832d764..093734033f5 100644
--- a/framework/src/main/java/org/checkerframework/framework/type/GenericAnnotatedTypeFactory.java
+++ b/framework/src/main/java/org/checkerframework/framework/type/GenericAnnotatedTypeFactory.java
@@ -2021,17 +2021,29 @@ public Store getEmptyStore() {
   }
 
   /**
-   * Returns the AnnotatedTypeFactory of the subchecker and copies the current visitor state to the
-   * sub-factory so that the types are computed properly. Because the visitor state is copied, call
-   * this method each time a subfactory is needed rather than store the returned subfactory in a
-   * field.
+   * Returns the type factory used by a subchecker. Returns null if no matching subchecker was found
+   * or if the type factory is null. The caller must know the exact checker class to request.
    *
-   * @see BaseTypeChecker#getTypeFactoryOfSubchecker(Class)
+   * <p>Because the visitor state is copied, call this method each time a subfactory is needed
+   * rather than store the returned subfactory in a field.
+   *
+   * @param subCheckerClass the exact class of the subchecker
+   * @param <T> the type of {@code subCheckerClass}'s {@link AnnotatedTypeFactory}
+   * @return the AnnotatedTypeFactory of the subchecker or null if no subchecker exists
    */
   @SuppressWarnings("TypeParameterUnusedInFormals") // Intentional abuse
-  public <T extends GenericAnnotatedTypeFactory<?, ?, ?, ?>, U extends BaseTypeChecker>
-      T getTypeFactoryOfSubchecker(Class<U> checkerClass) {
-    T subFactory = checker.getTypeFactoryOfSubchecker(checkerClass);
+  public <T extends GenericAnnotatedTypeFactory<?, ?, ?, ?>> @Nullable T getTypeFactoryOfSubchecker(
+      Class<? extends BaseTypeChecker> subCheckerClass) {
+    BaseTypeChecker subchecker = checker.getSubchecker(subCheckerClass);
+    if (subchecker == null) {
+      return null;
+    }
+
+    @SuppressWarnings(
+        "unchecked" // This might not be safe, but the caller of the method should use the correct
+    // type.
+    )
+    T subFactory = (T) subchecker.getTypeFactory();
     if (subFactory != null && subFactory.getVisitorState() != null) {
       // Copy the visitor state so that the types are computed properly.
       VisitorState subFactoryVisitorState = subFactory.getVisitorState();