Add clean up phase to remove annotations on local variables

annotator-core/src/main/java/edu/ucr/cs/riple/core/Annotator.java

@@ -35,8 +35,6 @@
 import edu.ucr.cs.riple.core.evaluators.VoidEvaluator;
 import edu.ucr.cs.riple.core.evaluators.suppliers.Supplier;
 import edu.ucr.cs.riple.core.evaluators.suppliers.TargetModuleSupplier;
-import edu.ucr.cs.riple.core.injectors.AnnotationInjector;
-import edu.ucr.cs.riple.core.injectors.PhysicalInjector;
 import edu.ucr.cs.riple.core.metadata.index.Fix;
 import edu.ucr.cs.riple.core.util.Utility;
 import java.util.Set;
@@ -49,8 +47,6 @@
  */
 public class Annotator {
 
-  /** Injector instance. */
-  private final AnnotationInjector injector;
   /** Annotator context. */
   public final Context context;
   /** Reports cache. */
@@ -62,7 +58,6 @@ public Annotator(Config config) {
     this.config = config;
     this.context = new Context(config);
     this.cache = new ReportCache(config);
-    this.injector = new PhysicalInjector(context);
   }
 
   /** Starts the annotating process consist of preprocess followed by the "annotate" phase. */
@@ -86,7 +81,7 @@ public void start() {
    */
   private void preprocess() {
     System.out.println("Preprocessing...");
-    context.checker.preprocess(injector);
+    context.checker.preprocess();
   }
 
   /** Performs iterations of inference/injection until no unseen fix is suggested. */
@@ -120,8 +115,9 @@ private void annotate() {
       }
     }
     if (config.suppressRemainingErrors) {
-      context.checker.suppressRemainingErrors(injector);
+      context.checker.suppressRemainingErrors();
     }
+    context.checker.cleanup();
     System.out.println("\nFinished annotating.");
     Utility.writeReports(context, cache.reports().stream().collect(ImmutableSet.toImmutableSet()));
   }
@@ -154,7 +150,7 @@ private void executeNextIteration(
             .filter(Report::approved)
             .flatMap(report -> config.chain ? report.tree.stream() : Stream.of(report.root))
             .collect(Collectors.toSet());
-    injector.injectFixes(selectedFixes);
+    context.injector.injectFixes(selectedFixes);
     // Update log.
     context.log.updateInjectedAnnotations(
         selectedFixes.stream().map(fix -> fix.change).collect(Collectors.toSet()));

annotator-core/src/main/java/edu/ucr/cs/riple/core/Context.java

@@ -27,6 +27,8 @@
 import com.google.common.collect.ImmutableSet;
 import edu.ucr.cs.riple.core.checkers.Checker;
 import edu.ucr.cs.riple.core.checkers.CheckerBaseClass;
+import edu.ucr.cs.riple.core.injectors.AnnotationInjector;
+import edu.ucr.cs.riple.core.injectors.PhysicalInjector;
 import edu.ucr.cs.riple.core.log.Log;
 import edu.ucr.cs.riple.core.metadata.index.Error;
 import edu.ucr.cs.riple.core.module.ModuleConfiguration;
@@ -63,6 +65,8 @@ public class Context {
   /** Checker instance. Used to execute checker specific tasks. */
   public final Checker<? extends Error> checker;
 
+  public final AnnotationInjector injector;
+
   /**
    * Builds context from command line arguments.
    *
@@ -76,6 +80,7 @@ public Context(Config config) {
     this.targetConfiguration = config.target;
     this.checker = CheckerBaseClass.getCheckerByName(config.checkerName, this);
     this.targetModuleInfo = new ModuleInfo(this, config.target, config.buildCommand);
+    this.injector = new PhysicalInjector(this);
     // Checker compatibility check must be after target module info is initialized.
     this.checker.verifyCheckerCompatibility();
   }

annotator-core/src/main/java/edu/ucr/cs/riple/core/checkers/Checker.java

@@ -25,7 +25,6 @@
 package edu.ucr.cs.riple.core.checkers;
 
 import com.google.common.collect.ImmutableSet;
-import edu.ucr.cs.riple.core.injectors.AnnotationInjector;
 import edu.ucr.cs.riple.core.metadata.index.Error;
 import edu.ucr.cs.riple.core.metadata.index.Fix;
 import edu.ucr.cs.riple.core.metadata.region.Region;
@@ -48,20 +47,11 @@ public interface Checker<T extends Error> {
    */
   Set<T> deserializeErrors(ModuleInfo module);
 
-  /**
-   * Suppresses remaining errors reported by the checker.
-   *
-   * @param injector Annotation injector to inject selected annotations.
-   */
-  void suppressRemainingErrors(AnnotationInjector injector);
+  /** Suppresses remaining errors reported by the checker. */
+  void suppressRemainingErrors();
 
-  /**
-   * Used to do any pre-processing steps before running the inference.
-   *
-   * @param injector Annotation injector, can be used to inject any annotations during the
-   *     pre-processing phase.
-   */
-  void preprocess(AnnotationInjector injector);
+  /** Used to do any pre-processing steps before running the inference. */
+  void preprocess();
 
   /**
    * Creates an {@link Error} instance from the given parameters.
@@ -95,4 +85,7 @@ T createError(
    *     build.
    */
   void prepareConfigFilesForBuild(ImmutableSet<ModuleConfiguration> configurations);
+
+  /** Used to do any post-processing steps after running the inference. */
+  void cleanup();
 }

annotator-core/src/main/java/edu/ucr/cs/riple/core/checkers/CheckerBaseClass.java

@@ -30,8 +30,12 @@
 import edu.ucr.cs.riple.core.checkers.ucrtaint.UCRTaint;
 import edu.ucr.cs.riple.core.metadata.index.Error;
 import edu.ucr.cs.riple.core.module.ModuleInfo;
+import edu.ucr.cs.riple.injector.changes.AddAnnotation;
+import edu.ucr.cs.riple.injector.changes.AddFullTypeMarkerAnnotation;
+import edu.ucr.cs.riple.injector.changes.RemoveAnnotation;
 import edu.ucr.cs.riple.injector.location.OnField;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /** Base class for all checker representations. */
 public abstract class CheckerBaseClass<T extends Error> implements Checker<T> {
@@ -82,4 +86,22 @@ public static Checker<?> getCheckerByName(String name, Context context) {
         throw new RuntimeException("Unknown checker name: " + name);
     }
   }
+
+  @Override
+  public void cleanup() {
+    // Remove all annotations added on local variables.
+    Set<RemoveAnnotation> annotations =
+        context.log.getInjectedAnnotations().stream()
+            .filter(addAnnotation -> addAnnotation.getLocation().isOnLocalVariable())
+            // Keep added type use annotations.
+            .map(
+                addAnnotation ->
+                    addAnnotation instanceof AddFullTypeMarkerAnnotation
+                        ? ((AddFullTypeMarkerAnnotation) addAnnotation).toDeclaration()
+                        : addAnnotation)
+            .map(AddAnnotation::getReverse)
+            .collect(Collectors.toSet());
+    // Remove
+    context.injector.removeAnnotations(annotations);
+  }
 }

annotator-core/src/main/java/edu/ucr/cs/riple/core/checkers/nullaway/NullAway.java

@@ -28,7 +28,6 @@
 import com.google.common.collect.ImmutableSet;
 import edu.ucr.cs.riple.core.Context;
 import edu.ucr.cs.riple.core.checkers.CheckerBaseClass;
-import edu.ucr.cs.riple.core.injectors.AnnotationInjector;
 import edu.ucr.cs.riple.core.metadata.field.FieldInitializationStore;
 import edu.ucr.cs.riple.core.metadata.index.Fix;
 import edu.ucr.cs.riple.core.metadata.region.Region;
@@ -219,11 +218,9 @@ protected ImmutableSet<Fix> generateFixesForUninitializedFields(
    *   <li>Explicit {@code Nullable} assignments to fields will be annotated as
    *       {@code @SuppressWarnings("NullAway")}.
    * </ul>
-   *
-   * @param injector Annotation injector to inject selected annotations.
    */
   @Override
-  public void suppressRemainingErrors(AnnotationInjector injector) {
+  public void suppressRemainingErrors() {
     // Collect regions with remaining errors.
     Utility.buildTarget(context);
     Set<NullAwayError> remainingErrors = deserializeErrors(context.targetModuleInfo);
@@ -332,7 +329,7 @@ public void suppressRemainingErrors(AnnotationInjector injector) {
                         fix.toField(), "SuppressWarnings", "NullAway.Init", false))
             .collect(Collectors.toSet());
     result.addAll(initializationSuppressWarningsAnnotations);
-    injector.injectAnnotations(result);
+    context.injector.injectAnnotations(result);
     // update log
     context.log.updateInjectedAnnotations(result);
     // Collect @NullUnmarked annotations on classes for any remaining error.
@@ -347,13 +344,13 @@ public void suppressRemainingErrors(AnnotationInjector injector) {
                         context.targetModuleInfo.getLocationOnClass(error.getRegion().clazz),
                         config.nullUnMarkedAnnotation))
             .collect(Collectors.toSet());
-    injector.injectAnnotations(nullUnMarkedAnnotations);
+    context.injector.injectAnnotations(nullUnMarkedAnnotations);
     // update log
     context.log.updateInjectedAnnotations(nullUnMarkedAnnotations);
   }
 
   @Override
-  public void preprocess(AnnotationInjector injector) {
+  public void preprocess() {
     // Collect @Initializer annotations. Heuristically, we add @Initializer on methods which writes
     // a @Nonnull value to more than one uninitialized field, and guarantees initialized fields are
     // nonnull at all exit paths.
@@ -370,7 +367,7 @@ public void preprocess(AnnotationInjector injector) {
             .map(onMethod -> new AddMarkerAnnotation(onMethod, config.initializerAnnot))
             .collect(Collectors.toSet());
     // Inject @Initializer annotations.
-    injector.injectAnnotations(initializers);
+    context.injector.injectAnnotations(initializers);
   }
 
   @Override

annotator-core/src/main/java/edu/ucr/cs/riple/core/checkers/ucrtaint/UCRTaint.java

@@ -27,7 +27,6 @@
 import com.google.common.collect.ImmutableSet;
 import edu.ucr.cs.riple.core.Context;
 import edu.ucr.cs.riple.core.checkers.CheckerBaseClass;
-import edu.ucr.cs.riple.core.injectors.AnnotationInjector;
 import edu.ucr.cs.riple.core.metadata.index.Fix;
 import edu.ucr.cs.riple.core.metadata.region.Region;
 import edu.ucr.cs.riple.core.module.ModuleConfiguration;
@@ -59,7 +58,7 @@ public UCRTaint(Context context) {
   }
 
   @Override
-  public void preprocess(AnnotationInjector injector) {}
+  public void preprocess() {}
 
   @Override
   public Set<UCRTaintError> deserializeErrors(ModuleInfo module) {
@@ -109,7 +108,7 @@ private UCRTaintError deserializeErrorFromJSON(JSONObject errorsJson, ModuleInfo
   }
 
   @Override
-  public void suppressRemainingErrors(AnnotationInjector injector) {
+  public void suppressRemainingErrors() {
     throw new RuntimeException(
         "Suppression for remaining errors is not supported for " + NAME + "yet!");
   }

injector/src/main/java/edu/ucr/cs/riple/injector/changes/AddTypeUseMarkerAnnotation.java → injector/src/main/java/edu/ucr/cs/riple/injector/changes/AddFullTypeMarkerAnnotation.java

@@ -49,10 +49,10 @@
  * output will be: {@code java.util.@Nullable Map<@Nullable String, java.lang.@Nullable String>
  * list;}
  */
-public class AddTypeUseMarkerAnnotation extends AddMarkerAnnotation {
+public class AddFullTypeMarkerAnnotation extends AnnotationChange implements AddAnnotation {
 
-  public AddTypeUseMarkerAnnotation(Location location, String annotation) {
-    super(location, annotation);
+  public AddFullTypeMarkerAnnotation(Location location, String annotation) {
+    super(location, new Name(annotation));
   }
 
   @Override
@@ -83,18 +83,18 @@ Modification computeTextModificationOn(T node) {
 
   @Override
   public RemoveAnnotation getReverse() {
-    return new RemoveTypeUseMarkerAnnotation(location, annotationName.fullName);
+    return new RemoveFullTypeMarkerAnnotation(location, annotationName.fullName);
   }
 
   @Override
   public boolean equals(Object other) {
     if (this == other) {
       return true;
     }
-    if (!(other instanceof AddTypeUseMarkerAnnotation)) {
+    if (!(other instanceof AddFullTypeMarkerAnnotation)) {
       return false;
     }
-    AddTypeUseMarkerAnnotation otherAdd = (AddTypeUseMarkerAnnotation) other;
+    AddFullTypeMarkerAnnotation otherAdd = (AddFullTypeMarkerAnnotation) other;
     return this.location.equals(otherAdd.location)
         && this.annotationName.equals(otherAdd.annotationName);
   }
@@ -126,4 +126,14 @@ private static Range findSimpleNameRangeInTypeName(Type type) {
     throw new RuntimeException(
         "Unexpected type to get range from: " + type + " : " + type.getClass());
   }
+
+  /**
+   * Converts this change to a declaration change. This is used to apply the change to the
+   * declaration only.
+   *
+   * @return a declaration change that adds the annotation to the declaration.
+   */
+  public AddMarkerAnnotation toDeclaration() {
+    return new AddMarkerAnnotation(location, annotationName.simpleName);
+  }
 }

injector/src/main/java/edu/ucr/cs/riple/injector/changes/RemoveTypeUseMarkerAnnotation.java → injector/src/main/java/edu/ucr/cs/riple/injector/changes/RemoveFullTypeMarkerAnnotation.java

@@ -47,9 +47,9 @@
  * {@code Map} and all the type arguments. The final output will be: {@code java.util.Map<String,
  * java.lang.String> list;}
  */
-public class RemoveTypeUseMarkerAnnotation extends RemoveMarkerAnnotation {
+public class RemoveFullTypeMarkerAnnotation extends RemoveMarkerAnnotation {
 
-  public RemoveTypeUseMarkerAnnotation(Location location, String annotation) {
+  public RemoveFullTypeMarkerAnnotation(Location location, String annotation) {
     super(location, annotation);
   }