[go_router] Adds GoRouterState to context (#2719)

diff --git a/packages/go_router/CHANGELOG.md b/packages/go_router/CHANGELOG.md
index 4d326cd..aea6f3a 100644
--- a/packages/go_router/CHANGELOG.md
+++ b/packages/go_router/CHANGELOG.md
@@ -1,6 +1,7 @@
-## NEXT
+## 5.1.2
 
-- Update README
+- Exposes uri and path parameters from GoRouter and fixes its notifications.
+- Updates README
 - Removes dynamic calls in examples.
 
 ## 5.1.1
diff --git a/packages/go_router/example/lib/shell_route.dart b/packages/go_router/example/lib/shell_route.dart
index 8c73c83..3076b3d 100644
--- a/packages/go_router/example/lib/shell_route.dart
+++ b/packages/go_router/example/lib/shell_route.dart
@@ -151,8 +151,7 @@
   }
 
   static int _calculateSelectedIndex(BuildContext context) {
-    final GoRouter route = GoRouter.of(context);
-    final String location = route.location;
+    final String location = GoRouterState.of(context).location;
     if (location.startsWith('/a')) {
       return 0;
     }
diff --git a/packages/go_router/lib/go_router.dart b/packages/go_router/lib/go_router.dart
index 09bfda9..ad2a74f 100644
--- a/packages/go_router/lib/go_router.dart
+++ b/packages/go_router/lib/go_router.dart
@@ -13,4 +13,5 @@
 export 'src/pages/custom_transition_page.dart';
 export 'src/route_data.dart' show GoRouteData, TypedGoRoute;
 export 'src/router.dart';
-export 'src/typedefs.dart' show GoRouterPageBuilder, GoRouterRedirect;
+export 'src/typedefs.dart'
+    show GoRouterPageBuilder, GoRouterRedirect, GoRouterWidgetBuilder;
diff --git a/packages/go_router/lib/src/builder.dart b/packages/go_router/lib/src/builder.dart
index 6170984..3afa6b4 100644
--- a/packages/go_router/lib/src/builder.dart
+++ b/packages/go_router/lib/src/builder.dart
@@ -47,6 +47,8 @@
   /// changes.
   final List<NavigatorObserver> observers;
 
+  final GoRouterStateRegistry _registry = GoRouterStateRegistry();
+
   /// Builds the top-level Navigator for the given [RouteMatchList].
   Widget build(
     BuildContext context,
@@ -59,17 +61,29 @@
       // empty box until then.
       return const SizedBox.shrink();
     }
-    try {
-      return tryBuild(
-          context, matchList, pop, routerNeglect, configuration.navigatorKey);
-    } on _RouteBuilderError catch (e) {
-      return _buildErrorNavigator(
-          context,
-          e,
-          Uri.parse(matchList.location.toString()),
-          pop,
-          configuration.navigatorKey);
-    }
+    return builderWithNav(
+      context,
+      Builder(
+        builder: (BuildContext context) {
+          try {
+            final Map<Page<Object?>, GoRouterState> newRegistry =
+                <Page<Object?>, GoRouterState>{};
+            final Widget result = tryBuild(context, matchList, pop,
+                routerNeglect, configuration.navigatorKey, newRegistry);
+            _registry.updateRegistry(newRegistry);
+            return GoRouterStateRegistryScope(
+                registry: _registry, child: result);
+          } on _RouteBuilderError catch (e) {
+            return _buildErrorNavigator(
+                context,
+                e,
+                Uri.parse(matchList.location.toString()),
+                pop,
+                configuration.navigatorKey);
+          }
+        },
+      ),
+    );
   }
 
   /// Builds the top-level Navigator by invoking the build method on each
@@ -83,21 +97,14 @@
     VoidCallback pop,
     bool routerNeglect,
     GlobalKey<NavigatorState> navigatorKey,
+    Map<Page<Object?>, GoRouterState> registry,
   ) {
     return builderWithNav(
       context,
-      GoRouterState(
-        configuration,
-        location: matchList.location.toString(),
-        name: null,
-        subloc: matchList.location.path,
-        queryParams: matchList.location.queryParameters,
-        queryParametersAll: matchList.location.queryParametersAll,
-        error: matchList.isError ? matchList.error : null,
-      ),
       _buildNavigator(
         pop,
-        buildPages(context, matchList, pop, routerNeglect, navigatorKey),
+        buildPages(
+            context, matchList, pop, routerNeglect, navigatorKey, registry),
         navigatorKey,
         observers: observers,
       ),
@@ -107,21 +114,22 @@
   /// Returns the top-level pages instead of the root navigator. Used for
   /// testing.
   @visibleForTesting
-  List<Page<dynamic>> buildPages(
+  List<Page<Object?>> buildPages(
       BuildContext context,
       RouteMatchList matchList,
       VoidCallback onPop,
       bool routerNeglect,
-      GlobalKey<NavigatorState> navigatorKey) {
+      GlobalKey<NavigatorState> navigatorKey,
+      Map<Page<Object?>, GoRouterState> registry) {
     try {
-      final Map<GlobalKey<NavigatorState>, List<Page<dynamic>>> keyToPage =
-          <GlobalKey<NavigatorState>, List<Page<dynamic>>>{};
+      final Map<GlobalKey<NavigatorState>, List<Page<Object?>>> keyToPage =
+          <GlobalKey<NavigatorState>, List<Page<Object?>>>{};
       final Map<String, String> params = <String, String>{};
       _buildRecursive(context, matchList, 0, onPop, routerNeglect, keyToPage,
-          params, navigatorKey);
+          params, navigatorKey, registry);
       return keyToPage[navigatorKey]!;
     } on _RouteBuilderError catch (e) {
-      return <Page<dynamic>>[
+      return <Page<Object?>>[
         _buildErrorPage(context, e, matchList.location),
       ];
     }
@@ -133,9 +141,10 @@
     int startIndex,
     VoidCallback pop,
     bool routerNeglect,
-    Map<GlobalKey<NavigatorState>, List<Page<dynamic>>> keyToPages,
+    Map<GlobalKey<NavigatorState>, List<Page<Object?>>> keyToPages,
     Map<String, String> params,
     GlobalKey<NavigatorState> navigatorKey,
+    Map<Page<Object?>, GoRouterState> registry,
   ) {
     if (startIndex >= matchList.matches.length) {
       return;
@@ -154,17 +163,17 @@
     };
     final GoRouterState state = buildState(match, newParams);
     if (route is GoRoute) {
-      final Page<dynamic> page = _buildPageForRoute(context, state, match);
-
+      final Page<Object?> page = _buildPageForRoute(context, state, match);
+      registry[page] = state;
       // If this GoRoute is for a different Navigator, add it to the
       // list of out of scope pages
       final GlobalKey<NavigatorState> goRouteNavKey =
           route.parentNavigatorKey ?? navigatorKey;
 
-      keyToPages.putIfAbsent(goRouteNavKey, () => <Page<dynamic>>[]).add(page);
+      keyToPages.putIfAbsent(goRouteNavKey, () => <Page<Object?>>[]).add(page);
 
       _buildRecursive(context, matchList, startIndex + 1, pop, routerNeglect,
-          keyToPages, newParams, navigatorKey);
+          keyToPages, newParams, navigatorKey, registry);
     } else if (route is ShellRoute) {
       // The key for the Navigator that will display this ShellRoute's page.
       final GlobalKey<NavigatorState> parentNavigatorKey = navigatorKey;
@@ -173,10 +182,10 @@
       final GlobalKey<NavigatorState> shellNavigatorKey = route.navigatorKey;
 
       // Add an entry for the parent navigator if none exists.
-      keyToPages.putIfAbsent(parentNavigatorKey, () => <Page<dynamic>>[]);
+      keyToPages.putIfAbsent(parentNavigatorKey, () => <Page<Object?>>[]);
 
       // Add an entry for the shell route's navigator
-      keyToPages.putIfAbsent(shellNavigatorKey, () => <Page<dynamic>>[]);
+      keyToPages.putIfAbsent(shellNavigatorKey, () => <Page<Object?>>[]);
 
       // Calling _buildRecursive can result in adding pages to the
       // parentNavigatorKey entry's list. Store the current length so
@@ -185,26 +194,26 @@
 
       // Build the remaining pages
       _buildRecursive(context, matchList, startIndex + 1, pop, routerNeglect,
-          keyToPages, newParams, shellNavigatorKey);
+          keyToPages, newParams, shellNavigatorKey, registry);
 
       // Build the Navigator
       final Widget child = _buildNavigator(
           pop, keyToPages[shellNavigatorKey]!, shellNavigatorKey);
 
       // Build the Page for this route
-      final Page<dynamic> page =
+      final Page<Object?> page =
           _buildPageForRoute(context, state, match, child: child);
-
+      registry[page] = state;
       // Place the ShellRoute's Page onto the list for the parent navigator.
       keyToPages
-          .putIfAbsent(parentNavigatorKey, () => <Page<dynamic>>[])
+          .putIfAbsent(parentNavigatorKey, () => <Page<Object?>>[])
           .insert(shellPageIdx, page);
     }
   }
 
   Navigator _buildNavigator(
     VoidCallback pop,
-    List<Page<dynamic>> pages,
+    List<Page<Object?>> pages,
     Key? navigatorKey, {
     List<NavigatorObserver> observers = const <NavigatorObserver>[],
   }) {
@@ -251,11 +260,11 @@
   }
 
   /// Builds a [Page] for [StackedRoute]
-  Page<dynamic> _buildPageForRoute(
+  Page<Object?> _buildPageForRoute(
       BuildContext context, GoRouterState state, RouteMatch match,
       {Widget? child}) {
     final RouteBase route = match.route;
-    Page<dynamic>? page;
+    Page<Object?>? page;
 
     if (route is GoRoute) {
       // Call the pageBuilder if it's non-null
@@ -277,8 +286,10 @@
 
     // Return the result of the route's builder() or pageBuilder()
     return page ??
-        buildPage(context, state,
-            _callRouteBuilder(context, state, match, childWidget: child));
+        // Uses a Builder to make sure its rebuild scope is limited to the page.
+        buildPage(context, state, Builder(builder: (BuildContext context) {
+          return _callRouteBuilder(context, state, match, childWidget: child);
+        }));
   }
 
   /// Calls the user-provided route builder from the [RouteMatch]'s [RouteBase].
@@ -356,7 +367,7 @@
 
   /// builds the page based on app type, i.e. MaterialApp vs. CupertinoApp
   @visibleForTesting
-  Page<dynamic> buildPage(
+  Page<Object?> buildPage(
     BuildContext context,
     GoRouterState state,
     Widget child,
@@ -393,7 +404,7 @@
       Uri uri, VoidCallback pop, GlobalKey<NavigatorState> navigatorKey) {
     return _buildNavigator(
       pop,
-      <Page<dynamic>>[
+      <Page<Object?>>[
         _buildErrorPage(context, e, uri),
       ],
       navigatorKey,
diff --git a/packages/go_router/lib/src/delegate.dart b/packages/go_router/lib/src/delegate.dart
index 9bcfa04..ae15383 100644
--- a/packages/go_router/lib/src/delegate.dart
+++ b/packages/go_router/lib/src/delegate.dart
@@ -78,12 +78,7 @@
 
     // Use the root navigator if no ShellRoute Navigators were found and didn't
     // pop
-    final NavigatorState? navigator = navigatorKey.currentState;
-
-    if (navigator == null) {
-      return SynchronousFuture<bool>(false);
-    }
-
+    final NavigatorState navigator = navigatorKey.currentState!;
     return navigator.maybePop();
   }
 
@@ -122,14 +117,15 @@
       final RouteMatch match = _matchList.matches[i];
       final RouteBase route = match.route;
       if (route is GoRoute && route.parentNavigatorKey != null) {
-        final bool canPop = route.parentNavigatorKey!.currentState!.canPop();
+        final bool canPop =
+            route.parentNavigatorKey!.currentState?.canPop() ?? false;
 
         // Continue if canPop is false.
         if (canPop) {
           return canPop;
         }
       } else if (route is ShellRoute) {
-        final bool canPop = route.navigatorKey.currentState!.canPop();
+        final bool canPop = route.navigatorKey.currentState?.canPop() ?? false;
 
         // Continue if canPop is false.
         if (canPop) {
@@ -183,6 +179,7 @@
   Future<void> setNewRoutePath(RouteMatchList configuration) {
     _matchList = configuration;
     assert(_matchList.isNotEmpty);
+    notifyListeners();
     // Use [SynchronousFuture] so that the initial url is processed
     // synchronously and remove unwanted initial animations on deep-linking
     return SynchronousFuture<void>(null);
diff --git a/packages/go_router/lib/src/misc/inherited_router.dart b/packages/go_router/lib/src/misc/inherited_router.dart
index abc5a9d..bd426d1 100644
--- a/packages/go_router/lib/src/misc/inherited_router.dart
+++ b/packages/go_router/lib/src/misc/inherited_router.dart
@@ -11,25 +11,17 @@
 ///
 /// Used for to find the current GoRouter in the widget tree. This is useful
 /// when routing from anywhere in your app.
-class InheritedGoRouter extends InheritedWidget {
+class InheritedGoRouter extends InheritedNotifier<GoRouter> {
   /// Default constructor for the inherited go router.
   const InheritedGoRouter({
     required super.child,
     required this.goRouter,
     super.key,
-  });
+  }) : super(notifier: goRouter);
 
   /// The [GoRouter] that is made available to the widget tree.
   final GoRouter goRouter;
 
-  /// Used by the Router architecture as part of the InheritedWidget.
-  @override
-  // ignore: prefer_expression_function_bodies
-  bool updateShouldNotify(covariant InheritedGoRouter oldWidget) {
-    // avoid rebuilding the widget tree if the router has not changed
-    return goRouter != oldWidget.goRouter;
-  }
-
   @override
   void debugFillProperties(DiagnosticPropertiesBuilder properties) {
     super.debugFillProperties(properties);
diff --git a/packages/go_router/lib/src/router.dart b/packages/go_router/lib/src/router.dart
index cc16088..9d78082 100644
--- a/packages/go_router/lib/src/router.dart
+++ b/packages/go_router/lib/src/router.dart
@@ -36,9 +36,7 @@
 ///  * [GoRoute], which provides APIs to define the routing table.
 ///  * [examples](https://github.com/flutter/packages/tree/main/packages/go_router/example),
 ///    which contains examples for different routing scenarios.
-class GoRouter extends ChangeNotifier
-    with NavigatorObserver
-    implements RouterConfig<RouteMatchList> {
+class GoRouter extends ChangeNotifier implements RouterConfig<RouteMatchList> {
   /// Default constructor to configure a GoRouter with a routes builder
   /// and an error page builder.
   ///
@@ -88,18 +86,14 @@
       routerNeglect: routerNeglect,
       observers: <NavigatorObserver>[
         ...observers ?? <NavigatorObserver>[],
-        this
       ],
       restorationScopeId: restorationScopeId,
       // wrap the returned Navigator to enable GoRouter.of(context).go() et al,
       // allowing the caller to wrap the navigator themselves
-      builderWithNav:
-          (BuildContext context, GoRouterState state, Navigator nav) =>
-              InheritedGoRouter(
-        goRouter: this,
-        child: nav,
-      ),
+      builderWithNav: (BuildContext context, Widget child) =>
+          InheritedGoRouter(goRouter: this, child: child),
     );
+    _routerDelegate.addListener(_handleStateMayChange);
 
     assert(() {
       log.info('setting initial location $initialLocation');
@@ -135,9 +129,23 @@
   @visibleForTesting
   RouteConfiguration get routeConfiguration => _routeConfiguration;
 
-  /// Get the current location.
-  String get location =>
-      _routerDelegate.currentConfiguration.location.toString();
+  /// Gets the current location.
+  // TODO(chunhtai): deprecates this once go_router_builder is migrated to
+  // GoRouterState.of.
+  String get location => _location;
+  String _location = '/';
+
+  /// Returns `true` if there is more than 1 page on the stack.
+  bool canPop() => _routerDelegate.canPop();
+
+  void _handleStateMayChange() {
+    final String newLocation =
+        _routerDelegate.currentConfiguration.location.toString();
+    if (_location != newLocation) {
+      _location = newLocation;
+      notifyListeners();
+    }
+  }
 
   /// Get a location from route name and parameters.
   /// This is useful for redirecting to a named location.
@@ -247,9 +255,6 @@
     );
   }
 
-  /// Returns `true` if there is more than 1 page on the stack.
-  bool canPop() => _routerDelegate.canPop();
-
   /// Pop the top page off the GoRouter's page stack.
   void pop() {
     assert(() {
@@ -276,29 +281,10 @@
     return inherited!.goRouter;
   }
 
-  /// The [Navigator] pushed `route`.
-  @override
-  void didPush(Route<dynamic> route, Route<dynamic>? previousRoute) =>
-      notifyListeners();
-
-  /// The [Navigator] popped `route`.
-  @override
-  void didPop(Route<dynamic> route, Route<dynamic>? previousRoute) =>
-      notifyListeners();
-
-  /// The [Navigator] removed `route`.
-  @override
-  void didRemove(Route<dynamic> route, Route<dynamic>? previousRoute) =>
-      notifyListeners();
-
-  /// The [Navigator] replaced `oldRoute` with `newRoute`.
-  @override
-  void didReplace({Route<dynamic>? newRoute, Route<dynamic>? oldRoute}) =>
-      notifyListeners();
-
   @override
   void dispose() {
     _routeInformationProvider.dispose();
+    _routerDelegate.removeListener(_handleStateMayChange);
     _routerDelegate.dispose();
     super.dispose();
   }
diff --git a/packages/go_router/lib/src/state.dart b/packages/go_router/lib/src/state.dart
index 7529c22..5197db7 100644
--- a/packages/go_router/lib/src/state.dart
+++ b/packages/go_router/lib/src/state.dart
@@ -2,13 +2,16 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-import 'package:flutter/foundation.dart';
+import 'package:flutter/widgets.dart';
 
+import '../go_router.dart';
 import 'configuration.dart';
+import 'misc/errors.dart';
 
 /// The route state during routing.
 ///
 /// The state contains parsed artifacts of the current URI.
+@immutable
 class GoRouterState {
   /// Default constructor for creating route state during routing.
   GoRouterState(
@@ -69,10 +72,63 @@
   /// A unique string key for this sub-route, e.g. ValueKey('/family/:fid')
   final ValueKey<String> pageKey;
 
+  /// Gets the [GoRouterState] from context.
+  ///
+  /// The returned [GoRouterState] will depends on which [GoRoute] or
+  /// [ShellRoute] the input `context` is in.
+  ///
+  /// This method only supports [GoRoute] and [ShellRoute] that generate
+  /// [ModalRoute]s. This is typically the case if one uses [GoRoute.builder],
+  /// [ShellRoute.builder], [CupertinoPage], [MaterialPage],
+  /// [CustomTransitionPage], or [NoTransitionPage].
+  ///
+  /// This method is fine to be called during [GoRoute.builder] or
+  /// [ShellRoute.builder].
+  ///
+  /// This method cannot be called during [GoRoute.pageBuilder] or
+  /// [ShellRoute.pageBuilder] since there is no [GoRouterState] to be
+  /// associated with.
+  ///
+  /// To access GoRouterState from a widget.
+  ///
+  /// ```
+  /// GoRoute(
+  ///   path: '/:id'
+  ///   builder: (_, __) => MyWidget(),
+  /// );
+  ///
+  /// class MyWidget extends StatelessWidget {
+  ///   @override
+  ///   Widget build(BuildContext context) {
+  ///     return Text('${GoRouterState.of(context).params['id']}');
+  ///   }
+  /// }
+  /// ```
+  static GoRouterState of(BuildContext context) {
+    final ModalRoute<Object?>? route = ModalRoute.of(context);
+    if (route == null) {
+      throw GoError('There is no modal route above the current context.');
+    }
+    final RouteSettings settings = route.settings;
+    if (settings is! Page<Object?>) {
+      throw GoError(
+          'The parent route must be a page route to have a GoRouterState');
+    }
+    final GoRouterStateRegistryScope? scope = context
+        .dependOnInheritedWidgetOfExactType<GoRouterStateRegistryScope>();
+    if (scope == null) {
+      throw GoError(
+          'There is no GoRouterStateRegistryScope above the current context.');
+    }
+    final GoRouterState state =
+        scope.notifier!._createPageRouteAssociation(settings, route);
+    return state;
+  }
+
   /// Get a location from route name and parameters.
   /// This is useful for redirecting to a named location.
-  // TODO(johnpryan): deprecate namedLocation API
-  // See https://github.com/flutter/flutter/issues/10772
+  @Deprecated(
+      'Uses GoRouter.of(context).routeInformationParser.namedLocation instead')
   String namedLocation(
     String name, {
     Map<String, String> params = const <String, String>{},
@@ -81,4 +137,123 @@
     return _configuration.namedLocation(name,
         params: params, queryParams: queryParams);
   }
+
+  @override
+  bool operator ==(Object other) {
+    return other is GoRouterState &&
+        other.location == location &&
+        other.subloc == subloc &&
+        other.name == name &&
+        other.path == path &&
+        other.fullpath == fullpath &&
+        other.params == params &&
+        other.queryParams == queryParams &&
+        other.queryParametersAll == queryParametersAll &&
+        other.extra == extra &&
+        other.error == error &&
+        other.pageKey == pageKey;
+  }
+
+  @override
+  int get hashCode => Object.hash(location, subloc, name, path, fullpath,
+      params, queryParams, queryParametersAll, extra, error, pageKey);
+}
+
+/// An inherited widget to host a [GoRouterStateRegistry] for the subtree.
+///
+/// Should not be used directly, consider using [GoRouterState.of] to access
+/// [GoRouterState] from the context.
+class GoRouterStateRegistryScope
+    extends InheritedNotifier<GoRouterStateRegistry> {
+  /// Creates a GoRouterStateRegistryScope.
+  const GoRouterStateRegistryScope({
+    super.key,
+    required GoRouterStateRegistry registry,
+    required super.child,
+  }) : super(notifier: registry);
+}
+
+/// A registry to record [GoRouterState] to [Page] relation.
+///
+/// Should not be used directly, consider using [GoRouterState.of] to access
+/// [GoRouterState] from the context.
+class GoRouterStateRegistry extends ChangeNotifier {
+  /// creates a [GoRouterStateRegistry].
+  GoRouterStateRegistry();
+
+  /// A [Map] that maps a [Page] to a [GoRouterState].
+  @visibleForTesting
+  final Map<Page<Object?>, GoRouterState> registry =
+      <Page<Object?>, GoRouterState>{};
+
+  final Map<Route<Object?>, Page<Object?>> _routePageAssociation =
+      <ModalRoute<Object?>, Page<Object?>>{};
+
+  GoRouterState _createPageRouteAssociation(
+      Page<Object?> page, ModalRoute<Object?> route) {
+    assert(route.settings == page);
+    assert(registry.containsKey(page));
+    final Page<Object?>? oldPage = _routePageAssociation[route];
+    if (oldPage == null) {
+      // This is a new association.
+      _routePageAssociation[route] = page;
+      // If there is an association, the registry relies on the route to remove
+      // entry from registry because it wants to preserve the GoRouterState
+      // until the route finishes the popping animations.
+      route.completed.then<void>((Object? result) {
+        // Can't use `page` directly because Route.settings may have changed during
+        // the lifetime of this route.
+        final Page<Object?> associatedPage =
+            _routePageAssociation.remove(route)!;
+        assert(registry.containsKey(associatedPage));
+        registry.remove(associatedPage);
+      });
+    } else if (oldPage != page) {
+      // Need to update the association to avoid memory leak.
+      _routePageAssociation[route] = page;
+      assert(registry.containsKey(oldPage));
+      registry.remove(oldPage);
+    }
+    assert(_routePageAssociation[route] == page);
+    return registry[page]!;
+  }
+
+  /// Updates this registry with new records.
+  void updateRegistry(Map<Page<Object?>, GoRouterState> newRegistry) {
+    bool shouldNotify = false;
+    final Set<Page<Object?>> pagesWithAssociation =
+        _routePageAssociation.values.toSet();
+    for (final MapEntry<Page<Object?>, GoRouterState> entry
+        in newRegistry.entries) {
+      final GoRouterState? existingState = registry[entry.key];
+      if (existingState != null) {
+        if (existingState != entry.value) {
+          shouldNotify =
+              shouldNotify || pagesWithAssociation.contains(entry.key);
+          registry[entry.key] = entry.value;
+        }
+        continue;
+      }
+      // Not in the _registry.
+      registry[entry.key] = entry.value;
+      // Adding or removing registry does not need to notify the listen since
+      // no one should be depending on them.
+    }
+    registry.removeWhere((Page<Object?> key, GoRouterState value) {
+      if (newRegistry.containsKey(key)) {
+        return false;
+      }
+      // For those that have page route association, it will be removed by the
+      // route future. Need to notify the listener so they can update the page
+      // route association if its page has changed.
+      if (pagesWithAssociation.contains(key)) {
+        shouldNotify = true;
+        return false;
+      }
+      return true;
+    });
+    if (shouldNotify) {
+      notifyListeners();
+    }
+  }
 }
diff --git a/packages/go_router/lib/src/typedefs.dart b/packages/go_router/lib/src/typedefs.dart
index f894488..2926994 100644
--- a/packages/go_router/lib/src/typedefs.dart
+++ b/packages/go_router/lib/src/typedefs.dart
@@ -44,8 +44,7 @@
 /// Signature of a go router builder function with navigator.
 typedef GoRouterBuilderWithNav = Widget Function(
   BuildContext context,
-  GoRouterState state,
-  Navigator navigator,
+  Widget child,
 );
 
 /// The signature of the redirect callback.
diff --git a/packages/go_router/pubspec.yaml b/packages/go_router/pubspec.yaml
index f2c4705..c315136 100644
--- a/packages/go_router/pubspec.yaml
+++ b/packages/go_router/pubspec.yaml
@@ -1,7 +1,7 @@
 name: go_router
 description: A declarative router for Flutter based on Navigation 2 supporting
   deep linking, data-driven routes and more
-version: 5.1.1
+version: 5.1.2
 repository: https://github.com/flutter/packages/tree/main/packages/go_router
 issue_tracker: https://github.com/flutter/flutter/issues?q=is%3Aissue+is%3Aopen+label%3A%22p%3A+go_router%22
 
diff --git a/packages/go_router/test/builder_test.dart b/packages/go_router/test/builder_test.dart
index 92eb17f..53c0d87 100644
--- a/packages/go_router/test/builder_test.dart
+++ b/packages/go_router/test/builder_test.dart
@@ -323,10 +323,9 @@
       configuration: configuration,
       builderWithNav: (
         BuildContext context,
-        GoRouterState state,
-        Navigator navigator,
+        Widget child,
       ) {
-        return navigator;
+        return child;
       },
       errorPageBuilder: (
         BuildContext context,
@@ -350,8 +349,8 @@
   @override
   Widget build(BuildContext context) {
     return MaterialApp(
-      home: builder.tryBuild(
-          context, matches, () {}, false, routeConfiguration.navigatorKey),
+      home: builder.tryBuild(context, matches, () {}, false,
+          routeConfiguration.navigatorKey, <Page<Object?>, GoRouterState>{}),
       // builder: (context, child) => ,
     );
   }
diff --git a/packages/go_router/test/go_router_state_test.dart b/packages/go_router/test/go_router_state_test.dart
new file mode 100644
index 0000000..c805aa9
--- /dev/null
+++ b/packages/go_router/test/go_router_state_test.dart
@@ -0,0 +1,153 @@
+// Copyright 2013 The Flutter Authors. 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:flutter/material.dart';
+import 'package:flutter_test/flutter_test.dart';
+import 'package:go_router/go_router.dart';
+import 'package:go_router/src/configuration.dart';
+
+import 'test_helpers.dart';
+
+void main() {
+  group('GoRouterState from context', () {
+    testWidgets('works in builder', (WidgetTester tester) async {
+      final List<GoRoute> routes = <GoRoute>[
+        GoRoute(
+            path: '/',
+            builder: (BuildContext context, _) {
+              final GoRouterState state = GoRouterState.of(context);
+              return Text('/ ${state.queryParams['p']}');
+            }),
+        GoRoute(
+            path: '/a',
+            builder: (BuildContext context, _) {
+              final GoRouterState state = GoRouterState.of(context);
+              return Text('/a ${state.queryParams['p']}');
+            }),
+      ];
+      final GoRouter router = await createRouter(routes, tester);
+      router.go('/?p=123');
+      await tester.pumpAndSettle();
+      expect(find.text('/ 123'), findsOneWidget);
+
+      router.go('/a?p=456');
+      await tester.pumpAndSettle();
+      expect(find.text('/a 456'), findsOneWidget);
+    });
+
+    testWidgets('works in subtree', (WidgetTester tester) async {
+      final List<GoRoute> routes = <GoRoute>[
+        GoRoute(
+            path: '/',
+            builder: (_, __) {
+              return Builder(builder: (BuildContext context) {
+                return Text(GoRouterState.of(context).location);
+              });
+            },
+            routes: <GoRoute>[
+              GoRoute(
+                  path: 'a',
+                  builder: (_, __) {
+                    return Builder(builder: (BuildContext context) {
+                      return Text(GoRouterState.of(context).location);
+                    });
+                  }),
+            ]),
+      ];
+      final GoRouter router = await createRouter(routes, tester);
+      router.go('/?p=123');
+      await tester.pumpAndSettle();
+      expect(find.text('/?p=123'), findsOneWidget);
+
+      router.go('/a');
+      await tester.pumpAndSettle();
+      expect(find.text('/a'), findsOneWidget);
+      // The query parameter is removed, so is the location in first page.
+      expect(find.text('/', skipOffstage: false), findsOneWidget);
+    });
+
+    testWidgets('registry retains GoRouterState for exiting route',
+        (WidgetTester tester) async {
+      final UniqueKey key = UniqueKey();
+      final List<GoRoute> routes = <GoRoute>[
+        GoRoute(
+            path: '/',
+            builder: (_, __) {
+              return Builder(builder: (BuildContext context) {
+                return Text(GoRouterState.of(context).location);
+              });
+            },
+            routes: <GoRoute>[
+              GoRoute(
+                  path: 'a',
+                  builder: (_, __) {
+                    return Builder(builder: (BuildContext context) {
+                      return Text(key: key, GoRouterState.of(context).location);
+                    });
+                  }),
+            ]),
+      ];
+      final GoRouter router =
+          await createRouter(routes, tester, initialLocation: '/a?p=123');
+      expect(tester.widget<Text>(find.byKey(key)).data, '/a?p=123');
+      final GoRouterStateRegistry registry = tester
+          .widget<GoRouterStateRegistryScope>(
+              find.byType(GoRouterStateRegistryScope))
+          .notifier!;
+      expect(registry.registry.length, 2);
+      router.go('/');
+      await tester.pump();
+      expect(registry.registry.length, 2);
+      // should retain the same location even if the location has changed.
+      expect(tester.widget<Text>(find.byKey(key)).data, '/a?p=123');
+
+      // Finish the pop animation.
+      await tester.pumpAndSettle();
+      expect(registry.registry.length, 1);
+      expect(find.byKey(key), findsNothing);
+    });
+
+    testWidgets('imperative pop clears out registry',
+        (WidgetTester tester) async {
+      final UniqueKey key = UniqueKey();
+      final GlobalKey<NavigatorState> nav = GlobalKey<NavigatorState>();
+      final List<GoRoute> routes = <GoRoute>[
+        GoRoute(
+            path: '/',
+            builder: (_, __) {
+              return Builder(builder: (BuildContext context) {
+                return Text(GoRouterState.of(context).location);
+              });
+            },
+            routes: <GoRoute>[
+              GoRoute(
+                  path: 'a',
+                  builder: (_, __) {
+                    return Builder(builder: (BuildContext context) {
+                      return Text(key: key, GoRouterState.of(context).location);
+                    });
+                  }),
+            ]),
+      ];
+      await createRouter(routes, tester,
+          initialLocation: '/a?p=123', navigatorKey: nav);
+      expect(tester.widget<Text>(find.byKey(key)).data, '/a?p=123');
+      final GoRouterStateRegistry registry = tester
+          .widget<GoRouterStateRegistryScope>(
+              find.byType(GoRouterStateRegistryScope))
+          .notifier!;
+      expect(registry.registry.length, 2);
+      nav.currentState!.pop();
+      await tester.pump();
+      expect(registry.registry.length, 2);
+      // should retain the same location even if the location has changed.
+      expect(tester.widget<Text>(find.byKey(key)).data, '/a?p=123');
+
+      // Finish the pop animation.
+      await tester.pumpAndSettle();
+      expect(registry.registry.length, 1);
+      expect(find.byKey(key), findsNothing);
+    });
+  });
+}
diff --git a/packages/go_router/test/go_router_test.dart b/packages/go_router/test/go_router_test.dart
index b884493..e10f11a 100644
--- a/packages/go_router/test/go_router_test.dart
+++ b/packages/go_router/test/go_router_test.dart
@@ -46,7 +46,7 @@
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
       expect(matches.first.fullpath, '/');
-      expect(router.screenFor(matches.first).runtimeType, HomeScreen);
+      expect(find.byType(HomeScreen), findsOneWidget);
     });
 
     testWidgets('If there is more than one route to match, use the first match',
@@ -61,7 +61,7 @@
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
       expect(matches.first.fullpath, '/');
-      expect(router.screenFor(matches.first).runtimeType, DummyScreen);
+      expect(find.byType(DummyScreen), findsOneWidget);
     });
 
     test('empty path', () {
@@ -120,7 +120,7 @@
       await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, TestErrorScreen);
+      expect(find.byType(TestErrorScreen), findsOneWidget);
     });
 
     testWidgets('match 2nd top level route', (WidgetTester tester) async {
@@ -137,10 +137,11 @@
 
       final GoRouter router = await createRouter(routes, tester);
       router.go('/login');
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
       expect(matches.first.subloc, '/login');
-      expect(router.screenFor(matches.first).runtimeType, LoginScreen);
+      expect(find.byType(LoginScreen), findsOneWidget);
     });
 
     testWidgets('match 2nd top level route with subroutes',
@@ -165,10 +166,11 @@
 
       final GoRouter router = await createRouter(routes, tester);
       router.go('/login');
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
       expect(matches.first.subloc, '/login');
-      expect(router.screenFor(matches.first).runtimeType, LoginScreen);
+      expect(find.byType(LoginScreen), findsOneWidget);
     });
 
     testWidgets('match top level route when location has trailing /',
@@ -188,10 +190,11 @@
 
       final GoRouter router = await createRouter(routes, tester);
       router.go('/login/');
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
       expect(matches.first.subloc, '/login');
-      expect(router.screenFor(matches.first).runtimeType, LoginScreen);
+      expect(find.byType(LoginScreen), findsOneWidget);
     });
 
     testWidgets('match top level route when location has trailing / (2)',
@@ -206,10 +209,11 @@
 
       final GoRouter router = await createRouter(routes, tester);
       router.go('/profile/');
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
       expect(matches.first.subloc, '/profile/foo');
-      expect(router.screenFor(matches.first).runtimeType, DummyScreen);
+      expect(find.byType(DummyScreen), findsOneWidget);
     });
 
     testWidgets('match top level route when location has trailing / (3)',
@@ -224,10 +228,47 @@
 
       final GoRouter router = await createRouter(routes, tester);
       router.go('/profile/?bar=baz');
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
       expect(matches.first.subloc, '/profile/foo');
-      expect(router.screenFor(matches.first).runtimeType, DummyScreen);
+      expect(find.byType(DummyScreen), findsOneWidget);
+    });
+
+    testWidgets('can access GoRouter parameters from builder',
+        (WidgetTester tester) async {
+      final List<GoRoute> routes = <GoRoute>[
+        GoRoute(path: '/', redirect: (_, __) => '/1'),
+        GoRoute(
+            path: '/:id',
+            builder: (BuildContext context, GoRouterState state) {
+              return Text(GoRouter.of(context).location);
+            }),
+      ];
+
+      final GoRouter router = await createRouter(routes, tester);
+      expect(find.text('/1'), findsOneWidget);
+      router.go('/123?id=456');
+      await tester.pumpAndSettle();
+      expect(find.text('/123?id=456'), findsOneWidget);
+    });
+
+    testWidgets('can access GoRouter parameters from error builder',
+        (WidgetTester tester) async {
+      final List<GoRoute> routes = <GoRoute>[
+        GoRoute(path: '/', builder: dummy),
+      ];
+
+      final GoRouter router = await createRouter(routes, tester,
+          errorBuilder: (BuildContext context, GoRouterState state) {
+        return Text(GoRouter.of(context).location);
+      });
+      router.go('/123?id=456');
+      await tester.pumpAndSettle();
+      expect(find.text('/123?id=456'), findsOneWidget);
+      router.go('/1234?id=456');
+      await tester.pumpAndSettle();
+      expect(find.text('/1234?id=456'), findsOneWidget);
     });
 
     testWidgets('match sub-route', (WidgetTester tester) async {
@@ -248,12 +289,13 @@
 
       final GoRouter router = await createRouter(routes, tester);
       router.go('/login');
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches.length, 2);
       expect(matches.first.subloc, '/');
-      expect(router.screenFor(matches.first).runtimeType, HomeScreen);
+      expect(find.byType(HomeScreen, skipOffstage: false), findsOneWidget);
       expect(matches[1].subloc, '/login');
-      expect(router.screenFor(matches[1]).runtimeType, LoginScreen);
+      expect(find.byType(LoginScreen), findsOneWidget);
     });
 
     testWidgets('match sub-routes', (WidgetTester tester) async {
@@ -289,39 +331,42 @@
         final List<RouteMatch> matches = router.routerDelegate.matches.matches;
         expect(matches, hasLength(1));
         expect(matches.first.fullpath, '/');
-        expect(router.screenFor(matches.first).runtimeType, HomeScreen);
+        expect(find.byType(HomeScreen), findsOneWidget);
       }
 
       router.go('/login');
+      await tester.pumpAndSettle();
       {
         final List<RouteMatch> matches = router.routerDelegate.matches.matches;
         expect(matches.length, 2);
         expect(matches.first.subloc, '/');
-        expect(router.screenFor(matches.first).runtimeType, HomeScreen);
+        expect(find.byType(HomeScreen, skipOffstage: false), findsOneWidget);
         expect(matches[1].subloc, '/login');
-        expect(router.screenFor(matches[1]).runtimeType, LoginScreen);
+        expect(find.byType(LoginScreen), findsOneWidget);
       }
 
       router.go('/family/f2');
+      await tester.pumpAndSettle();
       {
         final List<RouteMatch> matches = router.routerDelegate.matches.matches;
         expect(matches.length, 2);
         expect(matches.first.subloc, '/');
-        expect(router.screenFor(matches.first).runtimeType, HomeScreen);
+        expect(find.byType(HomeScreen, skipOffstage: false), findsOneWidget);
         expect(matches[1].subloc, '/family/f2');
-        expect(router.screenFor(matches[1]).runtimeType, FamilyScreen);
+        expect(find.byType(FamilyScreen), findsOneWidget);
       }
 
       router.go('/family/f2/person/p1');
+      await tester.pumpAndSettle();
       {
         final List<RouteMatch> matches = router.routerDelegate.matches.matches;
         expect(matches.length, 3);
         expect(matches.first.subloc, '/');
-        expect(router.screenFor(matches.first).runtimeType, HomeScreen);
+        expect(find.byType(HomeScreen, skipOffstage: false), findsOneWidget);
         expect(matches[1].subloc, '/family/f2');
-        expect(router.screenFor(matches[1]).runtimeType, FamilyScreen);
+        expect(find.byType(FamilyScreen, skipOffstage: false), findsOneWidget);
         expect(matches[2].subloc, '/family/f2/person/p1');
-        expect(router.screenFor(matches[2]).runtimeType, PersonScreen);
+        expect(find.byType(PersonScreen), findsOneWidget);
       }
     });
 
@@ -361,19 +406,22 @@
 
       final GoRouter router = await createRouter(routes, tester);
       router.go('/bar');
+      await tester.pumpAndSettle();
       List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(2));
-      expect(router.screenFor(matches[1]).runtimeType, Page1Screen);
+      expect(find.byType(Page1Screen), findsOneWidget);
 
       router.go('/foo/bar');
+      await tester.pumpAndSettle();
       matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(2));
-      expect(router.screenFor(matches[1]).runtimeType, FamilyScreen);
+      expect(find.byType(FamilyScreen), findsOneWidget);
 
       router.go('/foo');
+      await tester.pumpAndSettle();
       matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(2));
-      expect(router.screenFor(matches[1]).runtimeType, Page2Screen);
+      expect(find.byType(Page2Screen), findsOneWidget);
     });
 
     testWidgets('router state', (WidgetTester tester) async {
@@ -481,6 +529,7 @@
       final GoRouter router = await createRouter(routes, tester);
       const String loc = '/FaMiLy/f2';
       router.go(loc);
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
 
       // NOTE: match the lower case, since subloc is canonicalized to match the
@@ -489,7 +538,7 @@
       expect(router.location.toLowerCase(), loc.toLowerCase());
 
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, FamilyScreen);
+      expect(find.byType(FamilyScreen), findsOneWidget);
     });
 
     testWidgets(
@@ -504,9 +553,10 @@
 
       final GoRouter router = await createRouter(routes, tester);
       router.go('/user');
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, DummyScreen);
+      expect(find.byType(DummyScreen), findsOneWidget);
     });
 
     testWidgets('Handles the Android back button correctly',
@@ -1063,9 +1113,8 @@
       final GoRouter router = await createRouter(routes, tester);
       router.goNamed('person',
           params: <String, String>{'fid': 'f2', 'pid': 'p1'});
-
-      final List<RouteMatch> matches = router.routerDelegate.matches.matches;
-      expect(router.screenFor(matches.last).runtimeType, PersonScreen);
+      await tester.pumpAndSettle();
+      expect(find.byType(PersonScreen), findsOneWidget);
     });
 
     testWidgets('preserve path param spaces and slashes',
@@ -1087,10 +1136,11 @@
           .namedLocation('page1', params: <String, String>{'param1': param1});
       log.info('loc= $loc');
       router.go(loc);
+      await tester.pumpAndSettle();
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       log.info('param1= ${matches.first.decodedParams['param1']}');
-      expect(router.screenFor(matches.first).runtimeType, DummyScreen);
+      expect(find.byType(DummyScreen), findsOneWidget);
       expect(matches.first.decodedParams['param1'], param1);
     });
 
@@ -1112,9 +1162,9 @@
       final String loc = router.namedLocation('page1',
           queryParams: <String, String>{'param1': param1});
       router.go(loc);
-      await tester.pump();
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
-      expect(router.screenFor(matches.first).runtimeType, DummyScreen);
+      expect(find.byType(DummyScreen), findsOneWidget);
       expect(matches.first.queryParams['param1'], param1);
     });
   });
@@ -1337,10 +1387,10 @@
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, TestErrorScreen);
-      expect(
-          (router.screenFor(matches.first) as TestErrorScreen).ex, isNotNull);
-      log.info((router.screenFor(matches.first) as TestErrorScreen).ex);
+      expect(find.byType(TestErrorScreen), findsOneWidget);
+      final TestErrorScreen screen =
+          tester.widget<TestErrorScreen>(find.byType(TestErrorScreen));
+      expect(screen.ex, isNotNull);
     });
 
     testWidgets('route-level redirect loop', (WidgetTester tester) async {
@@ -1362,10 +1412,10 @@
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, TestErrorScreen);
-      expect(
-          (router.screenFor(matches.first) as TestErrorScreen).ex, isNotNull);
-      log.info((router.screenFor(matches.first) as TestErrorScreen).ex);
+      expect(find.byType(TestErrorScreen), findsOneWidget);
+      final TestErrorScreen screen =
+          tester.widget<TestErrorScreen>(find.byType(TestErrorScreen));
+      expect(screen.ex, isNotNull);
     });
 
     testWidgets('mixed redirect loop', (WidgetTester tester) async {
@@ -1384,10 +1434,10 @@
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, TestErrorScreen);
-      expect(
-          (router.screenFor(matches.first) as TestErrorScreen).ex, isNotNull);
-      log.info((router.screenFor(matches.first) as TestErrorScreen).ex);
+      expect(find.byType(TestErrorScreen), findsOneWidget);
+      final TestErrorScreen screen =
+          tester.widget<TestErrorScreen>(find.byType(TestErrorScreen));
+      expect(screen.ex, isNotNull);
     });
 
     testWidgets('top-level redirect loop w/ query params',
@@ -1405,10 +1455,10 @@
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, TestErrorScreen);
-      expect(
-          (router.screenFor(matches.first) as TestErrorScreen).ex, isNotNull);
-      log.info((router.screenFor(matches.first) as TestErrorScreen).ex);
+      expect(find.byType(TestErrorScreen), findsOneWidget);
+      final TestErrorScreen screen =
+          tester.widget<TestErrorScreen>(find.byType(TestErrorScreen));
+      expect(screen.ex, isNotNull);
     });
 
     testWidgets('expect null path/fullpath on top-level redirect',
@@ -1466,7 +1516,7 @@
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, LoginScreen);
+      expect(find.byType(LoginScreen), findsOneWidget);
     });
 
     testWidgets('route-level redirect state', (WidgetTester tester) async {
@@ -1495,7 +1545,7 @@
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, HomeScreen);
+      expect(find.byType(HomeScreen), findsOneWidget);
     });
 
     testWidgets('sub-sub-route-level redirect params',
@@ -1536,9 +1586,10 @@
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches.length, 3);
-      expect(router.screenFor(matches.first).runtimeType, HomeScreen);
-      expect(router.screenFor(matches[1]).runtimeType, FamilyScreen);
-      final PersonScreen page = router.screenFor(matches[2]) as PersonScreen;
+      expect(find.byType(HomeScreen, skipOffstage: false), findsOneWidget);
+      expect(find.byType(FamilyScreen, skipOffstage: false), findsOneWidget);
+      final PersonScreen page =
+          tester.widget<PersonScreen>(find.byType(PersonScreen));
       expect(page.fid, 'f2');
       expect(page.pid, 'p1');
     });
@@ -1554,10 +1605,10 @@
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
-      expect(router.screenFor(matches.first).runtimeType, TestErrorScreen);
-      expect(
-          (router.screenFor(matches.first) as TestErrorScreen).ex, isNotNull);
-      log.info((router.screenFor(matches.first) as TestErrorScreen).ex);
+      expect(find.byType(TestErrorScreen), findsOneWidget);
+      final TestErrorScreen screen =
+          tester.widget<TestErrorScreen>(find.byType(TestErrorScreen));
+      expect(screen.ex, isNotNull);
     });
 
     testWidgets('extra not null in redirect', (WidgetTester tester) async {
@@ -1752,11 +1803,12 @@
       for (final String fid in <String>['f2', 'F2']) {
         final String loc = '/family/$fid';
         router.go(loc);
+        await tester.pumpAndSettle();
         final List<RouteMatch> matches = router.routerDelegate.matches.matches;
 
         expect(router.location, loc);
         expect(matches, hasLength(1));
-        expect(router.screenFor(matches.first).runtimeType, FamilyScreen);
+        expect(find.byType(FamilyScreen), findsOneWidget);
         expect(matches.first.decodedParams['fid'], fid);
       }
     });
@@ -1780,11 +1832,12 @@
       for (final String fid in <String>['f2', 'F2']) {
         final String loc = '/family?fid=$fid';
         router.go(loc);
+        await tester.pumpAndSettle();
         final List<RouteMatch> matches = router.routerDelegate.matches.matches;
 
         expect(router.location, loc);
         expect(matches, hasLength(1));
-        expect(router.screenFor(matches.first).runtimeType, FamilyScreen);
+        expect(find.byType(FamilyScreen), findsOneWidget);
         expect(matches.first.queryParams['fid'], fid);
       }
     });
@@ -1805,10 +1858,11 @@
       final GoRouter router = await createRouter(routes, tester);
       final String loc = '/page1/${Uri.encodeComponent(param1)}';
       router.go(loc);
+      await tester.pumpAndSettle();
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       log.info('param1= ${matches.first.decodedParams['param1']}');
-      expect(router.screenFor(matches.first).runtimeType, DummyScreen);
+      expect(find.byType(DummyScreen), findsOneWidget);
       expect(matches.first.decodedParams['param1'], param1);
     });
 
@@ -1827,16 +1881,18 @@
 
       final GoRouter router = await createRouter(routes, tester);
       router.go('/page1?param1=$param1');
+      await tester.pumpAndSettle();
 
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
-      expect(router.screenFor(matches.first).runtimeType, DummyScreen);
+      expect(find.byType(DummyScreen), findsOneWidget);
       expect(matches.first.queryParams['param1'], param1);
 
       final String loc = '/page1?param1=${Uri.encodeQueryComponent(param1)}';
       router.go(loc);
+      await tester.pumpAndSettle();
 
       final List<RouteMatch> matches2 = router.routerDelegate.matches.matches;
-      expect(router.screenFor(matches2[0]).runtimeType, DummyScreen);
+      expect(find.byType(DummyScreen), findsOneWidget);
       expect(matches2[0].queryParams['param1'], param1);
     });
 
@@ -1879,7 +1935,7 @@
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
       expect(matches.first.fullpath, '/');
-      expect(router.screenFor(matches.first).runtimeType, HomeScreen);
+      expect(find.byType(HomeScreen), findsOneWidget);
     });
 
     testWidgets('duplicate path + query param', (WidgetTester tester) async {
@@ -1898,11 +1954,11 @@
       );
 
       router.go('/0?id=1');
-      await tester.pump();
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
       expect(matches, hasLength(1));
       expect(matches.first.fullpath, '/:id');
-      expect(router.screenFor(matches.first).runtimeType, HomeScreen);
+      expect(find.byType(HomeScreen), findsOneWidget);
     });
 
     testWidgets('push + query param', (WidgetTester tester) async {
@@ -1929,16 +1985,15 @@
       );
 
       router.go('/family?fid=f2');
-      await tester.pump();
+      await tester.pumpAndSettle();
       router.push('/person?fid=f2&pid=p1');
-      await tester.pump();
-      final FamilyScreen page1 =
-          router.screenFor(router.routerDelegate.matches.matches.first)
-              as FamilyScreen;
+      await tester.pumpAndSettle();
+      final FamilyScreen page1 = tester
+          .widget<FamilyScreen>(find.byType(FamilyScreen, skipOffstage: false));
       expect(page1.fid, 'f2');
 
-      final PersonScreen page2 = router
-          .screenFor(router.routerDelegate.matches.matches[1]) as PersonScreen;
+      final PersonScreen page2 =
+          tester.widget<PersonScreen>(find.byType(PersonScreen));
       expect(page2.fid, 'f2');
       expect(page2.pid, 'p1');
     });
@@ -1967,16 +2022,15 @@
       );
 
       router.go('/family', extra: <String, String>{'fid': 'f2'});
-      await tester.pump();
+      await tester.pumpAndSettle();
       router.push('/person', extra: <String, String>{'fid': 'f2', 'pid': 'p1'});
-      await tester.pump();
-      final FamilyScreen page1 =
-          router.screenFor(router.routerDelegate.matches.matches.first)
-              as FamilyScreen;
+      await tester.pumpAndSettle();
+      final FamilyScreen page1 = tester
+          .widget<FamilyScreen>(find.byType(FamilyScreen, skipOffstage: false));
       expect(page1.fid, 'f2');
 
-      final PersonScreen page2 = router
-          .screenFor(router.routerDelegate.matches.matches[1]) as PersonScreen;
+      final PersonScreen page2 =
+          tester.widget<PersonScreen>(find.byType(PersonScreen));
       expect(page2.fid, 'f2');
       expect(page2.pid, 'p1');
     });
@@ -2012,12 +2066,12 @@
       const String loc = '/family/$fid/person/$pid';
 
       router.push(loc);
-      await tester.pump();
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
 
       expect(router.location, loc);
       expect(matches, hasLength(2));
-      expect(router.screenFor(matches.last).runtimeType, PersonScreen);
+      expect(find.byType(PersonScreen), findsOneWidget);
       expect(matches.last.decodedParams['fid'], fid);
       expect(matches.last.decodedParams['pid'], pid);
     });
@@ -2059,13 +2113,13 @@
         'q1': 'v1',
         'q2': <String>['v2', 'v3'],
       });
-      await tester.pump();
+      await tester.pumpAndSettle();
       final List<RouteMatch> matches = router.routerDelegate.matches.matches;
 
       expect(matches, hasLength(1));
       expectLocationWithQueryParams(router.location);
       expect(
-        router.screenFor(matches.last),
+        tester.widget<DummyScreen>(find.byType(DummyScreen)),
         isA<DummyScreen>().having(
           (DummyScreen screen) => screen.queryParametersAll,
           'screen.queryParametersAll',
@@ -2109,13 +2163,62 @@
     final GoRouter router = await createRouter(routes, tester);
 
     router.go('/page?q1=v1&q2=v2&q2=v3');
-    await tester.pump();
+    await tester.pumpAndSettle();
     final List<RouteMatch> matches = router.routerDelegate.matches.matches;
 
     expect(matches, hasLength(1));
     expectLocationWithQueryParams(router.location);
     expect(
-      router.screenFor(matches.last),
+      tester.widget<DummyScreen>(find.byType(DummyScreen)),
+      isA<DummyScreen>().having(
+        (DummyScreen screen) => screen.queryParametersAll,
+        'screen.queryParametersAll',
+        queryParametersAll,
+      ),
+    );
+  });
+
+  testWidgets('goRouter should rebuild widget if ',
+      (WidgetTester tester) async {
+    const Map<String, dynamic> queryParametersAll = <String, List<dynamic>>{
+      'q1': <String>['v1'],
+      'q2': <String>['v2', 'v3'],
+    };
+    void expectLocationWithQueryParams(String location) {
+      final Uri uri = Uri.parse(location);
+      expect(uri.path, '/page');
+      expect(uri.queryParametersAll, queryParametersAll);
+    }
+
+    final List<GoRoute> routes = <GoRoute>[
+      GoRoute(
+        path: '/',
+        builder: (BuildContext context, GoRouterState state) =>
+            const HomeScreen(),
+      ),
+      GoRoute(
+        name: 'page',
+        path: '/page',
+        builder: (BuildContext context, GoRouterState state) {
+          expect(state.queryParametersAll, queryParametersAll);
+          expectLocationWithQueryParams(state.location);
+          return DummyScreen(
+            queryParametersAll: state.queryParametersAll,
+          );
+        },
+      ),
+    ];
+
+    final GoRouter router = await createRouter(routes, tester);
+
+    router.go('/page?q1=v1&q2=v2&q2=v3');
+    await tester.pumpAndSettle();
+    final List<RouteMatch> matches = router.routerDelegate.matches.matches;
+
+    expect(matches, hasLength(1));
+    expectLocationWithQueryParams(router.location);
+    expect(
+      tester.widget<DummyScreen>(find.byType(DummyScreen)),
       isA<DummyScreen>().having(
         (DummyScreen screen) => screen.queryParametersAll,
         'screen.queryParametersAll',
@@ -2415,46 +2518,6 @@
       await tester.pump();
     });
 
-    testWidgets('didPush notifies listeners', (WidgetTester tester) async {
-      await createGoRouter(tester)
-        ..addListener(expectAsync0(() {}))
-        ..didPush(
-          MaterialPageRoute<void>(builder: (_) => const Text('Current route')),
-          MaterialPageRoute<void>(builder: (_) => const Text('Previous route')),
-        );
-    });
-
-    testWidgets('didPop notifies listeners', (WidgetTester tester) async {
-      await createGoRouter(tester)
-        ..addListener(expectAsync0(() {}))
-        ..didPop(
-          MaterialPageRoute<void>(builder: (_) => const Text('Current route')),
-          MaterialPageRoute<void>(builder: (_) => const Text('Previous route')),
-        );
-    });
-
-    testWidgets('didRemove notifies listeners', (WidgetTester tester) async {
-      await createGoRouter(tester)
-        ..addListener(expectAsync0(() {}))
-        ..didRemove(
-          MaterialPageRoute<void>(builder: (_) => const Text('Current route')),
-          MaterialPageRoute<void>(builder: (_) => const Text('Previous route')),
-        );
-    });
-
-    testWidgets('didReplace notifies listeners', (WidgetTester tester) async {
-      await createGoRouter(tester)
-        ..addListener(expectAsync0(() {}))
-        ..didReplace(
-          newRoute: MaterialPageRoute<void>(
-            builder: (_) => const Text('Current route'),
-          ),
-          oldRoute: MaterialPageRoute<void>(
-            builder: (_) => const Text('Previous route'),
-          ),
-        );
-    });
-
     group('canPop', () {
       testWidgets(
         'It should return false if Navigator.canPop() returns false.',
diff --git a/packages/go_router/test/test_helpers.dart b/packages/go_router/test/test_helpers.dart
index 289cab4..be35707 100644
--- a/packages/go_router/test/test_helpers.dart
+++ b/packages/go_router/test/test_helpers.dart
@@ -9,8 +9,6 @@
 import 'package:flutter/src/foundation/diagnostics.dart';
 import 'package:flutter_test/flutter_test.dart';
 import 'package:go_router/go_router.dart';
-import 'package:go_router/src/match.dart';
-import 'package:go_router/src/matching.dart';
 
 Future<GoRouter> createGoRouter(WidgetTester tester) async {
   final GoRouter goRouter = GoRouter(
@@ -144,14 +142,16 @@
   String initialLocation = '/',
   int redirectLimit = 5,
   GlobalKey<NavigatorState>? navigatorKey,
+  GoRouterWidgetBuilder? errorBuilder,
 }) async {
   final GoRouter goRouter = GoRouter(
     routes: routes,
     redirect: redirect,
     initialLocation: initialLocation,
     redirectLimit: redirectLimit,
-    errorBuilder: (BuildContext context, GoRouterState state) =>
-        TestErrorScreen(state.error!),
+    errorBuilder: errorBuilder ??
+        (BuildContext context, GoRouterState state) =>
+            TestErrorScreen(state.error!),
     navigatorKey: navigatorKey,
   );
   await tester.pumpWidget(
@@ -219,26 +219,6 @@
 
 final GlobalKey<NavigatorState> navigatorKey = GlobalKey<NavigatorState>();
 
-extension Extension on GoRouter {
-  Page<dynamic> _pageFor(RouteMatch match) {
-    final RouteMatchList matchList = routerDelegate.matches;
-    final int i = matchList.matches.indexOf(match);
-    final List<Page<dynamic>> pages = routerDelegate.builder
-        .buildPages(
-          DummyBuildContext(),
-          matchList,
-          () {},
-          false,
-          navigatorKey,
-        )
-        .toList();
-    return pages[i];
-  }
-
-  Widget screenFor(RouteMatch match) =>
-      (_pageFor(match) as MaterialPage<void>).child;
-}
-
 class DummyBuildContext implements BuildContext {
   @override
   bool get debugDoingBuild => throw UnimplementedError();