[analyzer] Update ElementLocator2 to handle both local + top level variables

Change-Id: I3f84a381c4c103d478e74a28d7ebbe85ee46606c
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/391441
Reviewed-by: Brian Wilkerson <[email protected]>
Commit-Queue: Brian Wilkerson <[email protected]>
Reviewed-by: Konstantin Shcheglov <[email protected]>

Author: DanTup

pkg/analyzer/lib/src/dart/ast/element_locator.dart

@@ -559,6 +559,6 @@ class _ElementMapper2 extends GeneralizingAstVisitor<Element2> {
 
   @override
   Element2? visitVariableDeclaration(VariableDeclaration node) {
-    return node.declaredFragment?.element;
+    return node.declaredFragment?.element ?? node.declaredElement2;
   }
 }

pkg/analyzer/test/src/dart/ast/element_locator_test.dart

@@ -884,7 +884,18 @@ void f(int a) {
     expect(element, isLibraryElement2);
   }
 
-  test_locate_VariableDeclaration() async {
+  test_locate_VariableDeclaration_Local() async {
+    await resolveTestCode(r'''
+f() {
+  var x = 42;
+}
+''');
+    var node = findNode.variableDeclaration('x =');
+    var element = ElementLocator.locate2(node);
+    expect(element, isLocalVariableElement2);
+  }
+
+  test_locate_VariableDeclaration_TopLevel() async {
     await resolveTestCode('var x = 42;');
     var node = findNode.variableDeclaration('x =');
     var element = ElementLocator.locate2(node);

pkg/analyzer/test/util/element_type_matchers.dart

@@ -53,6 +53,8 @@ const isLocalFunctionElement = TypeMatcher<LocalFunctionElement>();
 
 const isLocalVariableElement = TypeMatcher<LocalVariableElement>();
 
+const isLocalVariableElement2 = TypeMatcher<LocalVariableElement2>();
+
 const isMethodElement = TypeMatcher<MethodElement>();
 
 const isMethodElement2 = TypeMatcher<MethodElement2>();