Implement augmentation for top level variable tests

Change-Id: Ic7c2728fb34c12eee6b79ab95bd258a8424a926c
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/450181
Reviewed-by: Konstantin Shcheglov <scheglov@google.com>
Reviewed-by: Samuel Rawlins <srawlins@google.com>
Commit-Queue: Konstantin Shcheglov <scheglov@google.com>

Author: fshcheglov

pkg/analyzer/lib/src/dart/element/element.dart

@@ -10257,6 +10257,16 @@ class TopLevelVariableElementImpl extends PropertyInducingElementImpl
   @trackedIncludedInId
   ElementKind get kind => ElementKind.TOP_LEVEL_VARIABLE;
 
+  @trackedDirectlyOpaque
+  TopLevelVariableFragmentImpl get lastFragment {
+    globalResultRequirements?.recordOpaqueApiUse(this, 'lastFragment');
+    var current = firstFragment;
+    while (current.nextFragment != null) {
+      current = current.nextFragment!;
+    }
+    return current;
+  }
+
   @Deprecated('Use library instead')
   @override
   @trackedIndirectly

pkg/analyzer/lib/src/summary2/element_builder.dart

@@ -739,15 +739,16 @@ class ElementBuilder {
 
   void _handleTopLevelVariableFragment(
     LibraryFragmentImpl libraryFragment,
-    FragmentImpl? lastVariableFragment,
+    FragmentImpl? lastFragment,
     TopLevelVariableFragmentImpl variableFragment,
   ) {
     assert(!variableFragment.isSynthetic);
     libraryFragment.addTopLevelVariable(variableFragment);
 
-    if (variableFragment.isAugmentation &&
-        lastVariableFragment is TopLevelVariableFragmentImpl) {
-      lastVariableFragment.addFragment(variableFragment);
+    var lastVariableElement = _topLevelVariableElement(lastFragment);
+
+    if (variableFragment.isAugmentation && lastVariableElement != null) {
+      lastVariableElement.lastFragment.addFragment(variableFragment);
       return;
     }