cassandra-trunk/src/java/org/apache/cassandra/repair/asymmetric/RangeMap.java
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

package org.apache.cassandra.repair.asymmetric;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NavigableMap;
import java.util.Objects;
import java.util.Set;
import java.util.TreeMap;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.AbstractIterator;

import org.apache.cassandra.dht.Range;
import org.apache.cassandra.dht.Token;

public class RangeMap<T> implements Map<Range<Token>, T>
{
    private static final Comparator<Range<Token>> comparator = Comparator.comparing((Range<Token> o) -> o.left);

    private final NavigableMap<Range<Token>, T> byStart;

    public RangeMap()
    {
        byStart = new TreeMap<>(comparator);
    }

    public int size()
    {
        return byStart.size();
    }

    public boolean isEmpty()
    {
        return byStart.isEmpty();
    }

    public boolean containsKey(Object key)
    {
        return byStart.containsKey(key);
    }

    public boolean containsValue(Object value)
    {
        return byStart.containsValue(value);
    }

    public T get(Object key)
    {
        return byStart.get(key);
    }

    public T put(Range<Token> key, T value)
    {
        assertNonIntersecting(key);
        return byStart.put(key, value);
    }

    private void assertNonIntersecting(Range<Token> range)
    {
        // todo: wraparound
        Range<Token> before = byStart.floorKey(range);
        Range<Token> after = byStart.ceilingKey(range);
        assert before == null || !before.intersects(range);
        assert after == null || !after.intersects(range);
    }

    public T remove(Object key)
    {
        return byStart.remove(key);
    }

    public void putAll(Map<? extends Range<Token>, ? extends T> m)
    {
        byStart.putAll(m);
    }

    public void clear()
    {
        byStart.clear();
    }

    public Set<Range<Token>> keySet()
    {
        return byStart.keySet();
    }

    public Collection<T> values()
    {
        return byStart.values();
    }

    public Set<Map.Entry<Range<Token>, T>> entrySet()
    {
        return byStart.entrySet();
    }

    /**
     * might return duplicate entries if range.isWrapAround()
     *
     * don't depend on the order of the entries returned
     */
    @VisibleForTesting
    Iterator<Map.Entry<Range<Token>, T>> intersectingEntryIterator(Range<Token> range)
    {
        return range.isWrapAround() ? new WrappingIntersectingIterator(range) : new IntersectingIterator(range);
    }

    public Set<Map.Entry<Range<Token>, T>> removeIntersecting(Range<Token> range)
    {
        Iterator<Map.Entry<Range<Token>, T>> iter = intersectingEntryIterator(range);
        Set<Map.Entry<Range<Token>, T>> intersecting = new HashSet<>();
        while (iter.hasNext())
        {
            Map.Entry<Range<Token>, T> entry = iter.next();
            intersecting.add(entry);
        }
        for (Map.Entry<Range<Token>, T> entry : intersecting)
            byStart.remove(entry.getKey());
        return intersecting;
    }

    private class WrappingIntersectingIterator extends AbstractIterator<Map.Entry<Range<Token>, T>>
    {
        private final Iterator<Iterator<Map.Entry<Range<Token>, T>>> iterators;
        private Iterator<Map.Entry<Range<Token>, T>> currentIter;

        public WrappingIntersectingIterator(Range<Token> range)
        {
            List<Iterator<Map.Entry<Range<Token>, T>>> iters = new ArrayList<>(2);
            for (Range<Token> unwrapped : range.unwrap())
                iters.add((new IntersectingIterator(unwrapped)));
            iterators = iters.iterator();
        }
        protected Map.Entry<Range<Token>, T> computeNext()
        {
            while (currentIter == null || !currentIter.hasNext())
            {
                if (!iterators.hasNext())
                    return endOfData();
                currentIter = iterators.next();
            }
            return currentIter.next();
        }
    }

    private class IntersectingIterator extends AbstractIterator<Map.Entry<Range<Token>, T>>
    {
        private final Iterator<Map.Entry<Range<Token>, T>> tailIterator;
        private final Range<Token> range;
        // since we guarantee no ranges overlap in byStart, we know the last entry is possibly the wrap around range
        private boolean shouldReturnLast = false;

        public IntersectingIterator(Range<Token> range)
        {
            Range<Token> startKey = byStart.floorKey(range);
            tailIterator = startKey == null ? byStart.entrySet().iterator() :
                                              byStart.tailMap(startKey, true).entrySet().iterator();
            Range<Token> last = byStart.isEmpty() ? null : byStart.lastKey();
            if (last != null && last.isWrapAround() && last.intersects(range))
                shouldReturnLast = true;
            this.range = range;
        }

        protected Map.Entry<Range<Token>, T> computeNext()
        {
            if (shouldReturnLast)
            {
                shouldReturnLast = false;
                return new Entry<>(byStart.lastEntry());
            }
            while (tailIterator.hasNext())
            {
                Entry<Range<Token>, T> candidateNext = new Entry<>(tailIterator.next());
                Range<Token> candidateRange = candidateNext.getKey();

                if (candidateRange.isWrapAround()) // we know we already returned any wrapping range
                    continue;

                if (candidateRange.left.compareTo(range.right) >= 0 && (!range.isWrapAround())) // range is unwrapped, but that means one range has right == min token and is still wrapping
                    return endOfData();

                if (range.left.compareTo(candidateRange.right) >= 0)
                    continue;

                return candidateNext;
            }
            return endOfData();
        }
    }

    public String toString()
    {
        return byStart.toString();
    }

    static class Entry<K, V> implements Map.Entry<K, V>
    {
        private final V v;
        private final K k;

        Entry(K key, V val)
        {
            this.k = key;
            this.v = val;
        }

        Entry(Map.Entry<K, V> toClone)
        {
            this(toClone.getKey(), toClone.getValue());
        }
        public K getKey()
        {
            return k;
        }

        public V getValue()
        {
            return v;
        }

        public V setValue(V value)
        {
            throw new UnsupportedOperationException();
        }

        public boolean equals(Object o)
        {
            if (this == o) return true;
            if (!(o instanceof Map.Entry)) return false;
            Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o;
            return v.equals(entry.getValue()) &&
                   k.equals(entry.getKey());
        }

        public int hashCode()
        {
            return Objects.hash(v, k);
        }

        public String toString()
        {
            return "Entry{" +
                   "v=" + v +
                   ", k=" + k +
                   '}';
        }
    }
}
