Initial support for PolyNull

nullaway/src/main/java/com/uber/nullaway/LibraryModels.java

@@ -139,6 +139,23 @@ default ImmutableSetMultimap<MethodRef, MethodRef> ensuresNonNullIfTrueMethodCal
     return ImmutableMap.of();
   }
 
+  /**
+   * Get the locations in library method signatures that have linked, polymorphic nullness.
+   *
+   * <p>At a call, NullAway introduces nullability inference variables at the modeled input
+   * locations and generates constraints from invocation arguments, including lambdas and method
+   * references. All input occurrences must infer the same nullness, which is then substituted at
+   * every modeled location. Explicit method type arguments seed the corresponding inference
+   * variables after substitution. A parameter index of {@code -1} denotes the return type. This
+   * model is used only in JSpecify mode, and the modeled method's enclosing class is expected to be
+   * modeled as {@code @NullMarked}.
+   *
+   * @return map from methods to signature locations with polymorphic nullness
+   */
+  default ImmutableSetMultimap<MethodRef, PolyNullLocation> polyNullLocations() {
+    return ImmutableSetMultimap.of();
+  }
+
   /**
    * Get the (className, type argument index) pairs for library classes where the generic type
    * variable has a {@code @Nullable} upper bound. Only used in JSpecify mode.
@@ -339,6 +356,23 @@ public String toString() {
     }
   }
 
+  /**
+   * A location within a method signature whose nullness is linked to other locations for the same
+   * method.
+   *
+   * @param parameterIndex zero-based parameter index, or {@code -1} for the return type
+   * @param typePath path within the parameter or return type; an empty path denotes its top level
+   */
+  public record PolyNullLocation(
+      int parameterIndex, ImmutableList<NestedAnnotationInfo.TypePathEntry> typePath) {
+
+    public PolyNullLocation {
+      if (parameterIndex < -1) {
+        throw new IllegalArgumentException("parameter index must be -1 or greater");
+      }
+    }
+  }
+
   /** Representation of a field as a qualified class name + a field name */
   record FieldRef(String enclosingClassName, String fieldName) {
 

nullaway/src/main/java/com/uber/nullaway/NullAway.java

@@ -3009,7 +3009,9 @@ private boolean mayBeNullMethodCall(
     if (Nullness.hasNullableAnnotation(exprSymbol, config)) {
       return true;
     }
-    if (config.isJSpecifyMode() && exprSymbol.getReturnType().getKind().equals(TypeKind.TYPEVAR)) {
+    if (config.isJSpecifyMode()
+        && (exprSymbol.getReturnType().getKind().equals(TypeKind.TYPEVAR)
+            || genericsChecks.hasPolyNullModel(exprSymbol, state))) {
       // It is important to pass a correct TreePath to getGenericReturnNullnessAtInvocation.  So, we
       // do a search under path to find invocationTree.  This shouldn't be too costly in the common
       // case, and it's important for correctness.

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

@@ -0,0 +1,215 @@
+package com.uber.nullaway.generics;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.errorprone.VisitorState;
+import com.sun.source.tree.MethodInvocationTree;
+import com.sun.tools.javac.code.Symbol;
+import com.sun.tools.javac.code.Type;
+import com.sun.tools.javac.util.ListBuffer;
+import com.uber.nullaway.LibraryModels.PolyNullLocation;
+import com.uber.nullaway.Nullness;
+import com.uber.nullaway.librarymodel.NestedTypePathUpdater;
+import java.util.ArrayList;
+import java.util.IdentityHashMap;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Objects;
+import javax.lang.model.element.Element;
+import org.jspecify.annotations.Nullable;
+
+/** Constructs and applies polymorphic-nullness constraints for modeled method locations. */
+final class PolyNullInference {
+
+  /** Diagnostic for incompatible constraints on modeled PolyNull locations. */
+  static final String INFERENCE_FAILURE_MESSAGE =
+      "inference failure: polymorphic nullness constrained to both @NonNull and @Nullable";
+
+  /** The result of resolving PolyNull for one invocation in a generic inference session. */
+  record PolyNullInferenceResult(@Nullable Nullness nullness) {}
+
+  /** The modeled method-type overlay and shared PolyNull variable for one call. */
+  record PolyNullInferenceContext(
+      Type.MethodType inferenceMethodType,
+      ImmutableList<PolyNullLocation> locations,
+      Type.TypeVar inferenceVariable) {
+
+    /** Returns whether the overlay contains at least one modeled parameter location. */
+    boolean hasInputLocations() {
+      return locations.stream().anyMatch(location -> location.parameterIndex() >= 0);
+    }
+
+    /** Returns whether the overlay contains at least one modeled return location. */
+    boolean hasReturnLocations() {
+      return locations.stream().anyMatch(location -> location.parameterIndex() == -1);
+    }
+  }
+
+  private PolyNullInference() {}
+
+  /** Applies a resolved PolyNull annotation to every modeled parameter and return location. */
+  @SuppressWarnings({"ReferenceEquality", "TypeEquals"}) // deliberate reference equality checks
+  static Type.MethodType applyToMethodType(
+      Type.MethodType methodType, ImmutableSet<PolyNullLocation> locations, Type annotationType) {
+    boolean changed = false;
+    ListBuffer<Type> updatedParameterTypes = new ListBuffer<>();
+    int parameterIndex = 0;
+    for (com.sun.tools.javac.util.List<Type> remaining = methodType.argtypes;
+        remaining.nonEmpty();
+        remaining = remaining.tail, parameterIndex++) {
+      Type parameterType = remaining.head;
+      Type updatedParameterType =
+          applyToType(parameterType, parameterIndex, locations, annotationType);
+      updatedParameterTypes.append(updatedParameterType);
+      changed |= updatedParameterType != parameterType;
+    }
+    Type returnType = methodType.restype;
+    Type updatedReturnType = applyToReturnType(returnType, locations, annotationType);
+    changed |= updatedReturnType != returnType;
+    return changed
+        ? new Type.MethodType(
+            updatedParameterTypes.toList(), updatedReturnType, methodType.thrown, methodType.tsym)
+        : methodType;
+  }
+
+  /** Applies a resolved PolyNull annotation to every modeled location within a return type. */
+  static Type applyToReturnType(
+      Type returnType, ImmutableSet<PolyNullLocation> locations, Type annotationType) {
+    return applyToType(returnType, -1, locations, annotationType);
+  }
+
+  /** Resolves all PolyNull contexts after a shared generic-inference solver run. */
+  static IdentityHashMap<MethodInvocationTree, PolyNullInferenceResult> resolveContexts(
+      IdentityHashMap<MethodInvocationTree, PolyNullInferenceContext> contexts,
+      Map<Element, ConstraintSolver.InferredNullability> solution) {
+    IdentityHashMap<MethodInvocationTree, PolyNullInferenceResult> results =
+        new IdentityHashMap<>();
+    for (Map.Entry<MethodInvocationTree, PolyNullInferenceContext> entry : contexts.entrySet()) {
+      results.put(
+          entry.getKey(), new PolyNullInferenceResult(resolveContext(entry.getValue(), solution)));
+    }
+    return results;
+  }
+
+  /** Resolves the shared PolyNull variable for one invocation. */
+  static Nullness resolveContext(
+      PolyNullInferenceContext inferenceContext,
+      Map<Element, ConstraintSolver.InferredNullability> solution) {
+    ConstraintSolver.InferredNullability inferred =
+        solution.getOrDefault(
+            inferenceContext.inferenceVariable().asElement(),
+            ConstraintSolver.InferredNullability.NONNULL);
+    return inferred == ConstraintSolver.InferredNullability.NULLABLE
+        ? Nullness.NULLABLE
+        : Nullness.NONNULL;
+  }
+
+  /** Creates a method-type overlay with one shared variable at every modeled PolyNull location. */
+  @SuppressWarnings({"ReferenceEquality", "TypeEquals"}) // deliberate reference equality checks
+  static PolyNullInferenceContext createContext(
+      Symbol.MethodSymbol methodSymbol,
+      Type.MethodType methodType,
+      ImmutableSet<PolyNullLocation> locations,
+      Type nullableAnnotationType,
+      VisitorState state) {
+    Map<Integer, java.util.List<PolyNullLocation>> locationsByParameter = new LinkedHashMap<>();
+    ImmutableList.Builder<PolyNullLocation> appliedLocations = ImmutableList.builder();
+    Type.TypeVar inferenceVariable =
+        createInferenceVariable(methodSymbol, 0, nullableAnnotationType, state);
+    for (PolyNullLocation location : locations) {
+      int parameterIndex = location.parameterIndex();
+      if (parameterIndex < 0 || parameterIndex >= methodType.argtypes.size()) {
+        continue;
+      }
+      locationsByParameter
+          .computeIfAbsent(parameterIndex, unused -> new ArrayList<>())
+          .add(location);
+    }
+    ListBuffer<Type> updatedParameterTypes = new ListBuffer<>();
+    int parameterIndex = 0;
+    for (com.sun.tools.javac.util.List<Type> remaining = methodType.argtypes;
+        remaining.nonEmpty();
+        remaining = remaining.tail, parameterIndex++) {
+      Type updated = remaining.head;
+      for (PolyNullLocation location :
+          locationsByParameter.getOrDefault(parameterIndex, java.util.List.of())) {
+        Type replaced =
+            NestedTypePathUpdater.replaceType(updated, location.typePath(), inferenceVariable);
+        if (replaced != updated) {
+          updated = replaced;
+          appliedLocations.add(location);
+        }
+      }
+      updatedParameterTypes.append(updated);
+    }
+    Type updatedReturnType = methodType.restype;
+    for (PolyNullLocation location : locations) {
+      if (location.parameterIndex() == -1) {
+        Type replaced =
+            NestedTypePathUpdater.replaceType(
+                updatedReturnType, location.typePath(), inferenceVariable);
+        if (replaced != updatedReturnType) {
+          updatedReturnType = replaced;
+          appliedLocations.add(location);
+        }
+      }
+    }
+    return new PolyNullInferenceContext(
+        new Type.MethodType(
+            updatedParameterTypes.toList(), updatedReturnType, methodType.thrown, methodType.tsym),
+        appliedLocations.build(),
+        inferenceVariable);
+  }
+
+  /** Adds a call-result subtype constraint using the invocation's shared PolyNull variable. */
+  static void addResultConstraints(
+      ConstraintSolver solver,
+      PolyNullInferenceContext inferenceContext,
+      Type targetType,
+      boolean assignedToLocal) {
+    if (!inferenceContext.hasReturnLocations()) {
+      return;
+    }
+    solver.addSubtypeConstraint(
+        inferenceContext.inferenceMethodType().getReturnType(), targetType, assignedToLocal);
+  }
+
+  /** Returns whether {@code typeVariable} is a PolyNull variable from one of {@code contexts}. */
+  static boolean containsInferenceVariable(
+      IdentityHashMap<MethodInvocationTree, PolyNullInferenceContext> contexts,
+      Element typeVariable) {
+    return contexts.values().stream()
+        .anyMatch(context -> Objects.equals(context.inferenceVariable().asElement(), typeVariable));
+  }
+
+  /** Creates the nullable-bounded synthetic type variable for one PolyNull invocation. */
+  private static Type.TypeVar createInferenceVariable(
+      Symbol.MethodSymbol methodSymbol,
+      int group,
+      Type nullableAnnotationType,
+      VisitorState state) {
+    Symbol.TypeVariableSymbol symbol =
+        new Symbol.TypeVariableSymbol(
+            0, state.getName("$PolyNull$" + group), Type.noType, methodSymbol);
+    Type nullableObject =
+        TypeSubstitutionUtils.typeWithAnnot(state.getSymtab().objectType, nullableAnnotationType);
+    Type.TypeVar variable = new Type.TypeVar(symbol, nullableObject, state.getSymtab().botType);
+    symbol.type = variable;
+    return variable;
+  }
+
+  /** Applies {@code annotationType} at modeled locations within one method-type component. */
+  private static Type applyToType(
+      Type type,
+      int parameterIndex,
+      ImmutableSet<PolyNullLocation> locations,
+      Type annotationType) {
+    Type updated = type;
+    for (PolyNullLocation location : locations) {
+      if (location.parameterIndex() == parameterIndex) {
+        updated = NestedTypePathUpdater.addAnnotation(updated, location.typePath(), annotationType);
+      }
+    }
+    return updated;
+  }
+}

nullaway/src/main/java/com/uber/nullaway/handlers/CompositeHandler.java

@@ -41,6 +41,7 @@
 import com.sun.tools.javac.code.Types;
 import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.ErrorMessage;
+import com.uber.nullaway.LibraryModels.PolyNullLocation;
 import com.uber.nullaway.MethodParameterNullness;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.Nullness;
@@ -376,6 +377,16 @@ public Type.MethodType onOverrideMethodType(
     return currentType;
   }
 
+  @Override
+  public ImmutableSet<PolyNullLocation> onGetPolyNullLocations(
+      Symbol.MethodSymbol methodSymbol, VisitorState state) {
+    ImmutableSet.Builder<PolyNullLocation> result = ImmutableSet.builder();
+    for (Handler h : handlers) {
+      result.addAll(h.onGetPolyNullLocations(methodSymbol, state));
+    }
+    return result.build();
+  }
+
   @Override
   public FieldSkipResult shouldSkipFieldInitializationCheck(
       Symbol.ClassSymbol classSymbol, Symbol fieldSymbol, VisitorState state) {

nullaway/src/main/java/com/uber/nullaway/handlers/Handler.java

@@ -38,6 +38,7 @@
 import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.ErrorMessage;
 import com.uber.nullaway.LibraryModels;
+import com.uber.nullaway.LibraryModels.PolyNullLocation;
 import com.uber.nullaway.MethodParameterNullness;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.Nullness;
@@ -505,6 +506,18 @@ default Type.MethodType onOverrideMethodType(
     return methodType;
   }
 
+  /**
+   * Returns modeled polymorphic-nullness locations for {@code methodSymbol}.
+   *
+   * <p>The generic type-checking machinery uses these locations when resolving an invocation. A
+   * handler should only provide model metadata here; substitution and call-site reasoning belong in
+   * {@code GenericsChecks}.
+   */
+  default ImmutableSet<PolyNullLocation> onGetPolyNullLocations(
+      Symbol.MethodSymbol methodSymbol, VisitorState state) {
+    return ImmutableSet.of();
+  }
+
   enum FieldSkipResult {
     /** do not skip the check */
     NO,