/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch 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.elasticsearch.search.scan;

import com.google.common.collect.Maps;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.*;
import org.apache.lucene.util.Bits;
import org.elasticsearch.common.lucene.docset.AllDocIdSet;
import org.elasticsearch.common.lucene.search.XFilteredQuery;
import org.elasticsearch.search.internal.SearchContext;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Map;

/**
 * The scan context allows to optimize readers we already processed during scanning. We do that by keeping track
 * of the count per reader, and if we are done with it, we no longer process it by using a filter that returns
 * null docIdSet for this reader.
 */
public class ScanContext {

    private final Map<IndexReader, ReaderState> readerStates = Maps.newHashMap();

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

    public TopDocs execute(SearchContext context) throws IOException {
        ScanCollector collector = new ScanCollector(readerStates, context.from(), context.size(), context.trackScores());
        Query query = new XFilteredQuery(context.query(), new ScanFilter(readerStates, collector));
        try {
            context.searcher().search(query, collector);
        } catch (ScanCollector.StopCollectingException e) {
            // all is well
        }
        return collector.topDocs();
    }

    static class ScanCollector extends Collector {

        private final Map<IndexReader, ReaderState> readerStates;

        private final int from;

        private final int to;

        private final ArrayList<ScoreDoc> docs;

        private final boolean trackScores;

        private Scorer scorer;

        private int docBase;

        private int counter;

        private IndexReader currentReader;
        private ReaderState readerState;

        ScanCollector(Map<IndexReader, ReaderState> readerStates, int from, int size, boolean trackScores) {
            this.readerStates = readerStates;
            this.from = from;
            this.to = from + size;
            this.trackScores = trackScores;
            this.docs = new ArrayList<>(size);
        }

        void incCounter(int count) {
            this.counter += count;
        }

        public TopDocs topDocs() {
            return new TopDocs(docs.size(), docs.toArray(new ScoreDoc[docs.size()]), 0f);
        }

        @Override
        public void setScorer(Scorer scorer) throws IOException {
            this.scorer = scorer;
        }

        @Override
        public void collect(int doc) throws IOException {
            if (counter >= from) {
                docs.add(new ScoreDoc(docBase + doc, trackScores ? scorer.score() : 0f));
            }
            readerState.count++;
            counter++;
            if (counter >= to) {
                throw StopCollectingException;
            }
        }

        @Override
        public void setNextReader(AtomicReaderContext context) throws IOException {
            // if we have a reader state, and we haven't registered one already, register it
            // we need to check in readersState since even when the filter return null, setNextReader is still
            // called for that reader (before)
            if (currentReader != null && !readerStates.containsKey(currentReader)) {
                assert readerState != null;
                readerState.done = true;
                readerStates.put(currentReader, readerState);
            }
            this.currentReader = context.reader();
            this.docBase = context.docBase;
            this.readerState = new ReaderState();
        }

        @Override
        public boolean acceptsDocsOutOfOrder() {
            return false;
        }

        public static final RuntimeException StopCollectingException = new StopCollectingException();

        static class StopCollectingException extends RuntimeException {
            @Override
            public Throwable fillInStackTrace() {
                return null;
            }
        }
    }

    public static class ScanFilter extends Filter {

        private final Map<IndexReader, ReaderState> readerStates;

        private final ScanCollector scanCollector;

        public ScanFilter(Map<IndexReader, ReaderState> readerStates, ScanCollector scanCollector) {
            this.readerStates = readerStates;
            this.scanCollector = scanCollector;
        }

        @Override
        public DocIdSet getDocIdSet(AtomicReaderContext context, Bits acceptedDocs) throws IOException {
            ReaderState readerState = readerStates.get(context.reader());
            if (readerState != null && readerState.done) {
                scanCollector.incCounter(readerState.count);
                return null;
            }
            return new AllDocIdSet(context.reader().maxDoc());
        }
    }

    static class ReaderState {
        public int count;
        public boolean done;
    }
}

@@ -827,47 +827,8 @@
 n;%0A%0A
-import com.google.common.collect.Maps;%0A
 impo
@@ -875,16 +875,16 @@
 ontext;%0A
+
 import o
@@ -1111,16 +1111,87 @@
 dQuery;%0A
+import org.elasticsearch.common.util.concurrent.ConcurrentCollections;%0A
 import o
@@ -1312,16 +1312,37 @@
 va.util.
+concurrent.Concurrent
 Map;%0A%0A/*
@@ -1643,32 +1643,42 @@
   private final 
+Concurrent
 Map%3CIndexReader,
@@ -1710,20 +1710,43 @@
 s = 
-Maps.newHash
+ConcurrentCollections.newConcurrent
 Map(
@@ -2382,32 +2382,42 @@
   private final 
+Concurrent
 Map%3CIndexReader,
@@ -2775,24 +2775,24 @@
 aderState;%0A%0A
-
         Scan
@@ -2801,16 +2801,26 @@
 llector(
+Concurrent
 Map%3CInde
@@ -5076,16 +5076,26 @@
 e final 
+Concurrent
 Map%3CInde
@@ -5209,16 +5209,26 @@
 nFilter(
+Concurrent
 Map%3CInde

Use ConcurrentHashMap instead of HashMap, because the readerStates is accessed by multiple threads during the entire scroll session.