Bug fix with type substitutions after inference (#1277)

Fixes #1267 

Due to possible type mutations, e.g., when substituting for bounds,
sometimes we see `TypeVar`s in types that are not the same object as in
the original method type for which we performed inference. Previously,
we wouldn't substitute properly in such cases. Here, we collect all
`TypeVar`s with the same `Symbol` observed during inference and
substitute for all of them, fixing the problem. This makes the Caffeine
integration test green again.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Improved nullability handling for generics: all matching occurrences
of a type variable are now annotated with a synthetic nullable marker,
ensuring consistent analysis and avoiding missed substitutions across
repeated type-variable occurrences.
* **Tests**
* Updated tests to reflect improved generic inference behavior; removed
prior suppression/workarounds.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: msridhar

nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java

@@ -657,15 +657,26 @@ private Type substituteInferredNullabilityForTypeVariables(
     for (Map.Entry<Element, ConstraintSolver.InferredNullability> entry :
         typeVarNullability.entrySet()) {
       if (entry.getValue() == NULLABLE) {
-        Type curTypeVar = (Type) entry.getKey().asType();
-        typeVars.append(curTypeVar);
-        inferredTypes.append(
-            TypeSubstitutionUtils.typeWithAnnot(
-                curTypeVar, GenericsChecks.getSyntheticNullAnnotType(state)));
+        // find all TypeVars occurring in targetType with the same symbol and substitute for those.
+        // we can have multiple such TypeVars due to previous substitutions that modified the type
+        // in some way, e.g., by changing its bounds
+        Element symbol = entry.getKey();
+        TypeVarWithSymbolCollector tvc = new TypeVarWithSymbolCollector(symbol);
+        targetType.accept(tvc, null);
+        for (Type.TypeVar tv : tvc.getMatches()) {
+          typeVars.append(tv);
+          inferredTypes.append(
+              TypeSubstitutionUtils.typeWithAnnot(tv, getSyntheticNullAnnotType(state)));
+        }
       }
     }
-    return TypeSubstitutionUtils.subst(
-        state.getTypes(), targetType, typeVars.toList(), inferredTypes.toList(), config);
+    com.sun.tools.javac.util.List<Type> typeVarsToReplace = typeVars.toList();
+    if (!typeVarsToReplace.isEmpty()) {
+      return TypeSubstitutionUtils.subst(
+          state.getTypes(), targetType, typeVarsToReplace, inferredTypes.toList(), config);
+    } else {
+      return targetType;
+    }
   }
 
   /**
@@ -1437,6 +1448,8 @@ public boolean isNullableAnnotated(Type type) {
     return Nullness.hasNullableAnnotation(type.getAnnotationMirrors().stream(), config);
   }
 
+  private @Nullable Type syntheticNullAnnotType;
+
   /**
    * Returns a "fake" {@link Type} object representing a synthetic {@code @Nullable} annotation.
    *
@@ -1450,12 +1463,15 @@ public boolean isNullableAnnotated(Type type) {
    *     Symtab}.
    * @return a fake {@code Type} for a synthetic {@code @Nullable} annotation.
    */
-  public static Type getSyntheticNullAnnotType(VisitorState state) {
-    Names names = Names.instance(state.context);
-    Symtab symtab = Symtab.instance(state.context);
-    Name name = names.fromString("nullaway.synthetic");
-    Symbol.PackageSymbol packageSymbol = new Symbol.PackageSymbol(name, symtab.noSymbol);
-    Name simpleName = names.fromString("Nullable");
-    return new Type.ErrorType(simpleName, packageSymbol, Type.noType);
+  private Type getSyntheticNullAnnotType(VisitorState state) {
+    if (syntheticNullAnnotType == null) {
+      Names names = Names.instance(state.context);
+      Symtab symtab = Symtab.instance(state.context);
+      Name name = names.fromString("nullaway.synthetic");
+      Symbol.PackageSymbol packageSymbol = new Symbol.PackageSymbol(name, symtab.noSymbol);
+      Name simpleName = names.fromString("Nullable");
+      syntheticNullAnnotType = new Type.ErrorType(simpleName, packageSymbol, Type.noType);
+    }
+    return syntheticNullAnnotType;
   }
 }

nullaway/src/main/java/com/uber/nullaway/generics/TypeVarWithSymbolCollector.java

@@ -0,0 +1,124 @@
+package com.uber.nullaway.generics;
+
+import com.sun.tools.javac.code.Type;
+import com.sun.tools.javac.code.Type.ArrayType;
+import com.sun.tools.javac.code.Type.CapturedType;
+import com.sun.tools.javac.code.Type.ClassType;
+import com.sun.tools.javac.code.Type.ForAll;
+import com.sun.tools.javac.code.Type.MethodType;
+import com.sun.tools.javac.code.Type.TypeVar;
+import com.sun.tools.javac.code.Type.WildcardType;
+import com.sun.tools.javac.code.Types;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.LinkedHashSet;
+import java.util.Set;
+import javax.lang.model.element.Element;
+import org.jspecify.annotations.Nullable;
+
+/**
+ * Visitor that collects every TypeVar whose symbol is the given Element.
+ *
+ * <p>Usage: TypeVarWithSymbolCollector v = new TypeVarWithSymbolCollector(elem);
+ * rootType.accept(v,null); Set<TypeVar> matches = v.getMatches();
+ *
+ * <p>Not safe to run multiple times; create a fresh visitor for each root type to scan.
+ */
+public final class TypeVarWithSymbolCollector extends Types.DefaultTypeVisitor<Void, Void> {
+
+  private final Element symbol;
+  private final Set<TypeVar> matches = new LinkedHashSet<>();
+
+  public TypeVarWithSymbolCollector(Element symbol) {
+    this.symbol = symbol;
+  }
+
+  private final Set<Type> seen = java.util.Collections.newSetFromMap(new IdentityHashMap<>());
+
+  /** Walk a (possibly null) type. */
+  private void scan(@Nullable Type t) {
+    if (t != null && seen.add(t)) {
+      t.accept(this, null);
+    }
+  }
+
+  /** Results (unmodifiable). */
+  public Set<TypeVar> getMatches() {
+    return Collections.unmodifiableSet(matches);
+  }
+
+  // ---- Core matching logic ----
+  @Override
+  public Void visitTypeVar(TypeVar t, Void p) {
+    if (t.tsym == symbol) {
+      matches.add(t);
+    }
+    scan(t.getUpperBound());
+    scan(t.getLowerBound());
+    return null;
+  }
+
+  // ---- Common container types ----
+  @Override
+  public Void visitClassType(ClassType t, Void p) {
+    for (Type arg : t.getTypeArguments()) {
+      scan(arg);
+    }
+    // For inner classes, the enclosing type may also carry args.
+    scan(t.getEnclosingType());
+    return null;
+  }
+
+  @Override
+  public Void visitArrayType(ArrayType t, Void p) {
+    scan(t.getComponentType());
+    return null;
+  }
+
+  @Override
+  public Void visitWildcardType(WildcardType t, Void p) {
+    scan(t.getExtendsBound());
+    scan(t.getSuperBound());
+    return null;
+  }
+
+  @Override
+  public Void visitCapturedType(CapturedType t, Void p) {
+    scan(t.getUpperBound());
+    scan(t.getLowerBound());
+    scan(t.wildcard);
+    return null;
+  }
+
+  // ---- Functional / executable types ----
+  @Override
+  public Void visitMethodType(MethodType t, Void p) {
+    scan(t.getReturnType());
+    for (Type pt : t.getParameterTypes()) {
+      scan(pt);
+    }
+    for (Type thrown : t.getThrownTypes()) {
+      scan(thrown);
+    }
+    return null;
+  }
+
+  @Override
+  public Void visitForAll(ForAll t, Void p) {
+    scan(t.qtype);
+    for (Type type : t.getTypeArguments()) {
+      scan(type);
+    }
+    return null;
+  }
+
+  // ---- Trivial / leaf types we don't need to descend into ----
+  @Override
+  public Void visitType(Type t, Void p) {
+    // Fallback: best-effort traversal via type arguments, if any.
+    for (Type arg : t.getTypeArguments()) {
+      scan(arg);
+    }
+    return null;
+  }
+}

nullaway/src/test/java/com/uber/nullaway/jspecify/GenericMethodTests.java

@@ -1110,9 +1110,7 @@ public void nullableWildcardFromCaffeine() {
             "        void test() {",
             "            JCacheLoaderAdapter<K, V> adapter = new JCacheLoaderAdapter<>();",
             "            caffeine.<K, @Nullable Expirable<V>>build(adapter);",
-            "            // TODO inference is buggy for this case; doesn't substitute inferred types properly",
-            "            // See https://github.com/uber/NullAway/issues/1267",
-            "            @SuppressWarnings(\"NullAway\")",
+            "            // also works with inference",
             "            Object o = caffeine.build(adapter);",
             "        }",
             "    }",