From 1a39ef83808e9d67e27ff5f36c6882e4015f79c3 Mon Sep 17 00:00:00 2001
From: Tobias Winkler <tobias.winkler@cs.rwth-aachen.de>
Date: Fri, 16 Oct 2020 11:04:12 +0200
Subject: [PATCH] minor extension to expressions/prism program

---
 src/storage/expressions.cpp | 1 +
 src/storage/prism.cpp       | 7 +++++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/storage/expressions.cpp b/src/storage/expressions.cpp
index 88a5176..4cc2941 100644
--- a/src/storage/expressions.cpp
+++ b/src/storage/expressions.cpp
@@ -32,6 +32,7 @@ void define_expressions(py::module& m) {
     // Variable
     py::class_<storm::expressions::Variable, std::shared_ptr<storm::expressions::Variable>>(m, "Variable", "Represents a variable")
         .def_property_readonly("name", &storm::expressions::Variable::getName, "Variable name")
+        .def_property_readonly("manager", &storm::expressions::Variable::getManager, "Variable manager")
         .def("has_boolean_type", &storm::expressions::Variable::hasBooleanType, "Check if the variable is of boolean type")
         .def("has_integer_type", &storm::expressions::Variable::hasIntegerType, "Check if the variable is of integer type")
         .def("has_rational_type", &storm::expressions::Variable::hasRationalType, "Check if the variable is of rational type")
diff --git a/src/storage/prism.cpp b/src/storage/prism.cpp
index aab935c..109ee8f 100644
--- a/src/storage/prism.cpp
+++ b/src/storage/prism.cpp
@@ -87,10 +87,12 @@ void define_prism(py::module& m) {
             ;
 
     py::class_<Update> update(m, "PrismUpdate", "An update in a Prism command");
-    update.def_property_readonly("assignments", [](Update const& update) {
+    update.def(py::init<uint_fast64_t, storm::expressions::Expression const&, std::vector<storm::prism::Assignment> const&>())
+           .def_property_readonly("assignments", [](Update const& update) {
                     return update.getAssignments();
                 }, "Assignments in the update")
            .def_property_readonly("probability_expression", &Update::getLikelihoodExpression, "The probability expression for this update")
+           .def_property_readonly("global_index", &Update::getGlobalIndex, "Retrieves the global index of the update, that is, a unique index over all modules")
            .def("substitute", &Update::substitute, "Substitutes all identifiers in the update according to the given map")
            .def("simplify", &Update::simplify, "Simplifies the update in various ways (also removes identity assignments)")
            .def("get_assignment", &Update::getAssignment, py::arg("variable_name"), "Retrieves a reference to the assignment for the variable with the given name")
@@ -99,7 +101,8 @@ void define_prism(py::module& m) {
            ;
 
     py::class_<Assignment> assignment(m, "PrismAssignment", "An assignment in prism");
-    assignment.def_property_readonly("variable", &Assignment::getVariable, "Variable that is updated")
+    assignment.def(py::init<storm::expressions::Variable const&, storm::expressions::Expression const&>())
+            .def_property_readonly("variable", &Assignment::getVariable, "Variable that is updated")
             .def_property_readonly("expression", &Assignment::getExpression, "Expression for the update")
             .def("__str__", &streamToString<Assignment>)
             ;