From 9de9db14a4d0db8d742ad86a06511807579416cc Mon Sep 17 00:00:00 2001
From: Jack Rosen <jrosen081@gmail.com>
Date: Fri, 12 Jan 2024 12:04:35 -0500
Subject: [PATCH] Work correctly with functions

---
 .../WhoopDIKitMacros/InjectableMacro.swift    |  7 ++++-
 Tests/WhoopDIKitTests/InjectableTests.swift   | 26 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/Sources/WhoopDIKitMacros/InjectableMacro.swift b/Sources/WhoopDIKitMacros/InjectableMacro.swift
index b190976..c833c00 100644
--- a/Sources/WhoopDIKitMacros/InjectableMacro.swift
+++ b/Sources/WhoopDIKitMacros/InjectableMacro.swift
@@ -200,6 +200,11 @@ extension VariableDeclSyntax {
     }
 
     var typeName: TypeSyntax? {
-        self.bindings.first?.typeAnnotation?.type.trimmed
+        guard let annotationType = self.bindings.first?.typeAnnotation?.type.trimmed else { return nil }
+        if (annotationType.is(FunctionTypeSyntax.self)) {
+            return "@escaping \(annotationType)"
+        } else {
+            return annotationType
+        }
     }
 }
diff --git a/Tests/WhoopDIKitTests/InjectableTests.swift b/Tests/WhoopDIKitTests/InjectableTests.swift
index f81bac6..4670850 100644
--- a/Tests/WhoopDIKitTests/InjectableTests.swift
+++ b/Tests/WhoopDIKitTests/InjectableTests.swift
@@ -105,4 +105,30 @@ final class InjectableTests: XCTestCase {
         """,
         macros: ["Injectable": InjectableMacro.self])
     }
+
+    func testInjectWithClosures() {
+        assertMacroExpansion(
+            """
+            @Injectable struct ClosureHolder {
+                let closure: () -> String
+            }
+            """,
+            expandedSource: """
+            struct ClosureHolder {
+                let closure: () -> String
+
+                internal static func inject() -> Self {
+                    Self.init(closure: WhoopDI.inject(nil))
+                }
+
+                internal init(closure: @escaping () -> String) {
+                    self.closure = closure
+                }
+            }
+            
+            extension ClosureHolder : Injectable {
+            }
+            """,
+            macros: ["Injectable": InjectableMacro.self])
+    }
 }