Fix attachments collection view sizing

Packages/ComposeUI/Sources/ComposeUI/Views/Attachments/AttachmentsSection.swift

@@ -13,17 +13,35 @@ import InstanceFeatures
 
 struct AttachmentsSection: View {
     @ObservedObject var draft: Draft
+    private let spacing: CGFloat = 8
+    private let minItemSize: CGFloat = 100
     
     var body: some View {
+        #if os(visionOS)
+        collectionView
+        #else
+        if #available(iOS 16.0, *) {
+            collectionView
+        } else {
+            LegacyCollectionViewSizingView {
+                collectionView
+            } computeHeight: { width in
+                WrappedCollectionView.totalHeight(width: width, minItemSize: minItemSize, spacing: spacing, items: draft.attachments.count + 1)
+            }
+        }
+        #endif
+    }
+    
+    private var collectionView: some View {
         WrappedCollectionView(
             draft: draft,
-            spacing: 8,
+            spacing: spacing,
-            minItemSize: 100
+            minItemSize: minItemSize
         )
         // Impose a minimum height, because otherwise it defaults to zero which prevents the collection
         // view from laying out, and leaving the intrinsic content size at zero too.
         // Add 4 to the minItemSize because otherwise drag-and-drop while reordering can alter the contentOffset by that much.
-        .frame(minHeight: 104)
+        .frame(minHeight: minItemSize + 4)
     }
     
     static func insertAttachments(in draft: Draft, at index: Int, itemProviders: [NSItemProvider]) {
@@ -40,6 +58,41 @@ struct AttachmentsSection: View {
     }
 }
 
+#if !os(visionOS)
+@available(iOS, obsoleted: 16.0)
+private struct LegacyCollectionViewSizingView<Content: View>: View {
+    @ViewBuilder let content: Content
+    let computeHeight: (CGFloat) -> CGFloat
+    @State private var width: CGFloat = 0
+
+    var body: some View {
+        let height = computeHeight(width)
+        
+        content
+            .frame(height: max(height, 10))
+            .overlay {
+                GeometryReader { proxy in
+                    Color.clear
+                        .preference(key: WidthPrefKey.self, value: proxy.size.width)
+                        .onPreferenceChange(WidthPrefKey.self) {
+                            width = $0
+                        }
+                }
+            }
+    }
+}
+
+private struct WidthPrefKey: PreferenceKey {
+    static var defaultValue: CGFloat { 0 }
+    static func reduce(value: inout CGFloat, nextValue: () -> CGFloat) {
+        let next = nextValue()
+        if next != 0 {
+            value = next
+        }
+    }
+}
+#endif
+
 // Use a UIViewControllerRepresentable so we have something from which to present the gallery VC.
 private struct WrappedCollectionView: UIViewControllerRepresentable {
     @ObservedObject var draft: Draft
@@ -59,6 +112,51 @@ private struct WrappedCollectionView: UIViewControllerRepresentable {
         }
         uiViewController.updateAttachments()
     }
+    
+    @available(iOS 16.0, *)
+    func sizeThatFits(_ proposal: ProposedViewSize, uiViewController: WrappedCollectionViewController, context: Context) -> CGSize? {
+        guard let width = proposal.width,
+              width.isFinite else {
+            return nil
+        }
+        let count = draft.attachments.count + 1
+        return CGSize(
+            width: width,
+            height: Self.totalHeight(width: width, minItemSize: minItemSize, spacing: spacing, items: count)
+        )
+    }
+
+    fileprivate static func itemSize(width: CGFloat, minItemSize: CGFloat, spacing: CGFloat) -> (CGFloat, Int) {
+        // The maximum item size is 2*minItemSize + spacing - 1,
+        // in the case where one item fits in the row but we are one pt short of
+        // adding a second item.
+        var itemSize = minItemSize
+        var fittingCount = floor((width + spacing) / (itemSize + spacing))
+        var usedSpaceForFittingCount = fittingCount * itemSize + (fittingCount - 1) * spacing
+        var remainingSpace = width - usedSpaceForFittingCount
+        if fittingCount == 0 {
+            return (0, 0)
+        } else if fittingCount == 1 && remainingSpace > minItemSize / 2 {
+            // If there's only one item that would fit at min size, and giving
+            // it the rest of the space would increase it by at least 50%,
+            // add a second item anywyas.
+            itemSize = (width - spacing) / 2
+            fittingCount = 2
+            usedSpaceForFittingCount = fittingCount * itemSize + (fittingCount - 1) * spacing
+            remainingSpace = width - usedSpaceForFittingCount
+        }
+        itemSize = itemSize + remainingSpace / fittingCount
+        return (itemSize, Int(fittingCount))
+    }
+    
+    fileprivate static func totalHeight(width: CGFloat, minItemSize: CGFloat, spacing: CGFloat, items: Int) -> CGFloat {
+        let (size, itemsPerRow) = itemSize(width: width, minItemSize: minItemSize, spacing: spacing)
+        guard itemsPerRow != 0 else {
+            return 0
+        }
+        let rows = ceil(Double(items) / Double(itemsPerRow))
+        return size * rows + spacing * (rows - 1)
+    }
 }
 
 private class WrappedCollectionViewController: UIViewController {
@@ -70,6 +168,10 @@ private class WrappedCollectionViewController: UIViewController {
     fileprivate var currentInteractiveMoveCell: HostingCollectionViewCell?
     fileprivate var addAttachment: ((DraftAttachment) -> Void)? = nil
     
+    var collectionView: UICollectionView {
+        view as! UICollectionView
+    }
+    
     init(spacing: CGFloat, minItemSize: CGFloat) {
         self.spacing = spacing
         self.minItemSize = minItemSize
@@ -82,7 +184,7 @@ private class WrappedCollectionViewController: UIViewController {
     
     override func loadView() {
         let layout = UICollectionViewCompositionalLayout { [unowned self] section, environment in
-            let (itemSize, itemsPerRow) = self.itemSize(width: environment.container.contentSize.width)
+            let (itemSize, itemsPerRow) = WrappedCollectionView.itemSize(width: environment.container.contentSize.width, minItemSize: minItemSize, spacing: spacing)
             
             let items = Array(repeating: NSCollectionLayoutItem(layoutSize: .init(widthDimension: .absolute(itemSize), heightDimension: .absolute(itemSize))), count: itemsPerRow)
             let group = NSCollectionLayoutGroup.horizontal(layoutSize: .init(widthDimension: .fractionalWidth(1), heightDimension: .absolute(itemSize)), subitems: items)
@@ -91,10 +193,12 @@ private class WrappedCollectionViewController: UIViewController {
             section.interGroupSpacing = spacing
             return section
         }
-        let attachmentCell = UICollectionView.CellRegistration<HostingCollectionViewCell, DraftAttachment> { cell, indexPath, attachment in
+        let attachmentCell = UICollectionView.CellRegistration<HostingCollectionViewCell, DraftAttachment> { [unowned self] cell, indexPath, attachment in
+            cell.containingViewController = self
             cell.setView(AttachmentCollectionViewCellView(attachment: attachment))
         }
         let addButtonCell = UICollectionView.CellRegistration<HostingCollectionViewCell, Bool> { [unowned self] cell, indexPath, item in
+            cell.containingViewController = self
             cell.setView(AddAttachmentButton(viewController: self, enabled: item))
         }
         let collectionView = IntrinsicContentSizeCollectionView(frame: .zero, collectionViewLayout: layout)
@@ -140,27 +244,6 @@ private class WrappedCollectionViewController: UIViewController {
         collectionView.addGestureRecognizer(longPressRecognizer)
     }
     
-    private func itemSize(width: CGFloat) -> (CGFloat, Int) {
-        // The maximum item size is 2*minItemSize + spacing - 1,
-        // in the case where one item fits in the row but we are one pt short of
-        // adding a second item.
-        var itemSize = minItemSize
-        var fittingCount = floor((width + spacing) / (itemSize + spacing))
-        var usedSpaceForFittingCount = fittingCount * itemSize + (fittingCount - 1) * spacing
-        var remainingSpace = width - usedSpaceForFittingCount
-        if fittingCount == 1 && remainingSpace > minItemSize / 2 {
-            // If there's only one item that would fit at min size, and giving
-            // it the rest of the space would increase it by at least 50%,
-            // add a second item anywyas.
-            itemSize = (width - spacing) / 2
-            fittingCount = 2
-            usedSpaceForFittingCount = fittingCount * itemSize + (fittingCount - 1) * spacing
-            remainingSpace = width - usedSpaceForFittingCount
-        }
-        itemSize = itemSize + remainingSpace / fittingCount
-        return (itemSize, Int(fittingCount))
-    }
-
     func updateAttachments() {
         var snapshot = NSDiffableDataSourceSnapshot<Section, Item>()
         snapshot.appendSections([.all])
@@ -300,6 +383,8 @@ private class HostingCollectionViewCell: UICollectionViewCell {
 #else
 @available(iOS, obsoleted: 16.0)
 private class HostingCollectionViewCell: UICollectionViewCell {
+    weak var containingViewController: UIViewController?
+    
     @available(iOS, obsoleted: 16.0)
     private var hostController: UIHostingController<AnyView>?
     private(set) var hostView: UIView?
@@ -324,11 +409,12 @@ private class HostingCollectionViewCell: UICollectionViewCell {
                 hostController.rootView = AnyView(view)
             } else {
                 let host = UIHostingController(rootView: AnyView(view))
+                containingViewController!.addChild(host)
                 host.view.frame = contentView.bounds
                 contentView.addSubview(host.view)
+                host.didMove(toParent: containingViewController!)
                 hostController = host
                 hostView = host.view
-                // drop the hosting controller on the floor and hope nothing bad happens
             }
         }
     }