Fix for filtering directives with protocol2.

Change-Id: I9f01bd6e2f3a8fdb31a0036e26412db7f38081b4
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/239223
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>
diff --git a/pkg/analysis_server/lib/src/services/completion/dart/fuzzy_filter_sort.dart b/pkg/analysis_server/lib/src/services/completion/dart/fuzzy_filter_sort.dart
index 59a5f61..b1247c7 100644
--- a/pkg/analysis_server/lib/src/services/completion/dart/fuzzy_filter_sort.dart
+++ b/pkg/analysis_server/lib/src/services/completion/dart/fuzzy_filter_sort.dart
@@ -17,7 +17,12 @@
   double score(CompletionSuggestionBuilder suggestion) {
     var suggestionTextToMatch = suggestion.completion;
 
-    if (suggestion.kind == CompletionSuggestionKind.NAMED_ARGUMENT) {
+    if (suggestion.kind == CompletionSuggestionKind.KEYWORD) {
+      var index = suggestionTextToMatch.indexOf(' ');
+      if (index != -1) {
+        suggestionTextToMatch = suggestionTextToMatch.substring(0, index);
+      }
+    } else if (suggestion.kind == CompletionSuggestionKind.NAMED_ARGUMENT) {
       var index = suggestionTextToMatch.indexOf(':');
       if (index != -1) {
         suggestionTextToMatch = suggestionTextToMatch.substring(0, index);
diff --git a/pkg/analysis_server/test/services/completion/dart/location/compilation_unit_test.dart b/pkg/analysis_server/test/services/completion/dart/location/compilation_unit_test.dart
new file mode 100644
index 0000000..d5130d5
--- /dev/null
+++ b/pkg/analysis_server/test/services/completion/dart/location/compilation_unit_test.dart
@@ -0,0 +1,83 @@
+// Copyright (c) 2022, 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 'package:analyzer_utilities/check/check.dart';
+import 'package:test_reflective_loader/test_reflective_loader.dart';
+
+import '../../../../client/completion_driver_test.dart';
+import '../completion_check.dart';
+
+void main() {
+  defineReflectiveSuite(() {
+    defineReflectiveTests(CompilationUnitTest1);
+    defineReflectiveTests(CompilationUnitTest2);
+  });
+}
+
+@reflectiveTest
+class CompilationUnitTest1 extends AbstractCompletionDriverTest
+    with CompilationUnitTestCases {
+  @override
+  TestingCompletionProtocol get protocol => TestingCompletionProtocol.version1;
+}
+
+@reflectiveTest
+class CompilationUnitTest2 extends AbstractCompletionDriverTest
+    with CompilationUnitTestCases {
+  @override
+  TestingCompletionProtocol get protocol => TestingCompletionProtocol.version2;
+}
+
+mixin CompilationUnitTestCases on AbstractCompletionDriverTest {
+  Future<void> test_definingUnit_export() async {
+    var response = await getTestCodeSuggestions('''
+exp^
+''');
+
+    check(response).suggestions.includesAll([
+      (suggestion) => suggestion
+        ..completion.isEqualTo("export '';")
+        ..kind.isKeyword
+        ..hasSelection(offset: 8),
+    ]);
+
+    if (isProtocolVersion2) {
+      check(response).suggestions.excludesAll([
+        (suggestion) => suggestion.completion.startsWith('import'),
+      ]);
+    }
+  }
+
+  Future<void> test_definingUnit_import() async {
+    var response = await getTestCodeSuggestions('''
+imp^
+''');
+
+    check(response).suggestions.includesAll([
+      (suggestion) => suggestion
+        ..completion.isEqualTo("import '';")
+        ..kind.isKeyword
+        ..hasSelection(offset: 8),
+    ]);
+
+    if (isProtocolVersion2) {
+      check(response).suggestions.excludesAll([
+        (suggestion) => suggestion.completion.startsWith('export'),
+      ]);
+    }
+  }
+
+  Future<void> test_definingUnit_part() async {
+    var response = await getTestCodeSuggestions('''
+par^
+''');
+
+    check(response).suggestions.includesAll([
+      (suggestion) => suggestion
+        ..completion.isEqualTo("part '';")
+        ..kind.isKeyword
+        ..hasSelection(offset: 6),
+    ]);
+  }
+}
diff --git a/pkg/analysis_server/test/services/completion/dart/location/test_all.dart b/pkg/analysis_server/test/services/completion/dart/location/test_all.dart
index 5fc5401..b0fe488 100644
--- a/pkg/analysis_server/test/services/completion/dart/location/test_all.dart
+++ b/pkg/analysis_server/test/services/completion/dart/location/test_all.dart
@@ -5,6 +5,7 @@
 import 'package:test_reflective_loader/test_reflective_loader.dart';
 
 import 'class_body_test.dart' as class_body;
+import 'compilation_unit_test.dart' as compilation_unit;
 import 'enum_constant_test.dart' as enum_constant;
 import 'enum_test.dart' as enum_;
 import 'field_formal_parameter_test.dart' as field_formal_parameter;
@@ -15,6 +16,7 @@
 void main() {
   defineReflectiveSuite(() {
     class_body.main();
+    compilation_unit.main();
     enum_constant.main();
     enum_.main();
     field_formal_parameter.main();