From 36e8359efaf730ebd805d1b1311462c44b0d91fc Mon Sep 17 00:00:00 2001
From: dehnert <dehnert@cs.rwth-aachen.de>
Date: Thu, 10 Sep 2015 22:50:38 +0200
Subject: [PATCH] added some useful functions to variable partition

Former-commit-id: b290129abb17db1f283de668df8105178d624861
---
 .../prism/menu_games/VariablePartition.cpp    | 37 +++++++++++++++----
 .../prism/menu_games/VariablePartition.h      | 21 +++++++++++
 2 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/src/storage/prism/menu_games/VariablePartition.cpp b/src/storage/prism/menu_games/VariablePartition.cpp
index c6536ba5b..53c0d2a61 100644
--- a/src/storage/prism/menu_games/VariablePartition.cpp
+++ b/src/storage/prism/menu_games/VariablePartition.cpp
@@ -1,5 +1,7 @@
 #include "src/storage/prism/menu_games/VariablePartition.h"
 
+#include <boost/algorithm/string/join.hpp>
+
 #include "src/utility/macros.h"
 
 namespace storm {
@@ -65,13 +67,8 @@ namespace storm {
                                 continue;
                             }
                             
-                            for (auto const& variable : variableBlocks[blockToMerge]) {
-                                variableBlocks[blockToKeep].insert(variable);
-                            }
-                            
-                            for (auto const& expression : expressionBlocks[blockToMerge]) {
-                                expressionBlocks[blockToKeep].insert(expression);
-                            }
+                            variableBlocks[blockToKeep].insert(variableBlocks[blockToMerge].begin(), variableBlocks[blockToMerge].end());
+                            expressionBlocks[blockToKeep].insert(expressionBlocks[blockToMerge].begin(), expressionBlocks[blockToMerge].end());
                         }
                     }
                     
@@ -121,6 +118,32 @@ namespace storm {
                 }
                 return result;
             }
+            
+            std::set<uint_fast64_t> const& VariablePartition::getExpressionsUsingVariable(storm::expressions::Variable const& variable) const {
+                STORM_LOG_ASSERT(this->relevantVariables.find(variable) != this->relevantVariables.end(), "Illegal variable '" << variable.getName() << "' for partition.");
+                return this->variableToExpressionsMapping.find(variable)->second;
+            }
+            
+            storm::expressions::Expression const& VariablePartition::getExpression(uint_fast64_t expressionIndex) const {
+                return this->expressions[expressionIndex];
+            }
+
+            std::ostream& operator<<(std::ostream& out, VariablePartition const& partition) {
+                std::vector<std::string> blocks;
+                for (auto const& block : partition.variableBlocks) {
+                    std::vector<std::string> variablesInBlock;
+                    for (auto const& variable : block) {
+                        variablesInBlock.push_back(variable.getName());
+                    }
+                    blocks.push_back("[" + boost::algorithm::join(variablesInBlock, ", ") + "]");
+                }
+                
+                out << "{";
+                out << boost::join(blocks, ", ");
+                out << "}";
+                return out;
+            }
+            
         }
     }
 }
\ No newline at end of file
diff --git a/src/storage/prism/menu_games/VariablePartition.h b/src/storage/prism/menu_games/VariablePartition.h
index 10b605cf8..46651fe8f 100644
--- a/src/storage/prism/menu_games/VariablePartition.h
+++ b/src/storage/prism/menu_games/VariablePartition.h
@@ -4,6 +4,7 @@
 #include <unordered_map>
 #include <set>
 #include <vector>
+#include <ostream>
 
 #include "src/storage/expressions/Variable.h"
 #include "src/storage/expressions/Expression.h"
@@ -103,6 +104,24 @@ namespace storm {
                  */
                 std::set<uint_fast64_t> getRelatedExpressions(std::set<storm::expressions::Variable> const& variables) const;
                 
+                /*!
+                 * Retrieves the indices of the expressions in which the given variable appears.
+                 *
+                 * @param variable The variable for which to retrieve the expressions.
+                 * @return The indices of all expressions using the given variable.
+                 */
+                std::set<uint_fast64_t> const& getExpressionsUsingVariable(storm::expressions::Variable const& variable) const;
+                
+                /*!
+                 * Retrieves the expression with the given index.
+                 *
+                 * @param expressionIndex The index of the expression to retrieve.
+                 * @return The corresponding expression.
+                 */
+                storm::expressions::Expression const& getExpression(uint_fast64_t expressionIndex) const;
+                
+                friend std::ostream& operator<<(std::ostream& out, VariablePartition const& partition);
+                
             private:
                 /*!
                  * Merges the blocks with the given indices.
@@ -130,6 +149,8 @@ namespace storm {
                 std::vector<storm::expressions::Expression> expressions;
             };
             
+            std::ostream& operator<<(std::ostream& out, VariablePartition const& partition);
+            
         }
     }
 }