klee,s2e: refactored state management

klee/include/klee/Executor.h

@@ -26,6 +26,7 @@
 #include "klee/Internal/Module/Cell.h"
 #include "klee/Internal/Module/KInstruction.h"
 #include "klee/Internal/Module/KModule.h"
+#include "klee/StateManager.h"
 
 namespace llvm {
 class BasicBlock;
@@ -94,8 +95,7 @@ class Executor {
     std::unique_ptr<ExternalDispatcher> m_externalDispatcher;
 
     // ======= Execution state management =======
-    Searcher *m_searcher;
-    StateSet m_states;
+    StateManager m_stateManager;
 
     void executeInstruction(ExecutionState &state, KInstruction *ki);
 
@@ -135,20 +135,13 @@ class Executor {
     virtual StatePair fork(ExecutionState &current, const ref<Expr> &condition, bool keepConditionTrueInCurrentState,
                            std::function<void(ExecutionStatePtr, const StatePair &)> onBeforeNotify) = 0;
 
-    // remove state from queue and delete
-    virtual void terminateState(ExecutionStatePtr state);
-
     virtual const llvm::Module *setModule(llvm::Module *module);
 
     static void reexecuteCurrentInstructionInForkedState(ExecutionStatePtr state, const StatePair &sp);
     static void skipCurrentInstructionInForkedState(ExecutionStatePtr state, const StatePair &sp);
 
-    /*** State accessor methods ***/
-    size_t getStatesCount() const {
-        return m_states.size();
-    }
-    const StateSet &getStates() {
-        return m_states;
+    const StateSet &states() const {
+        return m_stateManager.states();
     }
 
     KModulePtr getModule() const {

klee/include/klee/StateManager.h

@@ -0,0 +1,82 @@
+///
+/// Copyright (C) 2026, Vitaly Chipounov
+///
+/// Permission is hereby granted, free of charge, to any person obtaining a copy
+/// of this software and associated documentation files (the "Software"), to deal
+/// in the Software without restriction, including without limitation the rights
+/// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+/// copies of the Software, and to permit persons to whom the Software is
+/// furnished to do so, subject to the following conditions:
+///
+/// The above copyright notice and this permission notice shall be included in all
+/// copies or substantial portions of the Software.
+///
+/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+/// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+/// SOFTWARE.
+///
+
+#ifndef KLEE_STATEMANAGER_H
+#define KLEE_STATEMANAGER_H
+
+#include <klee/Searcher.h>
+
+namespace klee {
+
+class StateManager : public Searcher {
+private:
+    Searcher *m_searcher = nullptr;
+    StateSet m_states;
+
+public:
+    virtual ~StateManager() = default;
+
+    ExecutionStatePtr selectState() {
+        assert(m_searcher && "No searcher set");
+        auto ret = m_searcher->selectState();
+        assert(m_states.find(ret) != m_states.end() && "Selected state is not managed");
+        return ret;
+    }
+
+    void addState(klee::ExecutionStatePtr state) {
+        m_states.insert(state);
+
+        if (m_searcher) {
+            m_searcher->addState(state);
+        }
+    }
+
+    void removeState(klee::ExecutionStatePtr state) {
+        if (m_searcher) {
+            m_searcher->removeState(state);
+        }
+
+        m_states.erase(state);
+    }
+
+    bool empty() {
+        auto ret = m_states.empty();
+        assert(ret == m_states.empty());
+        return ret;
+    }
+
+    void setSearcher(Searcher *searcher) {
+        m_searcher = searcher;
+    }
+
+    Searcher *searcher() const {
+        return m_searcher;
+    }
+
+    const StateSet &states() const {
+        return m_states;
+    }
+};
+
+} // namespace klee
+
+#endif // KLEE_STATEMANAGER_H

klee/lib/Core/Executor.cpp

@@ -81,8 +81,7 @@ namespace klee {
 extern cl::opt<bool> UseExprSimplifier;
 } // namespace klee
 
-Executor::Executor(LLVMContext &context)
-    : m_kmodule(0), m_externalDispatcher(std::make_unique<ExternalDispatcher>()), m_searcher(0) {
+Executor::Executor(LLVMContext &context) : m_kmodule(0), m_externalDispatcher(std::make_unique<ExternalDispatcher>()) {
 }
 
 const Module *Executor::setModule(llvm::Module *module) {
@@ -1427,16 +1426,6 @@ void Executor::executeInstruction(ExecutionState &state, KInstruction *ki) {
     }
 }
 
-void Executor::terminateState(ExecutionStatePtr state) {
-    *klee::stats::completedPaths += 1;
-
-    if (m_searcher) {
-        m_searcher->removeState(state);
-    }
-
-    m_states.erase(state);
-}
-
 extern "C" {
 typedef uint64_t (*external_fcn_t)(...);
 }

libs2ecore/include/s2e/S2EExecutor.h

@@ -135,18 +135,8 @@ class S2EExecutor : public klee::Executor {
 
     static void doInterruptAll(int intno, int is_int, int error_code, uintptr_t next_eip, int is_hw);
 
-    /** Suspend the given state (does not kill it) */
-    bool suspendState(S2EExecutionState *state);
-
-    /** Puts back the previously suspended state in the queue */
-    bool resumeState(S2EExecutionState *state);
-
-    klee::Searcher *getSearcher() const {
-        return m_searcher;
-    }
-
     void setSearcher(klee::Searcher *s) {
-        m_searcher = s;
+        m_stateManager.setSearcher(s);
     }
 
     StatePair forkCondition(S2EExecutionState *state, klee::ref<klee::Expr> condition,

libs2ecore/src/S2EExecutor.cpp

@@ -418,7 +418,8 @@ S2EExecutor::S2EExecutor(S2E *s2e, TCGLLVMTranslator *translator)
     }
 #endif
 
-    m_searcher = constructUserSearcher();
+    auto searcher = constructUserSearcher();
+    m_stateManager.setSearcher(searcher);
 
     g_s2e_fork_on_symbolic_address = ForkOnSymbolicAddress;
     g_s2e_concretize_io_addresses = ConcretizeIoAddress;
@@ -470,10 +471,7 @@ S2EExecutionState *S2EExecutor::createInitialState() {
     state->m_active = true;
     state->setForking(EnableForking);
 
-    m_states.insert(state);
-    if (m_searcher) {
-        m_searcher->addState(state);
-    }
+    m_stateManager.addState(state);
 
 #define __DEFINE_EXT_OBJECT_RO(name)                                                  \
     {                                                                                 \
@@ -593,7 +591,8 @@ void S2EExecutor::computeNewStateGuids(std::unordered_map<ExecutionStatePtr, uin
 }
 
 void S2EExecutor::doLoadBalancing() {
-    if (m_states.size() < 2) {
+    auto states = m_stateManager.states();
+    if (states.size() < 2) {
         return;
     }
 
@@ -604,7 +603,7 @@ void S2EExecutor::doLoadBalancing() {
 
     std::vector<S2EExecutionStatePtr> allStates;
 
-    for (auto it : m_states) {
+    for (auto it : states) {
         S2EExecutionStatePtr s2estate = static_pointer_cast<S2EExecutionState>(it);
         if (!s2estate->isZombie() && !s2estate->isPinned()) {
             allStates.push_back(s2estate);
@@ -660,7 +659,7 @@ void S2EExecutor::doLoadBalancing() {
     for (auto state : allStates) {
         auto s2estate = static_pointer_cast<S2EExecutionState>(state);
         if (!currentSet.count(s2estate)) {
-            Executor::terminateState(s2estate);
+            m_stateManager.removeState(s2estate);
 
             // This is important if we kill the current state
             s2estate->zombify();
@@ -801,8 +800,8 @@ void S2EExecutor::doStateSwitch(S2EExecutionStatePtr oldState, S2EExecutionState
 ExecutionStatePtr S2EExecutor::selectSearcherState(S2EExecutionStatePtr state) {
     ExecutionStatePtr newState;
 
-    if (!m_searcher->empty()) {
-        newState = m_searcher->selectState();
+    if (!m_stateManager.empty()) {
+        newState = m_stateManager.selectState();
     }
 
     if (!newState) {
@@ -831,12 +830,6 @@ S2EExecutionStatePtr S2EExecutor::selectNextState(S2EExecutionStatePtr state) {
         return nullptr;
     }
 
-    // This assertion must go before the cast to S2EExecutionState.
-    // In case the searcher returns a bogus state, this allows
-    // spotting it immediately. The dynamic cast however, might cause
-    // memory corruptions.
-    assert(m_states.find(nstate) != m_states.end());
-
     auto newState = dynamic_pointer_cast<S2EExecutionState>(nstate);
 
     assert(newState);
@@ -1230,10 +1223,7 @@ S2EExecutor::StatePair S2EExecutor::fork(ExecutionState &current, const klee::re
     notifyFork(current, condition, ret);
 
     if (newState) {
-        if (m_searcher) {
-            m_searcher->addState(newState);
-        }
-        m_states.insert(newState);
+        m_stateManager.addState(newState);
         StateSet addedStates = {newState};
         auto s2eState = static_cast<S2EExecutionState *>(&current);
         m_s2e->getCorePlugin()->onUpdateStates.emit(s2eState, addedStates, StateSet());
@@ -1579,11 +1569,13 @@ void S2EExecutor::terminateState(klee::ExecutionStatePtr state, const std::strin
 }
 
 void S2EExecutor::terminateState(ExecutionStatePtr s) {
+    *klee::stats::completedPaths += 1;
+
     auto state = static_pointer_cast<S2EExecutionState>(s);
 
     m_s2e->getCorePlugin()->onStateKill.emit(state.get());
 
-    Executor::terminateState(state);
+    m_stateManager.removeState(state);
     state->zombify();
 
     g_s2e->getWarningsStream().flush();
@@ -1682,29 +1674,6 @@ void S2EExecutor::doInterruptAll(int intno, int is_int, int error_code, uintptr_
     g_s2e_state->setRunningExceptionEmulationCode(false);
 }
 
-/** Suspend the given state (does not kill it) */
-bool S2EExecutor::suspendState(S2EExecutionState *state) {
-    if (m_searcher) {
-        m_searcher->removeState(state);
-        size_t r = m_states.erase(state);
-        assert(r == 1);
-        return true;
-    }
-    return false;
-}
-
-bool S2EExecutor::resumeState(S2EExecutionState *state) {
-    if (m_searcher) {
-        if (m_states.find(state) != m_states.end()) {
-            return false;
-        }
-        m_states.insert(state);
-        m_searcher->addState(state);
-        return true;
-    }
-    return false;
-}
-
 S2ETranslationBlock *S2EExecutor::allocateS2ETb() {
     S2ETranslationBlockPtr se_tb(new S2ETranslationBlock);
     m_s2eTbs.insert(se_tb);
@@ -1784,7 +1753,7 @@ void s2e_set_tb_function(void *se_tb, void *llvmFunction) {
 }
 
 void s2e_flush_tb_cache() {
-    if (g_s2e && g_s2e->getExecutor()->getStatesCount() > 1) {
+    if (g_s2e && g_s2e->getExecutor()->states().size() > 1) {
         if (!FlushTBsOnStateSwitch) {
             g_s2e->getWarningsStream() << "Flushing TB cache with more than 1 state. Dangerous. Expect crashes.\n";
         }

libs2eplugins/src/s2e/Plugins/Core/BaseInstructions.cpp

@@ -860,7 +860,7 @@ void BaseInstructions::handleBuiltInOps(S2EExecutionState *state, uint64_t opcod
         }
 
         case BASE_S2E_STATE_COUNT: { /* Get number of active states */
-            target_ulong count = s2e()->getExecutor()->getStatesCount();
+            target_ulong count = s2e()->getExecutor()->states().size();
             state->regs()->write(CPU_OFFSET(regs[R_EAX]), &count, sizeof(count));
             break;
         }

libs2eplugins/src/s2e/Plugins/OSMonitors/Windows/WindowsCrashDumpGenerator.cpp

@@ -175,7 +175,7 @@ int WindowsCrashDumpInvoker::generateCrashDump(lua_State *L) {
 
         // Fetch the right state
         // XXX: Avoid linear search
-        const klee::StateSet &states = g_s2e->getExecutor()->getStates();
+        const klee::StateSet &states = g_s2e->getExecutor()->states();
         for (auto it : states) {
             auto ss = static_pointer_cast<S2EExecutionState>(it);
             if (ss->getID() == stateId) {

libs2eplugins/src/s2e/Plugins/PathLimiters/ResourceMonitor.cpp

@@ -154,19 +154,19 @@ void ResourceMonitor::updateMemoryUsage() {
 
 void ResourceMonitor::dropStates() {
     S2EExecutor *executor = s2e()->getExecutor();
-    assert(executor->getStatesCount() > 0 && "no states left to remove\n");
+    assert(executor->states().size() > 0 && "no states left to remove\n");
 
     // try to kill 5% of the states
-    size_t nrStatesToTerminate = (size_t) (0.05 * executor->getStatesCount());
+    size_t nrStatesToTerminate = (size_t) (0.05 * executor->states().size());
 
-    if (nrStatesToTerminate < 1 && executor->getStatesCount() > 0) {
+    if (nrStatesToTerminate < 1 && executor->states().size() > 0) {
         nrStatesToTerminate = 1; // kill at least one state
     }
 
     if (nrStatesToTerminate > 0) {
-        const klee::StateSet &states = s2e()->getExecutor()->getStates();
+        const klee::StateSet &states = s2e()->getExecutor()->states();
         getDebugStream() << "ResourceMonitor: will kill " << nrStatesToTerminate << " states, reason: out of memory"
-                         << " #states = " << executor->getStatesCount() << "\n";
+                         << " #states = " << executor->states().size() << "\n";
         getDebugStream() << "ResourceMonitor: limit - rss (MB) = " << ((m_cgroupMemLimit - m_rss) / 1024 / 1024)
                          << "\n";
 

libs2eplugins/src/s2e/Plugins/Searchers/SeedScheduler.cpp

@@ -131,7 +131,7 @@ void SeedScheduler::onProcessForkDecide(bool *proceed) {
     static unsigned previousAvailableSeedCount = 0;
     static time_point lastSeedFetchTime;
 
-    if (s2e()->getExecutor()->getStatesCount() >= 3) {
+    if (s2e()->getExecutor()->states().size() >= 3) {
         // We have plenty of states to load balance, so no need
         // to block instance forking.
         return;
@@ -247,7 +247,7 @@ void SeedScheduler::onNewBlockCovered(S2EExecutionState *state) {
 }
 
 void SeedScheduler::terminateIdleInstance() {
-    if (s2e()->getExecutor()->getStatesCount() > 1) {
+    if (s2e()->getExecutor()->states().size() > 1) {
         getDebugStream() << "idle detection: too many states\n";
         return;
     }
@@ -317,7 +317,7 @@ void SeedScheduler::processSeedStateMachine(time_point currentTime) {
         if (!foundBlocks && !foundCrashes) {
             m_explorationState = WAIT_FOR_NEW_SEEDS;
         } else if ((m_stateKilled || s2e()->getCurrentInstanceCount() > 1) &&
-                   (s2e()->getExecutor()->getStatesCount() == 1)) {
+                   (s2e()->getExecutor()->states().size() == 1)) {
             // The warm up phase terminates if no seedless states remain, i.e., there
             // is only state 0 remaining, in which case we have to wait for new seeds
             // as there is nothing else to do. We have to check for m_stateKilled because
@@ -340,7 +340,7 @@ void SeedScheduler::processSeedStateMachine(time_point currentTime) {
             /* Prioritize normal seeds if S2E couldn't find coverage on its own */
             m_seeds->enableSeeds(true);
             m_explorationState = WAIT_SEED_SCHEDULING;
-        } else if (s2e()->getExecutor()->getStatesCount() == 1) {
+        } else if (s2e()->getExecutor()->states().size() == 1) {
             /* Prioritize normal seeds if no other states are running */
             m_seeds->enableSeeds(true);
             m_explorationState = WAIT_SEED_SCHEDULING;
@@ -351,7 +351,7 @@ void SeedScheduler::processSeedStateMachine(time_point currentTime) {
 
     } else if (m_explorationState == WAIT_SEED_EXECUTION) {
         /* Give newly fetched seed some time to execute */
-        if (!recentSeedFetch || (s2e()->getExecutor()->getStatesCount() == 1)) {
+        if (!recentSeedFetch || (s2e()->getExecutor()->states().size() == 1)) {
             m_explorationState = WAIT_FOR_NEW_SEEDS;
         }
     }