From cbc53e1ff3e377e340c0d1d951f874001fc80a61 Mon Sep 17 00:00:00 2001
From: Joachim Klein <klein@tcs.inf.tu-dresden.de>
Date: Fri, 12 Oct 2018 14:26:08 +0200
Subject: [PATCH] imported patch common-explicit.StateValues.partition.patch

---
 prism/src/explicit/StateValues.java | 44 +++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/prism/src/explicit/StateValues.java b/prism/src/explicit/StateValues.java
index f50c7cd8..d983a128 100644
--- a/prism/src/explicit/StateValues.java
+++ b/prism/src/explicit/StateValues.java
@@ -33,6 +33,8 @@ import java.io.IOException;
 import java.util.BitSet;
 import java.util.List;
 import java.util.function.Predicate;
+import java.util.Map;
+import java.util.TreeMap;
 
 import parser.State;
 import parser.ast.ExpressionBinaryOp;
@@ -1841,6 +1843,48 @@ public class StateValues implements StateVector
 		}
 	}
 
+	/** Partition the state values by (boolean) value, returns a Map: Value -&gt; States-with-that-value */
+	public Map<Boolean, BitSet> partitionBool() {
+		if (!type.equals(TypeBool.getInstance()))
+			throw new UnsupportedOperationException("Can not invoke paritionBool() for StateValues with "+type+" elements.");
+		return this.<Boolean>partition();
+	}
+
+	/** Partition the state values by (double) value, returns a Map: Value -&gt; States-with-that-value */
+	public Map<Double, BitSet> partitionDouble() {
+		if (!type.equals(TypeDouble.getInstance()))
+			throw new UnsupportedOperationException("Can not invoke paritionDouble() for StateValues with "+type+" elements.");
+		return this.<Double>partition();
+	}
+
+	/** Partition the state values by (integer) value, returns a Map: Value -&gt; States-with-that-value */
+	public Map<Integer, BitSet> partitionInt() {
+		if (!type.equals(TypeInt.getInstance()))
+			throw new UnsupportedOperationException("Can not invoke paritionInt() for StateValues with "+type+" elements.");
+		return this.<Integer>partition();
+	}
+
+	@SuppressWarnings("unchecked")
+	/**
+	 * Helper: Partition the state values by value.
+	 * Returns a Map: Value -&gt; States-with-that-value
+	 * @param <T> the type of the values
+	 */
+	private <T extends Comparable<T>> Map<T, BitSet> partition() {
+		Map<T, BitSet> result = new TreeMap<T, BitSet>();
+
+		for (int i = 0; i < getSize(); i++) {
+			BitSet set = result.get(getValue(i));
+			if (set == null) {
+				set = new BitSet();
+				result.put((T) getValue(i), set);
+			}
+			set.set(i);
+		}
+
+		return result;
+	}
+
 	/**
 	 * Make a (deep) copy of this vector
 	 */