From f328e69dc4c986a87b1744597858c92567a2b622 Mon Sep 17 00:00:00 2001
From: TimQu <tim.quatmann@cs.rwth-aachen.de>
Date: Wed, 9 Sep 2020 13:23:29 +0200
Subject: [PATCH] Added extract function to FilteredRewardModel

---
 src/storm/utility/FilteredRewardModel.h | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/src/storm/utility/FilteredRewardModel.h b/src/storm/utility/FilteredRewardModel.h
index 5280cfd98..43d4deb2e 100644
--- a/src/storm/utility/FilteredRewardModel.h
+++ b/src/storm/utility/FilteredRewardModel.h
@@ -37,8 +37,27 @@ namespace storm {
             }
 
             RewardModelType const& get() const {
+                STORM_LOG_ASSERT(rewardModel, "tried to get a reward model but none is available. Was it extracted before?");
                 return *rewardModel;
             }
+            
+            /*!
+             * Extracts the reward model. After calling this, this object should not be queried anymore
+             * @return
+             */
+            RewardModelType extract() {
+                STORM_LOG_ASSERT(rewardModel, "tried to extract a reward model but none is available. Was it extracted already before?");
+                RewardModelType result;
+                if (localRewardModel) {
+                    result = std::move(localRewardModel.get());
+                } else {
+                    result = *rewardModel; // Creates a copy
+                }
+                rewardModel = nullptr;
+                return result;
+            }
+            
+            
         private:
                 std::unique_ptr<RewardModelType> localRewardModel;
                 RewardModelType const* rewardModel;