An implementation of the Cartesian product of ordered collections, in Java.

Cartesian.java

/* Copyright 2012 LinkedIn Corp.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 */

import com.google.common.base.Function;
import com.google.common.collect.Iterables;
import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;

/**
 * Implements the Cartesian product of ordered collections.
 * 
 * @author <a href="mailto:jmkristian@gmail.com?subject=Cartesian.product">John
 *         Kristian</a>
 */
public class Cartesian {
  /**
   * Generate the <a href="http://en.wikipedia.org/wiki/Cartesian_product">Cartesian
   * product</a> of the given axes. For axes [[a1, a2 ...], [b1, b2 ...], [c1, c2 ...]
   * ...] the product is [{a1, b1, c1 ...} ... {a1, b1, c2 ...} ... {a1, b2, c1 ...} ...
   * {aN, bN, cN ...}]. In other words, the results are generated in same order as these
   * nested loops:
   * 
   * <pre>
   * for (T a : [a1, a2 ...])
   *   for (T b : [b1, b2 ...])
   *     for (T c : [c1, c2 ...])
   *       ...
   *         result = new T[]{a, b, c ...};
   * </pre>
   * 
   * Each result is a new array of T, whose elements refer to the elements of the axes.
   * <p>
   * Don't change the axes while iterating over their product, as a rule. Changes to an
   * axis can affect the product or cause iteration to fail (which is usually bad). To
   * prevent this, you can pass clones of your axes to this method.
   * <p>
   * The implementation is lazy. This method iterates over the axes, and returns an
   * Iterable that contains a reference to each axis. Iterating over the product causes
   * iteration over each axis. Methods of each axis are called as late as practical. The
   * Iterator constructs just one result at a time and doesn't retain references to them.
   * For large axes, this uses much less memory than a collection of the results.
   */
  public static <T> Iterable<T[]> product(Class<T> resultType,
                                          Iterable<? extends Iterable<? extends T>> axes) {
    return new Product<T>(resultType, newArray(Iterable.class, axes));
  }

  /**
   * Like {@link #product(Class,Iterable) product}(resultType, Arrays.asList(axes)), but
   * slightly more efficient.
   */
  public static <T> Iterable<T[]> product(Class<T> resultType, Iterable<? extends T>... axes) {
    return new Product<T>(resultType, axes.clone());
  }

  /**
   * Like {@link #product(Class,Iterable) product}(T.class, axes), except each result is a
   * List instead of an array. So, the result element type can be generic (unlike an
   * array).
   */
  public static <T> Iterable<List<T>> product(Iterable<? extends Iterable<? extends T>> axes) {
    return asLists(product(Object.class, axes));
    // The internal data structures are untyped, but the result is type safe.
  }

  /**
   * Like {@link #product(Iterable) product}(Arrays.asList(axes)), but slightly more
   * efficient.
   */
  public static <T> Iterable<List<T>> product(Iterable<? extends T>... axes) {
    return asLists(product(Object.class, axes));
  }

  // Don't make this public. It's really dangerous.
  private static <T> Iterable<List<T>> asLists(Iterable<Object[]> arrays) {
    return Iterables.transform(arrays, new AsList<T>());
  }

  // Don't make this public. It's really dangerous.
  private static class AsList<T> implements Function<Object[], List<T>> {
    @Override
    @SuppressWarnings("unchecked")
    public List<T> apply(Object[] array) {
      return Arrays.asList((T[]) array);
    }
  }

  /** Create a generic array containing references to the given objects. */
  private static <T> T[] newArray(Class<? super T> elementType, Iterable<? extends T> from) {
    List<T> list = new ArrayList<T>();
    for (T f : from)
      list.add(f);
    return list.toArray(newArray(elementType, list.size()));
  }

  /** Create a generic array. */
  @SuppressWarnings("unchecked")
  private static <T> T[] newArray(Class<? super T> elementType, int length) {
    return (T[]) Array.newInstance(elementType, length);
  }

  private static class Product<T> implements Iterable<T[]> {
    private final Class<T> _resultType;
    private final Iterable<? extends T>[] _axes;

    /** Caution: the given array of axes is contained by reference, not cloned. */
    Product(Class<T> resultType, Iterable<? extends T>[] axes) {
      _resultType = resultType;
      _axes = axes;
    }

    @Override
    public Iterator<T[]> iterator() {
      if (_axes.length <= 0) // an edge case
        return Collections.singletonList(newArray(_resultType, 0)).iterator();
      return new ProductIterator<T>(_resultType, _axes);
    }

    @Override
    public String toString() {
      return "Cartesian.product(" + Arrays.toString(_axes) + ")";
    }

    private static class ProductIterator<T> implements Iterator<T[]> {
      private final Iterable<? extends T>[] _axes;
      private final Iterator<? extends T>[] _iterators; // one per axis
      private final T[] _result; // a copy of the last result
      /**
       * The minimum index such that this.next() will return an array that contains
       * _iterators[index].next(). There are some special sentinel values: NEW means this
       * is a freshly constructed iterator, DONE means all combinations have been
       * exhausted (so this.hasNext() == false) and _iterators.length means the value is
       * unknown (to be determined by this.hasNext).
       */
      private int _nextIndex = NEW;
      private static final int NEW = -2;
      private static final int DONE = -1;

      /** Caution: the given array of axes is contained by reference, not cloned. */
      ProductIterator(Class<T> resultType, Iterable<? extends T>[] axes) {
        _axes = axes;
        _iterators = Cartesian.<Iterator<? extends T>> newArray(Iterator.class, _axes.length);
        for (int a = 0; a < _axes.length; ++a) {
          _iterators[a] = axes[a].iterator();
        }
        _result = newArray(resultType, _iterators.length);
      }

      private void close() {
        _nextIndex = DONE;
        // Release references, to encourage garbage collection:
        Arrays.fill(_iterators, null);
        Arrays.fill(_result, null);
      }

      @Override
      public boolean hasNext() {
        if (_nextIndex == NEW) { // This is the first call to hasNext().
          _nextIndex = 0; // start here
          for (Iterator<? extends T> iter : _iterators) {
            if (!iter.hasNext()) {
              close(); // no combinations
              break;
            }
          }
        } else if (_nextIndex >= _iterators.length) {
          // This is the first call to hasNext() after next() returned a result.
          // Determine the _nextIndex to be used by next():
          for (_nextIndex = _iterators.length - 1; _nextIndex >= 0; --_nextIndex) {
            Iterator<? extends T> iter = _iterators[_nextIndex];
            if (iter.hasNext()) {
              break; // start here
            }
            if (_nextIndex == 0) { // All combinations have been generated.
              close();
              break;
            }
            // Repeat this axis, with the next value from the previous axis.
            iter = _axes[_nextIndex].iterator();
            _iterators[_nextIndex] = iter;
            if (!iter.hasNext()) { // Oops; this axis can't be repeated.
              close(); // no more combinations
              break;
            }
          }
        }
        return _nextIndex >= 0;
      }

      @Override
      public T[] next() {
        if (!hasNext())
          throw new NoSuchElementException("!hasNext");
        for (; _nextIndex < _iterators.length; ++_nextIndex) {
          _result[_nextIndex] = _iterators[_nextIndex].next();
        }
        return _result.clone();
      }

      @Override
      public void remove() {
        for (Iterator<? extends T> iter : _iterators) {
          iter.remove();
        }
      }

      @Override
      public String toString() {
        return "Cartesian.product(" + Arrays.toString(_axes) + ").iterator()";
      }
    }
  }
}

TestCartesianProduct.java

import static java.util.Arrays.asList;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.NoSuchElementException;
import org.testng.annotations.Test;

/**
 * Unit tests for Cartesian.product.
 * 
 * @author <a href="mailto:jmkristian@gmail.com?subject=TestCartesianProduct">John
 *         Kristian</a>
 */
public class TestCartesianProduct {
  /** Empty product (because at least one of the axes is empty). */
  @Test
  @SuppressWarnings("unchecked")
  public void test0() {
    for (Iterable<List<Object>> p : asList(Cartesian.product(asList()),
                                           Cartesian.product(Collections.emptyList(),
                                                             Collections.emptySet()),
                                           Cartesian.product(asList("a", "b"), asList()),
                                           Cartesian.product(asList(), asList("a", "b")),
                                           Cartesian.product(asList("a", "b"),
                                                             Collections.emptySet(),
                                                             asList(1, 2, 3)))) {
      assertDone(p.iterator());
    }
  }

  /** A single product (because all the axes have just one value). */
  @Test
  @SuppressWarnings("unchecked")
  public void test1() {
    for (Iterable<List<Object>> c : asList(cartesianProduct(Collections.<Iterable<BigDecimal>> emptySet()),
                                           cartesianProduct())) {
      Iterator<List<Object>> i = c.iterator();
      assertNext(i); // i.next() returns a zero-length array
      assertDone(i);
    }
    Iterator<List<Object>> i = cartesianProduct(asList("fred")).iterator();
    assertNext(i, "fred");
    assertDone(i);

    i = cartesianProduct(asList("alpha"), asList("bravo")).iterator();
    assertNext(i, "alpha", "bravo");
    assertDone(i);

    i = cartesianProduct(asList("A"), asList(1), Collections.singleton(Math.PI)).iterator();
    assertNext(i, "A", 1, Math.PI);
    assertDone(i);
  }

  /** Two products. */
  @Test
  @SuppressWarnings("unchecked")
  public void test2() {
    Iterator<List<Object>> i = cartesianProduct(asList(asList("A", "B"))).iterator();
    assertNext(i, "A");
    assertNext(i, "B");
    assertDone(i);

    i = cartesianProduct(asList("A", "B"), Collections.singleton(1)).iterator();
    assertNext(i, "A", 1);
    assertNext(i, "B", 1);
    assertDone(i);

    i = cartesianProduct(Collections.singleton("A"), asList(1, 2),
                         Collections.singletonList(Math.PI)).iterator();
    assertNext(i, "A", 1, Math.PI);
    assertNext(i, "A", 2, Math.PI);
    assertDone(i);
  }

  /** Multiple iterators over the same product. */
  @Test
  @SuppressWarnings("unchecked")
  public void testReiterable() {
    Iterable<List<Object>> c = cartesianProduct(asList("A"), asList(1, 2));
    Iterator<List<Object>> i1 = c.iterator();
    assertNext(i1, "A", 1);
    Iterator<List<Object>> i2 = c.iterator();
    assertNext(i1, "A", 2);
    assertNext(i2, "A", 1);
    assertDone(i1);
    assertNext(i2, "A", 2);
    assertDone(i2);
    i1 = c.iterator();
    assertNext(i1, "A", 1);
    assertNext(i1, "A", 2);
    assertDone(i1);
    assertDone(i2);
  }

  /** Different overloaded methods generate the same products. */
  @Test
  @SuppressWarnings("unchecked")
  public void testOverloads() {
    for (Iterable<List<Object>> c : asList(cartesianProduct(asList("A", "B"), asList(1, 2)),
                                           cartesianProduct(asList(asList("A", "B"), asList(1, 2))))) {
      Iterator<List<Object>> i = c.iterator();
      assertNext(i, "A", 1);
      assertNext(i, "A", 2);
      assertNext(i, "B", 1);
      assertNext(i, "B", 2);
      assertDone(i);
    }
  }

  /**
   * The results of Cartesian.product are distinct arrays, containing references to the
   * axis elements.
   */
  @Test
  @SuppressWarnings("unchecked")
  public void testResultsAreShallowClones() {
    final String v0 = "A";
    final Integer v10 = 1;
    final Integer v11 = 2;
    final Iterator<List<Object>> iter = cartesianProduct(asList(v0), asList(v10, v11, 3)).iterator();
    List<Object> r0 = iter.next();
    assertTrue(r0.get(0) == v0); // same object
    assertTrue(r0.get(1) == v10); // same object
    List<Object> r1 = iter.next();
    assertTrue(r1.get(0) == v0); // same object
    assertTrue(r1.get(1) == v11); // same object

    assertFalse(r0 == r1);
    assertTrue(r0.get(0) == r1.get(0));
    r0.set(0, "B");
    assertFalse(r0.get(0) == r1.get(0));

    List<Object> r3 = iter.next();
    assertTrue(r3.get(0) == v0); // same object
  }

  /** The results of Cartesian.product have the correct array element type. */
  @Test
  @SuppressWarnings("unchecked")
  public void testResultType() {
    { // result type is not a superclass
      Integer[] result = Cartesian.product(Integer.class, asList(asList(1), asList(2), asList(3)))
                                  .iterator().next();
      assertEquals(result, new Integer[] { 1, 2, 3 });
      assertEquals(result.getClass().getComponentType(), Integer.class);
      assertWrongType(result, Long.valueOf(1));

      result = Cartesian.product(Integer.class, Collections.<Iterable<Integer>> emptySet())
                        .iterator().next();
      assertEquals(result.length, 0);
      assertEquals(result.getClass().getComponentType(), Integer.class);
    }
    { // result type is not a subclass
      Number[] result = Cartesian.product(Number.class, asList(asList(1), asList(2), asList(3)))
                                 .iterator().next();
      assertEquals(result, new Integer[] { 1, 2, 3 });
      assertEquals(result.getClass().getComponentType(), Number.class);
      assertEquals(Cartesian.product(asList(1), asList(2), asList(3)).iterator().next(),
                   asList(result));
      result[0] = Long.valueOf(1);
      result[0] = Float.valueOf(1);

      result = Cartesian.product(Number.class, Collections.<Iterable<Integer>> emptySet())
                        .iterator().next();
      assertEquals(result.length, 0);
      assertEquals(result.getClass().getComponentType(), Number.class);
    }
    { // result is a generic type
      List<Integer> item = new ArrayList<Integer>(asList(1));
      Iterable<List<Integer>> axis = new ArrayList<List<Integer>>(asList(item));
      List<Integer>[] result = Cartesian.product(List.class, asList(axis)).iterator().next();
      assertEquals(result[0], item);
      assertEquals(result.getClass().getComponentType(), List.class);
      result[0] = new LinkedList<Integer>(asList(1));
      assertWrongType(result, new HashSet<Integer>(asList(1)));
    }
  }

  public static Iterable<List<Object>> cartesianProduct(Iterable<? extends Object>... axes) {
    return Cartesian.<Object> product(axes);
  }

  public static Iterable<List<Object>> cartesianProduct(Iterable<? extends Iterable<? extends Object>> axes) {
    return Cartesian.<Object> product(axes);
  }

  private static void assertNext(Iterator<List<Object>> subject, Object... expected)
      throws AssertionError {
    assertTrue(subject.hasNext(), subject + ".hasNext()");
    assertEquals(subject.next(), asList(expected), subject + ".next()");
  }

  private static void assertDone(Iterator<List<Object>> subject) throws AssertionError {
    assertFalse(subject.hasNext(), subject + ".hasNext()");
    try {
      subject.next();
      fail("expected NoSuchElementException");
    } catch (NoSuchElementException expected) {
    }
  }

  private static void assertWrongType(Object[] subject, Object element) {
    try {
      subject[0] = element;
      fail("expected " + ArrayStoreException.class.getName());
    } catch (ArrayStoreException expected) {
    }
  }
}