Fix a [_ViewportElement] RenderObjectChild update bug (#96377)

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

@@ -211,6 +211,9 @@ class _ViewportElement extends MultiChildRenderObjectElement {
   /// Creates an element that uses the given widget as its configuration.
   _ViewportElement(Viewport widget) : super(widget);
 
+  bool _doingMountOrUpdate = false;
+  int? _centerSlotIndex;
+
   @override
   Viewport get widget => super.widget as Viewport;
 
@@ -219,26 +222,67 @@ class _ViewportElement extends MultiChildRenderObjectElement {
 
   @override
   void mount(Element? parent, Object? newSlot) {
+    assert(!_doingMountOrUpdate);
+    _doingMountOrUpdate = true;
     super.mount(parent, newSlot);
     _updateCenter();
+    assert(_doingMountOrUpdate);
+    _doingMountOrUpdate = false;
   }
 
   @override
   void update(MultiChildRenderObjectWidget newWidget) {
+    assert(!_doingMountOrUpdate);
+    _doingMountOrUpdate = true;
     super.update(newWidget);
     _updateCenter();
+    assert(_doingMountOrUpdate);
+    _doingMountOrUpdate = false;
   }
 
   void _updateCenter() {
     // TODO(ianh): cache the keys to make this faster
     if (widget.center != null) {
-      renderObject.center = children.singleWhere(
-        (Element element) => element.widget.key == widget.center,
-      ).renderObject as RenderSliver?;
+      int elementIndex = 0;
+      for (final Element e in children) {
+        if (e.widget.key == widget.center) {
+          renderObject.center = e.renderObject as RenderSliver?;
+          break;
+        }
+        elementIndex++;
+      }
+      assert(elementIndex < children.length);
+      _centerSlotIndex = elementIndex;
     } else if (children.isNotEmpty) {
       renderObject.center = children.first.renderObject as RenderSliver?;
+      _centerSlotIndex = 0;
     } else {
       renderObject.center = null;
+      _centerSlotIndex = null;
+    }
+  }
+
+  @override
+  void insertRenderObjectChild(RenderObject child, IndexedSlot<Element?> slot) {
+    super.insertRenderObjectChild(child, slot);
+    // Once [mount]/[update] are done, the `renderObject.center` will be updated
+    // in [_updateCenter].
+    if (!_doingMountOrUpdate && slot.index == _centerSlotIndex) {
+      renderObject.center = child as RenderSliver?;
+    }
+  }
+
+  @override
+  void moveRenderObjectChild(RenderObject child, IndexedSlot<Element?> oldSlot, IndexedSlot<Element?> newSlot) {
+    super.moveRenderObjectChild(child, oldSlot, newSlot);
+    assert(_doingMountOrUpdate);
+  }
+
+  @override
+  void removeRenderObjectChild(RenderObject child, Object? slot) {
+    super.removeRenderObjectChild(child, slot);
+    if (!_doingMountOrUpdate && renderObject.center == child) {
+      renderObject.center = null;
     }
   }
 

packages/flutter/test/widgets/custom_scroll_view_test.dart

@@ -6,6 +6,81 @@ import 'package:flutter/widgets.dart';
 import 'package:flutter_test/flutter_test.dart';
 
 void main() {
+  // Regression test for https://github.com/flutter/flutter/issues/96024
+  testWidgets('CustomScrollView.center update test 1', (WidgetTester tester) async {
+    final Key centerKey = UniqueKey();
+    late StateSetter setState;
+    bool hasKey = false;
+    await tester.pumpWidget(Directionality(
+      textDirection: TextDirection.ltr,
+      child: CustomScrollView(
+        center: centerKey,
+        slivers: <Widget>[
+          const SliverToBoxAdapter(key: Key('a'), child: SizedBox(height: 100.0)),
+          StatefulBuilder(
+            key: centerKey,
+            builder: (BuildContext context, StateSetter setter) {
+              setState = setter;
+              if (hasKey) {
+                return const SliverToBoxAdapter(
+                  key: Key('b'),
+                  child: SizedBox(height: 100.0),
+                );
+              } else {
+                return const SliverToBoxAdapter(
+                  child: SizedBox(height: 100.0),
+                );
+              }
+            },
+          ),
+        ],
+      ),
+    ));
+    await tester.pumpAndSettle();
+
+    // Change the center key will trigger the old RenderObject remove and a new
+    // RenderObject insert.
+    setState(() {
+      hasKey = true;
+    });
+
+    await tester.pumpAndSettle();
+
+    // Pass without throw.
+  });
+
+  testWidgets('CustomScrollView.center update test 2', (WidgetTester tester) async {
+    const List<Widget> slivers1 = <Widget>[
+      SliverToBoxAdapter(key: Key('a'), child: SizedBox(height: 100.0)),
+      SliverToBoxAdapter(key: Key('b'), child: SizedBox(height: 100.0)),
+      SliverToBoxAdapter(key: Key('c'), child: SizedBox(height: 100.0)),
+    ];
+
+    const List<Widget> slivers2 = <Widget>[
+      SliverToBoxAdapter(key: Key('c'), child: SizedBox(height: 100.0)),
+      SliverToBoxAdapter(key: Key('d'), child: SizedBox(height: 100.0)),
+      SliverToBoxAdapter(key: Key('a'), child: SizedBox(height: 100.0)),
+    ];
+
+    Widget buildFrame(List<Widget> slivers, Key center) {
+      return Directionality(
+        textDirection: TextDirection.ltr,
+        child: CustomScrollView(
+          center: center,
+          slivers: slivers,
+        ),
+      );
+    }
+
+    await tester.pumpWidget(buildFrame(slivers1, const Key('b')));
+    await tester.pumpAndSettle();
+
+    await tester.pumpWidget(buildFrame(slivers2, const Key('d')));
+    await tester.pumpAndSettle();
+
+    // Pass without throw.
+  });
+
   testWidgets('CustomScrollView.center', (WidgetTester tester) async {
     await tester.pumpWidget(const Directionality(
       textDirection: TextDirection.ltr,