Do not eagerly allocate inherited widget caches when initializing element tree (#95596)

Author: Jonah Williams

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

@@ -395,6 +395,53 @@ abstract class Widget extends DiagnosticableTree {
   }
 }
 
+/// A cache of inherited elements with an optional parent cache.
+///
+/// When looking up an inherited element, this will look through parent
+/// caches until the element is found or the end is reached. When an element
+/// is found, it is cached at the first element where the search began.
+///
+/// The intention of this cache is to speed up the initial build of widget
+/// trees that contain a significant number of inherited widgets by deferring
+/// expensive map allocations until they are needed, and by only allocating
+/// in the "closest" hash map.
+///
+/// This will not cache `null` results if an inherited element is not found.
+@visibleForTesting
+class InheritedTreeCache {
+  /// Create a new [InheritedTreeCache] with an optional parent.
+  InheritedTreeCache([this._parent]);
+
+  final InheritedTreeCache? _parent;
+  final HashMap<Type, InheritedElement> _current = HashMap<Type, InheritedElement>();
+
+  /// Place the [element] in the cache under [type].
+  void operator []=(Type type, InheritedElement element) {
+    _current[type] = element;
+  }
+
+  /// Find the nearest [InheritedElement] of type [type], or `null` if it
+  /// cannot be found.
+  ///
+  /// This operation will also cache the inherited element to improve the
+  /// speed of future lookups.
+  InheritedElement? operator[](Type type) {
+    InheritedElement? potential = _current[type];
+    if (potential != null) {
+      return potential;
+    }
+    potential =  _parent?._lookupWithoutCaching(type);
+    if (potential != null) {
+      _current[type] = potential;
+    }
+    return potential;
+  }
+
+  InheritedElement? _lookupWithoutCaching(Type type) {
+    return _current[type] ?? _parent?._lookupWithoutCaching(type);
+  }
+}
+
 /// A widget that does not require mutable state.
 ///
 /// A stateless widget is a widget that describes part of the user interface by
@@ -3937,7 +3984,7 @@ abstract class Element extends DiagnosticableTree implements BuildContext {
       // implementation to decide whether to rebuild based on whether we had
       // dependencies here.
     }
-    _inheritedWidgets = null;
+    _inheritedLookup = null;
     _lifecycleState = _ElementLifecycle.inactive;
   }
 
@@ -4129,7 +4176,8 @@ abstract class Element extends DiagnosticableTree implements BuildContext {
     return null;
   }
 
-  Map<Type, InheritedElement>? _inheritedWidgets;
+  InheritedTreeCache? _inheritedLookup;
+
   Set<InheritedElement>? _dependencies;
   bool _hadUnsatisfiedDependencies = false;
 
@@ -4166,7 +4214,7 @@ abstract class Element extends DiagnosticableTree implements BuildContext {
   @override
   T? dependOnInheritedWidgetOfExactType<T extends InheritedWidget>({Object? aspect}) {
     assert(_debugCheckStateIsActiveForAncestorLookup());
-    final InheritedElement? ancestor = _inheritedWidgets == null ? null : _inheritedWidgets![T];
+    final InheritedElement? ancestor = _inheritedLookup?[T];
     if (ancestor != null) {
       return dependOnInheritedElement(ancestor, aspect: aspect) as T;
     }
@@ -4177,13 +4225,13 @@ abstract class Element extends DiagnosticableTree implements BuildContext {
   @override
   InheritedElement? getElementForInheritedWidgetOfExactType<T extends InheritedWidget>() {
     assert(_debugCheckStateIsActiveForAncestorLookup());
-    final InheritedElement? ancestor = _inheritedWidgets == null ? null : _inheritedWidgets![T];
+    final InheritedElement? ancestor = _inheritedLookup?[T];
     return ancestor;
   }
 
   void _updateInheritance() {
     assert(_lifecycleState == _ElementLifecycle.active);
-    _inheritedWidgets = _parent?._inheritedWidgets;
+    _inheritedLookup = _parent?._inheritedLookup;
   }
 
   @override
@@ -5191,12 +5239,8 @@ class InheritedElement extends ProxyElement {
   @override
   void _updateInheritance() {
     assert(_lifecycleState == _ElementLifecycle.active);
-    final Map<Type, InheritedElement>? incomingWidgets = _parent?._inheritedWidgets;
-    if (incomingWidgets != null)
-      _inheritedWidgets = HashMap<Type, InheritedElement>.of(incomingWidgets);
-    else
-      _inheritedWidgets = HashMap<Type, InheritedElement>();
-    _inheritedWidgets![widget.runtimeType] = this;
+    _inheritedLookup = InheritedTreeCache(_parent?._inheritedLookup)
+      ..[widget.runtimeType] = this;
   }
 
   @override

packages/flutter/test/widgets/inherited_tree_cache_test.dart

@@ -0,0 +1,79 @@
+// Copyright 2014 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/src/widgets/framework.dart';
+import 'package:flutter_test/flutter_test.dart';
+
+void main() {
+  testWidgets('InheritedTreeCache returns null if element is not found', (WidgetTester tester) async {
+    final InheritedTreeCache parent = InheritedTreeCache();
+
+    expect(parent[InheritedElementA], isNull);
+  });
+
+  testWidgets('InheritedTreeCache can look up element from parent', (WidgetTester tester) async {
+    final InheritedTreeCache parent = InheritedTreeCache();
+    final InheritedTreeCache child = InheritedTreeCache(parent);
+    final InheritedElementA elementA = InheritedElementA(const InheritedWidgetA());
+
+    parent[InheritedElementA] = elementA;
+
+    expect(child[InheritedElementA], elementA);
+  });
+
+  testWidgets('InheritedTreeCache can look up multiple elements from parent', (WidgetTester tester) async {
+    final InheritedTreeCache parent = InheritedTreeCache();
+    final InheritedTreeCache child = InheritedTreeCache(parent);
+    final InheritedElementA elementA = InheritedElementA(const InheritedWidgetA());
+    final InheritedElementA elementB = InheritedElementA(const InheritedWidgetB());
+
+    parent[InheritedElementA] = elementA;
+    parent[InheritedElementB] = elementB;
+
+    expect(child[InheritedElementA], elementA);
+    expect(child[InheritedElementB], elementB);
+  });
+
+  testWidgets('InheritedTreeCache does not cache nulls', (WidgetTester tester) async {
+    final InheritedTreeCache parent = InheritedTreeCache();
+    final InheritedTreeCache child = InheritedTreeCache(parent);
+    final InheritedElementA elementA = InheritedElementA(const InheritedWidgetA());
+
+    // First look up element that has not been cached.
+    expect(child[InheritedElementA], null);
+
+    // Then manually add element to parent.
+    parent[InheritedElementA] = elementA;
+
+    // Then the child should be able to find it.
+    expect(child[InheritedElementA], elementA);
+  });
+}
+
+class InheritedElementA extends InheritedElement {
+  InheritedElementA(InheritedWidget widget) : super(widget);
+}
+
+class InheritedWidgetA extends InheritedWidget {
+  const InheritedWidgetA({ Key? key }) : super(child: const SizedBox(), key: key);
+
+  @override
+  bool updateShouldNotify(covariant InheritedWidget oldWidget) {
+    return true;
+  }
+}
+
+class InheritedElementB extends InheritedElement {
+  InheritedElementB(InheritedWidget widget) : super(widget);
+}
+
+class InheritedWidgetB extends InheritedWidget {
+  const InheritedWidgetB({ Key? key }) : super(child: const SizedBox(), key: key);
+
+  @override
+  bool updateShouldNotify(covariant InheritedWidget oldWidget) {
+    return true;
+  }
+}