Remove non needed controllers in SegmentedButton. (#134064)

Author: polina-c

packages/flutter/lib/src/material/segmented_button.dart

@@ -236,12 +236,32 @@ class SegmentedButton<T> extends StatefulWidget {
   final Widget? selectedIcon;
 
   @override
-  State<SegmentedButton<T>> createState() => _SegmentedButtonState<T>();
+  State<SegmentedButton<T>> createState() => SegmentedButtonState<T>();
 }
 
-class _SegmentedButtonState<T> extends State<SegmentedButton<T>> {
+/// State for [SegmentedButton].
+@visibleForTesting
+class SegmentedButtonState<T> extends State<SegmentedButton<T>> {
   bool get _enabled => widget.onSelectionChanged != null;
-  final Map<ButtonSegment<T>, MaterialStatesController> _statesControllers = <ButtonSegment<T>, MaterialStatesController>{};
+
+  /// Controllers for the [ButtonSegment]s.
+  @visibleForTesting
+  final Map<ButtonSegment<T>, MaterialStatesController> statesControllers = <ButtonSegment<T>, MaterialStatesController>{};
+
+  @override
+  void didUpdateWidget(covariant SegmentedButton<T> oldWidget) {
+    super.didUpdateWidget(oldWidget);
+    if (oldWidget != widget) {
+      statesControllers.removeWhere((ButtonSegment<T> segment, MaterialStatesController controller) {
+        if (widget.segments.contains(segment)) {
+          return false;
+        } else {
+          controller.dispose();
+          return true;
+        }
+      });
+    }
+  }
 
   void _handleOnPressed(T segmentValue) {
     if (!_enabled) {
@@ -325,7 +345,7 @@ class _SegmentedButtonState<T> extends State<SegmentedButton<T>> {
         : segment.label != null
           ? segment.icon
           : null;
-      final MaterialStatesController controller = _statesControllers.putIfAbsent(segment, () => MaterialStatesController());
+      final MaterialStatesController controller = statesControllers.putIfAbsent(segment, () => MaterialStatesController());
       controller.value = <MaterialState>{
           if (segmentSelected) MaterialState.selected,
       };
@@ -391,7 +411,7 @@ class _SegmentedButtonState<T> extends State<SegmentedButton<T>> {
 
   @override
   void dispose() {
-    for (final MaterialStatesController controller in _statesControllers.values) {
+    for (final MaterialStatesController controller in statesControllers.values) {
       controller.dispose();
     }
     super.dispose();

packages/flutter/test/material/segmented_button_test.dart

@@ -21,6 +21,52 @@ Widget boilerplate({required Widget child}) {
 }
 
 void main() {
+  testWidgetsWithLeakTracking('SegmentedButton releases state controllers for deleted segments', (WidgetTester tester) async {
+    final ThemeData theme = ThemeData(useMaterial3: true);
+    final Key key = UniqueKey();
+
+    Widget buildApp(Widget button) {
+      return MaterialApp(
+        theme: theme,
+        home: Scaffold(
+          body: Center(
+            child: button,
+          ),
+        ),
+      );
+    }
+
+    await tester.pumpWidget(
+      buildApp(
+        SegmentedButton<int>(
+          key: key,
+          segments: const <ButtonSegment<int>>[
+            ButtonSegment<int>(value: 1, label: Text('1')),
+            ButtonSegment<int>(value: 2, label: Text('2')),
+          ],
+          selected: const <int>{2},
+        ),
+      ),
+    );
+
+    await tester.pumpWidget(
+      buildApp(
+        SegmentedButton<int>(
+          key: key,
+          segments: const <ButtonSegment<int>>[
+            ButtonSegment<int>(value: 2, label: Text('2')),
+            ButtonSegment<int>(value: 3, label: Text('3')),
+          ],
+          selected: const <int>{2},
+        ),
+      ),
+    );
+
+    final SegmentedButtonState<int> state = tester.state(find.byType(SegmentedButton<int>));
+    expect(state.statesControllers, hasLength(2));
+    expect(state.statesControllers.keys.first.value, 2);
+    expect(state.statesControllers.keys.last.value, 3);
+  });
 
   testWidgetsWithLeakTracking('SegmentedButton is built with Material of type MaterialType.transparency', (WidgetTester tester) async {
     final ThemeData theme = ThemeData(useMaterial3: true);