use stream and group by

soot-oss · Dec 7, 2024 · e0a5cad · e0a5cad
1 parent dec694a
commit e0a5cad
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 15 deletions.
diff --git a/sootup.core/src/main/java/sootup/core/model/Body.java b/sootup.core/src/main/java/sootup/core/model/Body.java
@@ -26,9 +26,11 @@
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.util.*;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import org.apache.commons.lang3.tuple.Pair;
 import sootup.core.graph.MutableBlockStmtGraph;
 import sootup.core.graph.MutableStmtGraph;
 import sootup.core.graph.StmtGraph;
@@ -497,21 +499,16 @@ public static Map<LValue, Collection<Stmt>> collectDefs(Collection<Stmt> stmts)
   }
 
   /**
-   * Collects all using statements of a Local from a list of statements
+   * Collects all using statements of a Values from a list of statements
    *
    * @param stmts The searched list of statements
-   * @return A map of Locals and their using statements
+   * @return A map of Values and their using statements
    */
-  public static Map<Value, Collection<Stmt>> collectUses(Collection<Stmt> stmts) {
-    Map<Value, Collection<Stmt>> allUses = new HashMap<>();
-    for (Stmt stmt : stmts) {
-      for (Iterator<Value> iterator = stmt.getUses().iterator(); iterator.hasNext(); ) {
-        Value value = iterator.next();
-        Collection<Stmt> localUses = allUses.computeIfAbsent(value, key -> new ArrayList<>());
-        localUses.add(stmt);
-        allUses.put(value, localUses);
-      }
-    }
-    return allUses;
+  public static Map<Value, List<Stmt>> collectUses(Collection<Stmt> stmts) {
+    return stmts.stream()
+        .flatMap(stmt -> stmt.getUses().map(value -> (Pair.of(value, stmt))))
+        .collect(
+            Collectors.groupingBy(
+                Pair::getLeft, Collectors.mapping(Pair::getRight, Collectors.toList())));
   }
 }
diff --git a/sootup.interceptors/src/main/java/sootup/interceptors/Aggregator.java b/sootup.interceptors/src/main/java/sootup/interceptors/Aggregator.java
@@ -69,7 +69,7 @@ public Aggregator(boolean dontAggregateFieldLocals) {
   public void interceptBody(@Nonnull Body.BodyBuilder builder, @Nonnull View view) {
     MutableStmtGraph graph = builder.getStmtGraph();
     List<Stmt> stmts = builder.getStmts();
-    Map<Value, Collection<Stmt>> usesMap = Body.collectUses(stmts);
+    Map<Value, List<Stmt>> usesMap = Body.collectUses(stmts);
 
     for (Stmt stmt : stmts) {
       if (!(stmt instanceof JAssignStmt)) {

diff --git a/sootup.interceptors/src/main/java/sootup/interceptors/DeadAssignmentEliminator.java b/sootup.interceptors/src/main/java/sootup/interceptors/DeadAssignmentEliminator.java
@@ -211,7 +211,7 @@ public void interceptBody(@Nonnull Body.BodyBuilder builder, @Nonnull View view)
       return;
     }
 
-    Map<Value, Collection<Stmt>> essentialUses = Body.collectUses(essentialStmts);
+    Map<Value, List<Stmt>> essentialUses = Body.collectUses(essentialStmts);
     // Eliminate dead assignments from invokes such as x = f(), where x is no longer used
     List<JAssignStmt> postProcess = new ArrayList<>();
     for (Stmt stmt : stmts) {