1 files changed, 130 insertions, 53 deletions
diff --git a/guava-testlib/src/com/google/common/collect/testing/google/SortedMapGenerators.java b/guava-testlib/src/com/google/common/collect/testing/google/SortedMapGenerators.java
index e37ae5f..6ce5a77 100644
--- a/guava-testlib/src/com/google/common/collect/testing/google/SortedMapGenerators.java
+++ b/guava-testlib/src/com/google/common/collect/testing/google/SortedMapGenerators.java
@@ -17,19 +17,21 @@
 package com.google.common.collect.testing.google;
 
 import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.collect.testing.Helpers.mapEntry;
 
 import com.google.common.annotations.GwtCompatible;
 import com.google.common.collect.ImmutableSortedMap;
-import com.google.common.collect.Ordering;
+import com.google.common.collect.ImmutableSortedMap.Builder;
 import com.google.common.collect.testing.SampleElements;
-import com.google.common.collect.testing.TestListGenerator;
-import com.google.common.collect.testing.TestStringListGenerator;
-import com.google.common.collect.testing.TestStringSortedMapGenerator;
+import com.google.common.collect.testing.TestCollectionGenerator;
+import com.google.common.collect.testing.TestMapEntrySetGenerator;
+import com.google.common.collect.testing.TestStringSetGenerator;
 
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map.Entry;
-import java.util.SortedMap;
+import java.util.Set;
 
 /**
  * Generators of sorted maps and derived collections.
@@ -38,85 +40,160 @@ import java.util.SortedMap;
  * @author Jesse Wilson
  * @author Jared Levy
  * @author Hayward Chan
- * @author Chris Povirk
- * @author Louis Wasserman
  */
 @GwtCompatible
 public class SortedMapGenerators {
-  public static class ImmutableSortedMapGenerator extends TestStringSortedMapGenerator {
-    @Override public SortedMap<String, String> create(Entry<String, String>[] entries) {
-      ImmutableSortedMap.Builder<String, String> builder = ImmutableSortedMap.naturalOrder();
-      for (Entry<String, String> entry : entries) {
-        checkNotNull(entry);
-        builder.put(entry.getKey(), entry.getValue());
+
+  public static final Comparator<Entry<String, String>> ENTRY_COMPARATOR
+      = new Comparator<Entry<String, String>>() {
+          @Override
+          public int compare(
+              Entry<String, String> o1, Entry<String, String> o2) {
+            return o1.getKey().compareTo(o2.getKey());
+          }
+  };
+
+  public static class ImmutableSortedMapKeySetGenerator
+      extends TestStringSetGenerator {
+
+    @Override protected Set<String> create(String[] elements) {
+      Builder<String, Integer> builder = ImmutableSortedMap.naturalOrder();
+      for (String key : elements) {
+        builder.put(key, 4);
       }
-      return builder.build();
+      return builder.build().keySet();
+    }
+    @Override public List<String> order(List<String> insertionOrder) {
+      Collections.sort(insertionOrder);
+      return insertionOrder;
     }
   }
 
-  public static class ImmutableSortedMapEntryListGenerator
-      implements TestListGenerator<Entry<String, Integer>> {
+  public static class ImmutableSortedMapValuesGenerator
+      implements TestCollectionGenerator<String> {
 
     @Override
-    public SampleElements<Entry<String, Integer>> samples() {
-      return new SampleElements<Entry<String, Integer>>(
-          mapEntry("foo", 5),
-          mapEntry("bar", 3),
-          mapEntry("baz", 17),
-          mapEntry("quux", 1),
-          mapEntry("toaster", -2));
+    public SampleElements<String> samples() {
+      return new SampleElements.Strings();
     }
 
-    @SuppressWarnings("unchecked")
     @Override
-    public Entry<String, Integer>[] createArray(int length) {
-      return new Entry[length];
+    public Collection<String> create(Object... elements) {
+      Builder<Integer, String> builder = ImmutableSortedMap.naturalOrder();
+      for (int i = 0; i < elements.length; i++) {
+        builder.put(i, toStringOrNull(elements[i]));
+      }
+      return builder.build().values();
     }
 
     @Override
-    public Iterable<Entry<String, Integer>> order(List<Entry<String, Integer>> insertionOrder) {
-      return new Ordering<Entry<String, Integer>>() {
-        @Override
-        public int compare(Entry<String, Integer> left, Entry<String, Integer> right) {
-          return left.getKey().compareTo(right.getKey());
-        }
-      }.sortedCopy(insertionOrder);
+    public String[] createArray(int length) {
+      return new String[length];
     }
 
     @Override
-    public List<Entry<String, Integer>> create(Object... elements) {
-      ImmutableSortedMap.Builder<String, Integer> builder = ImmutableSortedMap.naturalOrder();
-      for (Object o : elements) {
-        @SuppressWarnings("unchecked")
-        Entry<String, Integer> entry = (Entry<String, Integer>) o;
-        builder.put(entry);
+    public List<String> order(List<String> insertionOrder) {
+      return insertionOrder;
+    }
+  }
+
+  public static class ImmutableSortedMapSubMapEntryGenerator
+      extends TestMapEntrySetGenerator<String, String> {
+
+    public ImmutableSortedMapSubMapEntryGenerator() {
+      super(new SampleElements.Strings(), new SampleElements.Strings());
+    }
+
+    @Override public Set<Entry<String, String>> createFromEntries(
+        Entry<String, String>[] entries) {
+      Builder<String, String> builder = ImmutableSortedMap.naturalOrder();
+      builder.put(SampleElements.Strings.BEFORE_FIRST, "begin");
+      builder.put(SampleElements.Strings.AFTER_LAST, "end");
+      for (Entry<String, String> entry : entries) {
+        checkNotNull(entry);
+        builder.put(entry.getKey(), entry.getValue());
+      }
+      return builder.build().subMap(SampleElements.Strings.MIN_ELEMENT,
+          SampleElements.Strings.AFTER_LAST).entrySet();
+    }
+    @Override public List<Entry<String, String>> order(
+        List<Entry<String, String>> insertionOrder) {
+      Collections.sort(insertionOrder, ENTRY_COMPARATOR);
+      return insertionOrder;
+    }
+  }
+
+  public static class ImmutableSortedMapHeadMapKeySetGenerator
+      extends TestStringSetGenerator {
+    @Override protected Set<String> create(String[] elements) {
+      Builder<String, Integer> builder = ImmutableSortedMap.naturalOrder();
+      builder.put(SampleElements.Strings.AFTER_LAST, -1);
+      for (String key : elements) {
+        builder.put(key, 4);
       }
-      return builder.build().entrySet().asList();
+      return builder.build().headMap(
+          SampleElements.Strings.AFTER_LAST).keySet();
+    }
+    @Override public List<String> order(List<String> insertionOrder) {
+      Collections.sort(insertionOrder);
+      return insertionOrder;
     }
   }
 
-  public static class ImmutableSortedMapKeyListGenerator extends TestStringListGenerator {
-    @Override protected List<String> create(String[] elements) {
-      ImmutableSortedMap.Builder<String, Integer> builder = ImmutableSortedMap.naturalOrder();
+  public static class ImmutableSortedMapTailMapValuesGenerator
+      implements TestCollectionGenerator<String> {
+
+    @Override
+    public SampleElements<String> samples() {
+      return new SampleElements.Strings();
+    }
+
+    @Override
+    public Collection<String> create(Object... elements) {
+      Builder<Integer, String> builder = ImmutableSortedMap.naturalOrder();
+      builder.put(-1, "begin");
       for (int i = 0; i < elements.length; i++) {
-        builder.put(elements[i], i);
+        builder.put(i, toStringOrNull(elements[i]));
       }
-      return builder.build().keySet().asList();
+      return builder.build().tailMap(0).values();
+    }
+
+    @Override
+    public String[] createArray(int length) {
+      return new String[length];
     }
 
     @Override
     public List<String> order(List<String> insertionOrder) {
-      return Ordering.natural().sortedCopy(insertionOrder);
+      return insertionOrder;
     }
   }
 
-  public static class ImmutableSortedMapValueListGenerator extends TestStringListGenerator {
-    @Override protected List<String> create(String[] elements) {
-      ImmutableSortedMap.Builder<Integer, String> builder = ImmutableSortedMap.naturalOrder();
-      for (int i = 0; i < elements.length; i++) {
-        builder.put(i, elements[i]);
+  public static class ImmutableSortedMapEntrySetGenerator
+      extends TestMapEntrySetGenerator<String, String> {
+
+    public ImmutableSortedMapEntrySetGenerator() {
+      super(new SampleElements.Strings(), new SampleElements.Strings());
+    }
+
+    @Override public Set<Entry<String, String>> createFromEntries(
+        Entry<String, String>[] entries) {
+      Builder<String, String> builder = ImmutableSortedMap.naturalOrder();
+      for (Entry<String, String> entry : entries) {
+        checkNotNull(entry);
+        builder.put(entry.getKey(), entry.getValue());
       }
-      return builder.build().values().asList();
+      return builder.build().entrySet();
+    }
+
+    @Override public List<Entry<String, String>> order(
+        List<Entry<String, String>> insertionOrder) {
+      Collections.sort(insertionOrder, ENTRY_COMPARATOR);
+      return insertionOrder;
     }
   }
+
+  private static String toStringOrNull(Object o) {
+    return (o == null) ? null : o.toString();
+  }
 }