Remove intermediate SearchLookup classes (elastic#68052)

SearchLookup has two intermediate classes, DocMap and StoredFieldsLookup, that
are simple factories for their Leaf implementations. They are never accessed
outside SearchLookup, with the exception of two calls on DocMap that can be
easily refactored. This commit removes them, making SearchLookup.getLeafSearchLookup
directly responsible for creating the leaf lookups.

modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java

@@ -453,13 +453,13 @@ private static DoubleValuesSource getDocValueSource(String variable, SearchLooku
         }
 
         String fieldname = parts[1].text;
-        MappedFieldType fieldType = lookup.doc().fieldType(fieldname);
+        MappedFieldType fieldType = lookup.fieldType(fieldname);
 
         if (fieldType == null) {
             throw new ParseException("Field [" + fieldname + "] does not exist in mappings", 5);
         }
 
-        IndexFieldData<?> fieldData = lookup.doc().getForField(fieldType);
+        IndexFieldData<?> fieldData = lookup.getForField(fieldType);
         final DoubleValuesSource valueSource;
         if (fieldType instanceof GeoPointFieldType) {
             // geo

server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchDocValuesPhase.java

@@ -58,7 +58,7 @@ public FetchSubPhaseProcessor getProcessor(FetchContext context) {
             }
             ValueFetcher fetcher = new DocValueFetcher(
                 ft.docValueFormat(fieldAndFormat.format, null),
-                context.searchLookup().doc().getForField(ft)
+                context.searchLookup().getForField(ft)
             );
             fields.add(new DocValueField(fieldAndFormat.field, fetcher));
         }

server/src/main/java/org/elasticsearch/search/lookup/LeafStoredFieldsLookup.java

@@ -18,8 +18,9 @@
  */
 package org.elasticsearch.search.lookup;
 
-import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.StoredFieldVisitor;
 import org.elasticsearch.ElasticsearchParseException;
+import org.elasticsearch.common.CheckedBiConsumer;
 import org.elasticsearch.index.fieldvisitor.SingleFieldsVisitor;
 import org.elasticsearch.index.mapper.MappedFieldType;
 
@@ -38,13 +39,14 @@
 public class LeafStoredFieldsLookup implements Map<Object, Object> {
 
     private final Function<String, MappedFieldType> fieldTypeLookup;
-    private final LeafReader reader;
+    private final CheckedBiConsumer<Integer, StoredFieldVisitor, IOException> reader;
 
     private int docId = -1;
 
     private final Map<String, FieldLookup> cachedFieldData = new HashMap<>();
 
-    LeafStoredFieldsLookup(Function<String, MappedFieldType> fieldTypeLookup, LeafReader reader) {
+    LeafStoredFieldsLookup(Function<String, MappedFieldType> fieldTypeLookup,
+                           CheckedBiConsumer<Integer, StoredFieldVisitor, IOException> reader) {
         this.fieldTypeLookup = fieldTypeLookup;
         this.reader = reader;
     }
@@ -136,7 +138,7 @@ private FieldLookup loadFieldData(String name) {
             List<Object> values = new ArrayList<>(2);
             SingleFieldsVisitor visitor = new SingleFieldsVisitor(data.fieldType(), values);
             try {
-                reader.document(docId, visitor);
+                reader.accept(docId, visitor);
             } catch (IOException e) {
                 throw new ElasticsearchParseException("failed to load field [{}]", e, name);
             }

server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java

@@ -50,9 +50,7 @@ public class SearchLookup {
      * {@code b}'s chain will contain {@code ["a", "b"]}.
      */
     private final Set<String> fieldChain;
-    private final DocLookup docMap;
     private final SourceLookup sourceLookup;
-    private final StoredFieldsLookup storedFieldsLookup;
     private final Function<String, MappedFieldType> fieldTypeLookup;
     private final BiFunction<MappedFieldType, Supplier<SearchLookup>, IndexFieldData<?>> fieldDataLookup;
 
@@ -64,10 +62,7 @@ public SearchLookup(Function<String, MappedFieldType> fieldTypeLookup,
                         BiFunction<MappedFieldType, Supplier<SearchLookup>, IndexFieldData<?>> fieldDataLookup) {
         this.fieldTypeLookup = fieldTypeLookup;
         this.fieldChain = Collections.emptySet();
-        docMap = new DocLookup(fieldTypeLookup,
-            fieldType -> fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name())));
-        sourceLookup = new SourceLookup();
-        storedFieldsLookup = new StoredFieldsLookup(fieldTypeLookup);
+        this.sourceLookup = new SourceLookup();
         this.fieldDataLookup = fieldDataLookup;
     }
 
@@ -80,10 +75,7 @@ public SearchLookup(Function<String, MappedFieldType> fieldTypeLookup,
      */
     private SearchLookup(SearchLookup searchLookup, Set<String> fieldChain) {
         this.fieldChain = Collections.unmodifiableSet(fieldChain);
-        this.docMap = new DocLookup(searchLookup.fieldTypeLookup,
-            fieldType -> searchLookup.fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name())));
         this.sourceLookup = searchLookup.sourceLookup;
-        this.storedFieldsLookup = searchLookup.storedFieldsLookup;
         this.fieldTypeLookup = searchLookup.fieldTypeLookup;
         this.fieldDataLookup = searchLookup.fieldDataLookup;
     }
@@ -111,13 +103,17 @@ public final SearchLookup forkAndTrackFieldReferences(String field) {
 
     public LeafSearchLookup getLeafSearchLookup(LeafReaderContext context) {
         return new LeafSearchLookup(context,
-                docMap.getLeafDocLookup(context),
+                new LeafDocLookup(fieldTypeLookup, this::getForField, context),
                 sourceLookup,
-                storedFieldsLookup.getLeafFieldsLookup(context));
+                new LeafStoredFieldsLookup(fieldTypeLookup, (doc, visitor) -> context.reader().document(doc, visitor)));
     }
 
-    public DocLookup doc() {
-        return docMap;
+    public MappedFieldType fieldType(String fieldName) {
+        return fieldTypeLookup.apply(fieldName);
+    }
+
+    public IndexFieldData<?> getForField(MappedFieldType fieldType) {
+        return fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name()));
     }
 
     public SourceLookup source() {

server/src/test/java/org/elasticsearch/index/query/SearchExecutionContextTests.java

@@ -560,7 +560,7 @@ private static List<String> collect(String field, SearchExecutionContext searchE
                 if (randomBoolean()) {
                     indexFieldData = searchExecutionContext.getForField(fieldType);
                 } else {
-                    indexFieldData = searchExecutionContext.lookup().doc().getForField(fieldType);
+                    indexFieldData = searchExecutionContext.lookup().getForField(fieldType);
                 }
                 searcher.search(query, new Collector() {
                     @Override
@@ -579,7 +579,7 @@ public void setScorer(Scorable scorer) {}
                             public void collect(int doc) throws IOException {
                                 ScriptDocValues<?> scriptDocValues;
                                 if(randomBoolean()) {
-                                    LeafDocLookup leafDocLookup = searchExecutionContext.lookup().doc().getLeafDocLookup(context);
+                                    LeafDocLookup leafDocLookup = searchExecutionContext.lookup().getLeafSearchLookup(context).doc();
                                     leafDocLookup.setDocument(doc);
                                     scriptDocValues = leafDocLookup.get(field);
                                 } else {

server/src/test/java/org/elasticsearch/search/lookup/LeafStoredFieldsLookupTests.java

@@ -21,19 +21,14 @@
 import org.apache.lucene.index.DocValuesType;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.IndexOptions;
-import org.apache.lucene.index.LeafReader;
-import org.apache.lucene.index.StoredFieldVisitor;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.test.ESTestCase;
 import org.junit.Before;
 
 import java.util.Collections;
 import java.util.List;
 
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.anyObject;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -53,15 +48,8 @@ public void setUp() throws Exception {
         FieldInfo mockFieldInfo = new FieldInfo("field", 1, false, false, true,
             IndexOptions.NONE, DocValuesType.NONE, -1, Collections.emptyMap(), 0, 0, 0, false);
 
-        LeafReader leafReader = mock(LeafReader.class);
-        doAnswer(invocation -> {
-            Object[] args = invocation.getArguments();
-            StoredFieldVisitor visitor = (StoredFieldVisitor) args[1];
-            visitor.doubleField(mockFieldInfo, 2.718);
-            return null;
-        }).when(leafReader).document(anyInt(), any(StoredFieldVisitor.class));
-
-        fieldsLookup = new LeafStoredFieldsLookup(field -> field.equals("field") || field.equals("alias") ? fieldType : null, leafReader);
+        fieldsLookup = new LeafStoredFieldsLookup(field -> field.equals("field") || field.equals("alias") ? fieldType : null,
+            (doc, visitor) -> visitor.doubleField(mockFieldInfo, 2.718));
     }
 
     public void testBasicLookup() {

test/framework/src/main/java/org/elasticsearch/index/mapper/MapperTestCase.java

@@ -315,7 +315,7 @@ protected final List<?> fetchFromDocValues(MapperService mapperService, MappedFi
             iw.addDocument(mapperService.documentMapper().parse(source(b -> b.field(ft.name(), sourceValue))).rootDoc());
         }, iw -> {
             SearchLookup lookup = new SearchLookup(mapperService::fieldType, fieldDataLookup);
-            ValueFetcher valueFetcher = new DocValueFetcher(format, lookup.doc().getForField(ft));
+            ValueFetcher valueFetcher = new DocValueFetcher(format, lookup.getForField(ft));
             IndexSearcher searcher = newSearcher(iw);
             LeafReaderContext context = searcher.getIndexReader().leaves().get(0);
             lookup.source().setSegmentAndDocument(context, 0);

test/framework/src/main/java/org/elasticsearch/script/ScoreAccessor.java

@@ -20,14 +20,13 @@
 package org.elasticsearch.script;
 
 import org.apache.lucene.search.Scorable;
-import org.elasticsearch.search.lookup.DocLookup;
 
 import java.io.IOException;
 
 /**
  * A float encapsulation that dynamically accesses the score of a document.
  *
- * The provided {@link DocLookup} is used to retrieve the score
+ * The provided {@link Scorable} is used to retrieve the score
  * for the current document.
  */
 public final class ScoreAccessor extends Number {

x-pack/plugin/mapper-flattened/src/test/java/org/elasticsearch/index/mapper/FlattenedFieldLookupTests.java

@@ -163,7 +163,7 @@ public void testScriptDocValuesLookup() {
             }
             return null;
         }, fieldDataSupplier);
-        LeafDocLookup docLookup = searchLookup.doc().getLeafDocLookup(null);
+        LeafDocLookup docLookup = searchLookup.getLeafSearchLookup(null).doc();
 
         assertEquals(docValues1, docLookup.get("json.key1"));
         assertEquals(docValues2, docLookup.get("json.key2"));