made jit-builder ready to deal with compositions that do not mention all automata, more work on rewards

Former-commit-id: 2faeed82fa [formerly 60157bbec4]
Former-commit-id: ea62a094f1

src/adapters/Z3ExpressionAdapter.cpp

@@ -265,7 +265,6 @@ namespace storm {
                 return this->translateExpression(expression.getVariable());
             }
 
-
             z3::expr Z3ExpressionAdapter::createVariable(storm::expressions::Variable const& variable) {
                 z3::expr z3Variable(context);
                 if (variable.getType().isBooleanType()) {

src/builder/DdJaniModelBuilder.cpp

@@ -1660,7 +1660,11 @@ namespace storm {
         
         template <storm::dd::DdType Type, typename ValueType>
         storm::dd::Bdd<Type> computeInitialStates(storm::jani::Model const& model, CompositionVariables<Type, ValueType> const& variables) {
-            storm::dd::Bdd<Type> initialStates = variables.rowExpressionAdapter->translateExpression(model.getInitialStatesExpression(true)).toBdd();
+            std::vector<std::reference_wrapper<storm::jani::Automaton const>> allAutomata;
+            for (auto const& automaton : model.getAutomata()) {
+                allAutomata.push_back(automaton);
+            }
+            storm::dd::Bdd<Type> initialStates = variables.rowExpressionAdapter->translateExpression(model.getInitialStatesExpression(allAutomata)).toBdd();
             for (auto const& automaton : model.getAutomata()) {
                 storm::dd::Bdd<Type> initialLocationIndices = variables.manager->getBddZero();
                 for (auto const& locationIndex : automaton.getInitialLocationIndices()) {

src/builder/ExplicitModelBuilder.cpp

@@ -13,6 +13,8 @@
 #include "src/settings/modules/CoreSettings.h"
 #include "src/settings/modules/IOSettings.h"
 
+#include "src/builder/RewardModelBuilder.h"
+
 #include "src/generator/PrismNextStateGenerator.h"
 #include "src/generator/JaniNextStateGenerator.h"
 
@@ -28,65 +30,6 @@
 namespace storm {
     namespace builder {
                         
-        /*!
-         * A structure that is used to keep track of a reward model currently being built.
-         */
-        template <typename ValueType>
-        class RewardModelBuilder {
-        public:
-            RewardModelBuilder(storm::generator::RewardModelInformation const& rewardModelInformation) : rewardModelName(rewardModelInformation.getName()), stateRewards(rewardModelInformation.hasStateRewards()), stateActionRewards(rewardModelInformation.hasStateActionRewards()), stateRewardVector(), stateActionRewardVector() {
-                STORM_LOG_THROW(!rewardModelInformation.hasTransitionRewards(), storm::exceptions::InvalidArgumentException, "Unable to treat transition rewards.");
-            }
-            
-            storm::models::sparse::StandardRewardModel<ValueType> build(uint_fast64_t rowCount, uint_fast64_t columnCount, uint_fast64_t rowGroupCount) {
-                boost::optional<std::vector<ValueType>> optionalStateRewardVector;
-                if (hasStateRewards()) {
-                    stateRewardVector.resize(rowGroupCount);
-                    optionalStateRewardVector = std::move(stateRewardVector);
-                }
-                
-                boost::optional<std::vector<ValueType>> optionalStateActionRewardVector;
-                if (hasStateActionRewards()) {
-                    stateActionRewardVector.resize(rowCount);
-                    optionalStateActionRewardVector = std::move(stateActionRewardVector);
-                }
-                
-                return storm::models::sparse::StandardRewardModel<ValueType>(std::move(optionalStateRewardVector), std::move(optionalStateActionRewardVector));
-            }
-            
-            std::string const& getName() const {
-                return rewardModelName;
-            }
-            
-            void addStateReward(ValueType const& value) {
-                stateRewardVector.push_back(value);
-            }
-
-            void addStateActionReward(ValueType const& value) {
-                stateActionRewardVector.push_back(value);
-            }
-            
-            bool hasStateRewards() const {
-                return stateRewards;
-            }
-            
-            bool hasStateActionRewards() const {
-                return stateActionRewards;
-            }
-            
-        private:
-            std::string rewardModelName;
-
-            bool stateRewards;
-            bool stateActionRewards;
-            
-            // The state reward vector.
-            std::vector<ValueType> stateRewardVector;
-            
-            // The state-action reward vector.
-            std::vector<ValueType> stateActionRewardVector;
-        };
-                        
         template <typename ValueType, typename RewardModelType, typename StateType>
         ExplicitModelBuilder<ValueType, RewardModelType, StateType>::ModelComponents::ModelComponents() : transitionMatrix(), stateLabeling(), rewardModels(), choiceLabeling() {
             // Intentionally left empty.

src/builder/RewardModelBuilder.cpp

@@ -0,0 +1,66 @@
+#include "src/builder/RewardModelBuilder.h"
+
+#include "src/adapters/CarlAdapter.h"
+
+#include "src/models/sparse/StandardRewardModel.h"
+
+#include "src/utility/macros.h"
+#include "src/exceptions/InvalidArgumentException.h"
+
+namespace storm {
+    namespace builder {
+        
+        template <typename ValueType>
+        RewardModelBuilder<ValueType>::RewardModelBuilder(RewardModelInformation const& rewardModelInformation) : rewardModelName(rewardModelInformation.getName()), stateRewards(rewardModelInformation.hasStateRewards()), stateActionRewards(rewardModelInformation.hasStateActionRewards()), stateRewardVector(), stateActionRewardVector() {
+            STORM_LOG_THROW(!rewardModelInformation.hasTransitionRewards(), storm::exceptions::InvalidArgumentException, "Unable to treat transition rewards.");
+        }
+        
+        template <typename ValueType>
+        storm::models::sparse::StandardRewardModel<ValueType> RewardModelBuilder<ValueType>::build(uint_fast64_t rowCount, uint_fast64_t columnCount, uint_fast64_t rowGroupCount) {
+            boost::optional<std::vector<ValueType>> optionalStateRewardVector;
+            if (hasStateRewards()) {
+                stateRewardVector.resize(rowGroupCount);
+                optionalStateRewardVector = std::move(stateRewardVector);
+            }
+            
+            boost::optional<std::vector<ValueType>> optionalStateActionRewardVector;
+            if (hasStateActionRewards()) {
+                stateActionRewardVector.resize(rowCount);
+                optionalStateActionRewardVector = std::move(stateActionRewardVector);
+            }
+            
+            return storm::models::sparse::StandardRewardModel<ValueType>(std::move(optionalStateRewardVector), std::move(optionalStateActionRewardVector));
+        }
+        
+        template <typename ValueType>
+        std::string const& RewardModelBuilder<ValueType>::getName() const {
+            return rewardModelName;
+        }
+        
+        template <typename ValueType>
+        void RewardModelBuilder<ValueType>::addStateReward(ValueType const& value) {
+            stateRewardVector.push_back(value);
+        }
+        
+        template <typename ValueType>
+        void RewardModelBuilder<ValueType>::addStateActionReward(ValueType const& value) {
+            stateActionRewardVector.push_back(value);
+        }
+        
+        template <typename ValueType>
+        bool RewardModelBuilder<ValueType>::hasStateRewards() const {
+            return stateRewards;
+        }
+        
+        template <typename ValueType>
+        bool RewardModelBuilder<ValueType>::hasStateActionRewards() const {
+            return stateActionRewards;
+        }
+     
+        template class RewardModelBuilder<double>;
+        template class RewardModelBuilder<storm::RationalNumber>;
+        template class RewardModelBuilder<storm::RationalFunction>;
+        template class RewardModelBuilder<storm::Interval>;
+        
+    }
+}

src/builder/RewardModelBuilder.h

@@ -0,0 +1,51 @@
+#pragma once
+
+#include <vector>
+#include <cstdint>
+
+#include "src/builder/RewardModelInformation.h"
+
+namespace storm {
+    namespace models {
+        namespace sparse {
+            template <typename ValueType>
+            class StandardRewardModel;
+        }
+    }
+    namespace builder {
+        
+        /*!
+         * A structure that is used to keep track of a reward model currently being built.
+         */
+        template <typename ValueType>
+        class RewardModelBuilder {
+        public:
+            RewardModelBuilder(RewardModelInformation const& rewardModelInformation);
+            
+            storm::models::sparse::StandardRewardModel<ValueType> build(uint_fast64_t rowCount, uint_fast64_t columnCount, uint_fast64_t rowGroupCount);
+            
+            std::string const& getName() const;
+            
+            void addStateReward(ValueType const& value);
+            
+            void addStateActionReward(ValueType const& value);
+            
+            bool hasStateRewards() const;
+            
+            bool hasStateActionRewards() const;
+            
+        private:
+            std::string rewardModelName;
+            
+            bool stateRewards;
+            bool stateActionRewards;
+            
+            // The state reward vector.
+            std::vector<ValueType> stateRewardVector;
+            
+            // The state-action reward vector.
+            std::vector<ValueType> stateActionRewardVector;
+        };
+        
+    }
+}

src/builder/RewardModelInformation.cpp

@@ -0,0 +1,39 @@
+#include "src/builder/RewardModelInformation.h"
+
+namespace storm {
+    namespace builder {
+        
+        RewardModelInformation::RewardModelInformation(std::string const& name, bool stateRewards, bool stateActionRewards, bool transitionRewards) : name(name), stateRewards(stateRewards), stateActionRewards(stateActionRewards), transitionRewards(transitionRewards) {
+            // Intentionally left empty.
+        }
+        
+        std::string const& RewardModelInformation::getName() const {
+            return name;
+        }
+        
+        bool RewardModelInformation::hasStateRewards() const {
+            return stateRewards;
+        }
+        
+        bool RewardModelInformation::hasStateActionRewards() const {
+            return stateActionRewards;
+        }
+        
+        bool RewardModelInformation::hasTransitionRewards() const {
+            return transitionRewards;
+        }
+        
+        void RewardModelInformation::setHasStateRewards() {
+            stateRewards = true;
+        }
+        
+        void RewardModelInformation::setHasStateActionRewards() {
+            stateActionRewards = true;
+        }
+        
+        void RewardModelInformation::setHasTransitionRewards() {
+            transitionRewards = true;
+        }
+        
+    }
+}

src/builder/RewardModelInformation.h

@@ -0,0 +1,29 @@
+#pragma once
+
+#include <string>
+
+namespace storm {
+    namespace builder {
+        
+        class RewardModelInformation {
+        public:
+            RewardModelInformation(std::string const& name, bool stateRewards, bool stateActionRewards, bool transitionRewards);
+            
+            std::string const& getName() const;
+            bool hasStateRewards() const;
+            bool hasStateActionRewards() const;
+            bool hasTransitionRewards() const;
+            
+            void setHasStateRewards();
+            void setHasStateActionRewards();
+            void setHasTransitionRewards();
+            
+        private:
+            std::string name;
+            bool stateRewards;
+            bool stateActionRewards;
+            bool transitionRewards;
+        };
+        
+    }
+}

src/builder/jit/Choice.cpp

@@ -36,6 +36,21 @@ namespace storm {
                 distribution.divide(value);
             }
 
+            template <typename IndexType, typename ValueType>
+            void Choice<IndexType, ValueType>::addReward(ValueType const& value) {
+                rewards.push_back(value);
+            }
+            
+            template <typename IndexType, typename ValueType>
+            void Choice<IndexType, ValueType>::addRewards(std::vector<ValueType>&& values) {
+                rewards = std::move(values);
+            }
+            
+            template <typename IndexType, typename ValueType>
+            std::vector<ValueType> const& Choice<IndexType, ValueType>::getRewards() const {
+                return rewards;
+            }
+            
             template <typename IndexType, typename ValueType>
             void Choice<IndexType, ValueType>::compress() {
                 distribution.compress();

src/builder/jit/Choice.h

@@ -18,12 +18,34 @@ namespace storm {
                 Distribution<IndexType, ValueType> const& getDistribution() const;
                 void divideDistribution(ValueType const& value);
                 
+                /*!
+                 * Adds the given value to the reward associated with this choice.
+                 */
+                void addReward(ValueType const& value);
+                
+                /*!
+                 * Adds the given choices rewards to this choice.
+                 */
+                void addRewards(std::vector<ValueType>&& values);
+                
+                /*!
+                 * Retrieves the rewards for this choice.
+                 */
+                std::vector<ValueType> const& getRewards() const;
+                
+                /*!
+                 * Compresses the underlying distribution.
+                 */
                 void compress();
                 
             private:
                 Distribution<IndexType, ValueType>& getDistribution();
 
+                /// The distribution of this choice.
                 Distribution<IndexType, ValueType> distribution;
+                
+                /// The reward values associated with this choice.
+                std::vector<ValueType> rewards;
             };
             
         }

src/builder/jit/ExplicitJitJaniModelBuilder.cpp

@@ -38,15 +38,32 @@ namespace storm {
             
             template <typename ValueType, typename RewardModelType>
             ExplicitJitJaniModelBuilder<ValueType, RewardModelType>::ExplicitJitJaniModelBuilder(storm::jani::Model const& model, storm::builder::BuilderOptions const& options) : options(options), model(model), modelComponentsBuilder(model.getModelType()) {
+                // Register all transient variables as transient.
                 for (auto const& variable : this->model.getGlobalVariables().getTransientVariables()) {
                     transientVariables.insert(variable.getExpressionVariable());
                 }
 
-                for (auto& automaton : this->model.getAutomata()) {
+                // Construct vector of the automata to be put in parallel.
+                storm::jani::Composition const& topLevelComposition = model.getSystemComposition();
+                if (topLevelComposition.isAutomatonComposition()) {
+                    parallelAutomata.push_back(model.getAutomaton(topLevelComposition.asAutomatonComposition().getAutomatonName()));
+                } else {
+                    STORM_LOG_ASSERT(topLevelComposition.isParallelComposition(), "Expected parallel composition.");
+                    storm::jani::ParallelComposition const& parallelComposition = topLevelComposition.asParallelComposition();
+                    
+                    for (auto const& composition : parallelComposition.getSubcompositions()) {
+                        STORM_LOG_ASSERT(composition->isAutomatonComposition(), "Expected flat parallel composition.");
+                        parallelAutomata.push_back(model.getAutomaton(composition->asAutomatonComposition().getAutomatonName()));
+                    }
+                }
+                
                 // FIXME: just for testing purposes.
+                for (auto& automaton : this->model.getAutomata()) {
                     automaton.pushEdgeAssignmentsToDestinations();
+                }
                 
-                    automatonToLocationVariable.emplace(automaton.getName(), model.getManager().declareFreshIntegerVariable(false, automaton.getName() + "_"));
+                for (auto const& automaton : parallelAutomata) {
+                    automatonToLocationVariable.emplace(automaton.get().getName(), model.getManager().declareFreshIntegerVariable(false, automaton.get().getName() + "_"));
                 }
                 
                 // If the program still contains undefined constants and we are not in a parametric setting, assemble an appropriate error message.
@@ -656,7 +673,7 @@ namespace storm {
                 for (auto const& expression : rangeExpressions) {
                     solver->add(expression);
                 }
-                solver->add(model.getInitialStatesExpression(true));
+                solver->add(model.getInitialStatesExpression(parallelAutomata));
                 
                 // Proceed as long as the solver can still enumerate initial states.
                 while (solver->check() == storm::solver::SmtSolver::CheckResult::Sat) {
@@ -683,7 +700,8 @@ namespace storm {
                         storm::expressions::Expression localBlockingExpression = expressionVariable != model->getManager().integer(variableValue);
                         blockingExpression = blockingExpression.isInitialized() ? blockingExpression || localBlockingExpression : localBlockingExpression;
                     }
-                    for (auto const& automaton : this->model.getAutomata()) {
+                    for (auto const& automatonRef : parallelAutomata) {
+                        storm::jani::Automaton const& automaton = automatonRef.get();
                         for (auto const& variable : automaton.getVariables().getBooleanVariables()) {
                             storm::expressions::Variable const& expressionVariable = variable.getExpressionVariable();
                             bool variableValue = model->getBooleanValue(expressionVariable);
@@ -704,8 +722,8 @@ namespace storm {
                     
                     // Gather iterators to the initial locations of all the automata.
                     std::vector<std::set<uint64_t>::const_iterator> initialLocationsIterators;
-                    for (auto const& automaton : this->model.getAutomata()) {
-                        initialLocationsIterators.push_back(automaton.getInitialLocationIndices().cbegin());
+                    for (auto const& automaton : parallelAutomata) {
+                        initialLocationsIterators.push_back(automaton.get().getInitialLocationIndices().cbegin());
                     }
                     
                     // Now iterate through all combinations of initial locations.
@@ -713,7 +731,7 @@ namespace storm {
                         cpptempl::data_list completeAssignment(initialStateAssignment);
                         
                         for (uint64_t index = 0; index < initialLocationsIterators.size(); ++index) {
-                            storm::jani::Automaton const& automaton = this->model.getAutomata()[index];
+                            storm::jani::Automaton const& automaton = parallelAutomata[index].get();
                             if (automaton.getNumberOfLocations() > 1) {
                                 completeAssignment.push_back(generateLocationAssignment(automaton, *initialLocationsIterators[index]));
                             }
@@ -723,8 +741,8 @@ namespace storm {
                         uint64_t index = 0;
                         for (; index < initialLocationsIterators.size(); ++index) {
                             ++initialLocationsIterators[index];
-                            if (initialLocationsIterators[index] == this->model.getAutomata()[index].getInitialLocationIndices().cend()) {
-                                initialLocationsIterators[index] = this->model.getAutomata()[index].getInitialLocationIndices().cbegin();
+                            if (initialLocationsIterators[index] == parallelAutomata[index].get().getInitialLocationIndices().cend()) {
+                                initialLocationsIterators[index] = parallelAutomata[index].get().getInitialLocationIndices().cbegin();
                             } else {
                                 break;
                             }
@@ -857,7 +875,8 @@ namespace storm {
                         nonTransientRealVariables.push_back(newRealVariable);
                     }
                 }
-                for (auto const& automaton : model.getAutomata()) {
+                for (auto const& automatonRef : parallelAutomata) {
+                    storm::jani::Automaton const& automaton = automatonRef.get();
                     for (auto const& variable : automaton.getVariables().getBooleanVariables()) {
                         cpptempl::data_map newBooleanVariable = generateBooleanVariable(variable.asBooleanVariable());
                         if (variable.isTransient()) {
@@ -917,7 +936,8 @@ namespace storm {
             void ExplicitJitJaniModelBuilder<ValueType, RewardModelType>::generateLocations(cpptempl::data_map& modelData) {
                 cpptempl::data_list locations;
 
-                for (auto const& automaton : this->model.getAutomata()) {
+                for (auto const& automatonRef : parallelAutomata) {
+                    storm::jani::Automaton const& automaton = automatonRef.get();
                     cpptempl::data_map locationData;
                     uint64_t locationIndex = 0;
                     for (auto const& location : automaton.getLocations()) {
@@ -1254,13 +1274,7 @@ namespace storm {
                         ++edgeIndex;
                     }
                 } else {
-                    STORM_LOG_ASSERT(topLevelComposition.isParallelComposition(), "Expected parallel composition.");
                     storm::jani::ParallelComposition const& parallelComposition = topLevelComposition.asParallelComposition();
-#ifndef NDEBUG
-                    for (auto const& composition : parallelComposition.getSubcompositions()) {
-                        STORM_LOG_ASSERT(composition->isAutomatonComposition(), "Expected flat parallel composition.");
-                    }
-#endif
                     
                     std::vector<std::set<uint64_t>> synchronizingActions(parallelComposition.getNumberOfSubcompositions());
                     uint64_t synchronizationVectorIndex = 0;

src/builder/jit/ExplicitJitJaniModelBuilder.h

@@ -91,6 +91,7 @@ namespace storm {
                 
                 storm::builder::BuilderOptions options;
                 storm::jani::Model model;
+                std::vector<std::reference_wrapper<storm::jani::Automaton const>> parallelAutomata;
                 
                 ModelComponentsBuilder<IndexType, ValueType> modelComponentsBuilder;
                 typename ExplicitJitJaniModelBuilder<ValueType, RewardModelType>::ImportFunctionType jitBuilderGetFunction;

src/builder/jit/ModelComponentsBuilder.cpp

@@ -7,6 +7,8 @@
 #include "src/models/sparse/MarkovAutomaton.h"
 #include "src/models/sparse/StandardRewardModel.h"
 
+#include "src/builder/RewardModelBuilder.h"
+
 #include "src/settings/SettingsManager.h"
 #include "src/settings/modules/CoreSettings.h"
 
@@ -40,12 +42,30 @@ namespace storm {
                     transitionMatrixBuilder->newRowGroup(currentRow);
                 }
                 if (!behaviour.empty()) {
+                    // Add state reward entries.
+                    auto stateRewardIt = behaviour.getStateRewards().begin();
+                    for (auto& rewardModelBuilder : rewardModelBuilders) {
+                        if (rewardModelBuilder.hasStateRewards()) {
+                            rewardModelBuilder.addStateReward(*stateRewardIt);
+                            ++stateRewardIt;
+                        }
+                    }
+                    
                     for (auto const& choice : behaviour.getChoices()) {
                         // Add the elements to the transition matrix.
                         for (auto const& element : choice.getDistribution()) {
                             transitionMatrixBuilder->addNextValue(currentRow, element.getIndex(), element.getValue());
                         }
                         
+                        // Add state-action reward entries.
+                        auto stateActionRewardIt = choice.getRewards().begin();
+                        for (auto& rewardModelBuilder : rewardModelBuilders) {
+                            if (rewardModelBuilder.hasStateActionRewards()) {
+                                rewardModelBuilder.addStateActionReward(*stateActionRewardIt);
+                                ++stateActionRewardIt;
+                            }
+                        }
+                        
                         // Proceed to next row.
                         ++currentRow;
                     }
@@ -53,8 +73,21 @@ namespace storm {
                     if (behaviour.isExpanded() && dontFixDeadlocks) {
                         STORM_LOG_THROW(false, storm::exceptions::WrongFormatException, "Error while creating sparse matrix from JANI model: found deadlock state and fixing deadlocks was explicitly disabled.");
                     } else {
+                        // Add the self-loop in the transition matrix.
                         transitionMatrixBuilder->addNextValue(currentRow, currentRowGroup, storm::utility::one<ValueType>());
                         ++currentRow;
+                        
+                        // Add the appropriate entries for the reward models.
+                        auto stateRewardIt = behaviour.getStateRewards().begin();
+                        for (auto& rewardModelBuilder : rewardModelBuilders) {
+                            if (rewardModelBuilder.hasStateRewards()) {
+                                rewardModelBuilder.addStateReward(*stateRewardIt);
+                                ++stateRewardIt;
+                            }
+                            if (rewardModelBuilder.hasStateActionRewards()) {
+                                rewardModelBuilder.addStateActionReward(storm::utility::zero<ValueType>());
+                            }
+                        }
                     }
                 }
                 ++currentRowGroup;
@@ -62,23 +95,36 @@ namespace storm {
             
             template <typename IndexType, typename ValueType>
             storm::models::sparse::Model<ValueType, storm::models::sparse::StandardRewardModel<ValueType>>* ModelComponentsBuilder<IndexType, ValueType>::build(IndexType const& stateCount) {
+                storm::storage::SparseMatrix<ValueType> transitionMatrix = this->transitionMatrixBuilder->build();
+                
                 // Start by building the labeling object.
                 storm::models::sparse::StateLabeling stateLabeling(stateCount);
                 for (auto& label : labels) {
                     stateLabeling.addLabel(label.first, std::move(label.second));
                 }
                 
+                // Then build all reward models.
+                std::unordered_map<std::string, storm::models::sparse::StandardRewardModel<ValueType>> rewardModels;
+                for (auto& rewardModelBuilder : rewardModelBuilders) {
+                    rewardModels.emplace(rewardModelBuilder.getName(), rewardModelBuilder.build(transitionMatrix.getRowCount(), transitionMatrix.getColumnCount(), transitionMatrix.getRowGroupCount()));
+                }
+                
                 if (modelType == storm::jani::ModelType::DTMC) {
-                    return new storm::models::sparse::Dtmc<ValueType, storm::models::sparse::StandardRewardModel<ValueType>>(this->transitionMatrixBuilder->build(), std::move(stateLabeling));
+                    return new storm::models::sparse::Dtmc<ValueType, storm::models::sparse::StandardRewardModel<ValueType>>(std::move(transitionMatrix), std::move(stateLabeling), std::move(rewardModels));
                 } else if (modelType == storm::jani::ModelType::CTMC) {
-                    return new storm::models::sparse::Ctmc<ValueType, storm::models::sparse::StandardRewardModel<ValueType>>(this->transitionMatrixBuilder->build(), std::move(stateLabeling));
+                    return new storm::models::sparse::Ctmc<ValueType, storm::models::sparse::StandardRewardModel<ValueType>>(std::move(transitionMatrix), std::move(stateLabeling), std::move(rewardModels));
                 } else if (modelType == storm::jani::ModelType::MDP) {
-                    return new storm::models::sparse::Mdp<ValueType, storm::models::sparse::StandardRewardModel<ValueType>>(this->transitionMatrixBuilder->build(), std::move(stateLabeling));
+                    return new storm::models::sparse::Mdp<ValueType, storm::models::sparse::StandardRewardModel<ValueType>>(std::move(transitionMatrix), std::move(stateLabeling), std::move(rewardModels));
                 } else {
                     return nullptr;
                 }
             }
             
+            template <typename IndexType, typename ValueType>
+            void ModelComponentsBuilder<IndexType, ValueType>::registerRewardModel(RewardModelInformation const& rewardModelInformation) {
+                rewardModelBuilders.emplace_back(rewardModelInformation);
+            }
+            
             template <typename IndexType, typename ValueType>
             void ModelComponentsBuilder<IndexType, ValueType>::registerLabel(std::string const& name, IndexType const& stateCount) {
                 labels.emplace_back(name, storm::storage::BitVector(stateCount));

src/builder/jit/ModelComponentsBuilder.h

@@ -27,6 +27,11 @@ namespace storm {
     }
     
     namespace builder {
+        class RewardModelInformation;
+        
+        template<typename ValueType>
+        class RewardModelBuilder;
+        
         namespace jit {
             
             template <typename IndexType, typename ValueType>
@@ -39,6 +44,8 @@ namespace storm {
 
                 storm::models::sparse::Model<ValueType, storm::models::sparse::StandardRewardModel<ValueType>>* build(IndexType const& stateCount);
                 
+                void registerRewardModel(RewardModelInformation const& rewardModelInformation);
+                
                 void registerLabel(std::string const& name, IndexType const& stateCount);
                 void addLabel(IndexType const& stateId, IndexType const& labelIndex);
                 
@@ -51,6 +58,7 @@ namespace storm {
                 IndexType currentRowGroup;
                 IndexType currentRow;
                 std::unique_ptr<storm::storage::SparseMatrixBuilder<ValueType>> transitionMatrixBuilder;
+                std::vector<storm::builder::RewardModelBuilder<ValueType>> rewardModelBuilders;
                 std::vector<std::pair<std::string, storm::storage::BitVector>> labels;
             };
             

src/builder/jit/StateBehaviour.cpp

@@ -7,14 +7,13 @@ namespace storm {
         namespace jit {
             
             template <typename IndexType, typename ValueType>
-            StateBehaviour<IndexType, ValueType>::StateBehaviour() : compressed(true), expanded(false) {
+            StateBehaviour<IndexType, ValueType>::StateBehaviour() : expanded(false) {
                 // Intentionally left empty.
             }
             
             template <typename IndexType, typename ValueType>
             void StateBehaviour<IndexType, ValueType>::addChoice(Choice<IndexType, ValueType>&& choice) {
                 choices.emplace_back(std::move(choice));
-                
             }
             
             template <typename IndexType, typename ValueType>
@@ -28,6 +27,21 @@ namespace storm {
                 return choices;
             }
             
+            template <typename IndexType, typename ValueType>
+            void StateBehaviour<IndexType, ValueType>::addStateReward(ValueType const& stateReward) {
+                stateRewards.push_back(stateReward);
+            }
+            
+            template <typename IndexType, typename ValueType>
+            void StateBehaviour<IndexType, ValueType>::addStateRewards(std::vector<ValueType>&& stateRewards) {
+                this->stateRewards = std::move(stateRewards);
+            }
+
+            template <typename IndexType, typename ValueType>
+            std::vector<ValueType> const& StateBehaviour<IndexType, ValueType>::getStateRewards() const {
+                return stateRewards;
+            }
+            
             template <typename IndexType, typename ValueType>
             void StateBehaviour<IndexType, ValueType>::reduce(storm::jani::ModelType const& modelType) {
                 if (choices.size() > 1) {
@@ -52,16 +66,6 @@ namespace storm {
                 }
             }
             
-            template <typename IndexType, typename ValueType>
-            void StateBehaviour<IndexType, ValueType>::compress() {
-                if (!compressed) {
-                    for (auto& choice : choices) {
-                        choice.compress();
-                    }
-                    compressed = true;
-                }
-            }
-            
             template <typename IndexType, typename ValueType>
             bool StateBehaviour<IndexType, ValueType>::isExpanded() const {
                 return expanded;
@@ -86,7 +90,6 @@ namespace storm {
             void StateBehaviour<IndexType, ValueType>::clear() {
                 choices.clear();
                 expanded = false;
-                compressed = true;
             }
             
             template class StateBehaviour<uint32_t, double>;

src/builder/jit/StateBehaviour.h

@@ -16,11 +16,24 @@ namespace storm {
                 
                 void addChoice(Choice<IndexType, ValueType>&& choice);
                 Choice<IndexType, ValueType>& addChoice();
-                
                 ContainerType const& getChoices() const;
                 
+                /*!
+                 * Adds the given state reward to the behavior of the state.
+                 */
+                void addStateReward(ValueType const& stateReward);
+                
+                /*!
+                 * Adds the given state rewards to the behavior of the state.
+                 */
+                void addStateRewards(std::vector<ValueType>&& stateRewards);
+                
+                /*!
+                 * Retrieves the rewards for this state.
+                 */
+                std::vector<ValueType> const& getStateRewards() const;
+                
                 void reduce(storm::jani::ModelType const& modelType);
-                void compress();
                 
                 bool isExpanded() const;
                 void setExpanded();
@@ -31,6 +44,10 @@ namespace storm {
                 
             private:
                 ContainerType choices;
+                
+                // The state rewards (under the different, selected reward models) of the state.
+                std::vector<ValueType> stateRewards;
+                
                 bool compressed;
                 bool expanded;
             };

src/generator/JaniNextStateGenerator.cpp

@@ -166,11 +166,15 @@ namespace storm {
             storm::utility::solver::SmtSolverFactory factory;
             std::unique_ptr<storm::solver::SmtSolver> solver = factory.create(model.getExpressionManager());
             
-            std::vector<storm::expressions::Expression> rangeExpressions = model.getAllRangeExpressions();
+            std::vector<std::reference_wrapper<storm::jani::Automaton const>> allAutomata;
+            for (auto const& automaton : model.getAutomata()) {
+                allAutomata.push_back(automaton);
+            }
+            std::vector<storm::expressions::Expression> rangeExpressions = model.getAllRangeExpressions(allAutomata);
             for (auto const& expression : rangeExpressions) {
                 solver->add(expression);
             }
-            solver->add(model.getInitialStatesExpression(true));
+            solver->add(model.getInitialStatesExpression(allAutomata));
             
             // Proceed as long as the solver can still enumerate initial states.
             std::vector<StateType> initialStateIndices;
@@ -613,7 +617,7 @@ namespace storm {
         }
         
         template<typename ValueType, typename StateType>
-        RewardModelInformation JaniNextStateGenerator<ValueType, StateType>::getRewardModelInformation(uint64_t const& index) const {
+        storm::builder::RewardModelInformation JaniNextStateGenerator<ValueType, StateType>::getRewardModelInformation(uint64_t const& index) const {
             return rewardModelInformation[index];
         }
         

src/generator/JaniNextStateGenerator.h

@@ -22,7 +22,7 @@ namespace storm {
             virtual StateBehavior<ValueType, StateType> expand(StateToIdCallback const& stateToIdCallback) override;
             
             virtual std::size_t getNumberOfRewardModels() const override;
-            virtual RewardModelInformation getRewardModelInformation(uint64_t const& index) const override;
+            virtual storm::builder::RewardModelInformation getRewardModelInformation(uint64_t const& index) const override;
                         
             virtual storm::models::sparse::StateLabeling label(storm::storage::BitVectorHashMap<StateType> const& states, std::vector<StateType> const& initialStateIndices = {}, std::vector<StateType> const& deadlockStateIndices = {}) override;
             
@@ -113,7 +113,7 @@ namespace storm {
             std::vector<storm::expressions::Variable> rewardVariables;
             
             /// A vector storing information about the corresponding reward models (variables).
-            std::vector<RewardModelInformation> rewardModelInformation;
+            std::vector<storm::builder::RewardModelInformation> rewardModelInformation;
             
             // A flag that stores whether at least one of the selected reward models has state-action rewards.
             bool hasStateActionRewards;

src/generator/NextStateGenerator.cpp

@@ -15,38 +15,6 @@
 namespace storm {
     namespace generator {
                     
-        RewardModelInformation::RewardModelInformation(std::string const& name, bool stateRewards, bool stateActionRewards, bool transitionRewards) : name(name), stateRewards(stateRewards), stateActionRewards(stateActionRewards), transitionRewards(transitionRewards) {
-            // Intentionally left empty.
-        }
-        
-        std::string const& RewardModelInformation::getName() const {
-            return name;
-        }
-        
-        bool RewardModelInformation::hasStateRewards() const {
-            return stateRewards;
-        }
-        
-        bool RewardModelInformation::hasStateActionRewards() const {
-            return stateActionRewards;
-        }
-        
-        bool RewardModelInformation::hasTransitionRewards() const {
-            return transitionRewards;
-        }
-        
-        void RewardModelInformation::setHasStateRewards() {
-            stateRewards = true;
-        }
-        
-        void RewardModelInformation::setHasStateActionRewards() {
-            stateActionRewards = true;
-        }
-        
-        void RewardModelInformation::setHasTransitionRewards() {
-            transitionRewards = true;
-        }
-        
         template<typename ValueType, typename StateType>
         NextStateGenerator<ValueType, StateType>::NextStateGenerator(storm::expressions::ExpressionManager const& expressionManager, VariableInformation const& variableInformation, NextStateGeneratorOptions const& options) : options(options), expressionManager(expressionManager.getSharedPointer()), variableInformation(variableInformation), evaluator(nullptr), state(nullptr) {
             // Intentionally left empty.

src/generator/NextStateGenerator.h

@@ -11,6 +11,7 @@
 #include "src/storage/expressions/ExpressionEvaluator.h"
 
 #include "src/builder/BuilderOptions.h"
+#include "src/builder/RewardModelInformation.h"
 
 #include "src/generator/VariableInformation.h"
 #include "src/generator/CompressedState.h"
@@ -29,25 +30,6 @@ namespace storm {
             MA
         };
         
-        class RewardModelInformation {
-        public:
-            RewardModelInformation(std::string const& name, bool stateRewards, bool stateActionRewards, bool transitionRewards);
-            
-            std::string const& getName() const;
-            bool hasStateRewards() const;
-            bool hasStateActionRewards() const;
-            bool hasTransitionRewards() const;
-            
-            void setHasStateRewards();
-            void setHasStateActionRewards();
-            void setHasTransitionRewards();
-        private:
-            std::string name;
-            bool stateRewards;
-            bool stateActionRewards;
-            bool transitionRewards;
-        };
-        
         template<typename ValueType, typename StateType = uint32_t>
         class NextStateGenerator {
         public:
@@ -72,7 +54,7 @@ namespace storm {
             bool satisfies(storm::expressions::Expression const& expression) const;
             
             virtual std::size_t getNumberOfRewardModels() const = 0;
-            virtual RewardModelInformation getRewardModelInformation(uint64_t const& index) const = 0;
+            virtual storm::builder::RewardModelInformation getRewardModelInformation(uint64_t const& index) const = 0;
             
             storm::expressions::SimpleValuation toValuation(CompressedState const& state) const;
             

src/generator/PrismNextStateGenerator.cpp

@@ -578,9 +578,9 @@ namespace storm {
         }
         
         template<typename ValueType, typename StateType>
-        RewardModelInformation PrismNextStateGenerator<ValueType, StateType>::getRewardModelInformation(uint64_t const& index) const {
+        storm::builder::RewardModelInformation PrismNextStateGenerator<ValueType, StateType>::getRewardModelInformation(uint64_t const& index) const {
             storm::prism::RewardModel const& rewardModel = rewardModels[index].get();
-            return RewardModelInformation(rewardModel.getName(), rewardModel.hasStateRewards(), rewardModel.hasStateActionRewards(), rewardModel.hasTransitionRewards());
+            return storm::builder::RewardModelInformation(rewardModel.getName(), rewardModel.hasStateRewards(), rewardModel.hasStateActionRewards(), rewardModel.hasTransitionRewards());
         }
                 
         template class PrismNextStateGenerator<double>;

src/generator/PrismNextStateGenerator.h

@@ -23,7 +23,7 @@ namespace storm {
             virtual StateBehavior<ValueType, StateType> expand(StateToIdCallback const& stateToIdCallback) override;
 
             virtual std::size_t getNumberOfRewardModels() const override;
-            virtual RewardModelInformation getRewardModelInformation(uint64_t const& index) const override;
+            virtual storm::builder::RewardModelInformation getRewardModelInformation(uint64_t const& index) const override;
             
             virtual storm::models::sparse::StateLabeling label(storm::storage::BitVectorHashMap<StateType> const& states, std::vector<StateType> const& initialStateIndices = {}, std::vector<StateType> const& deadlockStateIndices = {}) override;
 

src/storage/jani/Model.cpp

@@ -396,7 +396,7 @@ namespace storm {
             return initialStatesRestriction;
         }
         
-        storm::expressions::Expression Model::getInitialStatesExpression(bool includeAutomataInitialStatesExpressions) const {
+        storm::expressions::Expression Model::getInitialStatesExpression(std::vector<std::reference_wrapper<storm::jani::Automaton const>> const& automata) const {
             // Start with the restriction of variables.
             storm::expressions::Expression result = initialStatesRestriction;
             
@@ -412,8 +412,8 @@ namespace storm {
             }
             
             // If we are to include the expressions for the automata, do so now.
-            if (includeAutomataInitialStatesExpressions) {
-                for (auto const& automaton : automata) {
+            for (auto const& automatonReference : automata) {
+                storm::jani::Automaton const& automaton = automatonReference.get();
                 if (!automaton.getVariables().empty()) {
                     storm::expressions::Expression automatonInitialStatesExpression = automaton.getInitialStatesExpression();
                     if (automatonInitialStatesExpression.isInitialized() && !automatonInitialStatesExpression.isTrue()) {
@@ -421,7 +421,6 @@ namespace storm {
                     }
                 }
             }
-            }
             return result;
         }
         
@@ -433,16 +432,23 @@ namespace storm {
             return this->getModelType() == ModelType::DTMC || this->getModelType() == ModelType::MDP;
         }
         
-        std::vector<storm::expressions::Expression> Model::getAllRangeExpressions() const {
+        std::vector<storm::expressions::Expression> Model::getAllRangeExpressions(std::vector<std::reference_wrapper<storm::jani::Automaton const>> const& automata) const {
             std::vector<storm::expressions::Expression> result;
             for (auto const& variable : this->getGlobalVariables().getBoundedIntegerVariables()) {
                 result.push_back(variable.getRangeExpression());
             }
             
-            for (auto const& automaton : automata) {
+            if (automata.empty()) {
+                for (auto const& automaton : this->getAutomata()) {
                     std::vector<storm::expressions::Expression> automatonRangeExpressions = automaton.getAllRangeExpressions();
                     result.insert(result.end(), automatonRangeExpressions.begin(), automatonRangeExpressions.end());
                 }
+            } else {
+                for (auto const& automaton : automata) {
+                    std::vector<storm::expressions::Expression> automatonRangeExpressions = automaton.get().getAllRangeExpressions();
+                    result.insert(result.end(), automatonRangeExpressions.begin(), automatonRangeExpressions.end());
+                }
+            }
             return result;
         }
         

src/storage/jani/Model.h

@@ -289,10 +289,9 @@ namespace storm {
             /*!
              * Retrieves the expression defining the legal initial values of the variables.
              *
-             * @param includeAutomataInitialStatesExpressions If set to true, the expression defines the legal initial
-             * states not only for the global variables but also for the variables of each automaton.
+             * @param automata The resulting expression will also characterize the legal initial states for these automata.
              */
-            storm::expressions::Expression getInitialStatesExpression(bool includeAutomataInitialStatesExpressions = false) const;
+            storm::expressions::Expression getInitialStatesExpression(std::vector<std::reference_wrapper<storm::jani::Automaton const>> const& automata = {}) const;
             
             /*!
              * Determines whether this model is a deterministic one in the sense that each state only has one choice.
@@ -306,8 +305,10 @@ namespace storm {
             
             /*!
              * Retrieves a list of expressions that characterize the legal values of the variables in this model.
+             *
+             * @param automata If provided only range expressions from these automata will be created.
              */
-            std::vector<storm::expressions::Expression> getAllRangeExpressions() const;
+            std::vector<storm::expressions::Expression> getAllRangeExpressions(std::vector<std::reference_wrapper<storm::jani::Automaton const>> const& automata = {}) const;
             
             /*!
              * Retrieves whether this model has the standard composition, that is it composes all automata in parallel