Fix for subscribing for notifications after analysis.

pkg/analysis_server/lib/src/analysis_server.dart

Index: pkg/analysis_server/lib/src/analysis_server.dart
diff --git a/pkg/analysis_server/lib/src/analysis_server.dart b/pkg/analysis_server/lib/src/analysis_server.dart
index da2510b0a0ce5f865f5cadeaea96a8085a984ad1..b525f8f949f7383b2d9ba50e55fb07d1a04a9ee9 100644
--- a/pkg/analysis_server/lib/src/analysis_server.dart
+++ b/pkg/analysis_server/lib/src/analysis_server.dart
@@ -398,25 +398,22 @@ class AnalysisServer {
           List<AnalysisError> errors = analysisContext.getErrors(source).errors;
           sendAnalysisNotificationErrors(this, file, errors);
         }
-        // TODO(scheglov)
-        // 1. implement resolveCompilationUnit()
-        // 2. Share "if (dartUnit != null)"
-        if (service == AnalysisService.HIGHLIGHTS) {
-          CompilationUnit dartUnit = test_getResolvedCompilationUnit(file);
+        // Dart unit notifications.
+        if (AnalysisEngine.isDartFileName(file)) {
+          CompilationUnit dartUnit = getResolvedCompilationUnitToResendNotification(file);
           if (dartUnit != null) {
-            sendAnalysisNotificationHighlights(this, file, dartUnit);
-          }
-        }
-        if (service == AnalysisService.NAVIGATION) {
-          CompilationUnit dartUnit = test_getResolvedCompilationUnit(file);
-          if (dartUnit != null) {
-            sendAnalysisNotificationNavigation(this, file, dartUnit);
-          }
-        }
-        if (service == AnalysisService.OUTLINE) {
-          CompilationUnit dartUnit = test_getResolvedCompilationUnit(file);
-          if (dartUnit != null) {
-            sendAnalysisNotificationOutline(this, file, dartUnit);
+            switch (service) {
+              case AnalysisService.HIGHLIGHTS:
+                sendAnalysisNotificationHighlights(this, file, dartUnit);
+                break;
+              case AnalysisService.NAVIGATION:
+                // TODO(scheglov) consider support for one unit in 2+ libraries
+                sendAnalysisNotificationNavigation(this, file, dartUnit);
+                break;
+              case AnalysisService.OUTLINE:
+                sendAnalysisNotificationOutline(this, file, dartUnit);
+                break;
+            }
           }
         }
       }
@@ -447,6 +444,33 @@ class AnalysisServer {
   }
 
   /**
+   * Returns the [CompilationUnit] of the Dart file with the given [path] that
+   * should be used to resend notifications for already resolved unit.
+   * Returns `null` if the file is not a part of any context, library has not
+   * been yet resolved, or any problem happened.
+   */
+  CompilationUnit getResolvedCompilationUnitToResendNotification(String path) {
+    // prepare AnalysisContext
+    AnalysisContext context = _getAnalysisContext(path);
+    if (context == null) {
+      return null;
+    }
+    // prepare sources
+    Source unitSource = _getSource(path);
+    List<Source> librarySources = context.getLibrariesContaining(unitSource);
+    if (librarySources.isEmpty) {
+      return null;
+    }
+    // if library has not been resolved yet, the unit will be resolved later
+    Source librarySource = librarySources[0];
+    if (context.getLibraryElement(librarySource) == null) {
+      return null;
+    }
+    // if library has been already resolved, resolve unit
+    return context.resolveCompilationUnit2(unitSource, librarySource);
+  }
+
+  /**
    * Return the [CompilationUnit] of the Dart file with the given [path].
    * Return `null` if the file is not a part of any context.
    */