AnalysisDriver.getUnitElement() should return a structure with a key.

pkg/analyzer/lib/src/dart/analysis/search.dart

 // Copyright (c) 2016, the Dart project authors.  Please see the AUTHORS file
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
 import 'dart:async';
 
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/ast/visitor.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/visitor.dart';
@@ skipping 29 matching lines @@
     List<Element> elements = <Element>[];
 
     void addElement(Element element) {
       if (!element.isSynthetic && element.displayName == name) {
         elements.add(element);
       }
     }
 
     List<String> files = await _driver.getFilesDefiningClassMemberName(name);
     for (String file in files) {
-      CompilationUnitElement unitElement = await _driver.getUnitElement(file);
-      if (unitElement != null) {
-        for (ClassElement clazz in unitElement.types) {
+      UnitElementResult unitResult = await _driver.getUnitElement(file);
+      if (unitResult != null) {
+        for (ClassElement clazz in unitResult.element.types) {
           clazz.accessors.forEach(addElement);
           clazz.fields.forEach(addElement);
           clazz.methods.forEach(addElement);
         }
       }
     }
     return elements;
   }
 
   /**
@@ skipping 62 matching lines @@
   Future<List<Element>> topLevelElements(RegExp regExp) async {
     List<Element> elements = <Element>[];
 
     void addElement(Element element) {
       if (!element.isSynthetic && regExp.hasMatch(element.displayName)) {
         elements.add(element);
       }
     }
 
     for (FileState file in _driver.fsState.knownFiles) {
-      CompilationUnitElement unitElement =
-          await _driver.getUnitElement(file.path);
-      if (unitElement != null) {
+      UnitElementResult unitResult = await _driver.getUnitElement(file.path);
+      if (unitResult != null) {
+        CompilationUnitElement unitElement = unitResult.element;
         unitElement.accessors.forEach(addElement);
         unitElement.enums.forEach(addElement);
         unitElement.functions.forEach(addElement);
         unitElement.functionTypeAliases.forEach(addElement);
         unitElement.topLevelVariables.forEach(addElement);
         unitElement.types.forEach(addElement);
       }
     }
     return elements;
   }
@@ skipping 16 matching lines @@
       if (index != null) {
         _IndexRequest request = new _IndexRequest(index);
         var fileResults = await request.getUnresolvedMemberReferences(
             name,
             const {
               IndexRelationKind.IS_READ_BY: SearchResultKind.READ,
               IndexRelationKind.IS_WRITTEN_BY: SearchResultKind.WRITE,
               IndexRelationKind.IS_READ_WRITTEN_BY: SearchResultKind.READ_WRITE,
               IndexRelationKind.IS_INVOKED_BY: SearchResultKind.INVOCATION
             },
-            () => _driver.getUnitElement(file));
+            () => _getUnitElement(file));
         results.addAll(fileResults);
       }
     }
 
     return results;
   }
 
   Future<Null> _addResults(List<SearchResult> results, Element element,
       Map<IndexRelationKind, SearchResultKind> relationToResultKind) async {
     // Prepare the element name.
@@ skipping 35 matching lines @@
   Future<Null> _addResultsInFile(
       List<SearchResult> results,
       Element element,
       Map<IndexRelationKind, SearchResultKind> relationToResultKind,
       String file) async {
     AnalysisDriverUnitIndex index = await _driver.getIndex(file);
     if (index != null) {
       _IndexRequest request = new _IndexRequest(index);
       int elementId = request.findElementId(element);
       if (elementId != -1) {
-        List<SearchResult> fileResults = await request.getRelations(elementId,
-            relationToResultKind, () => _driver.getUnitElement(file));
+        List<SearchResult> fileResults = await request.getRelations(
+            elementId, relationToResultKind, () => _getUnitElement(file));
         results.addAll(fileResults);
       }
     }
   }
 
+  Future<CompilationUnitElement> _getUnitElement(String file) async {
+    UnitElementResult result = await _driver.getUnitElement(file);
+    return result?.element;
+  }
+
   Future<List<SearchResult>> _searchReferences(Element element) async {
     List<SearchResult> results = <SearchResult>[];
     await _addResults(results, element,
         const {IndexRelationKind.IS_REFERENCED_BY: SearchResultKind.REFERENCE});
     return results;
   }
 
   Future<List<SearchResult>> _searchReferences_CompilationUnit(
       CompilationUnitElement element) async {
     String path = element.source.fullName;
@@ skipping 610 matching lines @@
   }
 
   void _addResult(AstNode node, SearchResultKind kind) {
     bool isQualified = node.parent is Label;
     Element enclosingElement =
         _getEnclosingElement(enclosingUnitElement, node.offset);
     results.add(new SearchResult._(
         enclosingElement, kind, node.offset, node.length, true, isQualified));
   }
 }