[PEx] Corrections to new backtracking logic

Author: aman-goel

Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/Schedule.java

@@ -54,7 +54,9 @@ public Choice getChoice(int idx) {
      * @param choiceNum Choice depth
      */
     public void removeChoicesAfter(int choiceNum) {
-        choices.subList(choiceNum + 1, choices.size()).clear();
+        if ((choiceNum + 1) < choices.size()) {
+            choices.subList(choiceNum + 1, choices.size()).clear();
+        }
     }
 
     /**
@@ -117,6 +119,9 @@ public PValue<?> getCurrentDataChoice(int idx) {
 
     public ScheduleChoice getScheduleChoiceAt(int choiceNum) {
         for (int i = choiceNum; i >= 0; i--) {
+            if (choiceNum >= choices.size()) {
+                continue;
+            }
             Choice c = choices.get(i);
             if (c instanceof ScheduleChoice scheduleChoice) {
                 return scheduleChoice;

Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/explicit/ExplicitSearchScheduler.java

@@ -455,26 +455,29 @@ private void endCurrTask() {
     }
 
     private void setChildTask(SearchUnit unit, int choiceNum, SearchTask parentTask, boolean isExact) {
-        SearchTask newTask = searchStrategy.createTask(choiceNum, parentTask);
+        SearchTask newTask = searchStrategy.createTask(parentTask);
 
-        for (int i = 0; i <= choiceNum; i++) {
-            newTask.addPrefixChoice(schedule.getChoice(i));
-        }
+        int maxChoiceNum = choiceNum;
 
         newTask.addSuffixSearchUnit(choiceNum, unit);
 
         if (!isExact) {
             for (int i: parentTask.getSearchUnitKeys(false)) {
                 if (i > choiceNum) {
+                    if (i > maxChoiceNum) {
+                        maxChoiceNum = i;
+                    }
                     newTask.addSuffixSearchUnit(i, parentTask.getSearchUnit(i));
                 }
             }
         }
 
+        for (int i = 0; i <= maxChoiceNum; i++) {
+            newTask.addPrefixChoice(schedule.getChoice(i));
+        }
+
         parentTask.addChild(newTask);
         searchStrategy.addNewTask(newTask);
-
-        assert (choiceNum >= parentTask.getCurrChoiceNumber());
     }
 
     /**
@@ -534,14 +537,13 @@ private void postIterationCleanup() {
             if (newStepNumber == 0) {
                 reset();
                 stepState.resetToZero();
-                schedule.clear();
             } else {
                 stepNumber = newStepNumber;
                 choiceNumber = scheduleChoice.getChoiceNumber();
                 stepState.setTo(scheduleChoice.getChoiceState());
                 assert (!PExplicitGlobal.getGlobalMachine(scheduleChoice.getCurrent()).getSendBuffer().isEmpty());
-                schedule.removeChoicesAfter(backtrackChoiceNumber);
             }
+            schedule.removeChoicesAfter(backtrackChoiceNumber);
             PExplicitLogger.logBacktrack(newStepNumber, cIdx, unit);
             return;
         }

Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/explicit/strategy/SearchStrategy.java

@@ -32,16 +32,16 @@ public abstract class SearchStrategy implements Serializable {
      */
     int currTaskStartIteration = 0;
 
-    public SearchTask createTask(int choiceNum, SearchTask parentTask) {
-        SearchTask newTask = new SearchTask(allTasks.size(), choiceNum, parentTask);
+    public SearchTask createTask(SearchTask parentTask) {
+        SearchTask newTask = new SearchTask(allTasks.size(), parentTask);
         allTasks.add(newTask);
         pendingTasks.add(newTask.getId());
         return newTask;
     }
 
     public void createFirstTask() {
         assert (allTasks.size() == 0);
-        SearchTask firstTask = createTask(0, null);
+        SearchTask firstTask = createTask(null);
         setCurrTask(firstTask);
     }
 

Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/explicit/strategy/SearchTask.java

@@ -16,15 +16,14 @@ public class SearchTask implements Serializable {
     @Getter
     private final List<SearchTask> children = new ArrayList<>();
     @Getter
-    private final int currChoiceNumber;
+    private int currChoiceNumber = 0;
     @Getter
     private final List<Choice> prefixChoices = new ArrayList<>();
     @Getter
     private final Map<Integer, SearchUnit> searchUnits = new HashMap<>();
 
-    public SearchTask(int id, int choiceNum, SearchTask parentTask) {
+    public SearchTask(int id, SearchTask parentTask) {
         this.id = id;
-        this.currChoiceNumber = choiceNum;
         this.parentTask = parentTask;
     }
 
@@ -47,6 +46,9 @@ public void addPrefixChoice(Choice choice) {
 
     public void addSuffixSearchUnit(int choiceNum, SearchUnit unit) {
         searchUnits.put(choiceNum, unit.transferUnit());
+        if (choiceNum > currChoiceNumber) {
+            currChoiceNumber = choiceNum;
+        }
     }
 
     @Override
@@ -72,19 +74,10 @@ public String toStringDetailed() {
         if (isInitialTask()) {
             return String.format("%s @0::0 (parent: null)", this);
         }
-        Choice c = prefixChoices.get(currChoiceNumber);
-        if (c instanceof ScheduleChoice scheduleChoice) {
-            return String.format("%s @%d::%d (parent: %s)",
-                    this,
-                    scheduleChoice.getStepNumber(),
-                    currChoiceNumber,
-                    parentTask);
-        } else {
-            return String.format("%s -::%d (parent: %s)",
-                    this,
-                    currChoiceNumber,
-                    parentTask);
-        }
+        return String.format("%s ?::%d (parent: %s)",
+                this,
+                currChoiceNumber,
+                parentTask);
     }
 
     public List<Integer> getSearchUnitKeys(boolean reversed) {