[PEx] Cleanup and minor corrections to recent changes to SearchTask

Author: aman-goel

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

@@ -48,15 +48,6 @@ public Choice getChoice(int idx) {
         return choices.get(idx);
     }
 
-    /**
-     * Clear choice at a choice depth
-     *
-     * @param idx Choice depth
-     */
-    public void clearChoice(int idx) {
-        choices.get(idx).clearCurrent();
-    }
-
     /**
      * Remove choices after a choice depth
      *
@@ -99,7 +90,7 @@ public void setDataChoice(int stepNum, int choiceNum, PValue<?> current) {
             choices.add(null);
         }
         assert (choiceNum < choices.size());
-        choices.set(choiceNum, new DataChoice(stepNum, choiceNum, current));
+        choices.set(choiceNum, new DataChoice(current));
     }
 
     /**

Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/choice/Choice.java

@@ -13,21 +13,8 @@ public abstract class Choice<T> implements Serializable {
     @Setter
     protected T current;
 
-    /**
-     * Step number
-     */
-    @Getter
-    protected int stepNumber = 0;
-    /**
-     * Choice number
-     */
-    @Getter
-    protected int choiceNumber = 0;
-
-    protected Choice(T c, int stepNum, int choiceNum) {
+    protected Choice(T c) {
         this.current = c;
-        this.stepNumber = stepNum;
-        this.choiceNumber = choiceNum;
     }
 
     /**

Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/choice/DataChoice.java

@@ -10,12 +10,12 @@ public class DataChoice extends Choice<PValue<?>> {
     /**
      * Constructor
      */
-    public DataChoice(int stepNum, int choiceNum, PValue<?> c) {
-        super(c, stepNum, choiceNum);
+    public DataChoice(PValue<?> c) {
+        super(c);
     }
 
     public Choice copyCurrent() {
-        return new DataChoice(this.stepNumber, this.choiceNumber, this.current);
+        return new DataChoice(this.current);
     }
 
     @Override

Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/choice/ScheduleChoice.java

@@ -5,19 +5,30 @@
 import pexplicit.runtime.machine.PMachineId;
 import pexplicit.runtime.scheduler.explicit.StepState;
 
-import java.util.ArrayList;
-import java.util.List;
-
 @Getter
 @Setter
 public class ScheduleChoice extends Choice<PMachineId> {
+    /**
+     * Step number
+     */
+    private int stepNumber = 0;
+    /**
+     * Choice number
+     */
+    private int choiceNumber = 0;
+
+    /**
+     * Protocol state at the schedule step
+     */
     private StepState choiceState = null;
 
     /**
      * Constructor
      */
     public ScheduleChoice(int stepNum, int choiceNum, PMachineId c, StepState s) {
-        super(c, stepNum, choiceNum);
+        super(c);
+        this.stepNumber = stepNum;
+        this.choiceNumber = choiceNum;
         this.choiceState = s;
     }
 

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

@@ -343,8 +343,11 @@ public PMachine getNextScheduleChoice() {
         schedule.setScheduleChoice(stepNumber, choiceNumber, result.getPid());
 
         // update search unit in search task
-        if (!choices.isEmpty())
+        if (choices.isEmpty()) {
+            searchStrategy.getCurrTask().clearSearchUnit(choiceNumber);
+        } else {
             searchStrategy.getCurrTask().setScheduleSearchUnit(choiceNumber, choices);
+        }
 
         // increment choice number
         choiceNumber++;
@@ -403,8 +406,11 @@ public PValue<?> getNextDataChoice(List<PValue<?>> input_choices) {
         schedule.setDataChoice(stepNumber, choiceNumber, result);
 
         // update search unit in search task
-        if (!choices.isEmpty())
+        if (choices.isEmpty()) {
+            searchStrategy.getCurrTask().clearSearchUnit(choiceNumber);
+        } else {
             searchStrategy.getCurrTask().setDataSearchUnit(choiceNumber, choices);
+        }
 
         // increment choice number
         choiceNumber++;

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

@@ -88,9 +88,11 @@ public SearchTask setNextTask() {
      */
     public int getNumPendingChoices() {
         int numUnexplored = 0;
+        SearchTask task = getCurrTask();
+        numUnexplored += task.getNumUnexploredChoices();
         for (Integer tid : pendingTasks) {
-            SearchTask task = getTask(tid);
-            numUnexplored += task.getNumUnexploredScheduleChoices() + task.getNumUnexploredDataChoices();
+            task = getTask(tid);
+            numUnexplored += task.getNumUnexploredChoices();
         }
         return numUnexplored;
     }
@@ -102,8 +104,10 @@ public int getNumPendingChoices() {
      */
     public int getNumPendingDataChoices() {
         int numUnexplored = 0;
+        SearchTask task = getCurrTask();
+        numUnexplored += task.getNumUnexploredDataChoices();
         for (Integer tid : pendingTasks) {
-            SearchTask task = getTask(tid);
+            task = getTask(tid);
             numUnexplored += task.getNumUnexploredDataChoices();
         }
         return numUnexplored;

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

@@ -1,7 +1,6 @@
 package pexplicit.runtime.scheduler.explicit.strategy;
 
 import lombok.Getter;
-import pexplicit.runtime.PExplicitGlobal;
 import pexplicit.runtime.machine.PMachineId;
 import pexplicit.runtime.scheduler.choice.*;
 import pexplicit.values.PValue;
@@ -22,10 +21,6 @@ public class SearchTask implements Serializable {
     private final List<Choice> prefixChoices = new ArrayList<>();
     @Getter
     private final Map<Integer, SearchUnit> searchUnits = new HashMap<>();
-    @Getter
-    private int numUnexploredScheduleChoices = 0;
-    @Getter
-    private int numUnexploredDataChoices = 0;
 
     public SearchTask(int id, int choiceNum, SearchTask parentTask) {
         this.id = id;
@@ -51,11 +46,6 @@ public void addPrefixChoice(Choice choice) {
     }
 
     public void addSuffixSearchUnit(int choiceNum, SearchUnit unit) {
-        if (unit instanceof ScheduleSearchUnit scheduleSearchUnit) {
-            numUnexploredScheduleChoices += scheduleSearchUnit.getUnexplored().size();
-        } else {
-            numUnexploredDataChoices += ((DataSearchUnit) unit).getUnexplored().size();
-        }
         searchUnits.put(choiceNum, unit.transferUnit());
     }
 
@@ -82,11 +72,19 @@ public String toStringDetailed() {
         if (isInitialTask()) {
             return String.format("%s @0::0 (parent: null)", this);
         }
-        return String.format("%s @%d::%d (parent: %s)",
-                this,
-                prefixChoices.get(currChoiceNumber).getStepNumber(),
-                currChoiceNumber,
-                parentTask);
+        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);
+        }
     }
 
     public List<Integer> getSearchUnitKeys(boolean reversed) {
@@ -202,10 +200,10 @@ public int getNumUnexploredDataChoices() {
     /**
      * Clear search unit at a choice depth
      *
-     * @param idx Choice depth
+     * @param choiceNum Choice depth
      */
-    public void clearSearchUnit(int idx) {
-        searchUnits.remove(idx);
+    public void clearSearchUnit(int choiceNum) {
+        searchUnits.remove(choiceNum);
     }
 
 }