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

I previously fixed this for variables, but functions are the same and can have elements but no fragments too.

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

Author: DanTup

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

@@ -401,7 +401,7 @@ class _ElementMapper2 extends GeneralizingAstVisitor<Element2> {
 
   @override
   Element2? visitFunctionDeclaration(FunctionDeclaration node) {
-    return node.declaredFragment?.element;
+    return node.declaredElement2 ?? node.declaredFragment?.element;
   }
 
   @override

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

@@ -586,7 +586,18 @@ enum E {
     expect(element, findElement2.extensionType('A'));
   }
 
-  test_locate_FunctionDeclaration() async {
+  test_locate_FunctionDeclaration_local() async {
+    await resolveTestCode(r'''
+void f() {
+  int g() => 3;
+}
+''');
+    var node = findNode.functionDeclaration('g');
+    var element = ElementLocator.locate2(node);
+    expect(element, isLocalFunctionElement);
+  }
+
+  test_locate_FunctionDeclaration_topLevel() async {
     await resolveTestCode('int f() => 3;');
     var node = findNode.functionDeclaration('f');
     var element = ElementLocator.locate2(node);