From 89fab7269aeaabfad4514bb5b9403dff299c063e Mon Sep 17 00:00:00 2001
From: Dave Parker <dave.x.parker@gmail.com>
Date: Tue, 22 Jul 2014 12:39:37 +0000
Subject: [PATCH] Bug fix in getSuccessorsIterator(s) in SubNondetModel (showed
 up as regression test failure in prism-games-heuristics-merge), plus required
 missing method getSuccessorsIterator(s,i) in NondetModel.

git-svn-id: https://www.prismmodelchecker.org/svn/prism/prism/trunk@8935 bbc10eb1-c90d-0410-af57-cb519fbb1720
---
 prism/src/explicit/MDPSimple.java       |  6 ++++++
 prism/src/explicit/MDPSparse.java       | 12 ++++++++++++
 prism/src/explicit/NondetModel.java     |  7 +++++--
 prism/src/explicit/STPGAbstrSimple.java | 12 ++++++++++++
 prism/src/explicit/SubNondetModel.java  | 16 +++++++++++++++-
 5 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/prism/src/explicit/MDPSimple.java b/prism/src/explicit/MDPSimple.java
index 5c9a75bd..333ed0e2 100644
--- a/prism/src/explicit/MDPSimple.java
+++ b/prism/src/explicit/MDPSimple.java
@@ -555,6 +555,12 @@ public class MDPSimple extends MDPExplicit implements NondetModelSimple
 		return trans.get(s).get(i).containsOneOf(set);
 	}
 
+	@Override
+	public Iterator<Integer> getSuccessorsIterator(final int s, final int i)
+	{
+		return trans.get(s).get(i).getSupport().iterator();
+	}
+
 	// Accessors (for MDP)
 
 	@Override
diff --git a/prism/src/explicit/MDPSparse.java b/prism/src/explicit/MDPSparse.java
index 7ee8e576..9c74b3a7 100644
--- a/prism/src/explicit/MDPSparse.java
+++ b/prism/src/explicit/MDPSparse.java
@@ -535,6 +535,18 @@ public class MDPSparse extends MDPExplicit
 		return false;
 	}
 
+	@Override
+	public Iterator<Integer> getSuccessorsIterator(final int s, final int i)
+	{
+		int start = choiceStarts[rowStarts[s]];
+		int end = choiceStarts[rowStarts[s + 1]];
+		List<Integer> succs = new ArrayList<Integer>();
+		for (int j = start; j < end; j++) {
+			succs.add(cols[j]);
+		}
+		return succs.iterator();
+	}
+
 	// Accessors (for MDP)
 
 	@Override
diff --git a/prism/src/explicit/NondetModel.java b/prism/src/explicit/NondetModel.java
index 3cc9a9cd..44758269 100644
--- a/prism/src/explicit/NondetModel.java
+++ b/prism/src/explicit/NondetModel.java
@@ -27,6 +27,7 @@
 package explicit;
 
 import java.util.BitSet;
+import java.util.Iterator;
 
 import prism.PrismLog;
 import strat.MDStrategy;
@@ -85,9 +86,11 @@ public interface NondetModel extends Model
 	public boolean someSuccessorsInSet(int s, int i, BitSet set);
 	
 	/**
-	 * Get an iterator over the transitions of state s and action i.
+	 * Get an iterator over the successor states from choice {@code i} of state {@code s}.
+	 * @param s The state
+	 * @param i Choice index
 	 */
-	//public Iterator<Entry<Integer, Double>> getTransitionsIterator(int s, int i);
+	public Iterator<Integer> getSuccessorsIterator(int s, int i);
 	
 	/**
 	 * Construct a model that is induced by applying strategy {@code strat} to this model.
diff --git a/prism/src/explicit/STPGAbstrSimple.java b/prism/src/explicit/STPGAbstrSimple.java
index 6e1e6c6c..0a67faa5 100644
--- a/prism/src/explicit/STPGAbstrSimple.java
+++ b/prism/src/explicit/STPGAbstrSimple.java
@@ -499,6 +499,18 @@ public class STPGAbstrSimple extends ModelExplicit implements STPG, NondetModelS
 		return trans.get(s).get(i).containsOneOf(set);
 	}
 	
+	@Override
+	public Iterator<Integer> getSuccessorsIterator(final int s, final int i)
+	{
+		// Need to build set to avoid duplicates
+		// So not necessarily the fastest method to access successors
+		HashSet<Integer> succs = new HashSet<Integer>();
+		for (Distribution distr : trans.get(s).get(i)) {
+			succs.addAll(distr.getSupport());
+		}
+		return succs.iterator();
+	}
+
 	// Accessors (for STPG)
 
 	@Override
diff --git a/prism/src/explicit/SubNondetModel.java b/prism/src/explicit/SubNondetModel.java
index d8773e39..1b2ab2a0 100644
--- a/prism/src/explicit/SubNondetModel.java
+++ b/prism/src/explicit/SubNondetModel.java
@@ -192,7 +192,7 @@ public class SubNondetModel implements NondetModel
 		List<Integer> succ = new ArrayList<Integer>();
 		for (int i = 0; i < model.getNumChoices(s); i++) {
 			if (actions.get(s).get(i)) {
-				Iterator<Integer> it = model.getSuccessorsIterator(s);
+				Iterator<Integer> it = model.getSuccessorsIterator(s, i);
 				while (it.hasNext()) {
 					int j = it.next();
 					succ.add(inverseTranslateState(j));
@@ -386,6 +386,20 @@ public class SubNondetModel implements NondetModel
 		return model.someSuccessorsInSet(s, i, set);
 	}
 
+	@Override
+	public Iterator<Integer> getSuccessorsIterator(int s, int i)
+	{
+		s = translateState(s);
+		i = translateAction(s, i);
+		List<Integer> succ = new ArrayList<Integer>();
+		Iterator<Integer> it = model.getSuccessorsIterator(s, i);
+		while (it.hasNext()) {
+			int j = it.next();
+			succ.add(inverseTranslateState(j));
+		}
+		return succ.iterator();
+	}
+
 	private BitSet translateSet(BitSet set)
 	{
 		BitSet translatedBitSet = new BitSet();