From b405a67b547dd7c1e27cb869a2d45c057df9bed2 Mon Sep 17 00:00:00 2001
From: dehnert <dehnert@cs.rwth-aachen.de>
Date: Mon, 5 Sep 2016 13:53:39 +0200
Subject: [PATCH] removed RewardIncrement. fixed PRISM to JANI converter

Former-commit-id: c189fa8e602b44ed01ad818855cc88d2afbbac6e [formerly 63dccbdb95335734d4e090de7e2588fe5c2715e3]
Former-commit-id: 36449defd0df173f97c975aa66dccbf5c6f22fda
---
 src/storage/jani/Automaton.cpp       | 23 +++++++------
 src/storage/jani/Automaton.h         |  8 ++---
 src/storage/jani/EdgeDestination.cpp | 48 +++++++++++++++++++++-------
 src/storage/jani/EdgeDestination.h   | 41 ++++++++++++++++--------
 src/storage/jani/Location.h          |  7 ++--
 src/storage/jani/Model.cpp           | 23 +++++++------
 src/storage/jani/Model.h             | 10 +++---
 src/storage/jani/RewardIncrement.cpp | 19 -----------
 src/storage/jani/RewardIncrement.h   | 39 ----------------------
 src/storage/jani/Variable.h          |  3 +-
 src/storage/jani/VariableSet.cpp     |  9 ++++--
 src/storage/jani/VariableSet.h       |  6 ++--
 src/storage/prism/Program.cpp        | 26 +++++++++------
 13 files changed, 130 insertions(+), 132 deletions(-)
 delete mode 100644 src/storage/jani/RewardIncrement.cpp
 delete mode 100644 src/storage/jani/RewardIncrement.h

diff --git a/src/storage/jani/Automaton.cpp b/src/storage/jani/Automaton.cpp
index 77deacdbf..29d0f838b 100644
--- a/src/storage/jani/Automaton.cpp
+++ b/src/storage/jani/Automaton.cpp
@@ -3,6 +3,7 @@
 #include "src/utility/macros.h"
 #include "src/exceptions/WrongFormatException.h"
 #include "src/exceptions/InvalidArgumentException.h"
+#include "src/exceptions/InvalidTypeException.h"
 
 namespace storm {
     namespace jani {
@@ -58,26 +59,28 @@ namespace storm {
             return name;
         }
 
-        void Automaton::addVariable(Variable const &variable) {
-            if(variable.isBooleanVariable()) {
+        Variable const& Automaton::addVariable(Variable const &variable) {
+            if (variable.isBooleanVariable()) {
                 return addBooleanVariable(variable.asBooleanVariable());
-            } else if(variable.isBoundedIntegerVariable()) {
+            } else if (variable.isBoundedIntegerVariable()) {
                 return addBoundedIntegerVariable(variable.asBoundedIntegerVariable());
-            } else if(variable.isUnboundedIntegerVariable()) {
+            } else if (variable.isUnboundedIntegerVariable()) {
                 return addUnboundedIntegerVariable(variable.asUnboundedIntegerVariable());
+            } else {
+                STORM_LOG_THROW(false, storm::exceptions::InvalidTypeException, "Variable has invalid type.");
             }
         }
         
-        void Automaton::addBooleanVariable(BooleanVariable const& variable) {
-            variables.addBooleanVariable(variable);
+        BooleanVariable const& Automaton::addBooleanVariable(BooleanVariable const& variable) {
+            return variables.addBooleanVariable(variable);
         }
         
-        void Automaton::addBoundedIntegerVariable(BoundedIntegerVariable const& variable) {
-            variables.addBoundedIntegerVariable(variable);
+        BoundedIntegerVariable const& Automaton::addBoundedIntegerVariable(BoundedIntegerVariable const& variable) {
+            return variables.addBoundedIntegerVariable(variable);
         }
 
-        void Automaton::addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable) {
-            variables.addUnboundedIntegerVariable(variable);
+        UnboundedIntegerVariable const& Automaton::addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable) {
+            return variables.addUnboundedIntegerVariable(variable);
         }
 
         VariableSet& Automaton::getVariables() {
diff --git a/src/storage/jani/Automaton.h b/src/storage/jani/Automaton.h
index 1fba060cb..3ea6cea4f 100644
--- a/src/storage/jani/Automaton.h
+++ b/src/storage/jani/Automaton.h
@@ -103,22 +103,22 @@ namespace storm {
             /*!
              * Adds the given variable to this automaton
              */
-            void addVariable(Variable const& variable);
+            Variable const& addVariable(Variable const& variable);
 
             /*!
              * Adds the given Boolean variable to this automaton.
              */
-            void addBooleanVariable(BooleanVariable const& variable);
+            BooleanVariable const& addBooleanVariable(BooleanVariable const& variable);
             
             /*!
              * Adds the given bounded integer variable to this automaton.
              */
-            void addBoundedIntegerVariable(BoundedIntegerVariable const& variable);
+            BoundedIntegerVariable const& addBoundedIntegerVariable(BoundedIntegerVariable const& variable);
             
             /*!
              * Adds the given unbounded integer variable to this automaton.
              */
-            void addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable);
+            UnboundedIntegerVariable const& addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable);
 
             /*!
              * Retrieves the variables of this automaton.
diff --git a/src/storage/jani/EdgeDestination.cpp b/src/storage/jani/EdgeDestination.cpp
index 1166488e3..c58a157e9 100644
--- a/src/storage/jani/EdgeDestination.cpp
+++ b/src/storage/jani/EdgeDestination.cpp
@@ -6,8 +6,17 @@
 namespace storm {
     namespace jani {
         
-        EdgeDestination::EdgeDestination(uint64_t locationIndex, storm::expressions::Expression const& probability, std::vector<Assignment> const& assignments, std::vector<RewardIncrement> const& rewardIncrements) : locationIndex(locationIndex), probability(probability), assignments(assignments), rewardIncrements(rewardIncrements) {
-            sortAssignments();
+        EdgeDestination::EdgeDestination(uint64_t locationIndex, storm::expressions::Expression const& probability, std::vector<Assignment> const& assignments) : locationIndex(locationIndex), probability(probability), assignments(assignments) {
+            for (auto const& assignment : assignments) {
+                if (assignment.isTransientAssignment()) {
+                    transientAssignments.push_back(assignment);
+                } else {
+                    nonTransientAssignments.push_back(assignment);
+                }
+            }
+            sortAssignments(this->assignments);
+            sortAssignments(transientAssignments);
+            sortAssignments(nonTransientAssignments);
         }
         
         void EdgeDestination::addAssignment(Assignment const& assignment) {
@@ -16,11 +25,15 @@ namespace storm {
                 STORM_LOG_THROW(oldAssignment.getExpressionVariable() != assignment.getExpressionVariable(), storm::exceptions::WrongFormatException, "Cannot add assignment '" << assignment << "', because another assignment '" << assignment << "' writes to the same target variable.");
             }
             assignments.push_back(assignment);
-            sortAssignments();
-        }
-        
-        void EdgeDestination::addRewardIncrement(RewardIncrement const& rewardIncrement) {
-            rewardIncrements.push_back(rewardIncrement);
+            sortAssignments(assignments);
+            
+            if (assignment.isTransientAssignment()) {
+                transientAssignments.push_back(assignment);
+                sortAssignments(transientAssignments);
+            } else {
+                nonTransientAssignments.push_back(assignment);
+                sortAssignments(nonTransientAssignments);
+            }
         }
         
         uint64_t EdgeDestination::getLocationIndex() const {
@@ -43,19 +56,30 @@ namespace storm {
             return assignments;
         }
         
-        std::vector<RewardIncrement> const& EdgeDestination::getRewardIncrements() const {
-            return rewardIncrements;
+        std::vector<Assignment>& EdgeDestination::getNonTransientAssignments() {
+            return nonTransientAssignments;
+        }
+        
+        std::vector<Assignment> const& EdgeDestination::getNonTransientAssignments() const {
+            return nonTransientAssignments;
+        }
+        
+        std::vector<Assignment>& EdgeDestination::getTransientAssignments() {
+            return transientAssignments;
         }
         
-        void EdgeDestination::sortAssignments() {
-            std::sort(this->assignments.begin(), this->assignments.end(), [] (storm::jani::Assignment const& assignment1, storm::jani::Assignment const& assignment2) {
+        std::vector<Assignment> const& EdgeDestination::getTransientAssignments() const {
+            return transientAssignments;
+        }
+        
+        void EdgeDestination::sortAssignments(std::vector<Assignment>& assignments) {
+            std::sort(assignments.begin(), assignments.end(), [] (storm::jani::Assignment const& assignment1, storm::jani::Assignment const& assignment2) {
                 bool smaller = assignment1.getExpressionVariable().getType().isBooleanType() && !assignment2.getExpressionVariable().getType().isBooleanType();
                 if (!smaller) {
                     smaller = assignment1.getExpressionVariable() < assignment2.getExpressionVariable();
                 }
                 return smaller;
             });
-
         }
         
     }
diff --git a/src/storage/jani/EdgeDestination.h b/src/storage/jani/EdgeDestination.h
index eb33bb7fc..d22e09ab3 100644
--- a/src/storage/jani/EdgeDestination.h
+++ b/src/storage/jani/EdgeDestination.h
@@ -5,7 +5,6 @@
 #include "src/storage/expressions/Expression.h"
 
 #include "src/storage/jani/Assignment.h"
-#include "src/storage/jani/RewardIncrement.h"
 
 namespace storm {
     namespace jani {
@@ -15,17 +14,12 @@ namespace storm {
             /*!
              * Creates a new edge destination.
              */
-            EdgeDestination(uint64_t locationIndex, storm::expressions::Expression const& probability, std::vector<Assignment> const& assignments = {}, std::vector<RewardIncrement> const& rewardIncrements = {});
+            EdgeDestination(uint64_t locationIndex, storm::expressions::Expression const& probability, std::vector<Assignment> const& assignments = {});
             
             /*!
              * Additionally performs the given assignment when choosing this destination.
              */
             void addAssignment(Assignment const& assignment);
-
-            /*!
-             * Additionally performs the given reward increment when choosing this destination.
-             */
-            void addRewardIncrement(RewardIncrement const& rewardIncrement);
             
             /*!
              * Retrieves the id of the destination location.
@@ -46,23 +40,38 @@ namespace storm {
              * Retrieves the assignments to make when choosing this destination.
              */
             std::vector<Assignment>& getAssignments();
-
+            
             /*!
              * Retrieves the assignments to make when choosing this destination.
              */
             std::vector<Assignment> const& getAssignments() const;
             
             /*!
-             * Retrieves the reward increments to make when choosing this destination.
+             * Retrieves the non-transient assignments to make when choosing this destination.
+             */
+            std::vector<Assignment>& getNonTransientAssignments();
+
+            /*!
+             * Retrieves the non-transient assignments to make when choosing this destination.
+             */
+            std::vector<Assignment> const& getNonTransientAssignments() const;
+
+            /*!
+             * Retrieves the non-transient assignments to make when choosing this destination.
              */
-            std::vector<RewardIncrement> const& getRewardIncrements() const;
+            std::vector<Assignment>& getTransientAssignments();
             
+            /*!
+             * Retrieves the non-transient assignments to make when choosing this destination.
+             */
+            std::vector<Assignment> const& getTransientAssignments() const;
+
         private:
             /*!
              * Sorts the assignments to make all assignments to boolean variables precede all others and order the
              * assignments within one variable group by the expression variables.
              */
-            void sortAssignments();
+            static void sortAssignments(std::vector<Assignment>& assignments);
             
             // The index of the destination location.
             uint64_t locationIndex;
@@ -72,9 +81,13 @@ namespace storm {
             
             // The assignments to make when choosing this destination.
             std::vector<Assignment> assignments;
-            
-            // The increments to rewards to make when choosing this destination.
-            std::vector<RewardIncrement> rewardIncrements;
+
+            // The assignments to make when choosing this destination.
+            std::vector<Assignment> nonTransientAssignments;
+
+            // The assignments to make when choosing this destination.
+            std::vector<Assignment> transientAssignments;
+
         };
         
     }
diff --git a/src/storage/jani/Location.h b/src/storage/jani/Location.h
index 349534d06..461b68e1f 100644
--- a/src/storage/jani/Location.h
+++ b/src/storage/jani/Location.h
@@ -25,7 +25,7 @@ namespace storm {
             std::string const& getName() const;
             
             /*!
-             *
+             * Retrieves the transient assignments of this location.
              */
             std::vector<Assignment> const& getTransientAssignments() const;
             
@@ -33,9 +33,12 @@ namespace storm {
              * Checks whether the location is valid, that is, whether the assignments are indeed all transient assignments.
              */
             void checkValid() const;
+            
         private:
-            // The name of the location.
+            /// The name of the location.
             std::string name;
+            
+            /// The transient assignments made in this location.
             std::vector<Assignment> transientAssignments;
         };
         
diff --git a/src/storage/jani/Model.cpp b/src/storage/jani/Model.cpp
index d76b02a07..8af4dc560 100644
--- a/src/storage/jani/Model.cpp
+++ b/src/storage/jani/Model.cpp
@@ -8,6 +8,7 @@
 #include "src/utility/macros.h"
 #include "src/exceptions/WrongFormatException.h"
 #include "src/exceptions/InvalidOperationException.h"
+#include "src/exceptions/InvalidTypeException.h"
 
 namespace storm {
     namespace jani {
@@ -101,26 +102,28 @@ namespace storm {
             return constants;
         }
 
-        void Model::addVariable(Variable const& variable) {
-            if(variable.isBooleanVariable()) {
+        Variable const& Model::addVariable(Variable const& variable) {
+            if (variable.isBooleanVariable()) {
                 return addBooleanVariable(variable.asBooleanVariable());
-            } else if(variable.isBoundedIntegerVariable()) {
+            } else if (variable.isBoundedIntegerVariable()) {
                 return addBoundedIntegerVariable(variable.asBoundedIntegerVariable());
-            } else if(variable.isUnboundedIntegerVariable()) {
+            } else if (variable.isUnboundedIntegerVariable()) {
                 return addUnboundedIntegerVariable(variable.asUnboundedIntegerVariable());
+            } else {
+                STORM_LOG_THROW(false, storm::exceptions::InvalidTypeException, "Variable has invalid type.");
             }
         }
 
-        void Model::addBooleanVariable(BooleanVariable const& variable) {
-            globalVariables.addBooleanVariable(variable);
+        BooleanVariable const& Model::addBooleanVariable(BooleanVariable const& variable) {
+            return globalVariables.addBooleanVariable(variable);
         }
         
-        void Model::addBoundedIntegerVariable(BoundedIntegerVariable const& variable) {
-            globalVariables.addBoundedIntegerVariable(variable);
+        BoundedIntegerVariable const& Model::addBoundedIntegerVariable(BoundedIntegerVariable const& variable) {
+            return globalVariables.addBoundedIntegerVariable(variable);
         }
         
-        void Model::addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable) {
-            globalVariables.addUnboundedIntegerVariable(variable);
+        UnboundedIntegerVariable const& Model::addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable) {
+            return globalVariables.addUnboundedIntegerVariable(variable);
         }
 
         VariableSet& Model::getGlobalVariables() {
diff --git a/src/storage/jani/Model.h b/src/storage/jani/Model.h
index 681229001..95163fda9 100644
--- a/src/storage/jani/Model.h
+++ b/src/storage/jani/Model.h
@@ -108,24 +108,24 @@ namespace storm {
             Constant const& getConstant(std::string const& name) const;
 
             /*!
-             * Adds the given variable to this model
+             * Adds the given variable to this model.
              */
-            void addVariable(Variable const& variable);
+            Variable const& addVariable(Variable const& variable);
 
             /*!
              * Adds the given boolean variable to this model.
              */
-            void addBooleanVariable(BooleanVariable const& variable);
+            BooleanVariable const& addBooleanVariable(BooleanVariable const& variable);
             
             /*!
              * Adds the given bounded integer variable to this model.
              */
-            void addBoundedIntegerVariable(BoundedIntegerVariable const& variable);
+            BoundedIntegerVariable const& addBoundedIntegerVariable(BoundedIntegerVariable const& variable);
             
             /*!
              * Adds the given unbounded integer variable to this model.
              */
-            void addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable);
+            UnboundedIntegerVariable const& addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable);
 
             /*!
              * Retrieves the variables of this automaton.
diff --git a/src/storage/jani/RewardIncrement.cpp b/src/storage/jani/RewardIncrement.cpp
deleted file mode 100644
index b82743d45..000000000
--- a/src/storage/jani/RewardIncrement.cpp
+++ /dev/null
@@ -1,19 +0,0 @@
-#include "src/storage/jani/RewardIncrement.h"
-
-namespace storm {
-    namespace jani {
-        
-        RewardIncrement::RewardIncrement(uint64_t rewardIndex, storm::expressions::Expression const& value) : rewardIndex(rewardIndex), value(value) {
-            // Intentionally left empty.
-        }
-        
-        uint64_t RewardIncrement::getRewardIndex() const {
-            return rewardIndex;
-        }
-        
-        storm::expressions::Expression const& RewardIncrement::getValue() const {
-            return value;
-        }
-        
-    }
-}
\ No newline at end of file
diff --git a/src/storage/jani/RewardIncrement.h b/src/storage/jani/RewardIncrement.h
deleted file mode 100644
index 383ce9fe0..000000000
--- a/src/storage/jani/RewardIncrement.h
+++ /dev/null
@@ -1,39 +0,0 @@
-#pragma once
-
-#include <cstdint>
-
-#include "src/storage/expressions/Expression.h"
-
-namespace storm {
-    namespace jani {
-        
-        class RewardIncrement {
-        public:
-            /*!
-             * Creates an increment of a reward (given by its index) by the given expression.
-             *
-             * @param rewardIndex The index of the reward to increment.
-             * @param value The expression defining the amount the reward is the incremented.
-             */
-            RewardIncrement(uint64_t rewardIndex, storm::expressions::Expression const& value);
-
-            /*!
-             * Retrieves the index of the reward to increment.
-             */
-            uint64_t getRewardIndex() const;
-            
-            /*!
-             * Retrieves the expression defining the amount by which the reward is to be incremented.
-             */
-            storm::expressions::Expression const& getValue() const;
-            
-        private:
-            // The index of the reward that is to be incremented.
-            uint64_t rewardIndex;
-            
-            // The expression defining the amount the reward is to be incremented.
-            storm::expressions::Expression value;
-        };
-        
-    }
-}
\ No newline at end of file
diff --git a/src/storage/jani/Variable.h b/src/storage/jani/Variable.h
index 884bf8d9b..d5ba1f20b 100644
--- a/src/storage/jani/Variable.h
+++ b/src/storage/jani/Variable.h
@@ -48,8 +48,7 @@ namespace storm {
              * @see hasInitExpression()
              */
             storm::expressions::Expression const& getInitExpression() const;
-
-                        
+            
             // Methods to determine the type of the variable.
             virtual bool isBooleanVariable() const;
             virtual bool isBoundedIntegerVariable() const;
diff --git a/src/storage/jani/VariableSet.cpp b/src/storage/jani/VariableSet.cpp
index 78cfab2cc..fe7ef300c 100644
--- a/src/storage/jani/VariableSet.cpp
+++ b/src/storage/jani/VariableSet.cpp
@@ -74,31 +74,34 @@ namespace storm {
             return detail::ConstVariables<UnboundedIntegerVariable>(unboundedIntegerVariables.begin(), unboundedIntegerVariables.end());
         }
 
-        void VariableSet::addBooleanVariable(BooleanVariable const& variable) {
+        BooleanVariable const& VariableSet::addBooleanVariable(BooleanVariable const& variable) {
             STORM_LOG_THROW(!this->hasVariable(variable.getName()), storm::exceptions::WrongFormatException, "Cannot add variable with name '" << variable.getName() << "', because a variable with that name already exists.");
             std::shared_ptr<BooleanVariable> newVariable = std::make_shared<BooleanVariable>(variable);
             variables.push_back(newVariable);
             booleanVariables.push_back(newVariable);
             nameToVariable.emplace(variable.getName(), variable.getExpressionVariable());
             variableToVariable.emplace(variable.getExpressionVariable(), newVariable);
+            return *newVariable;
         }
         
-        void VariableSet::addBoundedIntegerVariable(BoundedIntegerVariable const& variable) {
+        BoundedIntegerVariable const& VariableSet::addBoundedIntegerVariable(BoundedIntegerVariable const& variable) {
             STORM_LOG_THROW(!this->hasVariable(variable.getName()), storm::exceptions::WrongFormatException, "Cannot add variable with name '" << variable.getName() << "', because a variable with that name already exists.");
             std::shared_ptr<BoundedIntegerVariable> newVariable = std::make_shared<BoundedIntegerVariable>(variable);
             variables.push_back(newVariable);
             boundedIntegerVariables.push_back(newVariable);
             nameToVariable.emplace(variable.getName(), variable.getExpressionVariable());
             variableToVariable.emplace(variable.getExpressionVariable(), newVariable);
+            return *newVariable;
         }
         
-        void VariableSet::addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable) {
+        UnboundedIntegerVariable const& VariableSet::addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable) {
             STORM_LOG_THROW(!this->hasVariable(variable.getName()), storm::exceptions::WrongFormatException, "Cannot add variable with name '" << variable.getName() << "', because a variable with that name already exists.");
             std::shared_ptr<UnboundedIntegerVariable> newVariable = std::make_shared<UnboundedIntegerVariable>(variable);
             variables.push_back(newVariable);
             unboundedIntegerVariables.push_back(newVariable);
             nameToVariable.emplace(variable.getName(), variable.getExpressionVariable());
             variableToVariable.emplace(variable.getExpressionVariable(), newVariable);
+            return *newVariable;
         }
         
         bool VariableSet::hasVariable(std::string const& name) const {
diff --git a/src/storage/jani/VariableSet.h b/src/storage/jani/VariableSet.h
index 2dc5105e3..84e265ed1 100644
--- a/src/storage/jani/VariableSet.h
+++ b/src/storage/jani/VariableSet.h
@@ -100,17 +100,17 @@ namespace storm {
             /*!
              * Adds the given boolean variable to this set.
              */
-            void addBooleanVariable(BooleanVariable const& variable);
+            BooleanVariable const& addBooleanVariable(BooleanVariable const& variable);
 
             /*!
              * Adds the given bounded integer variable to this set.
              */
-            void addBoundedIntegerVariable(BoundedIntegerVariable const& variable);
+            BoundedIntegerVariable const& addBoundedIntegerVariable(BoundedIntegerVariable const& variable);
 
             /*!
              * Adds the given unbounded integer variable to this set.
              */
-            void addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable);
+            UnboundedIntegerVariable const& addUnboundedIntegerVariable(UnboundedIntegerVariable const& variable);
 
             /*!
              * Retrieves whether this variable set contains a variable with the given name.
diff --git a/src/storage/prism/Program.cpp b/src/storage/prism/Program.cpp
index bc50af09b..f946fbb98 100644
--- a/src/storage/prism/Program.cpp
+++ b/src/storage/prism/Program.cpp
@@ -1511,7 +1511,6 @@ namespace storm {
         }
         
         storm::jani::Model Program::toJani(bool allVariablesGlobal) const {
-#if 0
             // Start by creating an empty JANI model.
             storm::jani::ModelType modelType;
             switch (this->getModelType()) {
@@ -1535,14 +1534,20 @@ namespace storm {
                 janiModel.addConstant(storm::jani::Constant(constant.getName(), constant.getExpressionVariable(), constant.isDefined() ? boost::optional<storm::expressions::Expression>(constant.getExpression()) : boost::none));
             }
             
+            // Maintain a mapping from expression variables to JANI variables so we can fill in the correct objects when
+            // creating assignments.
+            std::map<storm::expressions::Variable, std::reference_wrapper<storm::jani::Variable const>> variableToVariableMap;
+            
             // Add all global variables of the PRISM program to the JANI model.
             for (auto const& variable : globalIntegerVariables) {
-                janiModel.addBoundedIntegerVariable(storm::jani::BoundedIntegerVariable(variable.getName(), variable.getExpressionVariable(), variable.getLowerBoundExpression(), variable.getUpperBoundExpression()));
+                storm::jani::BoundedIntegerVariable const& newVariable = janiModel.addBoundedIntegerVariable(storm::jani::BoundedIntegerVariable(variable.getName(), variable.getExpressionVariable(), variable.getLowerBoundExpression(), variable.getUpperBoundExpression()));
+                variableToVariableMap.emplace(variable.getExpressionVariable(), newVariable);
                 storm::expressions::Expression variableInitialExpression = variable.getExpressionVariable() == variable.getInitialValueExpression();
                 globalInitialStatesExpression = globalInitialStatesExpression.isInitialized() ? globalInitialStatesExpression && variableInitialExpression : variableInitialExpression;
             }
             for (auto const& variable : globalBooleanVariables) {
-                janiModel.addBooleanVariable(storm::jani::BooleanVariable(variable.getName(), variable.getExpressionVariable()));
+                storm::jani::BooleanVariable const& newVariable = janiModel.addBooleanVariable(storm::jani::BooleanVariable(variable.getName(), variable.getExpressionVariable()));
+                variableToVariableMap.emplace(variable.getExpressionVariable(), newVariable);
                 storm::expressions::Expression variableInitialExpression = storm::expressions::iff(variable.getExpressionVariable(), variable.getInitialValueExpression());
                 globalInitialStatesExpression = globalInitialStatesExpression.isInitialized() ? globalInitialStatesExpression && variableInitialExpression : variableInitialExpression;
             }
@@ -1591,12 +1596,14 @@ namespace storm {
                     std::set<uint_fast64_t> const& accessingModuleIndices = variablesToAccessingModuleIndices[variable.getExpressionVariable()];
                     // If there is exactly one module reading and writing the variable, we can make the variable local to this module.
                     if (!allVariablesGlobal && accessingModuleIndices.size() == 1) {
-                        automaton.addBoundedIntegerVariable(newIntegerVariable);
+                        storm::jani::BoundedIntegerVariable const& newVariable = automaton.addBoundedIntegerVariable(newIntegerVariable);
+                        variableToVariableMap.emplace(variable.getExpressionVariable(), newVariable);
                         storm::expressions::Expression variableInitialExpression = variable.getExpressionVariable() == variable.getInitialValueExpression();
                         initialStatesExpression = initialStatesExpression.isInitialized() ? initialStatesExpression && variableInitialExpression : variableInitialExpression;
                     } else if (!accessingModuleIndices.empty()) {
                         // Otherwise, we need to make it global.
-                        janiModel.addBoundedIntegerVariable(newIntegerVariable);
+                        storm::jani::BoundedIntegerVariable const& newVariable = janiModel.addBoundedIntegerVariable(newIntegerVariable);
+                        variableToVariableMap.emplace(variable.getExpressionVariable(), newVariable);
                         storm::expressions::Expression variableInitialExpression = variable.getExpressionVariable() == variable.getInitialValueExpression();
                         globalInitialStatesExpression = globalInitialStatesExpression.isInitialized() ? globalInitialStatesExpression && variableInitialExpression : variableInitialExpression;
                     }
@@ -1606,12 +1613,14 @@ namespace storm {
                     std::set<uint_fast64_t> const& accessingModuleIndices = variablesToAccessingModuleIndices[variable.getExpressionVariable()];
                     // If there is exactly one module reading and writing the variable, we can make the variable local to this module.
                     if (!allVariablesGlobal && accessingModuleIndices.size() == 1) {
-                        automaton.addBooleanVariable(newBooleanVariable);
+                        storm::jani::BooleanVariable const& newVariable = automaton.addBooleanVariable(newBooleanVariable);
+                        variableToVariableMap.emplace(variable.getExpressionVariable(), newVariable);
                         storm::expressions::Expression variableInitialExpression = storm::expressions::iff(variable.getExpressionVariable(), variable.getInitialValueExpression());
                         initialStatesExpression = initialStatesExpression.isInitialized() ? initialStatesExpression && variableInitialExpression : variableInitialExpression;
                     } else if (!accessingModuleIndices.empty()) {
                         // Otherwise, we need to make it global.
-                        janiModel.addBooleanVariable(newBooleanVariable);
+                        storm::jani::BooleanVariable const& newVariable = janiModel.addBooleanVariable(newBooleanVariable);
+                        variableToVariableMap.emplace(variable.getExpressionVariable(), newVariable);
                         storm::expressions::Expression variableInitialExpression = storm::expressions::iff(variable.getExpressionVariable(), variable.getInitialValueExpression());
                         globalInitialStatesExpression = globalInitialStatesExpression.isInitialized() ? globalInitialStatesExpression && variableInitialExpression : variableInitialExpression;
                     }
@@ -1640,7 +1649,7 @@ namespace storm {
                     for (auto const& update : command.getUpdates()) {
                         std::vector<storm::jani::Assignment> assignments;
                         for (auto const& assignment : update.getAssignments()) {
-                            assignments.push_back(storm::jani::Assignment(assignment.getVariable().getName(), assignment.getExpression()));
+                            assignments.push_back(storm::jani::Assignment(variableToVariableMap.at(assignment.getVariable()).get(), assignment.getExpression()));
                         }
                         
                         if (rateExpression) {
@@ -1669,7 +1678,6 @@ namespace storm {
             janiModel.finalize();
             
             return janiModel;
-#endif
         }
         
         std::ostream& operator<<(std::ostream& out, Program::ModelType const& type) {