In memory data structure

/*******************************************************************************
* To test this you either have to split this file into many different classes  *
* or do a global search / replace from 'public class' to 'public static class' *
* and place the entire code inside a new empty class (for quick testing only)  *
* The unit test below has to be in it's own .java file, or it won't run        *
*                                                                              *
* Questions, ideas, copyright: Sean Floyd <sean@mostlymagic.com>               *
*                                                                              *
* Inspired by this question:                                                   *
* http://stackoverflow.com/questions                                           *
*       /2874139/in-memory-data-structure-that-supports-boolean-querying       *
*******************************************************************************/



/***** Interfaces *************************************************************/

public interface Query<T> {
    Collection<T> get();
}

public interface DataStructure<K, V> {
    Query<V> and(Query<V> first, Query<V> second);

    Collection<K> keys();

    Query<V> not(Query<V> inner);

    Query<V> or(Query<V> first, Query<V> second);

    void put(V value, K... keys);

    void removeKey(K key);

    void removeValue(V value);

    Query<V> search(K key);

    Collection<V> values();

    Collection<V> values(K key);
}

/***** Implementation *********************************************************/

public class DataStructureImpl<K, V> implements DataStructure<K, V> {

    private final Map<K, Collection<V>> map = new HashMap<K, Collection<V>>();

    @Override
    public Query<V> and(final Query<V> first, final Query<V> second) {
        return new AndQuery<V>(first, second);
    }

    @Override
    public Collection<K> keys() {
        return new HashSet<K>(this.map.keySet());
    }

    @Override
    public Query<V> not(final Query<V> inner) {
        return new NotQuery<K, V>(inner, this);
    }

    @Override
    public Query<V> or(final Query<V> first, final Query<V> second) {
        return new OrQuery<V>(first, second);
    }

    @Override
    public void put(final V value, final K... keys) {
        for (final K key : keys) {
            if (!this.map.containsKey(key)) {
                this.map.put(key, new HashSet<V>());
            }
            this.map.get(key).add(value);
        }

    }

    @Override
    public void removeKey(final K key) {
        this.map.remove(key);
    }

    @Override
    public void removeValue(final V value) {
        final Collection<K> keysToRemove = new HashSet<K>();
        for (final Entry<K, Collection<V>> entry : this.map.entrySet()) {
            final Collection<V> rowValue = entry.getValue();
            if (rowValue.remove(rowValue) && rowValue.isEmpty()) {
                keysToRemove.add(entry.getKey());
            }
        }
        for (final K key : keysToRemove) {
            this.map.remove(key);
        }
    }

    @Override
    public Query<V> search(final K key) {
        return new SearchQuery<K, V>(key, this);
    }

    @Override
    public Collection<V> values() {
        final Set<V> set = new HashSet<V>();
        for (final Collection<V> row : this.map.values()) {
            set.addAll(row);
        }
        return set;
    }

    @Override
    public Collection<V> values(final K key) {
        return new HashSet<V>(this.map.containsKey(key) ? this.map.get(key)
                : Collections.<V> emptySet());
    }

}

public class AndQuery<V> implements Query<V> {

    private final Query<V> first;
    private final Query<V> second;

    public AndQuery(final Query<V> first, final Query<V> second) {
        this.first = first;
        this.second = second;
    }

    @Override
    public Collection<V> get() {
        final Collection<V> result = new ArrayList<V>(this.first.get());
        result.retainAll(this.second.get());
        return result;
    }

}




public class NotQuery<K, V> implements Query<V> {

    private final DataStructure<K, V> dataStructure;
    private final Query<V> inner;

    public NotQuery(final Query<V> inner,
            final DataStructure<K, V> dataStructure) {
        this.inner = inner;
        this.dataStructure = dataStructure;
    }

    @Override
    public Collection<V> get() {
        final Collection<V> values = this.dataStructure.values();
        values.removeAll(this.inner.get());
        return values;
    }

}

public class OrQuery<V> implements Query<V> {

    private final Query<V> first;
    private final Query<V> second;

    public OrQuery(final Query<V> first, final Query<V> second) {
        this.first = first;
        this.second = second;
    }

    @Override
    public Collection<V> get() {
        final Collection<V> result = new HashSet<V>(this.first.get());
        result.addAll(this.second.get());
        return result;
    }

}


public class SearchQuery<K, V> implements Query<V> {

    private final K key;
    private final DataStructure<K, V> dataStructure;

    public SearchQuery(final K key, final DataStructure<K, V> dataStructure) {
        this.key = key;
        this.dataStructure = dataStructure;
    }

    @Override
    public Collection<V> get() {
        return this.dataStructure.values(this.key);
    }

}

/***** JUnit Test *************************************************************/

public class DataStructureImplTest {
    enum Color {

        RED, BLUE, GREEN, BLACK, PINK, MAUVE, TURQUOISE, YELLOW, MAGENTA, GRAY, WHITE, BROWN;

    }

    enum Toon {

        TOM, JERRY, BUGS, DAFFY, MICKEY, DONALD, DAISY, GOOFY, WILEY;

    }

    private DataStructure<Color, Toon> data;

    @Before
    public void setup() {
        this.data = new DataStructureImpl<Color, Toon>();

        this.data.put(Toon.DONALD, Color.YELLOW, Color.BLUE, Color.WHITE);
        this.data.put(Toon.MICKEY, Color.BLACK, Color.WHITE, Color.RED);
        this.data.put(Toon.WILEY, Color.GRAY);
        this.data.put(Toon.TOM, Color.GRAY, Color.WHITE);
        this.data.put(Toon.JERRY, Color.BROWN);
        this.data.put(Toon.GOOFY, Color.BLACK, Color.WHITE, Color.RED);
        this.data.put(Toon.BUGS, Color.WHITE, Color.GRAY);
        this.data.put(Toon.DAISY, Color.PINK);
        this.data.put(Toon.DAFFY, Color.BLACK, Color.YELLOW);
    }

    @Test
    public void testCombinedQueries() {
        final Query<Toon> whiteQuery = this.data.search(Color.WHITE);
        final Query<Toon> blackQuery = this.data.search(Color.BLACK);
        final Query<Toon> redQuery = this.data.search(Color.RED);

        final Collection<Toon> blackOrWhite = this.data.or(blackQuery,
                whiteQuery).get();
        assertTrue(blackOrWhite.containsAll(EnumSet.of(Toon.DONALD,
                Toon.MICKEY, Toon.DAFFY)));
        assertFalse(blackOrWhite.contains(Toon.WILEY));

        final Collection<Toon> blackAndWhite = this.data.and(blackQuery,
                whiteQuery).get();
        assertTrue(blackAndWhite.containsAll(EnumSet
                .of(Toon.MICKEY, Toon.GOOFY)));
        assertFalse(blackAndWhite.contains(Toon.DONALD));

        final Collection<Toon> notWhite = this.data.not(whiteQuery).get();
        assertTrue(notWhite.containsAll(EnumSet.of(Toon.DAFFY, Toon.WILEY,
                Toon.JERRY)));
        assertFalse(notWhite.contains(Toon.TOM));

        final Collection<Toon> blackAndEitherWhiteOrNotRed = this.data.and(
                blackQuery, this.data.or(whiteQuery, this.data.not(redQuery)))
                .get();
        assertTrue(blackAndEitherWhiteOrNotRed.containsAll(EnumSet.of(
                Toon.MICKEY, Toon.GOOFY, Toon.DAFFY)));

    }

    @Test
    public void testSimpleQuery() {

        final Query<Toon> whiteQuery = this.data.search(Color.WHITE);
        assertTrue(whiteQuery.get().containsAll(
                EnumSet.of(Toon.DONALD, Toon.MICKEY)));
        assertFalse(whiteQuery.get().contains(Toon.WILEY));

    }
}