From 5a8db652be8d28604fd27ca1d413e8ef8b8200b3 Mon Sep 17 00:00:00 2001
From: Aman Goel <goelaman@amazon.com>
Date: Wed, 25 Oct 2023 14:12:01 -0700
Subject: [PATCH] [PCover] Optimize coverage stats tracking

---
 .../scheduler/search/SearchScheduler.java     |  6 ++-
 .../taskorchestration/BacktrackTask.java      | 45 +++++++++++++------
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/scheduler/search/SearchScheduler.java b/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/scheduler/search/SearchScheduler.java
index c8335aaa4..d97142e2d 100644
--- a/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/scheduler/search/SearchScheduler.java
+++ b/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/scheduler/search/SearchScheduler.java
@@ -390,6 +390,7 @@ private void setBacktrackTasks() {
       choice.clearBacktrack();
       if (i >= parentTask.getBacktrackChoiceDepth()) {
         parentTask.addPrefixChoice(i, choice);
+        parentTask.addPrefixCoverageStats(i, PSymGlobal.getCoverage().getPerChoiceDepthStats().get(i));
       }
     }
 
@@ -429,13 +430,14 @@ private void setBacktrackTaskAtDepth(
     newTask.setPrefixCoverage(prefixCoverage);
     newTask.setDepth(schedule.getChoice(backtrackChoiceDepth).getSchedulerDepth());
     newTask.setBacktrackChoiceDepth(backtrackChoiceDepth);
-    newTask.setPerChoiceDepthStats(PSymGlobal.getCoverage().getPerChoiceDepthStats());
     newTask.setParentTask(parentTask);
     newTask.setPriority();
     newTask.addSuffixChoice(schedule.getChoice(backtrackChoiceDepth));
+    newTask.addSuffixCoverageStats(PSymGlobal.getCoverage().getPerChoiceDepthStats().get(backtrackChoiceDepth));
     if (!isExact) {
       for (int i = backtrackChoiceDepth + 1; i < schedule.size(); i++) {
         newTask.addSuffixChoice(schedule.getChoice(i));
+        newTask.addSuffixCoverageStats(PSymGlobal.getCoverage().getPerChoiceDepthStats().get(i));
       }
     }
 
@@ -472,7 +474,7 @@ public BacktrackTask setNextBacktrackTask() throws InterruptedException {
     latestTask.getParentTask().cleanup();
 
     schedule.setChoices(latestTask.getAllChoices());
-    PSymGlobal.getCoverage().setPerChoiceDepthStats(latestTask.getPerChoiceDepthStats());
+    PSymGlobal.getCoverage().setPerChoiceDepthStats(latestTask.getAllPerChoiceDepthStats());
     return latestTask;
   }
 
diff --git a/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/scheduler/search/taskorchestration/BacktrackTask.java b/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/scheduler/search/taskorchestration/BacktrackTask.java
index 8354c6aeb..2c3f63ace 100644
--- a/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/scheduler/search/taskorchestration/BacktrackTask.java
+++ b/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/scheduler/search/taskorchestration/BacktrackTask.java
@@ -20,8 +20,12 @@ public class BacktrackTask implements Serializable {
   @Getter private List<Schedule.Choice> suffixChoices = new ArrayList<>();
 
   @Getter
-  private final List<CoverageStats.CoverageChoiceDepthStats> perChoiceDepthStats =
-      new ArrayList<>();
+  private final Map<Integer, CoverageStats.CoverageChoiceDepthStats> prefixPerChoiceDepthStats =
+      new HashMap<>();
+
+  @Getter
+  private final List<CoverageStats.CoverageChoiceDepthStats> suffixPerChoiceDepthStats =
+          new ArrayList<>();
 
   @Getter private final int id;
   @Getter private final List<BacktrackTask> children = new ArrayList<>();
@@ -81,14 +85,8 @@ public static BacktrackTask getNextTask() throws InterruptedException {
   public void cleanup() {
     numBacktracks = 0;
     numDataBacktracks = 0;
-    perChoiceDepthStats.clear();
+    suffixPerChoiceDepthStats.clear();
     suffixChoices.clear();
-    // clear task choices if all children have completed
-    for (BacktrackTask child: children) {
-      if (!child.completed) {
-        return;
-      }
-    }
   }
 
   public void addPrefixChoice(int cdepth, Schedule.Choice choice) {
@@ -97,6 +95,11 @@ public void addPrefixChoice(int cdepth, Schedule.Choice choice) {
     prefixChoices.put(cdepth, choice);
   }
 
+  public void addPrefixCoverageStats(int cdepth, CoverageStats.CoverageChoiceDepthStats stats) {
+    // TODO: check if we need copy here
+    prefixPerChoiceDepthStats.put(cdepth, stats);
+  }
+
   public void addSuffixChoice(Schedule.Choice choice) {
     // TODO: check if we need copy here
     suffixChoices.add(choice.getCopy());
@@ -108,6 +111,11 @@ public void addSuffixChoice(Schedule.Choice choice) {
     }
   }
 
+  public void addSuffixCoverageStats(CoverageStats.CoverageChoiceDepthStats stats) {
+    // TODO: check if we need copy here
+    suffixPerChoiceDepthStats.add(stats.getCopy());
+  }
+
   public List<Schedule.Choice> getAllChoices() {
     List<Schedule.Choice> result = new ArrayList<>(suffixChoices);
     BacktrackTask task = this;
@@ -126,11 +134,22 @@ public List<Schedule.Choice> getAllChoices() {
     return result;
   }
 
-  public void setPerChoiceDepthStats(List<CoverageStats.CoverageChoiceDepthStats> inputStats) {
-    assert (perChoiceDepthStats.isEmpty());
-    for (CoverageStats.CoverageChoiceDepthStats stat : inputStats) {
-      perChoiceDepthStats.add(stat.getCopy());
+  public List<CoverageStats.CoverageChoiceDepthStats> getAllPerChoiceDepthStats() {
+    List<CoverageStats.CoverageChoiceDepthStats> result = new ArrayList<>(suffixPerChoiceDepthStats);
+    BacktrackTask task = this;
+    int i = backtrackChoiceDepth-1;
+    while(i >= 0) {
+      CoverageStats.CoverageChoiceDepthStats c = task.getPrefixPerChoiceDepthStats().get(i);
+      if (c == null) {
+        assert (!task.isInitialTask());
+        task = task.getParentTask();
+      } else {
+        result.add(0, c);
+        i--;
+      }
     }
+    assert(result.size() == (suffixPerChoiceDepthStats.size() + backtrackChoiceDepth));
+    return result;
   }
 
   public boolean isInitialTask() {