From 97507a6cc2129dae852a01171bd5202edad98f05 Mon Sep 17 00:00:00 2001
From: Sergey Kazakov <kazaimazai@gmail.com>
Date: Sat, 15 Jun 2024 16:33:03 +0100
Subject: [PATCH] Segment layout sizes update flow improvement  (#7)

---
 Sources/SegmentedPicker/SegmentedPicker.swift | 38 +++++++++++++++----
 1 file changed, 30 insertions(+), 8 deletions(-)

diff --git a/Sources/SegmentedPicker/SegmentedPicker.swift b/Sources/SegmentedPicker/SegmentedPicker.swift
index a400bf8..55c5bd6 100644
--- a/Sources/SegmentedPicker/SegmentedPicker.swift
+++ b/Sources/SegmentedPicker/SegmentedPicker.swift
@@ -7,6 +7,7 @@
 
 import SwiftUI
 
+
 public struct SegmentedPicker<Element, Content, Selection>: View
     where
     Content: View,
@@ -14,7 +15,7 @@ public struct SegmentedPicker<Element, Content, Selection>: View
     
     public typealias Data = [Element]
     
-    @State private var frames: [Data.Index: CGRect] = [:]
+    @State private var segmentSizes: [Data.Index: CGSize] = [:]
     @Binding private var selectedIndex: Data.Index?
 
     private let data: Data
@@ -42,8 +43,8 @@ public struct SegmentedPicker<Element, Content, Selection>: View
 
             if let index = selectedIndex {
                 selection()
-                    .frame(width: selectionFrame(at: index).width,
-                           height: selectionFrame(at: index).height)
+                    .frame(width: selectionSize(at: index).width,
+                           height: selectionSize(at: index).height)
                     .alignmentGuide(.horizontalCenterAlignment) { dimensions in
                         dimensions[HorizontalAlignment.center]
                     }
@@ -56,10 +57,14 @@ public struct SegmentedPicker<Element, Content, Selection>: View
                     )
                     .buttonStyle(PlainButtonStyle())
                     .background(GeometryReader { proxy in
-                        Color.clear.onAppear {
-                            frames[index] = proxy.frame(in: .global)
-                        }
+                        Color.clear.preference(
+                            key: SegmentSizePreferenceKey.self,
+                            value: SegmentSize(index: index, size: proxy.size)
+                        )
                     })
+                    .onPreferenceChange(SegmentSizePreferenceKey.self) { segment in
+                        segmentSizes[segment.index] = segment.size
+                    }
                     .alignmentGuide(.horizontalCenterAlignment,
                                     isActive: selectedIndex == index) { dimensions in
                         dimensions[HorizontalAlignment.center]
@@ -69,7 +74,24 @@ public struct SegmentedPicker<Element, Content, Selection>: View
         }
     }
     
-    private func selectionFrame(at index: Data.Index) -> CGRect {
-        frames[index] ?? .zero
+    private func selectionSize(at index: Data.Index) -> CGSize {
+        segmentSizes[index] ?? .zero
+    }
+}
+
+private extension SegmentedPicker {
+    struct SegmentSize: Equatable {
+        let index: Int
+        let size: CGSize
+    }
+    
+    struct SegmentSizePreferenceKey: PreferenceKey {
+        static var defaultValue: SegmentSize { SegmentSize(index: .zero, size: .zero) }
+        
+        static func reduce(value: inout SegmentSize, 
+                           nextValue: () -> SegmentSize) {
+            
+            value = nextValue()
+        }
     }
 }