Implement Overlay.of with inherited widget (#174315)

<!--
Thanks for filing a pull request!
Reviewers are typically assigned within a week of filing a request.
To learn more about code review, see our documentation on Tree Hygiene:
https://github.com/flutter/flutter/blob/main/docs/contributing/Tree-hygiene.md
-->

for https://github.com/flutter/flutter/pull/174239

## Pre-launch Checklist

- [ ] I read the [Contributor Guide] and followed the process outlined
there for submitting PRs.
- [ ] I read the [Tree Hygiene] wiki page, which explains my
responsibilities.
- [ ] I read and followed the [Flutter Style Guide], including [Features
we expect every widget to implement].
- [ ] I signed the [CLA].
- [ ] I listed at least one issue that this PR fixes in the description
above.
- [ ] I updated/added relevant documentation (doc comments with `///`).
- [ ] I added new tests to check the change I am making, or this PR is
[test-exempt].
- [ ] I followed the [breaking change policy] and added [Data Driven
Fixes] where supported.
- [ ] All existing and new tests are passing.

If you need help, consider asking for advice on the #hackers-new channel
on [Discord].

**Note**: The Flutter team is currently trialing the use of [Gemini Code
Assist for
GitHub](https://developers.google.com/gemini-code-assist/docs/review-github-code).
Comments from the `gemini-code-assist` bot should not be taken as
authoritative feedback from the Flutter team. If you find its comments
useful you can update your code accordingly, but if you are unsure or
disagree with the feedback, please feel free to wait for a Flutter team
member's review for guidance on which automated comments should be
addressed.

<!-- Links -->
[Contributor Guide]:
https://github.com/flutter/flutter/blob/main/docs/contributing/Tree-hygiene.md#overview
[Tree Hygiene]:
https://github.com/flutter/flutter/blob/main/docs/contributing/Tree-hygiene.md
[test-exempt]:
https://github.com/flutter/flutter/blob/main/docs/contributing/Tree-hygiene.md#tests
[Flutter Style Guide]:
https://github.com/flutter/flutter/blob/main/docs/contributing/Style-guide-for-Flutter-repo.md
[Features we expect every widget to implement]:
https://github.com/flutter/flutter/blob/main/docs/contributing/Style-guide-for-Flutter-repo.md#features-we-expect-every-widget-to-implement
[CLA]: https://cla.developers.google.com/
[flutter/tests]: https://github.com/flutter/tests
[breaking change policy]:
https://github.com/flutter/flutter/blob/main/docs/contributing/Tree-hygiene.md#handling-breaking-changes
[Discord]:
https://github.com/flutter/flutter/blob/main/docs/contributing/Chat.md
[Data Driven Fixes]:
https://github.com/flutter/flutter/blob/main/docs/contributing/Data-driven-Fixes.md

packages/flutter/lib/src/widgets/framework.dart

@@ -2449,6 +2449,7 @@ abstract class BuildContext {
   /// Returns the nearest widget of the given [InheritedWidget] subclass `T` or
   /// null if an appropriate ancestor is not found.
   ///
+  /// {@template flutter.widgets.BuildContext.getInheritedWidgetOfExactType}
   /// This method does not introduce a dependency the way that the more typical
   /// [dependOnInheritedWidgetOfExactType] does, so this context will not be
   /// rebuilt if the [InheritedWidget] changes. This function is meant for those
@@ -2465,6 +2466,7 @@ abstract class BuildContext {
   /// that value is not cached and/or reused later.
   ///
   /// Calling this method is O(1) with a small constant factor.
+  /// {@endtemplate}
   T? getInheritedWidgetOfExactType<T extends InheritedWidget>();
 
   /// Obtains the element corresponding to the nearest widget of the given type `T`,

packages/flutter/lib/src/widgets/lookup_boundary.dart

@@ -105,6 +105,29 @@ class LookupBoundary extends InheritedWidget {
     return candidate.widget as T;
   }
 
+  /// Obtains the nearest widget of the given type `T` within the current
+  /// [LookupBoundary] of `context`, which must be the type of a concrete
+  /// [InheritedWidget] subclass.
+  ///
+  /// This method behaves exactly like
+  /// [BuildContext.getInheritedWidgetOfExactType], except it only
+  /// considers [InheritedWidget]s of the specified type `T` between the
+  /// provided [BuildContext] and its closest [LookupBoundary] ancestor.
+  /// [InheritedWidget]s past that [LookupBoundary] are invisible to this
+  /// method. The root of the tree is treated as an implicit lookup boundary.
+  ///
+  /// {@macro flutter.widgets.BuildContext.getInheritedWidgetOfExactType}
+  static T? getInheritedWidgetOfExactType<T extends InheritedWidget>(
+    BuildContext context, {
+    Object? aspect,
+  }) {
+    final InheritedElement? candidate = getElementForInheritedWidgetOfExactType<T>(context);
+    if (candidate == null) {
+      return null;
+    }
+    return candidate.widget as T;
+  }
+
   /// Obtains the element corresponding to the nearest widget of the given type
   /// `T` within the current [LookupBoundary] of `context`.
   ///

packages/flutter/lib/src/widgets/overlay.dart

@@ -421,7 +421,9 @@ class _OverlayEntryWidgetState extends State<_OverlayEntryWidget> {
       child: _RenderTheaterMarker(
         theater: _theater,
         overlayEntryWidgetState: this,
-        child: widget.entry.builder(context),
+        // Use a Builder so that the `widget.entry.builder` can have access to
+        // _RenderTheaterMarker.of
+        child: Builder(builder: widget.entry.builder),
       ),
     );
   }
@@ -536,11 +538,11 @@ class Overlay extends StatefulWidget {
   /// OverlayState overlay = Overlay.of(context);
   /// ```
   ///
+  /// {@template flutter.widgets.Overlay.of}
   /// If `rootOverlay` is set to true, the state from the furthest instance of
   /// this class is given instead. Useful for installing overlay entries above
   /// all subsequent instances of [Overlay].
-  ///
-  /// This method can be expensive (it walks the element tree).
+  /// {@endtemplate}
   ///
   /// See also:
   ///
@@ -599,21 +601,18 @@ class Overlay extends StatefulWidget {
   /// OverlayState? overlay = Overlay.maybeOf(context);
   /// ```
   ///
-  /// If `rootOverlay` is set to true, the state from the furthest instance of
-  /// this class is given instead. Useful for installing overlay entries above
-  /// all subsequent instances of [Overlay].
-  ///
-  /// This method can be expensive (it walks the element tree).
+  /// {@macro flutter.widgets.Overlay.of}
   ///
   /// See also:
   ///
   ///  * [Overlay.of] for a similar function that returns a non-nullable result
   ///    and throws if an [Overlay] is not found.
-
   static OverlayState? maybeOf(BuildContext context, {bool rootOverlay = false}) {
-    return rootOverlay
-        ? LookupBoundary.findRootAncestorStateOfType<OverlayState>(context)
-        : LookupBoundary.findAncestorStateOfType<OverlayState>(context);
+    return _RenderTheaterMarker.maybeOf(
+      context,
+      targetRootOverlay: rootOverlay,
+      createDependency: false,
+    )?.overlayEntryWidgetState.widget.overlayState;
   }
 
   @override
@@ -2162,18 +2161,7 @@ class _RenderTheaterMarker extends InheritedWidget {
   }
 
   static _RenderTheaterMarker of(BuildContext context, {bool targetRootOverlay = false}) {
-    final _RenderTheaterMarker? marker;
-    if (targetRootOverlay) {
-      final InheritedElement? ancestor = _rootRenderTheaterMarkerOf(
-        context.getElementForInheritedWidgetOfExactType<_RenderTheaterMarker>(),
-      );
-      assert(ancestor == null || ancestor.widget is _RenderTheaterMarker);
-      marker = ancestor != null
-          ? context.dependOnInheritedElement(ancestor) as _RenderTheaterMarker?
-          : null;
-    } else {
-      marker = context.dependOnInheritedWidgetOfExactType<_RenderTheaterMarker>();
-    }
+    final _RenderTheaterMarker? marker = maybeOf(context, targetRootOverlay: targetRootOverlay);
     if (marker != null) {
       return marker;
     }
@@ -2192,6 +2180,32 @@ class _RenderTheaterMarker extends InheritedWidget {
     ]);
   }
 
+  static _RenderTheaterMarker? maybeOf(
+    BuildContext context, {
+    bool targetRootOverlay = false,
+    bool createDependency = true,
+  }) {
+    if (targetRootOverlay) {
+      final InheritedElement? ancestor = _rootRenderTheaterMarkerOf(
+        LookupBoundary.getElementForInheritedWidgetOfExactType<_RenderTheaterMarker>(context),
+      );
+      assert(ancestor == null || ancestor.widget is _RenderTheaterMarker);
+      if (ancestor == null) {
+        return null;
+      }
+      if (createDependency) {
+        return context.dependOnInheritedElement(ancestor) as _RenderTheaterMarker;
+      }
+      return ancestor.widget as _RenderTheaterMarker;
+    }
+
+    if (createDependency) {
+      return LookupBoundary.dependOnInheritedWidgetOfExactType<_RenderTheaterMarker>(context);
+    }
+
+    return LookupBoundary.getInheritedWidgetOfExactType<_RenderTheaterMarker>(context);
+  }
+
   static InheritedElement? _rootRenderTheaterMarkerOf(InheritedElement? theaterMarkerElement) {
     assert(theaterMarkerElement == null || theaterMarkerElement.widget is _RenderTheaterMarker);
     if (theaterMarkerElement == null) {
@@ -2199,7 +2213,9 @@ class _RenderTheaterMarker extends InheritedWidget {
     }
     InheritedElement? ancestor;
     theaterMarkerElement.visitAncestorElements((Element element) {
-      ancestor = element.getElementForInheritedWidgetOfExactType<_RenderTheaterMarker>();
+      ancestor = LookupBoundary.getElementForInheritedWidgetOfExactType<_RenderTheaterMarker>(
+        element,
+      );
       return false;
     });
     return ancestor == null ? theaterMarkerElement : _rootRenderTheaterMarkerOf(ancestor);

packages/flutter/test/widgets/overlay_test.dart

@@ -206,6 +206,13 @@ void main() {
     );
   });
 
+  test('of method calls getElementForInheritedWidgetOfExactType', () async {
+    final FakeBuildContext context = FakeBuildContext();
+    expect(context.called, isFalse);
+    expect(Overlay.maybeOf(context), isNull);
+    expect(context.called, isTrue);
+  });
+
   testWidgets('insert top', (WidgetTester tester) async {
     final GlobalKey overlayKey = GlobalKey();
     final List<String> buildOrder = <String>[];
@@ -1956,5 +1963,105 @@ class StatefulTestState extends State<StatefulTestWidget> {
   }
 }
 
+class FakeBuildContext extends BuildContext {
+  bool called = false;
+
+  @override
+  bool get debugDoingBuild => throw UnimplementedError();
+
+  @override
+  InheritedWidget dependOnInheritedElement(InheritedElement ancestor, {Object? aspect}) {
+    throw UnimplementedError();
+  }
+
+  @override
+  T? dependOnInheritedWidgetOfExactType<T extends InheritedWidget>({Object? aspect}) {
+    throw UnimplementedError();
+  }
+
+  @override
+  DiagnosticsNode describeElement(
+    String name, {
+    DiagnosticsTreeStyle style = DiagnosticsTreeStyle.errorProperty,
+  }) {
+    throw UnimplementedError();
+  }
+
+  @override
+  List<DiagnosticsNode> describeMissingAncestor({required Type expectedAncestorType}) {
+    throw UnimplementedError();
+  }
+
+  @override
+  DiagnosticsNode describeOwnershipChain(String name) {
+    throw UnimplementedError();
+  }
+
+  @override
+  DiagnosticsNode describeWidget(
+    String name, {
+    DiagnosticsTreeStyle style = DiagnosticsTreeStyle.errorProperty,
+  }) {
+    throw UnimplementedError();
+  }
+
+  @override
+  void dispatchNotification(Notification notification) {}
+
+  @override
+  T? findAncestorRenderObjectOfType<T extends RenderObject>() {
+    throw UnimplementedError();
+  }
+
+  @override
+  T? findAncestorStateOfType<T extends State<StatefulWidget>>() {
+    throw UnimplementedError();
+  }
+
+  @override
+  T? findAncestorWidgetOfExactType<T extends Widget>() {
+    throw UnimplementedError();
+  }
+
+  @override
+  RenderObject? findRenderObject() {
+    throw UnimplementedError();
+  }
+
+  @override
+  T? findRootAncestorStateOfType<T extends State<StatefulWidget>>() {
+    throw UnimplementedError();
+  }
+
+  @override
+  InheritedElement? getElementForInheritedWidgetOfExactType<T extends InheritedWidget>() {
+    called = true;
+    return null;
+  }
+
+  @override
+  T? getInheritedWidgetOfExactType<T extends InheritedWidget>() {
+    throw UnimplementedError();
+  }
+
+  @override
+  bool get mounted => throw UnimplementedError();
+
+  @override
+  BuildOwner? get owner => throw UnimplementedError();
+
+  @override
+  Size? get size => throw UnimplementedError();
+
+  @override
+  void visitAncestorElements(ConditionalElementVisitor visitor) {}
+
+  @override
+  void visitChildElements(ElementVisitor visitor) {}
+
+  @override
+  Widget get widget => throw UnimplementedError();
+}
+
 /// This helper makes leak tracker forgiving the entry is not disposed.
 OverlayEntry _buildOverlayEntry(WidgetBuilder builder) => OverlayEntry(builder: builder);