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

@@ -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

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

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.

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

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();