Reland #49925 (#50071)

* Reland "Fix custom Elements that wants to decorate State.build (#49925)" (#50068)"

This reverts commit 62fdcb031945c23dacb481bd464cbea8ec9037d4.

Updates stack filtering for new frames in stateful element

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

@@ -276,10 +276,11 @@ mixin WidgetsBinding on BindingBase, ServicesBinding, SchedulerBinding, GestureB
     const PartialStackFrame elementUpdateChild = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'Element', method: 'updateChild');
     const PartialStackFrame elementRebuild = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'Element', method: 'rebuild');
     const PartialStackFrame componentElementPerformRebuild = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'ComponentElement', method: 'performRebuild');
-    const PartialStackFrame componentElementFristBuild = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'ComponentElement', method: '_firstBuild');
+    const PartialStackFrame componentElementFirstBuild = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'ComponentElement', method: '_firstBuild');
     const PartialStackFrame componentElementMount = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'ComponentElement', method: 'mount');
-    const PartialStackFrame statefulElementFristBuild = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'StatefulElement', method: '_firstBuild');
+    const PartialStackFrame statefulElementFirstBuild = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'StatefulElement', method: '_firstBuild');
     const PartialStackFrame singleChildMount = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'SingleChildRenderObjectElement', method: 'mount');
+    const PartialStackFrame statefulElementRebuild = PartialStackFrame(package: 'package:flutter/src/widgets/framework.dart', className: 'StatefulElement', method: 'performRebuild');
 
     const String replacementString = '...     Normal element mounting';
 
@@ -290,7 +291,7 @@ mixin WidgetsBinding on BindingBase, ServicesBinding, SchedulerBinding, GestureB
         elementUpdateChild,
         componentElementPerformRebuild,
         elementRebuild,
-        componentElementFristBuild,
+        componentElementFirstBuild,
         componentElementMount,
       ],
       replacement: replacementString,
@@ -300,7 +301,7 @@ mixin WidgetsBinding on BindingBase, ServicesBinding, SchedulerBinding, GestureB
         elementUpdateChild,
         componentElementPerformRebuild,
         elementRebuild,
-        componentElementFristBuild,
+        componentElementFirstBuild,
         componentElementMount,
       ],
       replacement: replacementString,
@@ -312,9 +313,10 @@ mixin WidgetsBinding on BindingBase, ServicesBinding, SchedulerBinding, GestureB
         elementInflateWidget,
         elementUpdateChild,
         componentElementPerformRebuild,
+        statefulElementRebuild,
         elementRebuild,
-        componentElementFristBuild,
+        componentElementFirstBuild,
-        statefulElementFristBuild,
+        statefulElementFirstBuild,
         componentElementMount,
       ],
       replacement: replacementString,
@@ -323,9 +325,10 @@ mixin WidgetsBinding on BindingBase, ServicesBinding, SchedulerBinding, GestureB
       frames: <PartialStackFrame>[
         elementUpdateChild,
         componentElementPerformRebuild,
+        statefulElementRebuild,
         elementRebuild,
-        componentElementFristBuild,
+        componentElementFirstBuild,
-        statefulElementFristBuild,
+        statefulElementFirstBuild,
         componentElementMount,
       ],
       replacement: replacementString,

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

@@ -4586,13 +4586,7 @@ class StatefulElement extends ComponentElement {
   }
 
   @override
-  Widget build() {
+  Widget build() => _state.build(this);
-    if (_didChangeDependencies) {
-      _state.didChangeDependencies();
-      _didChangeDependencies = false;
-    }
-    return state.build(this);
-  }
 
   /// The [State] instance associated with this location in the tree.
   ///
@@ -4642,6 +4636,15 @@ class StatefulElement extends ComponentElement {
     super._firstBuild();
   }
 
+  @override
+  void performRebuild() {
+    if (_didChangeDependencies) {
+      _state.didChangeDependencies();
+      _didChangeDependencies = false;
+    }
+    super.performRebuild();
+  }
+
   @override
   void update(StatefulWidget newWidget) {
     super.update(newWidget);

packages/flutter/test/widgets/framework_test.dart

@@ -1259,6 +1259,80 @@ void main() {
     expect(state.didChangeDependenciesCount, 3);
     expect(state.deactivatedCount, 2);
   });
+
+  testWidgets('StatefulElement subclass can decorate State.build', (WidgetTester tester) async {
+    bool isDidChangeDependenciesDecorated;
+    bool isBuildDecorated;
+
+    final Widget child = Decorate(
+      didChangeDependencies: (bool value) {
+        isDidChangeDependenciesDecorated = value;
+      },
+      build: (bool value) {
+        isBuildDecorated = value;
+      }
+    );
+
+    await tester.pumpWidget(Inherited(0, child: child));
+
+    expect(isBuildDecorated, isTrue);
+    expect(isDidChangeDependenciesDecorated, isFalse);
+
+    await tester.pumpWidget(Inherited(1, child: child));
+
+    expect(isBuildDecorated, isTrue);
+    expect(isDidChangeDependenciesDecorated, isFalse);
+  });
+}
+
+class Decorate extends StatefulWidget {
+  const Decorate({
+    Key key,
+    @required this.didChangeDependencies,
+    @required this.build
+  }) :
+    assert(didChangeDependencies != null),
+    assert(build != null),
+    super(key: key);
+
+  final void Function(bool isInBuild) didChangeDependencies;
+  final void Function(bool isInBuild) build;
+
+  @override
+  _DecorateState createState() => _DecorateState();
+
+  @override
+  DecorateElement createElement() => DecorateElement(this);
+}
+
+class DecorateElement extends StatefulElement {
+  DecorateElement(Decorate widget): super(widget);
+
+  bool isDecorated = false;
+
+  @override
+  Widget build() {
+    try {
+      isDecorated = true;
+      return super.build();
+    } finally {
+      isDecorated = false;
+    }
+  }
+}
+
+class _DecorateState extends State<Decorate> {
+  @override
+  void didChangeDependencies() {
+    super.didChangeDependencies();
+    widget.didChangeDependencies.call((context as DecorateElement).isDecorated);
+  }
+  @override
+  Widget build(covariant DecorateElement context) {
+    context.dependOnInheritedWidgetOfExactType<Inherited>();
+    widget.build.call(context.isDecorated);
+    return Container();
+  }
 }
 
 class NullChildTest extends Widget {