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) {
@@ -952,7 +953,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;
     }
@@ -996,7 +998,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);
   }
 
@@ -1440,7 +1442,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;
@@ -1813,7 +1815,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,27 +60,27 @@
   /**
    * 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 visitorState = 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;
     }
     Symbol.MethodSymbol closestOverriddenMethod =
-        NullabilityUtil.getClosestOverriddenMethod(methodSymbol, state.getTypes());
+        NullabilityUtil.getClosestOverriddenMethod(methodSymbol, visitorState.getTypes());
     if (closestOverriddenMethod == null) {
       return;
     }
@@ -90,8 +90,9 @@
     } else {
       fieldNames = Collections.emptySet();
     }
-    validateOverridingRules(fieldNames, analysis, state, tree, closestOverriddenMethod);
-    super.onMatchMethod(analysis, tree, state, methodSymbol);
+    validateOverridingRules(
+        fieldNames, methodAnalysisContext.analysis(), visitorState, tree, closestOverriddenMethod);
+    super.onMatchMethod(tree, methodAnalysisContext);
   }
 
   /**
@@ -117,9 +118,10 @@
    * 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,28 +139,27 @@
    * <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 visitorState = methodAnalysisContext.state();
+    NullAway analysis = methodAnalysisContext.analysis();
     if (content.isEmpty()) {
       // we should not allow useless annotations.
       message =
           "empty @"
               + annotName
               + " is the default precondition for every method, please remove it.";
-      state.reportMatch(
+      visitorState.reportMatch(
           analysis
               .getErrorBuilder()
               .createErrorDescription(
                   new ErrorMessage(ErrorMessage.MessageTypes.ANNOTATION_VALUE_INVALID, message),
                   tree,
                   analysis.buildDescription(tree),
-                  state,
+                  visitorState,
                   null));
       return false;
     } else {
@@ -172,22 +173,23 @@
                     + fieldName
                     + " is not supported";
 
-            state.reportMatch(
+            visitorState.reportMatch(
                 analysis
                     .getErrorBuilder()
                     .createErrorDescription(
                         new ErrorMessage(
                             ErrorMessage.MessageTypes.ANNOTATION_VALUE_INVALID, message),
                         tree,
                         analysis.buildDescription(tree),
-                        state,
+                        visitorState,
                         null));
             return false;
           } else {
             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,14 +199,15 @@
                   + fieldName
                   + " in class "
                   + classSymbol.getSimpleName();
-          state.reportMatch(
+
+          visitorState.reportMatch(
               analysis
                   .getErrorBuilder()
                   .createErrorDescription(
                       new ErrorMessage(ErrorMessage.MessageTypes.ANNOTATION_VALUE_INVALID, message),
                       tree,
                       analysis.buildDescription(tree),
-                      state,
+                      visitorState,
                       null));
           return false;
         }

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;
   }