Refactoring downstream dependency cache

annotator-core/src/main/java/edu/ucr/cs/riple/core/cache/downstream/DownstreamImpactCacheImpl.java

@@ -24,24 +24,23 @@
 
 package edu.ucr.cs.riple.core.cache.downstream;
 
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
-
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import edu.ucr.cs.riple.core.Context;
 import edu.ucr.cs.riple.core.Report;
 import edu.ucr.cs.riple.core.cache.BaseCache;
-import edu.ucr.cs.riple.core.cache.Impact;
 import edu.ucr.cs.riple.core.evaluators.suppliers.DownstreamDependencySupplier;
 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.method.MethodRecord;
 import edu.ucr.cs.riple.core.metadata.region.MethodRegionRegistry;
 import edu.ucr.cs.riple.injector.changes.AddMarkerAnnotation;
 import edu.ucr.cs.riple.injector.location.Location;
 import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.OptionalInt;
 import java.util.Set;
-import java.util.function.Function;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
 
@@ -50,7 +49,7 @@
  * once created, cannot be updated.
  */
 public class DownstreamImpactCacheImpl
-    extends BaseCache<DownstreamImpact, ImmutableMap<Location, DownstreamImpact>>
+    extends BaseCache<DownstreamImpact, Map<Location, DownstreamImpact>>
     implements DownstreamImpactCache {
 
   /** Annotator context instance. */
@@ -65,24 +64,28 @@ public class DownstreamImpactCacheImpl
    * @param context Annotator context.
    */
   public DownstreamImpactCacheImpl(Context context) {
-    super(
-        context
-            .targetModuleInfo
-            .getMethodRegistry()
-            .getPublicMethodsWithNonPrimitivesReturn()
-            .stream()
-            .map(
-                methodNode ->
-                    new DownstreamImpact(
-                        new Fix(
-                            new AddMarkerAnnotation(
-                                methodNode.location, context.config.nullableAnnot),
-                            "null",
-                            true)))
-            .collect(toImmutableMap(Impact::toLocation, Function.identity())));
+    super(new HashMap<>());
     this.context = context;
   }
 
+  /**
+   * Retrieves the set of locations that impact of making them {@code @Nullable} should be computed
+   * on downstream dependencies and stored in this cache.
+   *
+   * @param context Annotator context.
+   * @return Set of locations that impact of making them {@code @Nullable} should be computed on
+   *     downstream dependencies and stored in this cache.
+   */
+  private Set<Location> retrieveLocationsToCacheImpactsOnDownstreamDependencies(Context context) {
+    return context
+        .targetModuleInfo
+        .getMethodRegistry()
+        .getPublicMethodsWithNonPrimitivesReturn()
+        .stream()
+        .map(MethodRecord::toLocation)
+        .collect(Collectors.toSet());
+  }
+
   @Override
   public void analyzeDownstreamDependencies() {
     System.out.println("Analyzing downstream dependencies...");
@@ -91,12 +94,14 @@ public void analyzeDownstreamDependencies() {
     MethodRegionRegistry methodRegionRegistry = new MethodRegionRegistry(supplier.getModuleInfo());
     // Generate fixes corresponding methods.
     ImmutableSet<Fix> fixes =
-        store.values().stream()
+        retrieveLocationsToCacheImpactsOnDownstreamDependencies(context).stream()
             .filter(
                 input ->
                     !methodRegionRegistry
                         .getCallersOfMethod(input.toMethod().clazz, input.toMethod().method)
-                        .isEmpty()) // skip methods that are not called anywhere.
+                        // skip methods that are not called anywhere. This has a significant impact
+                        // on performance.
+                        .isEmpty())
             .map(
                 downstreamImpact ->
                     new Fix(
@@ -108,14 +113,12 @@ public void analyzeDownstreamDependencies() {
     DownstreamImpactEvaluator evaluator = new DownstreamImpactEvaluator(supplier);
     ImmutableSet<Report> reports = evaluator.evaluate(fixes);
     // Update method status based on the results.
-    this.store
-        .values()
-        .forEach(
-            methodImpact ->
-                reports.stream()
-                    .filter(input -> input.root.toMethod().equals(methodImpact.toMethod()))
-                    .findAny()
-                    .ifPresent(methodImpact::setStatus));
+    reports.forEach(
+        report -> {
+          DownstreamImpact impact = new DownstreamImpact(report.root);
+          impact.setStatus(report);
+          store.put(report.root.toLocation(), impact);
+        });
     System.out.println("Analyzing downstream dependencies completed!");
   }
 

annotator-core/src/main/java/edu/ucr/cs/riple/core/metadata/method/MethodRecord.java

@@ -25,6 +25,7 @@
 package edu.ucr.cs.riple.core.metadata.method;
 
 import com.google.common.collect.ImmutableSet;
+import edu.ucr.cs.riple.injector.location.Location;
 import edu.ucr.cs.riple.injector.location.OnMethod;
 import java.util.HashSet;
 import java.util.Objects;
@@ -179,6 +180,15 @@ private static boolean isNullableAnnotation(String annot) {
         || annot.endsWith(".CheckForNull");
   }
 
+  /**
+   * Returns location of the method.
+   *
+   * @return Location of the method.
+   */
+  public Location toLocation() {
+    return location;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {