[flutter] unify reassemble and fast reassemble (#84363)

packages/flutter/lib/src/foundation/binding.dart

@@ -68,6 +68,13 @@ abstract class BindingBase {
   static bool _debugInitialized = false;
   static bool _debugServiceExtensionsRegistered = false;
 
+  /// Additional configuration used by the framework during hot reload.
+  ///
+  /// See also:
+  ///
+  ///  * [DebugReassembleConfig], which describes the configuration.
+  static DebugReassembleConfig? debugReassembleConfig;
+
   /// The main window to which this binding is bound.
   ///
   /// A number of additional bindings are defined as extensions of
@@ -629,3 +636,23 @@ abstract class BindingBase {
 Future<void> _exitApplication() async {
   exit(0);
 }
+
+/// Additional configuration used for hot reload reassemble optimizations.
+///
+/// Do not extend, implement, or mixin this class. This may only be instantiated
+/// in debug mode.
+class DebugReassembleConfig {
+  /// Create a new [DebugReassembleConfig].
+  ///
+  /// Throws a [FlutterError] if this is called in profile or release mode.
+  DebugReassembleConfig({
+    this.widgetName,
+  }) {
+    if (!kDebugMode) {
+      throw FlutterError('Cannot instaniate DebugReassembleConfig in profile or release mode.');
+    }
+  }
+
+  /// The name of the widget that was modified, or `null` if the change was elsewhere.
+  final String? widgetName;
+}

packages/flutter/lib/src/rendering/binding.dart

@@ -474,12 +474,14 @@ mixin RendererBinding on BindingBase, ServicesBinding, SchedulerBinding, Gesture
   @override
   Future<void> performReassemble() async {
     await super.performReassemble();
+    if (BindingBase.debugReassembleConfig?.widgetName == null) {
       Timeline.startSync('Dirty Render Tree', arguments: timelineArgumentsIndicatingLandmarkEvent);
       try {
         renderView.reassemble();
       } finally {
         Timeline.finishSync();
       }
+    }
     scheduleWarmUpFrame();
     await endOfFrame;
   }

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

@@ -443,17 +443,16 @@ mixin WidgetsBinding on BindingBase, ServicesBinding, SchedulerBinding, GestureB
       registerServiceExtension(
         name: 'fastReassemble',
         callback: (Map<String, Object> params) async {
+          // This mirrors the implementation of the 'reassemble' callback registration
+          // in lib/src/foundation/binding.dart, but with the extra binding config used
+          // to skip some reassemble work.
           final String? className = params['className'] as String?;
-          void markElementsDirty(Element element) {
-            if (element.widget.runtimeType.toString() == className) {
-              element.markNeedsBuild();
-            }
-            element.visitChildElements(markElementsDirty);
-          }
-          if (renderViewElement != null) {
-            markElementsDirty(renderViewElement!);
+          BindingBase.debugReassembleConfig = DebugReassembleConfig(widgetName: className);
+          try {
+            await reassembleApplication();
+          } finally {
+            BindingBase.debugReassembleConfig = null;
           }
-          await endOfFrame;
           return <String, String>{'type': 'Success'};
         },
       );
@@ -502,7 +501,7 @@ mixin WidgetsBinding on BindingBase, ServicesBinding, SchedulerBinding, GestureB
 
   Future<void> _forceRebuild() {
     if (renderViewElement != null) {
-      buildOwner!.reassemble(renderViewElement!);
+      buildOwner!.reassemble(renderViewElement!, null);
       return endOfFrame;
     }
     return Future<void>.value();
@@ -956,8 +955,9 @@ mixin WidgetsBinding on BindingBase, ServicesBinding, SchedulerBinding, GestureB
       return true;
     }());
 
-    if (renderViewElement != null)
-      buildOwner!.reassemble(renderViewElement!);
+    if (renderViewElement != null) {
+      buildOwner!.reassemble(renderViewElement!, BindingBase.debugReassembleConfig);
+    }
     return super.performReassemble();
   }
 

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

@@ -2978,11 +2978,12 @@ class BuildOwner {
   /// changed implementations.
   ///
   /// This is expensive and should not be called except during development.
-  void reassemble(Element root) {
+  void reassemble(Element root, DebugReassembleConfig? reassembleConfig) {
     Timeline.startSync('Dirty Element Tree');
     try {
       assert(root._parent == null);
       assert(root.owner == this);
+      root._debugReassembleConfig = reassembleConfig;
       root.reassemble();
     } finally {
       Timeline.finishSync();
@@ -3049,6 +3050,7 @@ abstract class Element extends DiagnosticableTree implements BuildContext {
       _widget = widget;
 
   Element? _parent;
+  DebugReassembleConfig? _debugReassembleConfig;
 
   // Custom implementation of `operator ==` optimized for the ".of" pattern
   // used with `InheritedWidgets`.
@@ -3175,10 +3177,15 @@ abstract class Element extends DiagnosticableTree implements BuildContext {
   @mustCallSuper
   @protected
   void reassemble() {
+    if (_debugShouldReassemble(_debugReassembleConfig, _widget)) {
       markNeedsBuild();
+      _debugReassembleConfig = null;
+    }
     visitChildren((Element child) {
+      child._debugReassembleConfig = _debugReassembleConfig;
       child.reassemble();
     });
+    _debugReassembleConfig = null;
   }
 
   bool _debugIsInScope(Element target) {
@@ -4784,7 +4791,9 @@ class StatefulElement extends ComponentElement {
 
   @override
   void reassemble() {
+    if (_debugShouldReassemble(_debugReassembleConfig, _widget)) {
       state.reassemble();
+    }
     super.reassemble();
   }
 
@@ -6435,3 +6444,9 @@ class _NullWidget extends Widget {
   @override
   Element createElement() => throw UnimplementedError();
 }
+
+// Whether a [DebugReassembleConfig] indicates that an element holding [widget] can skip
+// a reassemble.
+bool _debugShouldReassemble(DebugReassembleConfig? config, Widget? widget) {
+  return config == null || config.widgetName == null || widget?.runtimeType.toString() == config.widgetName;
+}

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

@@ -910,7 +910,7 @@ mixin WidgetInspectorService {
   Future<void> forceRebuild() {
     final WidgetsBinding binding = WidgetsBinding.instance!;
     if (binding.renderViewElement != null) {
-      binding.buildOwner!.reassemble(binding.renderViewElement!);
+      binding.buildOwner!.reassemble(binding.renderViewElement!, null);
       return binding.endOfFrame;
     }
     return Future<void>.value();

packages/flutter/test/widgets/fast_reassemble_test.dart

+// 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/foundation.dart';
+import 'package:flutter/widgets.dart';
+import 'package:flutter_test/flutter_test.dart';
+
+void main() {
+  testWidgets('reassemble with a className only marks subtrees from the first matching element as dirty', (WidgetTester tester) async {
+    await tester.pumpWidget(
+      const Foo(Bar(Fizz(SizedBox())))
+    );
+
+    expect(Foo.count, 0);
+    expect(Bar.count, 0);
+    expect(Fizz.count, 0);
+
+    DebugReassembleConfig config = DebugReassembleConfig(widgetName: 'Bar');
+    WidgetsBinding.instance!.buildOwner!.reassemble(WidgetsBinding.instance!.renderViewElement!, config);
+
+    expect(Foo.count, 0);
+    expect(Bar.count, 1);
+    expect(Fizz.count, 1);
+
+    config = DebugReassembleConfig(widgetName: 'Fizz');
+    WidgetsBinding.instance!.buildOwner!.reassemble(WidgetsBinding.instance!.renderViewElement!, config);
+
+    expect(Foo.count, 0);
+    expect(Bar.count, 1);
+    expect(Fizz.count, 2);
+
+    config = DebugReassembleConfig(widgetName: 'NoMatch');
+    WidgetsBinding.instance!.buildOwner!.reassemble(WidgetsBinding.instance!.renderViewElement!, config);
+
+    expect(Foo.count, 0);
+    expect(Bar.count, 1);
+    expect(Fizz.count, 2);
+
+    config = DebugReassembleConfig();
+    WidgetsBinding.instance!.buildOwner!.reassemble(WidgetsBinding.instance!.renderViewElement!, config);
+
+    expect(Foo.count, 1);
+    expect(Bar.count, 2);
+    expect(Fizz.count, 3);
+
+    WidgetsBinding.instance!.buildOwner!.reassemble(WidgetsBinding.instance!.renderViewElement!, null);
+
+    expect(Foo.count, 2);
+    expect(Bar.count, 3);
+    expect(Fizz.count, 4);
+  });
+}
+
+class Foo extends StatefulWidget {
+  const Foo(this.child, {Key? key}) : super(key: key);
+
+  final Widget child;
+  static int count = 0;
+
+  @override
+  State<Foo> createState() => _FooState();
+}
+
+class _FooState extends State<Foo> {
+  @override
+  void reassemble() {
+    Foo.count += 1;
+    super.reassemble();
+  }
+
+  @override
+  Widget build(BuildContext context) {
+    return widget.child;
+  }
+}
+
+
+class Bar extends StatefulWidget {
+  const Bar(this.child, {Key? key}) : super(key: key);
+
+  final Widget child;
+  static int count = 0;
+
+  @override
+  State<Bar> createState() => _BarState();
+}
+
+class _BarState extends State<Bar> {
+  @override
+  void reassemble() {
+    Bar.count += 1;
+    super.reassemble();
+  }
+
+  @override
+  Widget build(BuildContext context) {
+    return widget.child;
+  }
+}
+
+class Fizz extends StatefulWidget {
+  const Fizz(this.child, {Key? key}) : super(key: key);
+
+  final Widget child;
+  static int count = 0;
+
+  @override
+  State<Fizz> createState() => _FizzState();
+}
+
+class _FizzState extends State<Fizz> {
+  @override
+  void reassemble() {
+    Fizz.count += 1;
+    super.reassemble();
+  }
+
+  @override
+  Widget build(BuildContext context) {
+    return widget.child;
+  }
+}

packages/flutter/test/widgets/widget_inspector_test_utils.dart

@@ -59,7 +59,7 @@ class TestWidgetInspectorService extends Object with WidgetInspectorService {
     final WidgetsBinding binding = WidgetsBinding.instance!;
 
     if (binding.renderViewElement != null) {
-      binding.buildOwner!.reassemble(binding.renderViewElement!);
+      binding.buildOwner!.reassemble(binding.renderViewElement!, null);
     }
   }