diff --git a/src/storage/prism/menu_games/VariablePartition.cpp b/src/storage/prism/menu_games/VariablePartition.cpp
index e69de29bb..c6536ba5b 100644
--- a/src/storage/prism/menu_games/VariablePartition.cpp
+++ b/src/storage/prism/menu_games/VariablePartition.cpp
@@ -0,0 +1,126 @@
+#include "src/storage/prism/menu_games/VariablePartition.h"
+
+#include "src/utility/macros.h"
+
+namespace storm {
+ namespace prism {
+ namespace menu_games {
+ VariablePartition::VariablePartition(std::set<storm::expressions::Variable> const& relevantVariables, std::vector<storm::expressions::Expression> const& expressions) : relevantVariables(relevantVariables), expressionBlocks(expressions.size()) {
+ for (auto const& expression : expressions) {
+ this->addExpression(expression);
+ }
+ }
+
+ bool VariablePartition::addExpression(storm::expressions::Expression const& expression) {
+ // Register the expression for all variables that appear in it.
+ std::set<storm::expressions::Variable> expressionVariables = expression.getVariables();
+ for (auto const& variable : expressionVariables) {
+ variableToExpressionsMapping[variable].insert(this->expressions.size());
+ }
+
+ // Add aexpression and relate all the appearing variables.
+ this->expressions.push_back(expression);
+ return this->relate(expressionVariables);
+ }
+
+ bool VariablePartition::areRelated(storm::expressions::Variable const& firstVariable, storm::expressions::Variable const& secondVariable) {
+ return getBlockIndexOfVariable(firstVariable) == getBlockIndexOfVariable(secondVariable);
+ }
+
+ bool VariablePartition::relate(storm::expressions::Variable const& firstVariable, storm::expressions::Variable const& secondVariable) {
+ return this->relate({firstVariable, secondVariable});
+ }
+
+ bool VariablePartition::relate(std::set<storm::expressions::Variable> const& variables) {
+ // Determine all blocks that need to be merged.
+ std::set<uint_fast64_t> blocksToMerge;
+ for (auto const& variable : variables) {
+ blocksToMerge.insert(getBlockIndexOfVariable(variable));
+ }
+
+ STORM_LOG_ASSERT(!blocksToMerge.empty(), "Found no blocks to merge.");
+
+ // If we found a single block only, there is nothing to do.
+ if (blocksToMerge.size() == 1) {
+ return false;
+ }
+
+ this->mergeBlocks(blocksToMerge);
+ return true;
+ }
+
+ void VariablePartition::mergeBlocks(std::set<uint_fast64_t> const& blocksToMerge) {
+ // Determine which block to keep (to merge the other blocks into).
+ uint_fast64_t blockToKeep = *blocksToMerge.begin();
+
+ // Merge all blocks into the block to keep.
+ std::vector<std::set<storm::expressions::Variable>> newVariableBlocks;
+ std::vector<std::set<uint_fast64_t>> newExpressionBlocks;
+ for (uint_fast64_t blockIndex = 0; blockIndex < variableBlocks.size(); ++blockIndex) {
+
+ // If the block is the one into which the others are to be merged, we do so.
+ if (blockIndex == blockToKeep) {
+ for (auto const& blockToMerge : blocksToMerge) {
+ if (blockToMerge == blockToKeep) {
+ continue;
+ }
+
+ for (auto const& variable : variableBlocks[blockToMerge]) {
+ variableBlocks[blockToKeep].insert(variable);
+ }
+
+ for (auto const& expression : expressionBlocks[blockToMerge]) {
+ expressionBlocks[blockToKeep].insert(expression);
+ }
+ }
+ }
+
+ // Adjust the mapping for all variables we are moving to the new block.
+ for (auto const& variable : variableBlocks[blockIndex]) {
+ variableToBlockMapping[variable] = newVariableBlocks.size();
+ }
+
+ // Move the current block to the new partition.
+ newVariableBlocks.emplace_back(std::move(variableBlocks[blockIndex]));
+ newExpressionBlocks.emplace_back(std::move(expressionBlocks[blockIndex]));
+ }
+ }
+
+ std::set<storm::expressions::Variable> const& VariablePartition::getBlockOfVariable(storm::expressions::Variable const& variable) const {
+ return variableBlocks[getBlockIndexOfVariable(variable)];
+ }
+
+ uint_fast64_t VariablePartition::getNumberOfBlocks() const {
+ return this->variableBlocks.size();
+ }
+
+ std::set<storm::expressions::Variable> const& VariablePartition::getVariableBlockWithIndex(uint_fast64_t blockIndex) const {
+ return this->variableBlocks[blockIndex];
+ }
+
+ uint_fast64_t VariablePartition::getBlockIndexOfVariable(storm::expressions::Variable const& variable) const {
+ STORM_LOG_ASSERT(this->relevantVariables.find(variable) != this->relevantVariables.end(), "Illegal variable '" << variable.getName() << "' for partition.");
+ return this->variableToBlockMapping.find(variable)->second;
+ }
+
+ std::set<uint_fast64_t> const& VariablePartition::getRelatedExpressions(storm::expressions::Variable const& variable) const {
+ return this->expressionBlocks[getBlockIndexOfVariable(variable)];
+ }
+
+ std::set<uint_fast64_t> VariablePartition::getRelatedExpressions(std::set<storm::expressions::Variable> const& variables) const {
+ // Start by determining the indices of all expression blocks that are related to any of the variables.
+ std::set<uint_fast64_t> relatedExpressionBlockIndices;
+ for (auto const& variable : variables) {
+ relatedExpressionBlockIndices.insert(getBlockIndexOfVariable(variable));
+ }
+
+ // Then join the expressions of these blocks and return the result.
+ std::set<uint_fast64_t> result;
+ for (auto const& blockIndex : relatedExpressionBlockIndices) {
+ result.insert(expressionBlocks[blockIndex].begin(), expressionBlocks[blockIndex].end());
+ }
+ return result;
+ }
+ }
+ }
+}
\ 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 e69de29bb..10b605cf8 100644
--- a/src/storage/prism/menu_games/VariablePartition.h
+++ b/src/storage/prism/menu_games/VariablePartition.h
@@ -0,0 +1,137 @@
+#ifndef STORM_STORAGE_PRISM_MENU_GAMES_VARIABLEPARTITION_H_
+#define STORM_STORAGE_PRISM_MENU_GAMES_VARIABLEPARTITION_H_
+
+#include <unordered_map>
+#include <set>
+#include <vector>
+
+#include "src/storage/expressions/Variable.h"
+#include "src/storage/expressions/Expression.h"
+
+namespace storm {
+ namespace prism {
+ namespace menu_games {
+
+ class VariablePartition{
+ public:
+ /*!
+ * Constructs a new variable partition.
+ *
+ * @param relevantVariables The variables of this partition.
+ * @param expressions The (initial) expressions that define the partition.
+ */
+ VariablePartition(std::set<storm::expressions::Variable> const& relevantVariables, std::vector<storm::expressions::Expression> const& expressions = std::vector<storm::expressions::Expression>());
+
+ /*!
+ * Adds the expression and therefore indirectly may cause blocks of variables to be merged.
+ *
+ * @param expression The expression to add.
+ * @return True iff the partition changed.
+ */
+ bool addExpression(storm::expressions::Expression const& expression);
+
+ /*!
+ * Retrieves whether the two given variables are in the same block of the partition.
+ *
+ * @param firstVariable The first variable.
+ * @param secondVariable The second variable.
+ * @return True iff the two variables are in the same block.
+ */
+ bool areRelated(storm::expressions::Variable const& firstVariable, storm::expressions::Variable const& secondVariable);
+
+ /*!
+ * Places the given variables in the same block of the partition and performs the implied merges.
+ *
+ * @param firstVariable The first variable.
+ * @param secondVariable The second variable.
+ * @return True iff the partition changed.
+ */
+ bool relate(storm::expressions::Variable const& firstVariable, storm::expressions::Variable const& secondVariable);
+
+ /*!
+ * Places the given variables in the same block of the partition and performs the implied merges.
+ *
+ * @param variables The variables to relate.
+ * @return True iff the partition changed.
+ */
+ bool relate(std::set<storm::expressions::Variable> const& variables);
+
+ /*!
+ * Retrieves the block of related variables of the given variable.
+ *
+ * @param variable The variable whose block to retrieve.
+ * @return The block of the variable.
+ */
+ std::set<storm::expressions::Variable> const& getBlockOfVariable(storm::expressions::Variable const& variable) const;
+
+ /*!
+ * Retrieves the block index of the given variable.
+ *
+ * @param variable The variable for which to retrieve the block.
+ * @return The block index of the given variable.
+ */
+ uint_fast64_t getBlockIndexOfVariable(storm::expressions::Variable const& variable) const;
+
+ /*!
+ * Retrieves the number of blocks of the varible partition.
+ *
+ * @return The number of blocks in this partition.
+ */
+ uint_fast64_t getNumberOfBlocks() const;
+
+ /*!
+ * Retrieves the block with the given index.
+ *
+ * @param blockIndex The index of the block to retrieve.
+ * @return The block with the given index.
+ */
+ std::set<storm::expressions::Variable> const& getVariableBlockWithIndex(uint_fast64_t blockIndex) const;
+
+ /*!
+ * Retrieves the indices of the expressions related to the given variable.
+ *
+ * @param variable The variable for which to retrieve the related expressions.
+ * @return The related expressions.
+ */
+ std::set<uint_fast64_t> const& getRelatedExpressions(storm::expressions::Variable const& variable) const;
+
+ /*!
+ * Retrieves the indices of the expressions related to any of the given variables.
+ *
+ * @param variables The variables for which to retrieve the related expressions.
+ * @return The related expressions.
+ */
+ std::set<uint_fast64_t> getRelatedExpressions(std::set<storm::expressions::Variable> const& variables) const;
+
+ private:
+ /*!
+ * Merges the blocks with the given indices.
+ *
+ * @param blocksToMerge The indices of the blocks to merge.
+ */
+ void mergeBlocks(std::set<uint_fast64_t> const& blocksToMerge);
+
+ // The set of variables relevant for this partition.
+ std::set<storm::expressions::Variable> relevantVariables;
+
+ // A mapping from variables to their blocks.
+ std::unordered_map<storm::expressions::Variable, uint_fast64_t> variableToBlockMapping;
+
+ // The variable blocks of the partition.
+ std::vector<std::set<storm::expressions::Variable>> variableBlocks;
+
+ // The expression blocks of the partition.
+ std::vector<std::set<uint_fast64_t>> expressionBlocks;
+
+ // A mapping from variables to the indices of all expressions they appear in.
+ std::unordered_map<storm::expressions::Variable, std::set<uint_fast64_t>> variableToExpressionsMapping;
+
+ // The vector of all expressions.
+ std::vector<storm::expressions::Expression> expressions;
+ };
+
+ }
+ }
+}
+
+#endif /* STORM_STORAGE_PRISM_MENU_GAMES_VARIABLEPARTITION_H_ */
\ No newline at end of file