From 3861c502b4563468a0c9b9f9e7198b66bc262e45 Mon Sep 17 00:00:00 2001
From: Sebastian Junges <sebastian.junges@gmail.com>
Date: Mon, 24 Aug 2020 22:25:00 -0700
Subject: [PATCH] memory unfolder can label states with the memory state

---
 src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp | 12 +++++++++++-
 src/storm-pomdp/transformer/PomdpMemoryUnfolder.h   |  4 +++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp b/src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp
index 086a1e389..6687c30db 100644
--- a/src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp
+++ b/src/storm-pomdp/transformer/PomdpMemoryUnfolder.cpp
@@ -11,7 +11,8 @@ namespace storm {
 
             
             template<typename ValueType>
-            PomdpMemoryUnfolder<ValueType>::PomdpMemoryUnfolder(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory) : pomdp(pomdp), memory(memory) {
+            PomdpMemoryUnfolder<ValueType>::PomdpMemoryUnfolder(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory, bool addMemoryLabels)
+            : pomdp(pomdp), memory(memory), addMemoryLabels(addMemoryLabels) {
                 // intentionally left empty
             }
 
@@ -94,6 +95,15 @@ namespace storm {
                     }
                     labeling.addLabel(labelName, std::move(newStates));
                 }
+                if (addMemoryLabels) {
+                    for (uint64_t memState = 0; memState < memory.getNumberOfStates(); ++memState) {
+                        storm::storage::BitVector newStates(pomdp.getNumberOfStates() * memory.getNumberOfStates(), false);
+                        for (uint64_t modelState = 0; modelState < pomdp.getNumberOfStates(); ++modelState) {
+                            newStates.set(getUnfoldingState(modelState, memState));
+                        }
+                        labeling.addLabel("memstate_"+std::to_string(memState), newStates);
+                    }
+                }
                 return labeling;
             }
         
diff --git a/src/storm-pomdp/transformer/PomdpMemoryUnfolder.h b/src/storm-pomdp/transformer/PomdpMemoryUnfolder.h
index 69b18559a..41f8d0526 100644
--- a/src/storm-pomdp/transformer/PomdpMemoryUnfolder.h
+++ b/src/storm-pomdp/transformer/PomdpMemoryUnfolder.h
@@ -12,7 +12,7 @@ namespace storm {
 
         public:
             
-            PomdpMemoryUnfolder(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory);
+            PomdpMemoryUnfolder(storm::models::sparse::Pomdp<ValueType> const& pomdp, storm::storage::PomdpMemory const& memory, bool addMemoryLabels = false);
             
             std::shared_ptr<storm::models::sparse::Pomdp<ValueType>> transform() const;
 
@@ -33,6 +33,8 @@ namespace storm {
             
             storm::models::sparse::Pomdp<ValueType> const& pomdp;
             storm::storage::PomdpMemory const& memory;
+
+            bool addMemoryLabels;
         };
     }
 }
\ No newline at end of file