Use visitor for applying changes with locations

injector/src/main/java/edu/ucr/cs/riple/injector/Injector.java

@@ -29,6 +29,7 @@
 import com.github.javaparser.ast.ImportDeclaration;
 import edu.ucr.cs.riple.injector.changes.AddAnnotation;
 import edu.ucr.cs.riple.injector.changes.Change;
+import edu.ucr.cs.riple.injector.changes.ChangeVisitor;
 import edu.ucr.cs.riple.injector.changes.RemoveAnnotation;
 import edu.ucr.cs.riple.injector.modifications.Modification;
 import edu.ucr.cs.riple.injector.offsets.FileOffsetStore;
@@ -62,11 +63,12 @@ public <T extends Change> Set<FileOffsetStore> start(Set<T> changes) {
           } catch (IOException exception) {
             return;
           }
+          ChangeVisitor visitor = new ChangeVisitor(tree);
           Set<Modification> modifications = new HashSet<>();
           Set<ImportDeclaration> imports = new HashSet<>();
           for (Change change : changeList) {
             try {
-              Modification modification = change.translate(tree);
+              Modification modification = visitor.computeModification(change);
               if (modification != null) {
                 modifications.add(modification);
                 if (change instanceof AddAnnotation) {

injector/src/main/java/edu/ucr/cs/riple/injector/changes/Change.java

@@ -22,7 +22,6 @@
 
 package edu.ucr.cs.riple.injector.changes;
 
-import com.github.javaparser.ast.CompilationUnit;
 import com.github.javaparser.ast.nodeTypes.NodeWithAnnotations;
 import com.github.javaparser.ast.nodeTypes.NodeWithRange;
 import edu.ucr.cs.riple.injector.Helper;
@@ -48,18 +47,6 @@ public Change(Location location, String annotation) {
     this.annotationSimpleName = Helper.simpleName(annotation);
   }
 
-  /**
-   * Translate the change to a text modification in the source file.
-   *
-   * @param tree Compilation unit tree instance.
-   * @return A text modification instance if the translation is successful, otherwise {@code null}
-   *     will be returned.
-   */
-  @Nullable
-  public Modification translate(CompilationUnit tree) {
-    return this.location.apply(tree, this);
-  }
-
   /**
    * Visits the given node and translates the change to a text modification.
    *

injector/src/main/java/edu/ucr/cs/riple/injector/changes/ChangeVisitor.java

@@ -0,0 +1,187 @@
+/*
+ * Copyright (c) 2023 University of California, Riverside.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package edu.ucr.cs.riple.injector.changes;
+
+import static edu.ucr.cs.riple.injector.location.OnClass.isAnonymousClassFlatName;
+
+import com.github.javaparser.ast.CompilationUnit;
+import com.github.javaparser.ast.Node;
+import com.github.javaparser.ast.NodeList;
+import com.github.javaparser.ast.body.BodyDeclaration;
+import com.github.javaparser.ast.body.Parameter;
+import com.github.javaparser.ast.body.VariableDeclarator;
+import com.github.javaparser.utils.Pair;
+import edu.ucr.cs.riple.injector.Helper;
+import edu.ucr.cs.riple.injector.exceptions.TargetClassNotFound;
+import edu.ucr.cs.riple.injector.location.LocationVisitor;
+import edu.ucr.cs.riple.injector.location.OnClass;
+import edu.ucr.cs.riple.injector.location.OnField;
+import edu.ucr.cs.riple.injector.location.OnMethod;
+import edu.ucr.cs.riple.injector.location.OnParameter;
+import edu.ucr.cs.riple.injector.modifications.Modification;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
+import javax.annotation.Nullable;
+
+/**
+ * A visitor for computing the required {@link Modification} to a compilation unit on a specified
+ * location for the requested change.
+ */
+public class ChangeVisitor
+    implements LocationVisitor<Modification, Pair<NodeList<BodyDeclaration<?>>, Change>> {
+
+  /** Compilation unit which the changes will be applied. */
+  private final CompilationUnit cu;
+
+  public ChangeVisitor(CompilationUnit cu) {
+    this.cu = cu;
+  }
+
+  @Override
+  @Nullable
+  public Modification visitMethod(
+      OnMethod onMethod, Pair<NodeList<BodyDeclaration<?>>, Change> pair) {
+    final AtomicReference<Modification> ans = new AtomicReference<>();
+    final NodeList<BodyDeclaration<?>> members = pair.a;
+    final Change change = pair.b;
+    members.forEach(
+        bodyDeclaration ->
+            bodyDeclaration.ifCallableDeclaration(
+                callableDeclaration -> {
+                  if (ans.get() != null) {
+                    // already found the member.
+                    return;
+                  }
+                  if (onMethod.matchesCallableDeclaration(callableDeclaration)) {
+                    ans.set(change.visit(callableDeclaration));
+                  }
+                }));
+    if (ans.get() == null) {
+      members.forEach(
+          bodyDeclaration ->
+              bodyDeclaration.ifAnnotationMemberDeclaration(
+                  annotationMemberDeclaration -> {
+                    if (annotationMemberDeclaration
+                        .getNameAsString()
+                        .equals(Helper.extractCallableName(onMethod.method))) {
+                      ans.set(change.visit(annotationMemberDeclaration));
+                    }
+                  }));
+    }
+    return ans.get();
+  }
+
+  @Override
+  @Nullable
+  public Modification visitField(OnField onField, Pair<NodeList<BodyDeclaration<?>>, Change> pair) {
+    final AtomicReference<Modification> ans = new AtomicReference<>();
+    final NodeList<BodyDeclaration<?>> members = pair.a;
+    final Change change = pair.b;
+    members.forEach(
+        bodyDeclaration ->
+            bodyDeclaration.ifFieldDeclaration(
+                fieldDeclaration -> {
+                  if (ans.get() != null) {
+                    // already found the member.
+                    return;
+                  }
+                  NodeList<VariableDeclarator> vars =
+                      fieldDeclaration.asFieldDeclaration().getVariables();
+                  for (VariableDeclarator v : vars) {
+                    if (onField.variables.contains(v.getName().toString())) {
+                      ans.set(change.visit(fieldDeclaration));
+                      break;
+                    }
+                  }
+                }));
+    return ans.get();
+  }
+
+  @Override
+  @Nullable
+  public Modification visitParameter(
+      OnParameter onParameter, Pair<NodeList<BodyDeclaration<?>>, Change> pair) {
+    final AtomicReference<Modification> ans = new AtomicReference<>();
+    final NodeList<BodyDeclaration<?>> members = pair.a;
+    final Change change = pair.b;
+    members.forEach(
+        bodyDeclaration ->
+            bodyDeclaration.ifCallableDeclaration(
+                callableDeclaration -> {
+                  if (ans.get() != null) {
+                    // already found the member.
+                    return;
+                  }
+                  if (onParameter.matchesCallableDeclaration(callableDeclaration)) {
+                    NodeList<?> params = callableDeclaration.getParameters();
+                    if (onParameter.index < params.size()) {
+                      if (params.get(onParameter.index) != null) {
+                        Node param = params.get(onParameter.index);
+                        if (param instanceof Parameter) {
+                          ans.set(change.visit((Parameter) param));
+                        }
+                      }
+                    }
+                  }
+                }));
+    return ans.get();
+  }
+
+  @Override
+  @Nullable
+  public Modification visitClass(OnClass onClass, Pair<NodeList<BodyDeclaration<?>>, Change> pair) {
+    final NodeList<BodyDeclaration<?>> members = pair.a;
+    final Change change = pair.b;
+    if (isAnonymousClassFlatName(change.location.clazz)) {
+      return null;
+    }
+    // Get the enclosing class of the members
+    Optional<Node> optionalClass = members.getParentNode();
+    if (optionalClass.isEmpty() || !(optionalClass.get() instanceof BodyDeclaration<?>)) {
+      return null;
+    }
+    return change.visit(((BodyDeclaration<?>) optionalClass.get()));
+  }
+
+  /**
+   * Computes the required {@link Modification} that should be applied to the compilation unit for
+   * the given change.
+   *
+   * @param change the change to apply.
+   * @return the modification that should be applied.
+   */
+  @Nullable
+  public Modification computeModification(Change change) {
+    NodeList<BodyDeclaration<?>> members;
+    try {
+      members = Helper.getTypeDeclarationMembersByFlatName(cu, change.location.clazz);
+      if (members == null) {
+        return null;
+      }
+      return change.location.accept(this, new Pair<>(members, change));
+    } catch (TargetClassNotFound notFound) {
+      System.err.println(notFound.getMessage());
+      return null;
+    }
+  }
+}

injector/src/main/java/edu/ucr/cs/riple/injector/location/Location.java

@@ -24,15 +24,9 @@
 
 package edu.ucr.cs.riple.injector.location;
 
-import com.github.javaparser.ast.CompilationUnit;
-import com.github.javaparser.ast.NodeList;
-import com.github.javaparser.ast.body.BodyDeclaration;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Sets;
 import edu.ucr.cs.riple.injector.Helper;
-import edu.ucr.cs.riple.injector.changes.Change;
-import edu.ucr.cs.riple.injector.exceptions.TargetClassNotFound;
-import edu.ucr.cs.riple.injector.modifications.Modification;
 import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.Objects;
@@ -104,36 +98,6 @@ public static Location createLocationFromArrayInfo(String[] values) {
     throw new RuntimeException("Cannot reach this statement, values: " + Arrays.toString(values));
   }
 
-  /**
-   * Applies the change to the element in this location.
-   *
-   * @param members The list of members of the enclosing class of the target element.
-   * @param change The change to be applied on the target element.
-   * @return The modification that should be applied on the source file.
-   */
-  @Nullable
-  protected abstract Modification applyToMember(
-      NodeList<BodyDeclaration<?>> members, Change change);
-
-  /**
-   * Applies the change to the target element on the given compilation unit tree.
-   *
-   * @param tree CompilationUnit Tree to locate the target element.
-   * @param change Change to be applied on the target element.
-   * @return The modification that should be applied on the source file.
-   */
-  @Nullable
-  public Modification apply(CompilationUnit tree, Change change) {
-    NodeList<BodyDeclaration<?>> clazz;
-    try {
-      clazz = Helper.getTypeDeclarationMembersByFlatName(tree, this.clazz);
-    } catch (TargetClassNotFound notFound) {
-      System.err.println(notFound.getMessage());
-      return null;
-    }
-    return applyToMember(clazz, change);
-  }
-
   /**
    * If this location is of kind {@link LocationKind#METHOD}, calls the consumer on the location.
    *

injector/src/main/java/edu/ucr/cs/riple/injector/location/OnClass.java

@@ -24,16 +24,9 @@
 
 package edu.ucr.cs.riple.injector.location;
 
-import com.github.javaparser.ast.Node;
-import com.github.javaparser.ast.NodeList;
-import com.github.javaparser.ast.body.BodyDeclaration;
 import edu.ucr.cs.riple.injector.Helper;
-import edu.ucr.cs.riple.injector.changes.Change;
-import edu.ucr.cs.riple.injector.modifications.Modification;
 import java.nio.file.Path;
-import java.util.Optional;
 import java.util.regex.Pattern;
-import javax.annotation.Nullable;
 
 /** Represents a location for class element. This location is used to apply changes to a class. */
 public class OnClass extends Location {
@@ -52,20 +45,6 @@ public OnClass(String path, String clazz) {
     this(Helper.deserializePath(path), clazz);
   }
 
-  @Override
-  @Nullable
-  protected Modification applyToMember(NodeList<BodyDeclaration<?>> members, Change change) {
-    if (isAnonymousClassFlatName(change.location.clazz)) {
-      return null;
-    }
-    // Get the enclosing class of the members
-    Optional<Node> optionalClass = members.getParentNode();
-    if (optionalClass.isEmpty() || !(optionalClass.get() instanceof BodyDeclaration<?>)) {
-      return null;
-    }
-    return change.visit(((BodyDeclaration<?>) optionalClass.get()));
-  }
-
   /**
    * Checks if flat name is for an anonymous class.
    *

injector/src/main/java/edu/ucr/cs/riple/injector/location/OnField.java

@@ -24,19 +24,12 @@
 
 package edu.ucr.cs.riple.injector.location;
 
-import com.github.javaparser.ast.NodeList;
-import com.github.javaparser.ast.body.BodyDeclaration;
-import com.github.javaparser.ast.body.VariableDeclarator;
 import edu.ucr.cs.riple.injector.Helper;
-import edu.ucr.cs.riple.injector.changes.Change;
-import edu.ucr.cs.riple.injector.modifications.Modification;
 import java.nio.file.Path;
 import java.util.Collections;
 import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
-import javax.annotation.Nullable;
 
 /**
  * Represents a location for field element. This location is used to apply changes to a class field.
@@ -66,30 +59,6 @@ public OnField(String path, String clazz, Set<String> variables) {
     this(Helper.deserializePath(path), clazz, variables);
   }
 
-  @Override
-  @Nullable
-  protected Modification applyToMember(NodeList<BodyDeclaration<?>> members, Change change) {
-    final AtomicReference<Modification> ans = new AtomicReference<>();
-    members.forEach(
-        bodyDeclaration ->
-            bodyDeclaration.ifFieldDeclaration(
-                fieldDeclaration -> {
-                  if (ans.get() != null) {
-                    // already found the member.
-                    return;
-                  }
-                  NodeList<VariableDeclarator> vars =
-                      fieldDeclaration.asFieldDeclaration().getVariables();
-                  for (VariableDeclarator v : vars) {
-                    if (variables.contains(v.getName().toString())) {
-                      ans.set(change.visit(fieldDeclaration));
-                      break;
-                    }
-                  }
-                }));
-    return ans.get();
-  }
-
   @Override
   public void ifField(Consumer<OnField> consumer) {
     consumer.accept(this);