From 256cb82157692196488cf892f101d60c6787a8f4 Mon Sep 17 00:00:00 2001
From: Thomas Knoll <thomas.knolł@student.tugraz.at>
Date: Mon, 7 Aug 2023 16:04:02 +0200
Subject: [PATCH] added functions for querying choice map in shields

---
 src/storm/storage/PostScheduler.cpp |  6 ++++++
 src/storm/storage/PostScheduler.h   | 10 ++++++++++
 src/storm/storage/PreScheduler.cpp  |  8 ++++++++
 src/storm/storage/PreScheduler.h    |  8 ++++++++
 4 files changed, 32 insertions(+)

diff --git a/src/storm/storage/PostScheduler.cpp b/src/storm/storage/PostScheduler.cpp
index 11d0a7aed..3131e21c9 100644
--- a/src/storm/storage/PostScheduler.cpp
+++ b/src/storm/storage/PostScheduler.cpp
@@ -40,6 +40,12 @@ namespace storm {
             schedulerChoiceMapping[memoryState][modelState] = choice;
         }
 
+        template <typename ValueType>
+        PostSchedulerChoice<ValueType> const& PostScheduler<ValueType>::getChoice(uint_fast64_t modelState, uint_fast64_t memoryState) const {
+            STORM_LOG_ASSERT(modelState < schedulerChoiceMapping[memoryState].size(), "Illegal model state index");
+            return schedulerChoiceMapping[memoryState][modelState];
+        }
+
         template <typename ValueType>
         bool PostScheduler<ValueType>::isDeterministicScheduler() const {
             return true;
diff --git a/src/storm/storage/PostScheduler.h b/src/storm/storage/PostScheduler.h
index 99fc3c1b8..44fcb6abd 100644
--- a/src/storm/storage/PostScheduler.h
+++ b/src/storm/storage/PostScheduler.h
@@ -37,6 +37,15 @@ namespace storm {
              */
             void setChoice(PostSchedulerChoice<ValueType> const& newChoice, 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.
+             */
+            PostSchedulerChoice<ValueType> const& getChoice(uint_fast64_t modelState, uint_fast64_t memoryState = 0) const;
+
+
             /*!
              * Is the scheduler defined on the states indicated by the selected-states bitvector?
              */
@@ -97,3 +106,4 @@ namespace storm {
         };
     }
 }
+    
\ No newline at end of file
diff --git a/src/storm/storage/PreScheduler.cpp b/src/storm/storage/PreScheduler.cpp
index bb09a81f1..d78400731 100644
--- a/src/storm/storage/PreScheduler.cpp
+++ b/src/storm/storage/PreScheduler.cpp
@@ -40,6 +40,14 @@ namespace storm {
             schedulerChoice = choice;
         }
 
+        template <typename ValueType>
+        PreSchedulerChoice<ValueType> const& PreScheduler<ValueType>::getChoice(uint_fast64_t modelState, uint_fast64_t memoryState) const {
+            STORM_LOG_ASSERT(memoryState < getNumberOfMemoryStates(), "Illegal memory state index");
+            STORM_LOG_ASSERT(modelState < schedulerChoices[memoryState].size(), "Illegal model state index");
+            return schedulerChoices[memoryState][modelState];
+        }
+
+
         template <typename ValueType>
         void PreScheduler<ValueType>::printToStream(std::ostream& out, std::shared_ptr<storm::logic::ShieldExpression const> shieldingExpression, std::shared_ptr<storm::models::sparse::Model<ValueType>> model, bool skipUniqueChoices) const {
             STORM_LOG_THROW(model == nullptr || model->getNumberOfStates() == this->schedulerChoices.front().size(), storm::exceptions::InvalidOperationException, "The given model is not compatible with this scheduler.");
diff --git a/src/storm/storage/PreScheduler.h b/src/storm/storage/PreScheduler.h
index 5cfa2e8cb..6277c5051 100644
--- a/src/storm/storage/PreScheduler.h
+++ b/src/storm/storage/PreScheduler.h
@@ -31,6 +31,14 @@ namespace storm {
 
             void setChoice(PreSchedulerChoice<ValueType> const& choice, uint_fast64_t modelState, uint_fast64_t memoryState);
 
+            /*!
+             * 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.
+             */
+            PreSchedulerChoice<ValueType> const& getChoice(uint_fast64_t modelState, uint_fast64_t memoryState = 0) const;
+
             /*!
              * Prints the scheduler to the given output stream.
              */