From c18dcdc6bdf0355cf0e19f8428a6127e9af567ff Mon Sep 17 00:00:00 2001
From: Adam Boyd <adamboyd50@gmail.com>
Date: Fri, 26 Jan 2024 16:13:55 -0800
Subject: [PATCH] Optional Instantiator and ForwardingInstantiator Support
 (#45)

Co-authored-by: Dan Federman <dfed@me.com>
---
 .../Generators/ScopeGenerator.swift           | 25 ++++++++++++-
 Sources/SafeDICore/Models/Property.swift      | 36 ++++++++++++------
 Tests/SafeDIToolTests/SafeDIToolTests.swift   | 37 +++++++++++++++++++
 3 files changed, 84 insertions(+), 14 deletions(-)

diff --git a/Sources/SafeDICore/Generators/ScopeGenerator.swift b/Sources/SafeDICore/Generators/ScopeGenerator.swift
index 57a35b0b..4a86a251 100644
--- a/Sources/SafeDICore/Generators/ScopeGenerator.swift
+++ b/Sources/SafeDICore/Generators/ScopeGenerator.swift
@@ -165,8 +165,14 @@ actor ScopeGenerator: CustomStringConvertible {
                     switch property.propertyType {
                     case .instantiator, .forwardingInstantiator:
                         isConstant = false
-                        propertyDeclaration = "let \(property.label)"
-                        leadingConcreteTypeName = property.typeDescription.asSource
+                        let typeDescription = property.typeDescription.asSource
+                        let unwrappedTypeDescription = property.typeDescription.unwrappedTypeDescription.asSource
+                        if typeDescription == unwrappedTypeDescription {
+                            propertyDeclaration = "let \(property.label)"
+                        } else {
+                            propertyDeclaration = "let \(property.label): \(typeDescription)"
+                        }
+                        leadingConcreteTypeName = unwrappedTypeDescription
                     case .constant:
                         isConstant = true
                         if concreteTypeName == property.typeDescription.asSource {
@@ -331,3 +337,18 @@ extension Instantiable {
             ) ?? "/* @Instantiable type is incorrectly configured. Fix errors from @Instantiable macro to fix this error. */"
     }
 }
+
+// MARK: TypeDescription
+
+extension TypeDescription {
+    fileprivate var unwrappedTypeDescription: TypeDescription {
+        switch self {
+        case
+            let .optional(type),
+            let .implicitlyUnwrappedOptional(type):
+            return type.unwrappedTypeDescription
+        case .any, .array, .attributed, .closure, .composition, .dictionary, .metatype, .nested, .simple, .some, .tuple, .unknown, .void:
+            return self
+        }
+    }
+}
diff --git a/Sources/SafeDICore/Models/Property.swift b/Sources/SafeDICore/Models/Property.swift
index 8c281043..9cff1425 100644
--- a/Sources/SafeDICore/Models/Property.swift
+++ b/Sources/SafeDICore/Models/Property.swift
@@ -115,18 +115,7 @@ public struct Property: Codable, Hashable, Comparable, Sendable {
     }
 
     var propertyType: PropertyType {
-        switch typeDescription {
-        case let .simple(name, _):
-            if name == Dependency.instantiatorType {
-                return .instantiator
-            } else if name == Dependency.forwardingInstantiatorType {
-                return .forwardingInstantiator
-            } else {
-                return .constant
-            }
-        case .any, .array, .attributed, .closure, .composition, .dictionary, .implicitlyUnwrappedOptional, .metatype, .nested, .optional, .some, .tuple, .unknown, .void:
-            return .constant
-        }
+        typeDescription.propertyType
     }
 
     var generics: [TypeDescription] {
@@ -152,3 +141,26 @@ public struct Property: Codable, Hashable, Comparable, Sendable {
         case forwardingInstantiator
     }
 }
+
+// MARK: TypeDescription
+
+extension TypeDescription {
+    fileprivate var propertyType: Property.PropertyType {
+        switch self {
+        case let .simple(name, _):
+            if name == Dependency.instantiatorType {
+                return .instantiator
+            } else if name == Dependency.forwardingInstantiatorType {
+                return .forwardingInstantiator
+            } else {
+                return .constant
+            }
+        case
+            let .optional(type),
+            let .implicitlyUnwrappedOptional(type):
+            return type.propertyType
+        case .any, .array, .attributed, .closure, .composition, .dictionary, .metatype, .nested, .some, .tuple, .unknown, .void:
+            return .constant
+        }
+    }
+}
diff --git a/Tests/SafeDIToolTests/SafeDIToolTests.swift b/Tests/SafeDIToolTests/SafeDIToolTests.swift
index d32bf5ec..29760beb 100644
--- a/Tests/SafeDIToolTests/SafeDIToolTests.swift
+++ b/Tests/SafeDIToolTests/SafeDIToolTests.swift
@@ -3575,6 +3575,43 @@ final class SafeDIToolTests: XCTestCase {
         )
     }
 
+    func test_run_writesConvenienceExtensionOnRootOfTree_whenRootPropertyWithOptionalInstantiator() async throws {
+        let output = try await executeSystemUnderTest(
+            swiftFileContent: [
+                """
+                @Instantiable
+                public final class Root {
+                    @Instantiated
+                    let childBuilder: Instantiator<Child>?
+                }
+                """,
+                """
+                @Instantiable
+                public final class Child {}
+                """,
+            ],
+            buildDependencyTreeOutput: true
+        )
+
+        XCTAssertEqual(
+            try XCTUnwrap(output.dependencyTree),
+            """
+            // This file was generated by the SafeDIGenerateDependencyTree build tool plugin.
+            // Any modifications made to this file will be overwritten on subsequent builds.
+            // Please refrain from editing this file directly.
+
+            extension Root {
+                public convenience init() {
+                    let childBuilder: Instantiator<Child>? = Instantiator<Child> {
+                        Child()
+                    }
+                    self.init(childBuilder: childBuilder)
+                }
+            }
+            """
+        )
+    }
+
     // MARK: Error Tests
 
     func test_run_onCodeWithPropertyWithUnknownFulfilledType_throwsError() async {