Table macro fix for primary key trivia (#51)

* wip

* wip

Author: stephencelis

Sources/StructuredQueriesMacros/TableMacro.swift

@@ -34,7 +34,9 @@ extension TableMacro: ExtensionMacro {
     var diagnostics: [Diagnostic] = []
 
     // NB: A compiler bug prevents us from applying the '@_Draft' macro directly
-    var draftBindings: [(PatternBindingSyntax, queryOutputType: TypeSyntax?)] = []
+    var draftBindings: [
+      (PatternBindingSyntax, queryOutputType: TypeSyntax?, optionalize: Bool)
+    ] = []
     // NB: End of workaround
 
     var draftProperties: [DeclSyntax] = []
@@ -223,12 +225,8 @@ extension TableMacro: ExtensionMacro {
         )
       }
 
-      // NB: A compiled bug prevents us from applying the '@_Draft' macro directly
-      if identifier == primaryKey?.identifier {
-        draftBindings.append((binding.optionalized(), columnQueryOutputType))
-      } else {
-        draftBindings.append((binding, columnQueryOutputType))
-      }
+      // NB: A compiler bug prevents us from applying the '@_Draft' macro directly
+      draftBindings.append((binding, columnQueryOutputType, identifier == primaryKey?.identifier))
       // NB: End of workaround
 
       var assignedType: String? {
@@ -459,8 +457,12 @@ extension TableMacro: ExtensionMacro {
       .compactMap(\.memberBlock.members.trimmed)
       var memberwiseArguments: [PatternBindingSyntax] = []
       var memberwiseAssignments: [TokenSyntax] = []
-      for (binding, queryOutputType) in draftBindings {
-        let argument = binding.trimmed.annotated(queryOutputType).rewritten(selfRewriter)
+      for (binding, queryOutputType, optionalize) in draftBindings {
+        var argument = binding.trimmed
+        if optionalize {
+          argument = argument.optionalized()
+        }
+        argument = argument.annotated(queryOutputType).rewritten(selfRewriter)
         if argument.typeAnnotation == nil {
           let identifier =
             (argument.pattern.as(IdentifierPatternSyntax.self)?.identifier.trimmedDescription)

Tests/StructuredQueriesMacrosTests/TableMacroTests.swift

@@ -1703,5 +1703,75 @@ extension SnapshotTests {
         """#
       }
     }
+
+    @Test func commentAfterOptionalID() {
+      assertMacro {
+        """
+        @Table
+        struct Reminder {
+          let id: Int?  // TODO: Migrate to UUID
+          var title = ""
+        }
+        """
+      } expansion: {
+        #"""
+        struct Reminder {
+          let id: Int?  // TODO: Migrate to UUID
+          var title = ""
+        }
+
+        extension Reminder: StructuredQueriesCore.Table, StructuredQueriesCore.PrimaryKeyedTable {
+          public struct TableColumns: StructuredQueriesCore.TableDefinition, StructuredQueriesCore.PrimaryKeyedTableDefinition {
+            public typealias QueryValue = Reminder
+            public let id = StructuredQueriesCore.TableColumn<QueryValue, Int?>("id", keyPath: \QueryValue.id)
+            public let title = StructuredQueriesCore.TableColumn<QueryValue, Swift.String>("title", keyPath: \QueryValue.title, default: "")
+            public var primaryKey: StructuredQueriesCore.TableColumn<QueryValue, Int?> {
+              self.id
+            }
+            public static var allColumns: [any StructuredQueriesCore.TableColumnExpression] {
+              [QueryValue.columns.id, QueryValue.columns.title]
+            }
+          }
+          public struct Draft: StructuredQueriesCore.TableDraft {
+            public typealias PrimaryTable = Reminder
+            @Column(primaryKey: false)
+            let id: Int?
+            var title = ""
+            public struct TableColumns: StructuredQueriesCore.TableDefinition {
+              public typealias QueryValue = Reminder.Draft
+              public let id = StructuredQueriesCore.TableColumn<QueryValue, Int?>("id", keyPath: \QueryValue.id)
+              public let title = StructuredQueriesCore.TableColumn<QueryValue, Swift.String>("title", keyPath: \QueryValue.title, default: "")
+              public static var allColumns: [any StructuredQueriesCore.TableColumnExpression] {
+                [QueryValue.columns.id, QueryValue.columns.title]
+              }
+            }
+            public static let columns = TableColumns()
+            public static let tableName = Reminder.tableName
+            public init(decoder: inout some StructuredQueriesCore.QueryDecoder) throws {
+              self.id = try decoder.decode(Int.self)
+              self.title = try decoder.decode(Swift.String.self) ?? ""
+            }
+            public init(_ other: Reminder) {
+              self.id = other.id
+              self.title = other.title
+            }
+            public init(
+              id: Int? = nil,
+              title: Swift.String = ""
+            ) {
+              self.id = id
+              self.title = title
+            }
+          }
+          public static let columns = TableColumns()
+          public static let tableName = "reminders"
+          public init(decoder: inout some StructuredQueriesCore.QueryDecoder) throws {
+            self.id = try decoder.decode(Int.self)
+            self.title = try decoder.decode(Swift.String.self) ?? ""
+          }
+        }
+        """#
+      }
+    }
   }
 }