Optimize top-level element declarations search.

pkg/analysis_server/lib/src/services/index/store/split_store.dart

Index: pkg/analysis_server/lib/src/services/index/store/split_store.dart
diff --git a/pkg/analysis_server/lib/src/services/index/store/split_store.dart b/pkg/analysis_server/lib/src/services/index/store/split_store.dart
index cb2e123d3d98996212863de4a254212620875710..dbd963e7dff35265b8ae45e7cccd320a7b56c928 100644
--- a/pkg/analysis_server/lib/src/services/index/store/split_store.dart
+++ b/pkg/analysis_server/lib/src/services/index/store/split_store.dart
@@ -8,6 +8,7 @@ import 'dart:async';
 import 'dart:collection';
 import 'dart:typed_data';
 
+import 'package:analysis_server/src/analysis_server.dart';
 import 'package:analysis_server/src/services/index/index.dart';
 import 'package:analysis_server/src/services/index/index_store.dart';
 import 'package:analysis_server/src/services/index/store/codec.dart';
@@ -16,7 +17,24 @@ import 'package:analyzer/src/generated/element.dart';
 import 'package:analyzer/src/generated/engine.dart';
 import 'package:analyzer/src/generated/java_engine.dart';
 import 'package:analyzer/src/generated/source.dart';
-import 'package:analysis_server/src/analysis_server.dart';
+
+
+class _TopElementData {
+  final String name;
+  final int elementId;
+
+  factory _TopElementData(ElementCodec elementCodec, Element element) {
+    return new _TopElementData._(
+        element.name, elementCodec.encode(element, false));
+  }
+
+  _TopElementData._(this.name, this.elementId);
+
+  Element getElement(AnalysisContext context, ElementCodec elementCodec) {
+    return elementCodec.decode(context, elementId);
+  }
+}
+
 
 /**
  * A manager for files content.
@@ -507,13 +525,13 @@ class SplitIndexStore implements IndexStore {
   ContextCodec _contextCodec;
 
   /**
-   * Information about "universe" elements.
+   * Information about top-level / "universe" elements.
    * We need to keep them together to avoid loading of all index nodes.
    *
-   * Order of keys: contextId, nodeId, Relationship.
+   * Order of keys: contextId, nodeId.
    */
-  Map<int, Map<int, Map<Relationship, List<LocationData>>>> _contextNodeRelations =
-      new HashMap<int, Map<int, Map<Relationship, List<LocationData>>>>();
+  Map<int, Map<int, List<_TopElementData>>> _topDeclarations =
+      new Map<int, Map<int, List<_TopElementData>>>();
 
   /**
    * The mapping of library [Source] to the [Source]s of part units.
@@ -646,7 +664,7 @@ class SplitIndexStore implements IndexStore {
     _currentNode = _nodeManager.newNode(context);
     _currentContextId = _contextCodec.encode(context);
     // remove Universe information for the current node
-    for (Map<int, dynamic> nodeRelations in _contextNodeRelations.values) {
+    for (Map<int, dynamic> nodeRelations in _topDeclarations.values) {
       nodeRelations.remove(_currentNodeNameId);
     }
     // done
@@ -675,7 +693,7 @@ class SplitIndexStore implements IndexStore {
 
   @override
   void clear() {
-    _contextNodeRelations.clear();
+    _topDeclarations.clear();
     _nodeManager.clear();
     _relToNameMap.clear();
   }
@@ -693,11 +711,6 @@ class SplitIndexStore implements IndexStore {
 
   Future<List<Location>> getRelationships(Element element,
       Relationship relationship) {
-    // special support for UniverseElement
-    if (identical(element, UniverseElement.INSTANCE)) {
-      List<Location> locations = _getRelationshipsUniverse(relationship);
-      return new Future.value(locations);
-    }
     // prepare node names
     List<int> nodeNameIds;
     {
@@ -733,6 +746,26 @@ class SplitIndexStore implements IndexStore {
     });
   }
 
+  List<Element> getTopDeclarations(ElementNameFilter nameFilter) {
+    List<Element> elements = <Element>[];
+    _topDeclarations.forEach((contextId, contextLocations) {
+      AnalysisContext context = _contextCodec.decode(contextId);
+      if (context != null) {
+        for (List<_TopElementData> topDataList in contextLocations.values) {
+          for (_TopElementData topData in topDataList) {
+            if (nameFilter(topData.name)) {
+              Element element = topData.getElement(context, _elementCodec);
+              if (element != null) {
+                elements.add(element);
+              }
+            }
+          }
+        }
+      }
+    });
+    return elements;
+  }
+
   /**
    * Returns all relations with [Element]s with the given [name].
    */
@@ -781,16 +814,29 @@ class SplitIndexStore implements IndexStore {
     if (location == null) {
       return;
     }
-    // special support for UniverseElement
-    if (identical(element, UniverseElement.INSTANCE)) {
-      _recordRelationshipUniverse(relationship, location);
-      return;
-    }
     // other elements
     _recordNodeNameForElement(element, relationship);
     _currentNode.recordRelationship(element, relationship, location);
   }
 
+  void recordTopDeclaration(Element element) {
+    // in current context
+    Map<int, List<_TopElementData>> nodeDeclarations =
+        _topDeclarations[_currentContextId];
+    if (nodeDeclarations == null) {
+      nodeDeclarations = new Map<int, List<_TopElementData>>();
+      _topDeclarations[_currentContextId] = nodeDeclarations;
+    }
+    // in current node
+    List<_TopElementData> declarations = nodeDeclarations[_currentNodeNameId];
+    if (declarations == null) {
+      declarations = <_TopElementData>[];
+      nodeDeclarations[_currentNodeNameId] = declarations;
+    }
+    // record LocationData
+    declarations.add(new _TopElementData(_elementCodec, element));
+  }
+
   @override
   void removeContext(AnalysisContext context) {
     if (context == null) {
@@ -801,7 +847,7 @@ class SplitIndexStore implements IndexStore {
     // remove context information
     _contextToLibraryToUnits.remove(context);
     _contextToUnitToLibraries.remove(context);
-    _contextNodeRelations.remove(_contextCodec.encode(context));
+    _topDeclarations.remove(_contextCodec.encode(context));
     // remove context from codec
     _contextCodec.remove(context);
   }
@@ -862,29 +908,6 @@ class SplitIndexStore implements IndexStore {
     }
   }
 
-  List<Location> _getRelationshipsUniverse(Relationship relationship) {
-    List<Location> locations = <Location>[];
-    _contextNodeRelations.forEach((contextId, contextRelations) {
-      AnalysisContext context = _contextCodec.decode(contextId);
-      if (context != null) {
-        for (Map<Relationship, List<LocationData>> nodeRelations in
-            contextRelations.values) {
-          List<LocationData> nodeLocations = nodeRelations[relationship];
-          if (nodeLocations != null) {
-            for (LocationData locationData in nodeLocations) {
-              Location location =
-                  locationData.getLocation(context, _elementCodec);
-              if (location != null) {
-                locations.add(location);
-              }
-            }
-          }
-        }
-      }
-    });
-    return locations;
-  }
-
   void _recordLibraryWithUnit(AnalysisContext context, Source library,
       Source unit) {
     Map<Source, Set<Source>> libraryToUnits = _contextToLibraryToUnits[context];
@@ -910,32 +933,6 @@ class SplitIndexStore implements IndexStore {
     nameToNodeNames.add(nameId, _currentNodeNameId);
   }
 
-  void _recordRelationshipUniverse(Relationship relationship,
-      Location location) {
-    // in current context
-    Map<int, Map<Relationship, List<LocationData>>> nodeRelations =
-        _contextNodeRelations[_currentContextId];
-    if (nodeRelations == null) {
-      nodeRelations = new HashMap<int, Map<Relationship, List<LocationData>>>();
-      _contextNodeRelations[_currentContextId] = nodeRelations;
-    }
-    // in current node
-    Map<Relationship, List<LocationData>> relations =
-        nodeRelations[_currentNodeNameId];
-    if (relations == null) {
-      relations = new HashMap<Relationship, List<LocationData>>();
-      nodeRelations[_currentNodeNameId] = relations;
-    }
-    // for the given relationship
-    List<LocationData> locations = relations[relationship];
-    if (locations == null) {
-      locations = <LocationData>[];
-      relations[relationship] = locations;
-    }
-    // record LocationData
-    locations.add(new LocationData.forObject(_elementCodec, location));
-  }
-
   void _recordUnitInLibrary(AnalysisContext context, Source library,
       Source unit) {
     Map<Source, Set<Source>> unitToLibraries =
@@ -970,7 +967,7 @@ class SplitIndexStore implements IndexStore {
     // remove universe relations
     {
       int contextId = _contextCodec.encode(context);
-      Map<int, Object> nodeRelations = _contextNodeRelations[contextId];
+      Map<int, dynamic> nodeRelations = _topDeclarations[contextId];
       if (nodeRelations != null) {
         nodeRelations.remove(nodeNameId);
       }