Refactored data clumps with the help of LLMs (research project)

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

@@ -98,6 +98,7 @@
 import com.uber.nullaway.generics.GenericsChecks;
 import com.uber.nullaway.handlers.Handler;
 import com.uber.nullaway.handlers.Handlers;
+import com.uber.nullaway.handlers.MethodAnalysisContext;
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
@@ -384,7 +385,7 @@ public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState
       return Description.NO_MATCH;
     }
     Symbol.MethodSymbol methodSymbol = getSymbolForMethodInvocation(tree);
-    handler.onMatchMethodInvocation(this, tree, state, methodSymbol);
+    handler.onMatchMethodInvocation(tree, new MethodAnalysisContext(this, state, methodSymbol));
     // assuming this list does not include the receiver
     List<? extends ExpressionTree> actualParams = tree.getArguments();
     return handleInvocation(tree, state, methodSymbol, actualParams);
@@ -644,7 +645,7 @@ public Description matchMethod(MethodTree tree, VisitorState state) {
     // overridden method (if overridden method is in an annotated
     // package)
     Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(tree);
-    handler.onMatchMethod(this, tree, state, methodSymbol);
+    handler.onMatchMethod(tree, new MethodAnalysisContext(this, state, methodSymbol));
     boolean isOverriding = ASTHelpers.hasAnnotation(methodSymbol, "java.lang.Override", state);
     boolean exhaustiveOverride = config.exhaustiveOverride();
     if (isOverriding || !exhaustiveOverride) {
@@ -957,7 +958,8 @@ public Description matchLambdaExpression(LambdaExpressionTree tree, VisitorState
     // we need to update environment mapping before running the handler, as some handlers
     // (like Rx nullability) run dataflow analysis
     updateEnvironmentMapping(state.getPath(), state);
-    handler.onMatchLambdaExpression(this, tree, state, funcInterfaceMethod);
+    handler.onMatchLambdaExpression(
+        tree, new MethodAnalysisContext(this, state, funcInterfaceMethod));
     if (codeAnnotationInfo.isSymbolUnannotated(funcInterfaceMethod, config, handler)) {
       return Description.NO_MATCH;
     }
@@ -1001,7 +1003,7 @@ public Description matchMemberReference(MemberReferenceTree tree, VisitorState s
     Symbol.MethodSymbol referencedMethod = ASTHelpers.getSymbol(tree);
     Symbol.MethodSymbol funcInterfaceSymbol =
         NullabilityUtil.getFunctionalInterfaceMethod(tree, state.getTypes());
-    handler.onMatchMethodReference(this, tree, state, referencedMethod);
+    handler.onMatchMethodReference(tree, new MethodAnalysisContext(this, state, referencedMethod));
     return checkOverriding(funcInterfaceSymbol, referencedMethod, tree, state);
   }
 
@@ -1445,7 +1447,7 @@ private boolean okToReadBeforeInitialized(TreePath path, VisitorState state) {
       } else {
         castToNonNullArg =
             handler.castToNonNullArgumentPositionsForMethod(
-                this, state, methodSymbol, arguments, null);
+                arguments, null, new MethodAnalysisContext(this, state, methodSymbol));
       }
       if (castToNonNullArg != null && leaf.equals(arguments.get(castToNonNullArg))) {
         return true;
@@ -1818,7 +1820,7 @@ private Description checkCastToNonNullTakesNullable(
     } else {
       castToNonNullPosition =
           handler.castToNonNullArgumentPositionsForMethod(
-              this, state, methodSymbol, actualParams, null);
+              actualParams, null, new MethodAnalysisContext(this, state, methodSymbol));
     }
     if (castToNonNullPosition != null) {
       ExpressionTree actual = actualParams.get(castToNonNullPosition);

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

@@ -60,22 +60,22 @@ protected AbstractFieldContractHandler(String annotName) {
   /**
    * Verifies that the method being processed adheres to the annotation specifications.
    *
-   * @param analysis NullAway instance.
    * @param tree Method tree under processing.
-   * @param state Error Prone {@link VisitorState}.
-   * @param methodSymbol Processing method symbol.
+   * @param methodAnalysisContext The MethodAnalysisContext object
    */
   @Override
-  public void onMatchMethod(
-      NullAway analysis, MethodTree tree, VisitorState state, Symbol.MethodSymbol methodSymbol) {
+  public void onMatchMethod(MethodTree tree, MethodAnalysisContext methodAnalysisContext) {
+
+    Symbol.MethodSymbol methodSymbol = methodAnalysisContext.methodSymbol();
+    VisitorState state = methodAnalysisContext.state();
     Set<String> annotationContent =
         NullabilityUtil.getAnnotationValueArray(methodSymbol, annotName, false);
     boolean isAnnotated = annotationContent != null;
     boolean isValid =
         isAnnotated
             && validateAnnotationSyntax(
-                castToNonNull(annotationContent), analysis, tree, state, methodSymbol)
-            && validateAnnotationSemantics(analysis, state, tree, methodSymbol);
+                castToNonNull(annotationContent), tree, methodAnalysisContext)
+            && validateAnnotationSemantics(tree, methodAnalysisContext);
     if (isAnnotated && !isValid) {
       return;
     }
@@ -90,8 +90,9 @@ && validateAnnotationSyntax(
     } else {
       fieldNames = Collections.emptySet();
     }
-    validateOverridingRules(fieldNames, analysis, state, tree, closestOverriddenMethod);
-    super.onMatchMethod(analysis, tree, state, methodSymbol);
+    validateOverridingRules(
+        fieldNames, methodAnalysisContext.analysis(), state, tree, closestOverriddenMethod);
+    super.onMatchMethod(tree, methodAnalysisContext);
   }
 
   /**
@@ -117,9 +118,10 @@ protected abstract void validateOverridingRules(
    * Validates that a method implementation matches the semantics of the annotation.
    *
    * @return Returns true, if the annotation conforms to the semantic rules.
+   * @param methodAnalysisContext The MethodAnalysisContext object
    */
   protected abstract boolean validateAnnotationSemantics(
-      NullAway analysis, VisitorState state, MethodTree tree, Symbol.MethodSymbol methodSymbol);
+      MethodTree tree, MethodAnalysisContext methodAnalysisContext);
 
   /**
    * Validates whether the parameter inside annotation conforms to the syntax rules. Parameters must
@@ -137,14 +139,13 @@ protected abstract boolean validateAnnotationSemantics(
    * <p>
    *
    * @return Returns true, if the annotation conforms to the syntax rules.
+   * @param methodAnalysisContext The MethodAnalysisContext object
    */
   protected boolean validateAnnotationSyntax(
-      Set<String> content,
-      NullAway analysis,
-      MethodTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol) {
+      Set<String> content, MethodTree tree, MethodAnalysisContext methodAnalysisContext) {
     String message;
+    VisitorState state = methodAnalysisContext.state();
+    NullAway analysis = methodAnalysisContext.analysis();
     if (content.isEmpty()) {
       // we should not allow useless annotations.
       message =
@@ -187,7 +188,8 @@ protected boolean validateAnnotationSyntax(
             fieldName = fieldName.substring(fieldName.lastIndexOf(".") + 1);
           }
         }
-        Symbol.ClassSymbol classSymbol = castToNonNull(ASTHelpers.enclosingClass(methodSymbol));
+        Symbol.ClassSymbol classSymbol =
+            castToNonNull(ASTHelpers.enclosingClass(methodAnalysisContext.methodSymbol()));
         VariableElement field = getInstanceFieldOfClass(classSymbol, fieldName);
         if (field == null) {
           message =
@@ -197,6 +199,7 @@ protected boolean validateAnnotationSyntax(
                   + fieldName
                   + " in class "
                   + classSymbol.getSimpleName();
+
           state.reportMatch(
               analysis
                   .getErrorBuilder()

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

@@ -73,35 +73,25 @@ public void onMatchTopLevelClass(
   }
 
   @Override
-  public void onMatchMethod(
-      NullAway analysis, MethodTree tree, VisitorState state, Symbol.MethodSymbol methodSymbol) {
+  public void onMatchMethod(MethodTree tree, MethodAnalysisContext methodAnalysisContext) {
     // NoOp
   }
 
   @Override
   public void onMatchMethodInvocation(
-      NullAway analysis,
-      MethodInvocationTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol) {
+      MethodInvocationTree tree, MethodAnalysisContext methodAnalysisContext) {
     // NoOp
   }
 
   @Override
   public void onMatchLambdaExpression(
-      NullAway analysis,
-      LambdaExpressionTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol) {
+      LambdaExpressionTree tree, MethodAnalysisContext methodAnalysisContext) {
     // NoOp
   }
 
   @Override
   public void onMatchMethodReference(
-      NullAway analysis,
-      MemberReferenceTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol) {
+      MemberReferenceTree tree, MethodAnalysisContext methodAnalysisContext) {
     // NoOp
   }
 
@@ -238,11 +228,9 @@ public MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
   @Override
   @Nullable
   public Integer castToNonNullArgumentPositionsForMethod(
-      NullAway analysis,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol,
       List<? extends ExpressionTree> actualParams,
-      @Nullable Integer previousArgumentPosition) {
+      @Nullable Integer previousArgumentPosition,
+      MethodAnalysisContext methodAnalysisContext) {
     // NoOp
     return previousArgumentPosition;
   }

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

@@ -80,43 +80,33 @@ public void onMatchTopLevelClass(
   }
 
   @Override
-  public void onMatchMethod(
-      NullAway analysis, MethodTree tree, VisitorState state, Symbol.MethodSymbol methodSymbol) {
+  public void onMatchMethod(MethodTree tree, MethodAnalysisContext methodAnalysisContext) {
     for (Handler h : handlers) {
-      h.onMatchMethod(analysis, tree, state, methodSymbol);
+      h.onMatchMethod(tree, methodAnalysisContext);
     }
   }
 
   @Override
   public void onMatchLambdaExpression(
-      NullAway analysis,
-      LambdaExpressionTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol) {
+      LambdaExpressionTree tree, MethodAnalysisContext methodAnalysisContext) {
     for (Handler h : handlers) {
-      h.onMatchLambdaExpression(analysis, tree, state, methodSymbol);
+      h.onMatchLambdaExpression(tree, methodAnalysisContext);
     }
   }
 
   @Override
   public void onMatchMethodReference(
-      NullAway analysis,
-      MemberReferenceTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol) {
+      MemberReferenceTree tree, MethodAnalysisContext methodAnalysisContext) {
     for (Handler h : handlers) {
-      h.onMatchMethodReference(analysis, tree, state, methodSymbol);
+      h.onMatchMethodReference(tree, methodAnalysisContext);
     }
   }
 
   @Override
   public void onMatchMethodInvocation(
-      NullAway analysis,
-      MethodInvocationTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol) {
+      MethodInvocationTree tree, MethodAnalysisContext methodAnalysisContext) {
     for (Handler h : handlers) {
-      h.onMatchMethodInvocation(analysis, tree, state, methodSymbol);
+      h.onMatchMethodInvocation(tree, methodAnalysisContext);
     }
   }
 
@@ -310,15 +300,13 @@ public MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
   @Override
   @Nullable
   public Integer castToNonNullArgumentPositionsForMethod(
-      NullAway analysis,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol,
       List<? extends ExpressionTree> actualParams,
-      @Nullable Integer previousArgumentPosition) {
+      @Nullable Integer previousArgumentPosition,
+      MethodAnalysisContext methodAnalysisContext) {
     for (Handler h : handlers) {
       previousArgumentPosition =
           h.castToNonNullArgumentPositionsForMethod(
-              analysis, state, methodSymbol, actualParams, previousArgumentPosition);
+              actualParams, previousArgumentPosition, methodAnalysisContext);
     }
     return previousArgumentPosition;
   }

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

@@ -78,55 +78,37 @@ void onMatchTopLevelClass(
   /**
    * Called when NullAway first matches a particular method node.
    *
-   * @param analysis A reference to the running NullAway analysis.
    * @param tree The AST node for the method being matched.
-   * @param state The current visitor state.
-   * @param methodSymbol The method symbol for the method being matched.
+   * @param methodAnalysisContext The MethodAnalysisContext object
    */
-  void onMatchMethod(
-      NullAway analysis, MethodTree tree, VisitorState state, Symbol.MethodSymbol methodSymbol);
+  void onMatchMethod(MethodTree tree, MethodAnalysisContext methodAnalysisContext);
 
   /**
    * Called when NullAway first matches a particular method call-site.
    *
-   * @param analysis A reference to the running NullAway analysis.
    * @param tree The AST node for the method invocation (call-site) being matched.
-   * @param state The current visitor state.
-   * @param methodSymbol The method symbol for the method being called.
+   * @param methodAnalysisContext The MethodAnalysisContext object
    */
   void onMatchMethodInvocation(
-      NullAway analysis,
-      MethodInvocationTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol);
+      MethodInvocationTree tree, MethodAnalysisContext methodAnalysisContext);
 
   /**
    * Called when NullAway first matches a particular lambda expression.
    *
-   * @param analysis A reference to the running NullAway analysis.
    * @param tree The AST node for the lambda expression being matched.
-   * @param state The current visitor state.
-   * @param methodSymbol The method symbol for the functional interface of the lambda being matched.
+   * @param methodAnalysisContext The MethodAnalysisContext object
    */
   void onMatchLambdaExpression(
-      NullAway analysis,
-      LambdaExpressionTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol);
+      LambdaExpressionTree tree, MethodAnalysisContext methodAnalysisContext);
 
   /**
    * Called when NullAway first matches a particular method reference expression
    *
-   * @param analysis A reference to the running NullAway analysis.
    * @param tree The AST node for the method reference expression being matched.
-   * @param state The current visitor state.
-   * @param methodSymbol The method symbol for the reference being matched.
+   * @param methodAnalysisContext The MethodAnalysisContext object
    */
   void onMatchMethodReference(
-      NullAway analysis,
-      MemberReferenceTree tree,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol);
+      MemberReferenceTree tree, MethodAnalysisContext methodAnalysisContext);
 
   /**
    * Called when NullAway first matches a return statement.
@@ -386,24 +368,20 @@ MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
    * <p>See {@link LibraryModels#castToNonNullMethods()} for more information about general
    * configuration of <code>castToNonNull</code> methods.
    *
-   * @param analysis A reference to the running NullAway analysis.
-   * @param state The current visitor state.
-   * @param methodSymbol The method symbol for the potential castToNonNull method.
    * @param actualParams The actual parameters from the invocation node
    * @param previousArgumentPosition The result computed by the previous handler in the chain, if
    *     any.
    * @return The index of the parameter for which the method should act as a cast (if any). This
    *     value can be set only once through the full chain of handlers, with each handler deciding
    *     whether to propagate or override the value previousArgumentPosition passed by the previous
    *     handler in the chain.
+   * @param methodAnalysisContext The MethodAnalysisContext object
    */
   @Nullable
   Integer castToNonNullArgumentPositionsForMethod(
-      NullAway analysis,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol,
       List<? extends ExpressionTree> actualParams,
-      @Nullable Integer previousArgumentPosition);
+      @Nullable Integer previousArgumentPosition,
+      MethodAnalysisContext methodAnalysisContext);
 
   /**
    * Method to override the nullability of the upper bound for a generic type variable on a class.

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

@@ -186,12 +186,12 @@ public boolean onOverrideMayBeNullExpr(
   @Override
   @Nullable
   public Integer castToNonNullArgumentPositionsForMethod(
-      NullAway analysis,
-      VisitorState state,
-      Symbol.MethodSymbol methodSymbol,
       List<? extends ExpressionTree> actualParams,
-      @Nullable Integer previousArgumentPosition) {
-    OptimizedLibraryModels optLibraryModels = getOptLibraryModels(state.context);
+      @Nullable Integer previousArgumentPosition,
+      MethodAnalysisContext methodAnalysisContext) {
+    Symbol.MethodSymbol methodSymbol = methodAnalysisContext.methodSymbol();
+    OptimizedLibraryModels optLibraryModels =
+        getOptLibraryModels(methodAnalysisContext.state().context);
     ImmutableSet<Integer> newPositions = optLibraryModels.castToNonNullMethod(methodSymbol);
     if (newPositions.size() > 1) {
       // Library models sanity check