Handle disposal of a HeroState while a hero is animating (#5189)

Fixes https://github.com/flutter/flutter/issues/5178

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

@@ -4,6 +4,8 @@
 
 import 'dart:collection';
 
+import 'package:meta/meta.dart';
+
 import 'basic.dart';
 import 'binding.dart';
 import 'framework.dart';
@@ -156,6 +158,7 @@ class HeroState extends State<Hero> implements HeroHandle {
 
   GlobalKey _key = new GlobalKey();
   Size _placeholderSize;
+  VoidCallback _disposeCallback;
 
   @override
   bool get alwaysAnimate => config.alwaysAnimate;
@@ -202,6 +205,13 @@ class HeroState extends State<Hero> implements HeroHandle {
       _setChild(null);
   }
 
+  @override
+  void dispose() {
+    if (_disposeCallback != null)
+      _disposeCallback();
+    super.dispose();
+  }
+
   @override
   Widget build(BuildContext context) {
     if (_placeholderSize != null) {
@@ -231,6 +241,12 @@ class _HeroQuestState implements HeroHandle {
     this.currentTurns
   }) {
     assert(tag != null);
+
+    for (HeroState state in sourceStates)
+      state._disposeCallback = () => sourceStates.remove(state);
+
+    if (targetState != null)
+      targetState._disposeCallback = _handleTargetStateDispose;
   }
 
   final Object tag;
@@ -238,12 +254,14 @@ class _HeroQuestState implements HeroHandle {
   final Widget child;
   final Set<HeroState> sourceStates;
   final Rect animationArea;
-  final Rect targetRect;
-  final int targetTurns;
-  final HeroState targetState;
+  Rect targetRect;
+  int targetTurns;
+  HeroState targetState;
   final RectTween currentRect;
   final Tween<double> currentTurns;
 
+  OverlayEntry overlayEntry;
+
   @override
   bool get alwaysAnimate => true;
 
@@ -257,6 +275,10 @@ class _HeroQuestState implements HeroHandle {
     Set<HeroState> states = sourceStates;
     if (targetState != null)
       states = states.union(new HashSet<HeroState>.from(<HeroState>[targetState]));
+
+    for (HeroState state in states)
+      state._disposeCallback = null;
+
     return new _HeroManifest(
       key: key,
       config: child,
@@ -266,6 +288,13 @@ class _HeroQuestState implements HeroHandle {
     );
   }
 
+  void _handleTargetStateDispose() {
+    targetState = null;
+    targetTurns = 0;
+    targetRect = targetRect.center & Size.zero;
+    WidgetsBinding.instance.addPostFrameCallback((Duration d) => overlayEntry.markNeedsBuild());
+  }
+
   Widget build(BuildContext context, Animation<double> animation) {
     return new RelativePositionedTransition(
       rect: currentRect.animate(animation),
@@ -279,6 +308,17 @@ class _HeroQuestState implements HeroHandle {
       )
     );
   }
+
+  @mustCallSuper
+  void dispose() {
+    overlayEntry = null;
+
+    for (HeroState state in sourceStates)
+      state._disposeCallback = null;
+
+    if (targetState != null)
+      targetState._disposeCallback = null;
+  }
 }
 
 class _HeroMatch {
@@ -366,6 +406,8 @@ class HeroParty {
     }
 
     assert(!_heroes.any((_HeroQuestState hero) => !hero.taken));
+    for (_HeroQuestState hero in _heroes)
+      hero.dispose();
     _heroes = _newHeroes;
   }
 
@@ -393,6 +435,7 @@ class HeroParty {
           hero.targetState._setChild(hero.key);
         for (HeroState source in hero.sourceStates)
           source._resetChild();
+        hero.dispose();
       }
       _heroes.clear();
       _clearCurrentAnimation();
@@ -476,14 +519,15 @@ class HeroController extends NavigatorObserver {
     _overlayEntries.clear();
   }
 
-  void _addHeroToOverlay(Widget hero, Object tag, OverlayState overlay) {
-    OverlayEntry entry = new OverlayEntry(builder: (_) => hero);
+  OverlayEntry _addHeroToOverlay(WidgetBuilder hero, Object tag, OverlayState overlay) {
+    OverlayEntry entry = new OverlayEntry(builder: hero);
     assert(_animation.status != AnimationStatus.dismissed && _animation.status != AnimationStatus.completed);
     if (_animation.status == AnimationStatus.forward)
       _to.insertHeroOverlayEntry(entry, tag, overlay);
     else
       _from.insertHeroOverlayEntry(entry, tag, overlay);
     _overlayEntries.add(entry);
+    return entry;
   }
 
   Set<Key> _getMostValuableKeys() {
@@ -523,8 +567,11 @@ class HeroController extends NavigatorObserver {
       curve: curve
     ));
     for (_HeroQuestState hero in _party._heroes) {
-      Widget widget = hero.build(navigator.context, animation);
-      _addHeroToOverlay(widget, hero.tag, navigator.overlay);
+      hero.overlayEntry = _addHeroToOverlay(
+        (BuildContext context) => hero.build(navigator.context, animation),
+        hero.tag,
+        navigator.overlay
+      );
     }
   }
 }

packages/flutter/test/widget/heroes_test.dart

@@ -40,6 +40,18 @@ class ThreeRoute extends MaterialPageRoute<Null> {
   });
 }
 
+class MutatingRoute extends MaterialPageRoute<Null> {
+  MutatingRoute() : super(builder: (BuildContext context) {
+    return new Hero(tag: 'a', child: new Text('MutatingRoute'), key: new UniqueKey());
+  });
+
+  void markNeedsBuild() {
+    setState(() {
+      // Trigger a rebuild
+    });
+  }
+}
+
 void main() {
   testWidgets('Heroes animate', (WidgetTester tester) async {
 
@@ -143,4 +155,25 @@ void main() {
     expect(find.byKey(thirdKey), isOnStage);
     expect(find.byKey(thirdKey), isInCard);
   });
+
+  testWidgets('Heroes animate', (WidgetTester tester) async {
+    MutatingRoute route = new MutatingRoute();
+
+    await tester.pumpWidget(new MaterialApp(
+      home: new Material(child: new Block(children: <Widget>[
+        new Hero(tag: 'a', child: new Text('foo')),
+        new Builder(builder: (BuildContext context) {
+          return new FlatButton(child: new Text('two'), onPressed: () => Navigator.push(context, route));
+        })
+      ]))
+    ));
+
+    await tester.tap(find.text('two'));
+    await tester.pump(new Duration(milliseconds: 10));
+
+    route.markNeedsBuild();
+
+    await tester.pump(new Duration(milliseconds: 10));
+    await tester.pump(new Duration(seconds: 1));
+  });
 }