From 2005eb7e73e3a78131126eb452d6743a529ec4d6 Mon Sep 17 00:00:00 2001
From: gereon <gereon.kremer@rwth-aachen.de>
Date: Thu, 14 Mar 2013 20:20:16 +0100
Subject: [PATCH] Added getter routines, so we can retrieve the reward models

---
 src/ir/Program.cpp          | 12 ++++++++++++
 src/ir/Program.h            | 13 +++++++++++++
 src/ir/RewardModel.cpp      |  8 ++++++++
 src/ir/RewardModel.h        | 12 ++++++++++++
 src/ir/StateReward.cpp      |  7 +++++++
 src/ir/StateReward.h        |  8 ++++++++
 src/ir/TransitionReward.cpp |  8 ++++++++
 src/ir/TransitionReward.h   |  6 ++++++
 8 files changed, 74 insertions(+)

diff --git a/src/ir/Program.cpp b/src/ir/Program.cpp
index 7a11f31d2..691951f48 100644
--- a/src/ir/Program.cpp
+++ b/src/ir/Program.cpp
@@ -98,6 +98,18 @@ std::shared_ptr<std::set<uint_fast64_t>> const Program::getModulesByAction(std::
 	}
 }
 
+storm::ir::RewardModel Program::getRewardModel(std::string const & name) const {
+	auto it = this->rewards.find(name);
+	if (it == this->rewards.end()) {
+		// throw some exception here...
+	} else {
+		return it->second;
+	}
+}
+
+std::map<std::string, std::shared_ptr<storm::ir::expressions::BaseExpression>> Program::getLabels() const {
+	return this->labels;
+}
 
 } // namespace ir
 
diff --git a/src/ir/Program.h b/src/ir/Program.h
index 3ce3f3af9..b6daecdd7 100644
--- a/src/ir/Program.h
+++ b/src/ir/Program.h
@@ -93,6 +93,19 @@ public:
 	 */
 	std::shared_ptr<std::set<uint_fast64_t>> const getModulesByAction(std::string const& action) const;
 
+	/*!
+	 * Retrieve reward model with given name.
+	 * @param name Name of the reward model.
+	 * @return Reward model with given name.
+	 */
+	storm::ir::RewardModel getRewardModel(std::string const & name) const;
+
+	/*!
+	 * Retrieves all labels.
+	 * @return All labels.
+	 */
+	std::map<std::string, std::shared_ptr<storm::ir::expressions::BaseExpression>> getLabels() const;
+
 private:
 	// The type of the model.
 	ModelType modelType;
diff --git a/src/ir/RewardModel.cpp b/src/ir/RewardModel.cpp
index d8549db6e..bd3ce1e16 100644
--- a/src/ir/RewardModel.cpp
+++ b/src/ir/RewardModel.cpp
@@ -37,6 +37,14 @@ std::string RewardModel::toString() const {
 	return result.str();
 }
 
+std::vector<storm::ir::StateReward> RewardModel::getStateRewards() const {
+	return this->stateRewards;
+}
+
+std::vector<storm::ir::TransitionReward> RewardModel::getTransitionRewards() const {
+	return this->transitionRewards;
+}
+
 } // namespace ir
 
 } // namespace storm
diff --git a/src/ir/RewardModel.h b/src/ir/RewardModel.h
index fcf8dcf01..5cd7f2767 100644
--- a/src/ir/RewardModel.h
+++ b/src/ir/RewardModel.h
@@ -42,6 +42,18 @@ public:
 	 */
 	std::string toString() const;
 
+	/*!
+	 * Retrieve state rewards.
+	 * @return State rewards.
+	 */
+	std::vector<storm::ir::StateReward> getStateRewards() const;
+
+	/*!
+	 * Retrieve transition rewards.
+	 * @return Transition rewards.
+	 */
+	std::vector<storm::ir::TransitionReward> getTransitionRewards() const;
+
 private:
 	// The name of the reward model.
 	std::string rewardModelName;
diff --git a/src/ir/StateReward.cpp b/src/ir/StateReward.cpp
index 28ed4b72f..93bf738db 100644
--- a/src/ir/StateReward.cpp
+++ b/src/ir/StateReward.cpp
@@ -30,6 +30,13 @@ std::string StateReward::toString() const {
 	return result.str();
 }
 
+double StateReward::getReward(std::pair<std::vector<bool>, std::vector<int_fast64_t>> const * state) const {
+	if (this->statePredicate->getValueAsBool(state)) {
+		return this->rewardValue->getValueAsDouble(state);
+	}
+	return 0;
+}
+
 } // namespace ir
 
 } // namespace storm
diff --git a/src/ir/StateReward.h b/src/ir/StateReward.h
index 0a127fd0c..638b42ed7 100644
--- a/src/ir/StateReward.h
+++ b/src/ir/StateReward.h
@@ -42,6 +42,14 @@ public:
 	 */
 	std::string toString() const;
 
+	/*!
+	 * Returns the reward for the given state.
+	 * It the state fulfills the predicate, the reward value is returned, zero otherwise.
+	 * @param state State object.
+	 * @return Reward for given state.
+	 */
+	double getReward(std::pair<std::vector<bool>, std::vector<int_fast64_t>> const * state) const;
+
 private:
 	// The predicate that characterizes the states that obtain this reward.
 	std::shared_ptr<storm::ir::expressions::BaseExpression> statePredicate;
diff --git a/src/ir/TransitionReward.cpp b/src/ir/TransitionReward.cpp
index 6052f8b37..680df77e6 100644
--- a/src/ir/TransitionReward.cpp
+++ b/src/ir/TransitionReward.cpp
@@ -30,6 +30,14 @@ std::string TransitionReward::toString() const {
 	return result.str();
 }
 
+double TransitionReward::getReward(std::string const & label, std::pair<std::vector<bool>, std::vector<int_fast64_t>> const * state) const {
+	if (this->commandName != label) return 0;
+	if (this->statePredicate->getValueAsBool(state)) {
+		return this->rewardValue->getValueAsDouble(state);
+	}
+	return 0;
+}
+
 } // namespace ir
 
 } // namespace storm
diff --git a/src/ir/TransitionReward.h b/src/ir/TransitionReward.h
index 07e0084e0..89480109a 100644
--- a/src/ir/TransitionReward.h
+++ b/src/ir/TransitionReward.h
@@ -43,6 +43,12 @@ public:
 	 */
 	std::string toString() const;
 
+	/*!
+	 * Retrieves reward for given transition.
+	 * Returns reward value if source state fulfills predicate and the transition is labeled correctly, zero otherwise.
+	 */
+	double getReward(std::string const & label, std::pair<std::vector<bool>, std::vector<int_fast64_t>> const * state) const;
+
 private:
 	// The name of the command this transition-based reward is attached to.
 	std::string commandName;