From 7c6953255898f8a6d5c62c852be5b3c70e8e75f8 Mon Sep 17 00:00:00 2001
From: Manu Sridharan <msridhar@gmail.com>
Date: Wed, 30 Oct 2024 16:30:46 -0700
Subject: [PATCH] cleanup logic

---
 .../nullaway/generics/GenericsChecks.java     | 25 +++++++------------
 1 file changed, 9 insertions(+), 16 deletions(-)

diff --git a/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java b/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java
index 86337f2bd9..ddfcadab90 100644
--- a/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java
+++ b/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java
@@ -799,22 +799,15 @@ public static Nullness getGenericReturnNullnessAtInvocation(
       List<? extends Tree> typeArgumentTrees = tree.getTypeArguments();
       com.sun.tools.javac.util.List<Type> explicitTypeArgs =
           convertTreesToTypes(typeArgumentTrees); // Convert to Type objects
-
-      Type methodType = invokedMethodSymbol.type;
-      Type substitutedReturnType = null;
-      if (methodType instanceof Type.ForAll) {
-        Type.ForAll forAllType = (Type.ForAll) methodType;
-
-        // Extract the underlying MethodType (the actual signature)
-        Type.MethodType methodTypeInsideForAll = (Type.MethodType) forAllType.qtype;
-
-        // Substitute the argument and return types within the MethodType
-        substitutedReturnType =
-            state
-                .getTypes()
-                .subst(methodTypeInsideForAll.restype, forAllType.tvars, explicitTypeArgs);
-      }
-
+      Type.ForAll forAllType = (Type.ForAll) invokedMethodSymbol.type;
+      // Extract the underlying MethodType (the actual signature)
+      Type.MethodType methodTypeInsideForAll = (Type.MethodType) forAllType.asMethodType();
+      // Substitute type arguments inside the return type
+      // NOTE: if the return type it not a type variable, this is a noop
+      Type substitutedReturnType =
+          state
+              .getTypes()
+              .subst(methodTypeInsideForAll.restype, forAllType.tvars, explicitTypeArgs);
       if (substitutedReturnType != null
           && Objects.equals(getTypeNullness(substitutedReturnType, config), Nullness.NULLABLE)) {
         return Nullness.NULLABLE;