From 5b8f3cd4d6762e184cf7141fcb1df1639f2d80e6 Mon Sep 17 00:00:00 2001
From: Dave Parker <dave.x.parker@gmail.com>
Date: Wed, 24 Jun 2015 09:39:17 +0000
Subject: [PATCH] Add iterator to get actions from a DTMC (not usually stored,
 yet), defaulting them to null, and display these from
 exportToPrismExplicitTra().

git-svn-id: https://www.prismmodelchecker.org/svn/prism/prism/trunk@10076 bbc10eb1-c90d-0410-af57-cb519fbb1720
---
 prism/src/explicit/DTMC.java         |  6 +++
 prism/src/explicit/DTMCExplicit.java | 72 +++++++++++++++++++++++++---
 2 files changed, 71 insertions(+), 7 deletions(-)

diff --git a/prism/src/explicit/DTMC.java b/prism/src/explicit/DTMC.java
index 28721c6d..012d143c 100644
--- a/prism/src/explicit/DTMC.java
+++ b/prism/src/explicit/DTMC.java
@@ -29,6 +29,7 @@ package explicit;
 import java.util.*;
 import java.util.Map.Entry;
 
+import prism.Pair;
 import explicit.rewards.MCRewards;
 
 /**
@@ -46,6 +47,11 @@ public interface DTMC extends Model
 	 */
 	public Iterator<Entry<Integer, Double>> getTransitionsIterator(int s);
 
+	/**
+	 * Get an iterator over the transitions from state s, with their attached actions if present.
+	 */
+	public Iterator<Entry<Integer, Pair<Double, Object>>> getTransitionsAndActionsIterator(int s);
+
 	/**
 	 * Perform a single step of precomputation algorithm Prob0, i.e., for states i in {@code subset},
 	 * set bit i of {@code result} iff there is a transition to a state in {@code u}.
diff --git a/prism/src/explicit/DTMCExplicit.java b/prism/src/explicit/DTMCExplicit.java
index 65b1949a..299f6c1e 100644
--- a/prism/src/explicit/DTMCExplicit.java
+++ b/prism/src/explicit/DTMCExplicit.java
@@ -32,11 +32,12 @@ import java.util.BitSet;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.TreeMap;
+import java.util.Map.Entry;
 
 import common.IterableStateSet;
-
 import explicit.rewards.MCRewards;
 import prism.ModelType;
+import prism.Pair;
 import prism.PrismException;
 import prism.PrismLog;
 import prism.PrismUtils;
@@ -58,21 +59,25 @@ public abstract class DTMCExplicit extends ModelExplicit implements DTMC
 	public void exportToPrismExplicitTra(PrismLog out)
 	{
 		int i;
-		TreeMap<Integer, Double> sorted;
+		TreeMap<Integer, Pair<Double, Object>> sorted;
 		// Output transitions to .tra file
 		out.print(numStates + " " + getNumTransitions() + "\n");
-		sorted = new TreeMap<Integer, Double>();
+		sorted = new TreeMap<Integer, Pair<Double, Object>>();
 		for (i = 0; i < numStates; i++) {
 			// Extract transitions and sort by destination state index (to match PRISM-exported files)
-			Iterator<Map.Entry<Integer, Double>> iter = getTransitionsIterator(i);
+			Iterator<Map.Entry<Integer,Pair<Double, Object>>> iter = getTransitionsAndActionsIterator(i);
 			while (iter.hasNext()) {
-				Map.Entry<Integer, Double> e = iter.next();
+				Map.Entry<Integer, Pair<Double, Object>> e = iter.next();
 				sorted.put(e.getKey(), e.getValue());
 			}
 			// Print out (sorted) transitions
-			for (Map.Entry<Integer, Double> e : sorted.entrySet()) {
+			for (Map.Entry<Integer, Pair<Double, Object>> e : sorted.entrySet()) {
 				// Note use of PrismUtils.formatDouble to match PRISM-exported files
-				out.print(i + " " + e.getKey() + " " + PrismUtils.formatDouble(e.getValue()) + "\n");
+				out.print(i + " " + e.getKey() + " " + PrismUtils.formatDouble(e.getValue().first));
+				Object action = e.getValue().second; 
+				if (action != null && !"".equals(action))
+					out.print(" " + action);
+				out.print("\n");
 			}
 			sorted.clear();
 		}
@@ -132,6 +137,13 @@ public abstract class DTMCExplicit extends ModelExplicit implements DTMC
 	
 	// Accessors (for DTMC)
 	
+	@Override
+	public Iterator<Entry<Integer, Pair<Double, Object>>> getTransitionsAndActionsIterator(int s)
+	{
+		// Default implementation: extend iterator, setting all actions to null
+		return new AddDefaultActionToTransitionsIterator(getTransitionsIterator(s), null);
+	}
+
 	@Override
 	public void mvMult(double vect[], double result[], BitSet subset, boolean complement)
 	{
@@ -169,4 +181,50 @@ public abstract class DTMCExplicit extends ModelExplicit implements DTMC
 			result[s] = mvMultRewSingle(s, vect, mcRewards);
 		}
 	}
+	
+	public class AddDefaultActionToTransitionsIterator implements Iterator<Map.Entry<Integer,Pair<Double,Object>>>
+	{
+		private Iterator<Entry<Integer, Double>> transIter;
+		private Object defaultAction;
+		private Entry<Integer, Double> next;
+
+		public AddDefaultActionToTransitionsIterator(Iterator<Entry<Integer, Double>> transIter, Object defaultAction)
+		{
+			this.transIter = transIter;
+			this.defaultAction = defaultAction;
+		}
+
+		@Override
+		public Entry<Integer, Pair<Double, Object>> next()
+		{
+			next = transIter.next();
+			Entry<Integer, Pair<Double, Object>> next2 = new Entry<Integer, Pair<Double, Object>>()
+			{
+				@Override
+				public Pair<Double, Object> setValue(Pair<Double, Object> value)
+				{
+					return null; // read-only
+				}
+
+				@Override
+				public Pair<Double, Object> getValue()
+				{
+					return new Pair<Double, Object>(next.getValue(), defaultAction);
+				}
+
+				@Override
+				public Integer getKey()
+				{
+					return next.getKey();
+				}
+			};
+			return next2;
+		}
+
+		@Override
+		public boolean hasNext()
+		{
+			return transIter.hasNext();
+		}
+	}
 }