From 8ec0dbc11f2eb1f4a2f2ffd248bedb5a60b3f3bd Mon Sep 17 00:00:00 2001
From: Aman Goel <goelaman@amazon.com>
Date: Thu, 18 Apr 2024 21:21:58 +0000
Subject: [PATCH] [PExplicit] Minor refactoring

---
 .../runtime/scheduler/Scheduler.java          |  3 +
 .../explicit/ExplicitSearchScheduler.java     | 56 +++++++++++--------
 .../scheduler/replay/ReplayScheduler.java     |  3 -
 3 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/Scheduler.java b/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/Scheduler.java
index 2eab10684..60c4ddc9d 100644
--- a/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/Scheduler.java
+++ b/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/Scheduler.java
@@ -221,6 +221,9 @@ public void executeStep(PMachine sender) {
             schedule.setStepNumber(schedule.getStepNumber() + 1);
         }
 
+        // reset number of logs in current step
+        stepNumLogs = 0;
+
         // log start step
         PExplicitLogger.logStartStep(schedule.getStepNumber(), sender, msg);
 
diff --git a/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/explicit/ExplicitSearchScheduler.java b/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/explicit/ExplicitSearchScheduler.java
index fce1981a8..b91112234 100644
--- a/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/explicit/ExplicitSearchScheduler.java
+++ b/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/explicit/ExplicitSearchScheduler.java
@@ -143,6 +143,37 @@ protected void runStep() throws TimeoutException {
         TimeMonitor.checkTimeout();
         MemoryMonitor.checkMemout();
 
+        // check if we can skip the remaining schedule
+        if (skipRemainingSchedule()) {
+            scheduleTerminated = false;
+            skipLiveness = true;
+            isDoneStepping = true;
+            PExplicitLogger.logFinishedIteration(schedule.getStepNumber());
+            return;
+        }
+
+        // get a scheduling choice as sender machine
+        PMachine sender = getNextScheduleChoice();
+
+        if (sender == null) {
+            // done with this schedule
+            scheduleTerminated = true;
+            skipLiveness = false;
+            isDoneStepping = true;
+            PExplicitLogger.logFinishedIteration(schedule.getStepNumber());
+            return;
+        }
+
+        // execute a step from message in the sender queue
+        executeStep(sender);
+    }
+
+    /**
+     * Check if the remaining schedule can be skipped if current state is already in state cache
+     *
+     * @return true if remaining schedule can be skipped
+     */
+    boolean skipRemainingSchedule() {
         if (PExplicitGlobal.getConfig().getStateCachingMode() != StateCachingMode.None) {
             // perform state caching only if beyond backtrack choice number
             if (schedule.getChoiceNumber() > backtrackChoiceNumber) {
@@ -167,33 +198,12 @@ protected void runStep() throws TimeoutException {
                         Assert.cycle(false, "Cycle detected: Infinite loop found due to revisiting a state multiple times in the same schedule");
                     } else {
                         // done with this schedule
-                        scheduleTerminated = false;
-                        skipLiveness = true;
-                        isDoneStepping = true;
-                        PExplicitLogger.logFinishedIteration(schedule.getStepNumber());
-                        return;
+                        return true;
                     }
                 }
             }
         }
-
-        // reset number of logs in current step
-        stepNumLogs = 0;
-
-        // get a scheduling choice as sender machine
-        PMachine sender = getNextScheduleChoice();
-
-        if (sender == null) {
-            // done with this schedule
-            scheduleTerminated = true;
-            skipLiveness = false;
-            isDoneStepping = true;
-            PExplicitLogger.logFinishedIteration(schedule.getStepNumber());
-            return;
-        }
-
-        // execute a step from message in the sender queue
-        executeStep(sender);
+        return false;
     }
 
     /**
diff --git a/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/replay/ReplayScheduler.java b/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/replay/ReplayScheduler.java
index c2b157182..610e2348a 100644
--- a/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/replay/ReplayScheduler.java
+++ b/Src/PRuntimes/PExplicitRuntime/src/main/java/pexplicit/runtime/scheduler/replay/ReplayScheduler.java
@@ -55,9 +55,6 @@ protected void runIteration() throws TimeoutException {
 
     @Override
     protected void runStep() throws TimeoutException {
-        // reset number of logs in current step
-        stepNumLogs = 0;
-
         // get a scheduling choice as sender machine
         PMachine sender = getNextScheduleChoice();