refactored safety shield choices and scheduler

src/storm/shields/PostSafetyShield.cpp

@@ -43,28 +43,28 @@ namespace tempest {
                     ValueType maxProbability = *(choice_it + maxProbabilityIndex);
                     if(!relative && !choiceFilter(maxProbability, maxProbability, this->shieldingExpression->getValue())) {
                         STORM_LOG_WARN("No shielding action possible with absolute comparison for state with index " << state);
-                        shield.setChoice(0, storm::storage::Distribution<ValueType, IndexType>(), state);
+                        shield.setChoice(storm::storage::PostSchedulerChoice<ValueType>(), state, 0);
                         choice_it += rowGroupSize;
                         continue;
                     }
+                    storm::storage::PostSchedulerChoice<ValueType> choiceMapping;
                     for(uint choice = 0; choice < rowGroupSize; choice++, choice_it++) {
-                        storm::storage::Distribution<ValueType, IndexType> actionDistribution;
                         if(choiceFilter(*choice_it, maxProbability, this->shieldingExpression->getValue())) {
-                            actionDistribution.addProbability(choice, 1);
+                            choiceMapping.addChoice(choice, choice);
                         } else {
-                            actionDistribution.addProbability(maxProbabilityIndex, 1);
+                            choiceMapping.addChoice(choice, maxProbabilityIndex);
                         }
-                        shield.setChoice(choice, storm::storage::SchedulerChoice<ValueType>(actionDistribution), state);
                     }
+                    shield.setChoice(choiceMapping, state, 0);
                 } else {
-                    shield.setChoice(0, storm::storage::Distribution<ValueType, IndexType>(), state);
+                    shield.setChoice(storm::storage::PostSchedulerChoice<ValueType>(), state, 0);
                     choice_it += rowGroupSize;
                 }
             }
             return shield;
         }
 
-        // Explicitly instantiate appropriate
+        // Explicitly instantiate appropriate classes
         template class PostSafetyShield<double, typename storm::storage::SparseMatrix<double>::index_type>;
 #ifdef STORM_HAVE_CARL
         template class PostSafetyShield<storm::RationalNumber, typename storm::storage::SparseMatrix<storm::RationalNumber>::index_type>;

src/storm/shields/PreSafetyShield.cpp

@@ -39,29 +39,30 @@ namespace tempest {
             for(uint state = 0; state < this->rowGroupIndices.size() - 1; state++) {
                 uint rowGroupSize = this->rowGroupIndices[state + 1] - this->rowGroupIndices[state];
                 if(this->relevantStates.get(state)) {
-                    storm::storage::Distribution<ValueType, IndexType> actionDistribution;
+                    storm::storage::PreSchedulerChoice<ValueType> enabledChoices;
                     ValueType maxProbability = *std::max_element(choice_it, choice_it + rowGroupSize);
                     if(!relative && !choiceFilter(maxProbability, maxProbability, this->shieldingExpression->getValue())) {
                         STORM_LOG_WARN("No shielding action possible with absolute comparison for state with index " << state);
-                        shield.setChoice(storm::storage::Distribution<ValueType, IndexType>(), state);
+                        shield.setChoice(storm::storage::PreSchedulerChoice<ValueType>(), state, 0);
                         choice_it += rowGroupSize;
                         continue;
                     }
                     for(uint choice = 0; choice < rowGroupSize; choice++, choice_it++) {
                         if(choiceFilter(*choice_it, maxProbability, this->shieldingExpression->getValue())) {
-                            actionDistribution.addProbability(choice, *choice_it);
+                            enabledChoices.addChoice(choice, *choice_it);
                         }
                     }
-                    shield.setChoice(storm::storage::SchedulerChoice<ValueType>(actionDistribution), state);
+                    shield.setChoice(enabledChoices, state, 0);
 
                 } else {
-                    shield.setChoice(storm::storage::Distribution<ValueType, IndexType>(), state);
+                    shield.setChoice(storm::storage::PreSchedulerChoice<ValueType>(), state, 0);
                     choice_it += rowGroupSize;
                 }
+
             }
             return shield;
         }
-        // Explicitly instantiate appropriate
+        // Explicitly instantiate appropriate classes
         template class PreSafetyShield<double, typename storm::storage::SparseMatrix<double>::index_type>;
 #ifdef STORM_HAVE_CARL
         template class PreSafetyShield<storm::RationalNumber, typename storm::storage::SparseMatrix<storm::RationalNumber>::index_type>;

src/storm/storage/PostScheduler.cpp

@@ -9,14 +9,10 @@ namespace storm {
     namespace storage {
 
         template <typename ValueType>
-        PostScheduler<ValueType>::PostScheduler(uint_fast64_t numberOfModelStates, std::vector<uint_fast64_t> numberOfChoicesPerState, boost::optional<storm::storage::MemoryStructure> const& memoryStructure) : Scheduler<ValueType>(numberOfModelStates, memoryStructure) {
-            STORM_LOG_DEBUG(numberOfChoicesPerState.size() << " " << numberOfModelStates);
-            STORM_LOG_ASSERT(numberOfChoicesPerState.size() == numberOfModelStates, "Need to know amount of choices per model state");
+        PostScheduler<ValueType>::PostScheduler(uint_fast64_t numberOfModelStates, std::vector<uint_fast64_t> numberOfChoicesPerState, boost::optional<storm::storage::MemoryStructure> const& memoryStructure) : memoryStructure(memoryStructure) {
+            //STORM_LOG_ASSERT(numberOfChoicesPerState.size() == numberOfModelStates, "Need to know amount of choices per model state");
             uint_fast64_t numOfMemoryStates = memoryStructure ? memoryStructure->getNumberOfStates() : 1;
-            schedulerChoiceMapping = std::vector<std::vector<std::vector<SchedulerChoice<ValueType>>>>(numOfMemoryStates, std::vector<std::vector<SchedulerChoice<ValueType>>>(numberOfModelStates));
-            for(uint state = 0; state < numberOfModelStates; state++) {
-                schedulerChoiceMapping[0][state].resize(numberOfChoicesPerState[state]);
-            }
+            schedulerChoiceMapping = std::vector<std::vector<PostSchedulerChoice<ValueType>>>(numOfMemoryStates, std::vector<PostSchedulerChoice<ValueType>>(numberOfModelStates));
             numberOfChoices = 0;
             for(std::vector<uint_fast64_t>::iterator it = numberOfChoicesPerState.begin(); it != numberOfChoicesPerState.end(); ++it)
                 numberOfChoices += *it;
@@ -25,13 +21,10 @@ namespace storm {
         }
 
         template <typename ValueType>
-        PostScheduler<ValueType>::PostScheduler(uint_fast64_t numberOfModelStates, std::vector<uint_fast64_t> numberOfChoicesPerState, boost::optional<storm::storage::MemoryStructure>&& memoryStructure) : Scheduler<ValueType>(numberOfModelStates, std::move(memoryStructure)) {
+        PostScheduler<ValueType>::PostScheduler(uint_fast64_t numberOfModelStates, std::vector<uint_fast64_t> numberOfChoicesPerState, boost::optional<storm::storage::MemoryStructure>&& memoryStructure) : memoryStructure(std::move(memoryStructure)) {
             STORM_LOG_ASSERT(numberOfChoicesPerState.size() == numberOfModelStates, "Need to know amount of choices per model state");
             uint_fast64_t numOfMemoryStates = memoryStructure ? memoryStructure->getNumberOfStates() : 1;
-            schedulerChoiceMapping = std::vector<std::vector<std::vector<SchedulerChoice<ValueType>>>>(numOfMemoryStates, std::vector<std::vector<SchedulerChoice<ValueType>>>(numberOfModelStates));
-            for(uint state = 0; state < numberOfModelStates; state++) {
-                schedulerChoiceMapping[0][state].resize(numberOfChoicesPerState[state]);
-            }
+            schedulerChoiceMapping = std::vector<std::vector<PostSchedulerChoice<ValueType>>>(numOfMemoryStates, std::vector<PostSchedulerChoice<ValueType>>(numberOfModelStates));
             numberOfChoices = 0;
             for(std::vector<uint_fast64_t>::iterator it = numberOfChoicesPerState.begin(); it != numberOfChoicesPerState.end(); ++it)
                 numberOfChoices += *it;
@@ -40,21 +33,13 @@ namespace storm {
         }
 
         template <typename ValueType>
-        void PostScheduler<ValueType>::setChoice(OldChoice const& oldChoice, SchedulerChoice<ValueType> const& newChoice, uint_fast64_t modelState, uint_fast64_t memoryState) {
+        void PostScheduler<ValueType>::setChoice(PostSchedulerChoice<ValueType> const& choice, uint_fast64_t modelState, uint_fast64_t memoryState) {
             STORM_LOG_ASSERT(memoryState == 0, "Currently we do not support PostScheduler with memory");
             STORM_LOG_ASSERT(modelState < schedulerChoiceMapping[memoryState].size(), "Illegal model state index");
 
-            schedulerChoiceMapping[memoryState][modelState][oldChoice] = newChoice;
-        }
-
-        template <typename ValueType>
-        SchedulerChoice<ValueType> const& PostScheduler<ValueType>::getChoice(uint_fast64_t modelState, OldChoice oldChoice, uint_fast64_t memoryState)  {
-            STORM_LOG_ASSERT(memoryState < this->getNumberOfMemoryStates(), "Illegal memory state index");
-            STORM_LOG_ASSERT(modelState < schedulerChoiceMapping[memoryState].size(), "Illegal model state index");
-            return schedulerChoiceMapping[memoryState][modelState][oldChoice];
+            schedulerChoiceMapping[memoryState][modelState] = choice;
         }
 
-
         template <typename ValueType>
         bool PostScheduler<ValueType>::isDeterministicScheduler() const {
             return true;
@@ -89,7 +74,11 @@ namespace storm {
             out << ":" << std::endl;
             uint_fast64_t numOfSkippedStatesWithUniqueChoice = 0;
             for (uint_fast64_t state = 0; state < schedulerChoiceMapping.front().size(); ++state) {
+                PostSchedulerChoice<ValueType> const& choices = schedulerChoiceMapping[0][state];
+                if(choices.isEmpty() && !printUndefinedChoices) continue;
+
                 std::stringstream stateString;
+
                 // Print the state info
                 if (stateValuationsGiven) {
                     stateString << std::setw(widthOfStates)  << (std::to_string(state) + ": " + model->getStateValuations().getStateInfo(state));
@@ -98,42 +87,37 @@ namespace storm {
                 }
                 stateString << "    ";
 
+
                 bool firstChoiceIndex = true;
-                for(uint choiceIndex = 0; choiceIndex < schedulerChoiceMapping[0][state].size(); choiceIndex++) {
-                    SchedulerChoice<ValueType> const& choice = schedulerChoiceMapping[0][state][choiceIndex];
+                for(auto const& choiceMap : choices.getChoiceMap()) {
                     if(firstChoiceIndex) {
                         firstChoiceIndex = false;
                     } else {
                         stateString << ";    ";
                     }
-                    if (choice.isDefined()) {
-                        auto choiceProbPair = *(choice.getChoiceAsDistribution().begin());
+
                     if(choiceLabelsGiven) {
-                            auto choiceLabels = model->getChoiceLabeling().getLabelsOfChoice(model->getTransitionMatrix().getRowGroupIndices()[state] + choiceIndex);
-                            stateString << std::to_string(choiceIndex) << " {" << boost::join(choiceLabels, ", ") << "}: ";
+                        auto choiceLabels = model->getChoiceLabeling().getLabelsOfChoice(model->getTransitionMatrix().getRowGroupIndices()[state] + std::get<0>(choiceMap));
+                        stateString << std::to_string(std::get<0>(choiceMap)) << " {" << boost::join(choiceLabels, ", ") << "}: ";
                     } else {
-                            stateString << std::to_string(choiceIndex) << ": ";
+                        stateString << std::to_string(std::get<0>(choiceMap)) << ": ";
                     }
-                        //stateString << choiceProbPair.second << ": (";
+
                     if (choiceOriginsGiven) {
-                            stateString << model->getChoiceOrigins()->getChoiceInfo(model->getTransitionMatrix().getRowGroupIndices()[state] + choiceProbPair.first);
+                        stateString << model->getChoiceOrigins()->getChoiceInfo(model->getTransitionMatrix().getRowGroupIndices()[state] + std::get<1>(choiceMap));
                     } else {
-                            stateString << choiceProbPair.first;
+                        stateString << std::to_string(std::get<1>(choiceMap));
                     }
                     if (choiceLabelsGiven) {
-                            auto choiceLabels = model->getChoiceLabeling().getLabelsOfChoice(model->getTransitionMatrix().getRowGroupIndices()[state] + choiceProbPair.first);
+                        auto choiceLabels = model->getChoiceLabeling().getLabelsOfChoice(model->getTransitionMatrix().getRowGroupIndices()[state] + std::get<1>(choiceMap));
                         stateString << " {" << boost::join(choiceLabels, ", ") << "}";
                     }
 
-                    } else {
-                        if(!this->printUndefinedChoices) goto skipStatesWithUndefinedChoices;
-                        stateString << "undefined.";
-                    }
                     // Todo: print memory updates
                 }
                 out << stateString.str() << std::endl;
                 // jump to label if we find one undefined choice.
-                skipStatesWithUndefinedChoices:;
+                //skipStatesWithUndefinedChoices:;
             }
             out << "___________________________________________________________________" << std::endl;
         }

src/storm/storage/PostScheduler.h

@@ -2,7 +2,7 @@
 
 #include <cstdint>
 #include <map>
-#include "storm/storage/SchedulerChoice.h"
+#include "storm/storage/PostSchedulerChoice.h"
 #include "storm/storage/Scheduler.h"
 #include "storm/logic/ShieldExpression.h"
 
@@ -16,7 +16,7 @@ namespace storm {
          * A Choice can be undefined, deterministic
          */
         template <typename ValueType>
-        class PostScheduler : public Scheduler<ValueType> {
+        class PostScheduler {
         public:
             typedef uint_fast64_t OldChoice;
             /*!
@@ -35,7 +35,7 @@ namespace storm {
              * @param modelState The state of the model for which to set the choice.
              * @param memoryState The state of the memoryStructure for which to set the choice.
              */
-            void setChoice(OldChoice const& oldChoice, SchedulerChoice<ValueType> const& newChoice, uint_fast64_t modelState, uint_fast64_t memoryState = 0);
+            void setChoice(PostSchedulerChoice<ValueType> const& newChoice, uint_fast64_t modelState, uint_fast64_t memoryState = 0);
 
             /*!
              * Is the scheduler defined on the states indicated by the selected-states bitvector?
@@ -50,14 +50,6 @@ namespace storm {
              */
             void clearChoice(uint_fast64_t modelState, uint_fast64_t memoryState = 0);
 
-            /*!
-             * Gets the choice defined by the scheduler for the given model and memory state.
-             *
-             * @param state The state for which to get the choice.
-             * @param memoryState the memory state which we consider.
-             */
-            SchedulerChoice<ValueType> const& getChoice(uint_fast64_t modelState, OldChoice oldChoice, uint_fast64_t memoryState = 0) ;
-
             /*!
              * Compute the Action Support: A bit vector that indicates all actions that are selected with positive probability in some memory state
              */
@@ -92,8 +84,14 @@ namespace storm {
              */
             void printToStream(std::ostream& out, std::shared_ptr<storm::logic::ShieldExpression const> shieldingExpression, std::shared_ptr<storm::models::sparse::Model<ValueType>> model = nullptr, bool skipUniqueChoices = false) const;
         private:
-            std::vector<std::vector<std::vector<SchedulerChoice<ValueType>>>> schedulerChoiceMapping;
+            boost::optional<storm::storage::MemoryStructure> memoryStructure;
+            std::vector<std::vector<PostSchedulerChoice<ValueType>>> schedulerChoiceMapping;
+
+            bool printUndefinedChoices = false;
 
+            uint_fast64_t numOfUndefinedChoices;
+            uint_fast64_t numOfDeterministicChoices;
+            uint_fast64_t numOfDontCareStates;
             std::vector<uint_fast64_t> numberOfChoicesPerState;
             uint_fast64_t numberOfChoices;
         };

src/storm/storage/PostSchedulerChoice.cpp

@@ -0,0 +1,63 @@
+#include "storm/storage/PostSchedulerChoice.h"
+
+#include "storm/utility/constants.h"
+#include "storm/utility/macros.h"
+
+#include "storm/exceptions/InvalidOperationException.h"
+#include "storm/adapters/RationalFunctionAdapter.h"
+#include "storm/adapters/RationalNumberAdapter.h"
+
+namespace storm {
+    namespace storage {
+
+        template <typename ValueType>
+        PostSchedulerChoice<ValueType>::PostSchedulerChoice() {
+            // Intentionally left empty
+        }
+
+        template <typename ValueType>
+        void PostSchedulerChoice<ValueType>::addChoice(uint_fast64_t oldChoiceIndex, uint_fast64_t newChoiceIndex) {
+            choiceMap.emplace_back(oldChoiceIndex, newChoiceIndex);
+        }
+
+        template <typename ValueType>
+        std::vector<std::tuple<uint_fast64_t, uint_fast64_t>> const& PostSchedulerChoice<ValueType>::getChoiceMap() const {
+            return choiceMap;
+        }
+
+        template <typename ValueType>
+        std::tuple<uint_fast64_t, uint_fast64_t> const& PostSchedulerChoice<ValueType>::getChoice(uint_fast64_t choiceIndex) const {
+            return choiceMap.at(choiceIndex);
+        }
+
+        template <typename ValueType>
+        bool PostSchedulerChoice<ValueType>::isEmpty() const {
+            return choiceMap.size() == 0;
+        }
+
+        template <typename ValueType>
+        std::ostream& operator<<(std::ostream& out, PostSchedulerChoice<ValueType> const& schedulerChoice) {
+            if (!schedulerChoice.isEmpty()) {
+                bool firstChoice = true;
+                for(auto const& choice : schedulerChoice.getChoiceMap()) {
+                    if(firstChoice) firstChoice = false;
+                    else out << ", ";
+                    out << std::get<0>(choice) << " -> " << std::get<1>(choice);
+                }
+            } else {
+              out << "undefined";
+            }
+            return out;
+        }
+
+        template class PostSchedulerChoice<double>;
+        template std::ostream& operator<<(std::ostream& out, PostSchedulerChoice<double> const& schedulerChoice);
+        template class PostSchedulerChoice<float>;
+        template std::ostream& operator<<(std::ostream& out, PostSchedulerChoice<float> const& schedulerChoice);
+        template class PostSchedulerChoice<storm::RationalNumber>;
+        template std::ostream& operator<<(std::ostream& out, PostSchedulerChoice<storm::RationalNumber> const& schedulerChoice);
+        template class PostSchedulerChoice<storm::RationalFunction>;
+        template std::ostream& operator<<(std::ostream& out, PostSchedulerChoice<storm::RationalFunction> const& schedulerChoice);
+
+    }
+}

src/storm/storage/PostSchedulerChoice.h

@@ -0,0 +1,46 @@
+#pragma once
+
+#include "storm/utility/constants.h"
+
+namespace storm {
+    namespace storage {
+
+        template <typename ValueType>
+        class PostSchedulerChoice {
+
+        public:
+
+            /*!
+             * Creates an undefined scheduler choice
+             */
+            PostSchedulerChoice();
+
+            /*
+             *
+             */
+            void addChoice(uint_fast64_t oldChoiceIndex, uint_fast64_t newChoiceIndex);
+
+            /*
+             *
+             */
+            bool isEmpty() const;
+
+            /*
+             *
+             */
+            std::vector<std::tuple<uint_fast64_t, uint_fast64_t>> const& getChoiceMap() const;
+
+            /*
+             *
+             */
+            std::tuple<uint_fast64_t, uint_fast64_t> const& getChoice(uint_fast64_t choiceIndex) const;
+
+        private:
+            //std::vector<std::tuple<uint_fast64_t, storm::storage::Distribution<ValueType, uint_fast64_t>>> choiceMap;
+            std::vector<std::tuple<uint_fast64_t, uint_fast64_t>> choiceMap;
+        };
+
+        template<typename ValueType>
+        std::ostream& operator<<(std::ostream& out, PostSchedulerChoice<ValueType> const& schedulerChoice);
+    }
+}

src/storm/storage/PreScheduler.cpp

@@ -8,11 +8,36 @@
 namespace storm {
     namespace storage {
         template <typename ValueType>
-        PreScheduler<ValueType>::PreScheduler(uint_fast64_t numberOfModelStates, boost::optional<storm::storage::MemoryStructure> const& memoryStructure) : Scheduler<ValueType>(numberOfModelStates, memoryStructure) {
+        PreScheduler<ValueType>::PreScheduler(uint_fast64_t numberOfModelStates, boost::optional<storm::storage::MemoryStructure> const& memoryStructure) : memoryStructure(memoryStructure) {
+            uint_fast64_t numOfMemoryStates = memoryStructure ? memoryStructure->getNumberOfStates() : 1;
+            schedulerChoices = std::vector<std::vector<PreSchedulerChoice<ValueType>>>(numOfMemoryStates, std::vector<PreSchedulerChoice<ValueType>>(numberOfModelStates));
+            //dontCareStates =  std::vector<storm::storage::BitVector>(numOfMemoryStates, storm::storage::BitVector(numberOfModelStates, false));
+            numOfUndefinedChoices = numOfMemoryStates * numberOfModelStates;
+            numOfDeterministicChoices = 0;
+            numOfDontCareStates = 0;
         }
 
         template <typename ValueType>
-        PreScheduler<ValueType>::PreScheduler(uint_fast64_t numberOfModelStates, boost::optional<storm::storage::MemoryStructure>&& memoryStructure) : Scheduler<ValueType>(numberOfModelStates, std::move(memoryStructure)) {
+        PreScheduler<ValueType>::PreScheduler(uint_fast64_t numberOfModelStates, boost::optional<storm::storage::MemoryStructure>&& memoryStructure) : memoryStructure(std::move(memoryStructure)) {
+        }
+
+        template <typename ValueType>
+        bool PreScheduler<ValueType>::isMemorylessScheduler() const {
+            return getNumberOfMemoryStates() == 1;
+        }
+
+        template <typename ValueType>
+        uint_fast64_t PreScheduler<ValueType>::getNumberOfMemoryStates() const {
+            return memoryStructure ? memoryStructure->getNumberOfStates() : 1;
+        }
+
+        template <typename ValueType>
+        void PreScheduler<ValueType>::setChoice(PreSchedulerChoice<ValueType> const& choice, uint_fast64_t modelState, uint_fast64_t memoryState) {
+            STORM_LOG_ASSERT(memoryState < this->getNumberOfMemoryStates(), "Illegal memory state index");
+            STORM_LOG_ASSERT(modelState < this->schedulerChoices[memoryState].size(), "Illegal model state index");
+
+            auto& schedulerChoice = schedulerChoices[memoryState][modelState];
+            schedulerChoice = choice;
         }
 
         template <typename ValueType>
@@ -67,23 +92,23 @@ namespace storm {
                     }
 
                     // Print choice info
-                    SchedulerChoice<ValueType> const& choice = this->schedulerChoices[memoryState][state];
-                    if (choice.isDefined()) {
+                    PreSchedulerChoice<ValueType> const& choices = this->schedulerChoices[memoryState][state];
+                    if (!choices.isEmpty()) {
                         bool firstChoice = true;
-                        for (auto const& choiceProbPair : choice.getChoiceAsDistribution()) {
+                        for (auto const& choiceProbPair : choices.getChoiceMap()) {
                             if (firstChoice) {
                                 firstChoice = false;
                             } else {
                                 stateString << ";    ";
                             }
-                            stateString << choiceProbPair.second << ": (";
+                            stateString << std::get<0>(choiceProbPair) << ": (";
                             if (choiceOriginsGiven) {
-                                stateString << model->getChoiceOrigins()->getChoiceInfo(model->getTransitionMatrix().getRowGroupIndices()[state] + choiceProbPair.first);
+                                stateString << model->getChoiceOrigins()->getChoiceInfo(model->getTransitionMatrix().getRowGroupIndices()[state] + std::get<1>(choiceProbPair));
                             } else {
-                                stateString << choiceProbPair.first;
+                                stateString << std::get<1>(choiceProbPair);
                             }
                             if (choiceLabelsGiven) {
-                                auto choiceLabels = model->getChoiceLabeling().getLabelsOfChoice(model->getTransitionMatrix().getRowGroupIndices()[state] + choiceProbPair.first);
+                                auto choiceLabels = model->getChoiceLabeling().getLabelsOfChoice(model->getTransitionMatrix().getRowGroupIndices()[state] + std::get<1>(choiceProbPair));
                                 stateString << " {" << boost::join(choiceLabels, ", ") << "}";
                             }
                             stateString << ")";

src/storm/storage/PreScheduler.h

@@ -3,7 +3,7 @@
 #include <cstdint>
 #include <map>
 #include <memory>
-#include "storm/storage/SchedulerChoice.h"
+#include "storm/storage/PreSchedulerChoice.h"
 #include "storm/storage/Scheduler.h"
 #include "storm/logic/ShieldExpression.h"
 
@@ -12,14 +12,11 @@ namespace storm {
 
         /*
          * TODO needs obvious changes in all comment blocks
-         * This class defines which action is chosen in a particular state of a non-deterministic model. More concretely, a scheduler maps a state s to i
-         * if the scheduler takes the i-th action available in s (i.e. the choices are relative to the states).
-         * A Choice can be undefined, deterministic
          */
         template <typename ValueType>
-        class PreScheduler : public Scheduler<ValueType> {
+        class PreScheduler {
         public:
-            typedef uint_fast64_t OldChoice;
+
             /*!
              * Initializes a scheduler for the given number of model states.
              *
@@ -29,14 +26,25 @@ namespace storm {
             PreScheduler(uint_fast64_t numberOfModelStates, boost::optional<storm::storage::MemoryStructure> const& memoryStructure = boost::none);
             PreScheduler(uint_fast64_t numberOfModelStates, boost::optional<storm::storage::MemoryStructure>&& memoryStructure);
 
+            bool isMemorylessScheduler() const;
+            uint_fast64_t getNumberOfMemoryStates() const;
+
+            void setChoice(PreSchedulerChoice<ValueType> const& choice, uint_fast64_t modelState, uint_fast64_t memoryState);
+
             /*!
              * Prints the scheduler to the given output stream.
-             * @param out The output stream
-             * @param model If given, provides additional information for printing (e.g., displaying the state valuations instead of state indices)
-             * @param skipUniqueChoices If true, the (unique) choice for deterministic states (i.e., states with only one enabled choice) is not printed explicitly.
-             *                          Requires a model to be given.
              */
             void printToStream(std::ostream& out, std::shared_ptr<storm::logic::ShieldExpression const> shieldingExpression, std::shared_ptr<storm::models::sparse::Model<ValueType>> model = nullptr, bool skipUniqueChoices = false) const;
+
+        private:
+            boost::optional<storm::storage::MemoryStructure> memoryStructure;
+            std::vector<std::vector<PreSchedulerChoice<ValueType>>> schedulerChoices;
+
+            bool printUndefinedChoices = false;
+
+            uint_fast64_t numOfUndefinedChoices;
+            uint_fast64_t numOfDeterministicChoices;
+            uint_fast64_t numOfDontCareStates;
         };
     }
 }

src/storm/storage/PreSchedulerChoice.cpp

@@ -0,0 +1,56 @@
+#include "storm/storage/PreSchedulerChoice.h"
+
+#include "storm/utility/constants.h"
+#include "storm/utility/macros.h"
+
+#include "storm/exceptions/InvalidOperationException.h"
+#include "storm/adapters/RationalFunctionAdapter.h"
+#include "storm/adapters/RationalNumberAdapter.h"
+
+namespace storm {
+    namespace storage {
+
+        template <typename ValueType>
+        PreSchedulerChoice<ValueType>::PreSchedulerChoice() {
+            // Intentionally left empty
+        }
+
+        template <typename ValueType>
+        void PreSchedulerChoice<ValueType>::addChoice(uint_fast64_t choiceIndex, ValueType probToSatisfy) {
+            choiceMap.emplace_back(probToSatisfy, choiceIndex);
+        }
+
+        template <typename ValueType>
+        std::vector<std::tuple<ValueType, uint_fast64_t>> const& PreSchedulerChoice<ValueType>::getChoiceMap() const {
+            return choiceMap;
+        }
+
+        template <typename ValueType>
+        bool PreSchedulerChoice<ValueType>::isEmpty() const {
+            return choiceMap.size() == 0;
+        }
+
+        template <typename ValueType>
+        std::ostream& operator<<(std::ostream& out, PreSchedulerChoice<ValueType> const& schedulerChoice) {
+            out << schedulerChoice.getChoiceMap().size();
+            if (!schedulerChoice.isEmpty()) {
+                for(auto const& choice : schedulerChoice.getChoiceMap()) {
+                    out << std::get<0>(choice) << ": " << std::get<1>(choice);
+                }
+            } else {
+              out << "undefined";
+            }
+            return out;
+        }
+
+        template class PreSchedulerChoice<double>;
+        template std::ostream& operator<<(std::ostream& out, PreSchedulerChoice<double> const& schedulerChoice);
+        template class PreSchedulerChoice<float>;
+        template std::ostream& operator<<(std::ostream& out, PreSchedulerChoice<float> const& schedulerChoice);
+        template class PreSchedulerChoice<storm::RationalNumber>;
+        template std::ostream& operator<<(std::ostream& out, PreSchedulerChoice<storm::RationalNumber> const& schedulerChoice);
+        template class PreSchedulerChoice<storm::RationalFunction>;
+        template std::ostream& operator<<(std::ostream& out, PreSchedulerChoice<storm::RationalFunction> const& schedulerChoice);
+
+    }
+}

src/storm/storage/PreSchedulerChoice.h

@@ -0,0 +1,42 @@
+#pragma once
+
+#include "storm/utility/constants.h"
+
+namespace storm {
+    namespace storage {
+
+        template <typename ValueType>
+        class PreSchedulerChoice {
+
+        public:
+
+            /*!
+             * Creates an undefined scheduler choice
+             */
+            PreSchedulerChoice();
+
+            /*
+             *
+             */
+            void addChoice(uint_fast64_t choiceIndex, ValueType probToSatisfy);
+
+            /*
+             *
+             */
+            bool isEmpty() const;
+
+            /*
+             *
+             */
+            std::vector<std::tuple<ValueType, uint_fast64_t>> const& getChoiceMap() const;
+
+        private:
+            // For now we only consider shields with deterministic choices.
+            //std::map<ValueType, storm::storage::Distribution<ValueType, uint_fast64_t>> choiceMap;
+            std::vector<std::tuple<ValueType, uint_fast64_t>> choiceMap;
+        };
+
+        template<typename ValueType>
+        std::ostream& operator<<(std::ostream& out, PreSchedulerChoice<ValueType> const& schedulerChoice);
+    }
+}