From 91611318a3d216fbe68f665e9b18a0b8ed5258ee Mon Sep 17 00:00:00 2001
From: foyoodo <foyoodo@gmail.com>
Date: Tue, 17 Sep 2024 17:28:06 +0800
Subject: [PATCH] Draft for generic support

---
 Sources/Swift/FlexLayout.swift             | 26 ++++++++++++----------
 Sources/Swift/Impl/UIView+FlexLayout.swift | 13 +++++++----
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/Sources/Swift/FlexLayout.swift b/Sources/Swift/FlexLayout.swift
index bf8cf93d..937a7dd3 100644
--- a/Sources/Swift/FlexLayout.swift
+++ b/Sources/Swift/FlexLayout.swift
@@ -26,7 +26,7 @@ import FlexLayoutYogaKit
     label.flex.margin(10)
  ```
  */
-public final class Flex {
+public final class Flex<Base: UIView> {
     
     //
     // MARK: Properties
@@ -35,9 +35,11 @@ public final class Flex {
     /**
      Flex items's UIView.
     */
-    public private(set) weak var view: UIView?
+    public private(set) weak var base: Base?
     private let yoga: YGLayout
     
+    public var view: UIView? { base }
+    
     /**
      Item natural size, considering only properties of the view itself. Independent of the item frame.
      */
@@ -45,10 +47,10 @@ public final class Flex {
         return yoga.intrinsicSize
     }
     
-    init(view: UIView) {
-        self.view = view
-        self.yoga = view.yoga
-        
+    init(_ base: Base) {
+        self.base = base
+        self.yoga = base.yoga
+
         // Enable flexbox and overwrite Yoga default values.
         yoga.isEnabled = true
     }
@@ -66,7 +68,7 @@ public final class Flex {
      - Returns: The flex interface corresponding to the added view.
      */
     @discardableResult
-    public func addItem() -> Flex {
+    public func addItem() -> Flex<UIView> {
         let view = UIView()
         return addItem(view)
     }
@@ -80,8 +82,8 @@ public final class Flex {
      - Returns: The flex interface corresponding to the added view.
      */
     @discardableResult
-    public func addItem(_ view: UIView) -> Flex {
-        if let host = self.view {
+    public func addItem<Item: UIView>(_ view: Item) -> Flex<Item> {
+        if let host = base {
             host.addSubview(view)
             return view.flex
         } else {
@@ -1278,7 +1280,7 @@ public final class Flex {
     */
     @discardableResult
     public func backgroundColor(_ color: UIColor) -> Flex {
-        if let host = self.view {
+        if let host = base {
             host.backgroundColor = color
             return self
         } else {
@@ -1295,7 +1297,7 @@ public final class Flex {
      */
     @discardableResult
     public func cornerRadius(_ value: CGFloat) -> Flex {
-        if let host = self.view {
+        if let host = base {
             host.layer.cornerRadius = value
             return self
         } else {
@@ -1313,7 +1315,7 @@ public final class Flex {
      */
     @discardableResult
     public func border(_ width: CGFloat, _ color: UIColor) -> Flex {
-        if let host = self.view {
+        if let host = base {
             host.layer.borderWidth = width
             host.layer.borderColor = color.cgColor
             return self
diff --git a/Sources/Swift/Impl/UIView+FlexLayout.swift b/Sources/Swift/Impl/UIView+FlexLayout.swift
index cb8ec2c0..756586cf 100644
--- a/Sources/Swift/Impl/UIView+FlexLayout.swift
+++ b/Sources/Swift/Impl/UIView+FlexLayout.swift
@@ -16,12 +16,15 @@ import UIKit
 
 private var flexLayoutAssociatedObjectHandle = 72_399_923
 
-extension UIView {
-    public var flex: Flex {
-        if let flex = objc_getAssociatedObject(self, &flexLayoutAssociatedObjectHandle) as? Flex {
+public protocol FlexLayoutCompatible: UIView { }
+
+extension FlexLayoutCompatible {
+
+    public var flex: Flex<Self> {
+        if let flex = objc_getAssociatedObject(self, &flexLayoutAssociatedObjectHandle) as? Flex<Self> {
             return flex
         } else {
-            let flex = Flex(view: self)
+            let flex = Flex(self)
             objc_setAssociatedObject(self, &flexLayoutAssociatedObjectHandle, flex, objc_AssociationPolicy.OBJC_ASSOCIATION_RETAIN_NONATOMIC)
             return flex
         }
@@ -31,3 +34,5 @@ extension UIView {
         (objc_getAssociatedObject(self, &flexLayoutAssociatedObjectHandle) as? Flex) != nil
     }
 }
+
+extension UIView: FlexLayoutCompatible { }