Merge pull request #2008 from timll/develop

Allow dynamic, thread-safe loading of new method bodies and updating dependent data structures
soot-oss · Oct 12, 2023 · e09d471 · e09d471
2 parents ac0348c + 31dccba
commit e09d471
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 8 deletions.
diff --git a/src/main/java/soot/SootMethod.java b/src/main/java/soot/SootMethod.java
@@ -30,6 +30,7 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.StringTokenizer;
+import java.util.function.Consumer;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -415,6 +416,21 @@ public synchronized void setActiveBody(Body body) {
    * get retrieve its active body. Please call {@link SootClass#setApplicationClass()} on the relevant class.
    */
   public Body retrieveActiveBody() {
+    return retrieveActiveBody((b) -> { });
+  }
+
+  /**
+   * Returns the active body if present, else constructs an active body, calls the consumer and returns the body afterward.
+   *
+   * If you called Scene.v().loadClassAndSupport() for a class yourself, it will not be an application class, so you cannot
+   * get retrieve its active body. Please call {@link SootClass#setApplicationClass()} on the relevant class.
+   *
+   * @param consumer Consumer that takes in the body of the method. The consumer is only invoked if the current
+   *                 invocation constructs a new body and is guaranteed to terminate before the body is available
+   *                 to other threads.
+   * @return active body of the method
+   */
+  public Body retrieveActiveBody(Consumer<Body> consumer) {
     // Retrieve the active body so thread changes do not affect the
     // synchronization between if the body exists and the returned body.
     // This is a quick check just in case the activeBody exists.
@@ -445,6 +461,11 @@ public Body retrieveActiveBody() {
 
       // Method sources are not expected to be thread safe
       activeBody = ms.getBody(this, "jb");
+
+      // Call the consumer such that clients can update any data structures, caches, etc.
+      // atomically before the body is available to other threads.
+      consumer.accept(activeBody);
+
       setActiveBody(activeBody);
 
       // If configured, we drop the method source to save memory

diff --git a/src/main/java/soot/jimple/toolkits/ide/icfg/AbstractJimpleBasedICFG.java b/src/main/java/soot/jimple/toolkits/ide/icfg/AbstractJimpleBasedICFG.java
@@ -25,18 +25,17 @@
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
 
-import heros.DontSynchronize;
 import heros.SynchronizedBy;
 import heros.solver.IDESolver;
 
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
-import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 import soot.Body;
 import soot.PatchingChain;
@@ -53,7 +52,7 @@ public abstract class AbstractJimpleBasedICFG implements BiDiInterproceduralCFG<
 
   protected final boolean enableExceptions;
 
-  @DontSynchronize("written by single thread; read afterwards")
+  @SynchronizedBy("thread-safe data structure")
   private final Map<Unit, Body> unitToOwner = createUnitToOwnerMap();
 
   @SynchronizedBy("by use of synchronized LoadingCache class")
@@ -87,8 +86,13 @@ public AbstractJimpleBasedICFG() {
     this(true);
   }
 
+  /**
+   * Creates a new map used for the unitToOwner map. Must be thread-safe.
+   *
+   * @return a new thread-safe map
+   */
   protected Map<Unit, Body> createUnitToOwnerMap() {
-    return new LinkedHashMap<Unit, Body>();
+    return new ConcurrentHashMap<>();
   }
 
   public AbstractJimpleBasedICFG(boolean enableExceptions) {
@@ -242,10 +246,14 @@ public Set<Unit> getCallsFromWithin(SootMethod m) {
   public void initializeUnitToOwner(SootMethod m) {
     if (m.hasActiveBody()) {
       Body b = m.getActiveBody();
-      PatchingChain<Unit> units = b.getUnits();
-      for (Unit unit : units) {
-        unitToOwner.put(unit, b);
-      }
+      initializeUnitToOwner(b);
+    }
+  }
+
+  public void initializeUnitToOwner(Body b) {
+    PatchingChain<Unit> units = b.getUnits();
+    for (Unit unit : units) {
+      unitToOwner.put(unit, b);
     }
   }