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 + #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 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 blocks; + for (auto const& block : partition.variableBlocks) { + std::vector 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 #include #include +#include #include "src/storage/expressions/Variable.h" #include "src/storage/expressions/Expression.h" @@ -103,6 +104,24 @@ namespace storm { */ std::set getRelatedExpressions(std::set 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 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 expressions; }; + std::ostream& operator<<(std::ostream& out, VariablePartition const& partition); + } } }