diff --git a/pkg/linter/analyzer_use_new_elements.txt b/pkg/linter/analyzer_use_new_elements.txt
index d78cd8f..cc5b1ba 100644
--- a/pkg/linter/analyzer_use_new_elements.txt
+++ b/pkg/linter/analyzer_use_new_elements.txt
@@ -71,6 +71,7 @@
 lib/src/rules/file_names.dart
 lib/src/rules/flutter_style_todos.dart
 lib/src/rules/hash_and_equals.dart
+lib/src/rules/implementation_imports.dart
 lib/src/rules/implicit_call_tearoffs.dart
 lib/src/rules/implicit_reopen.dart
 lib/src/rules/invalid_case_patterns.dart
@@ -163,6 +164,7 @@
 lib/src/rules/type_annotate_public_apis.dart
 lib/src/rules/type_literal_in_constant_pattern.dart
 lib/src/rules/unintended_html_in_doc_comment.dart
+lib/src/rules/unnecessary_await_in_return.dart
 lib/src/rules/unnecessary_brace_in_string_interps.dart
 lib/src/rules/unnecessary_breaks.dart
 lib/src/rules/unnecessary_const.dart
diff --git a/pkg/linter/lib/src/rules/implementation_imports.dart b/pkg/linter/lib/src/rules/implementation_imports.dart
index a0d4dc1..5a894bb 100644
--- a/pkg/linter/lib/src/rules/implementation_imports.dart
+++ b/pkg/linter/lib/src/rules/implementation_imports.dart
@@ -4,8 +4,10 @@
 
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/ast/visitor.dart';
+import 'package:analyzer/dart/element/element2.dart';
 
 import '../analyzer.dart';
+import '../extensions.dart';
 
 const _desc = r"Don't import implementation files from another package.";
 
@@ -46,30 +48,32 @@
   @override
   void registerNodeProcessors(
       NodeLintRegistry registry, LinterContext context) {
-    var visitor = _Visitor(this);
+    var libraryUri = context.libraryElement2?.uri;
+    if (libraryUri == null) return;
+
+    // If the source URI is not a `package` URI bail out.
+    if (!isPackage(libraryUri)) return;
+
+    var visitor = _Visitor(this, libraryUri);
     registry.addImportDirective(this, visitor);
   }
 }
 
 class _Visitor extends SimpleAstVisitor<void> {
   final LintRule rule;
+  final Uri sourceUri;
 
-  _Visitor(this.rule);
+  _Visitor(this.rule, this.sourceUri);
 
   @override
   void visitImportDirective(ImportDirective node) {
-    var importUri = node.element?.importedLibrary?.source.uri;
-    var sourceUri = node.element?.source.uri;
+    Uri? importUri;
+    if (node.libraryImport?.uri case DirectiveUriWithSource importedLibrary) {
+      importUri = importedLibrary.source.uri;
+    }
 
     // Test for 'package:*/src/'.
-    if (!isImplementation(importUri)) {
-      return;
-    }
-
-    // If the source URI is not a `package` URI bail out.
-    if (!isPackage(sourceUri)) {
-      return;
-    }
+    if (!isImplementation(importUri)) return;
 
     if (!samePackage(importUri, sourceUri)) {
       rule.reportLint(node.uri);