From 28ca9857886bf0ce8959be77d98c1f6e4cc6562c Mon Sep 17 00:00:00 2001
From: Lennard Sprong <lennard@buku.io>
Date: Tue, 26 Nov 2024 13:07:53 +0100
Subject: [PATCH] feature: Emit Identifiable conformance on SelectionSets

---
 .../SelectionSetTemplateTests.swift           | 111 ++++++++++++++++++
 .../Templates/SelectionSetTemplate.swift      |   5 +-
 .../Sources/IR/IR+ComputedSelectionSet.swift  |  16 +++
 3 files changed, 131 insertions(+), 1 deletion(-)

diff --git a/Tests/ApolloCodegenTests/CodeGeneration/Templates/SelectionSet/SelectionSetTemplateTests.swift b/Tests/ApolloCodegenTests/CodeGeneration/Templates/SelectionSet/SelectionSetTemplateTests.swift
index e71e7780a..36ad1cf9b 100644
--- a/Tests/ApolloCodegenTests/CodeGeneration/Templates/SelectionSet/SelectionSetTemplateTests.swift
+++ b/Tests/ApolloCodegenTests/CodeGeneration/Templates/SelectionSet/SelectionSetTemplateTests.swift
@@ -125,6 +125,117 @@ class SelectionSetTemplateTests: XCTestCase {
     // then
     expect(String(actual.reversed())).to(equalLineByLine("}", ignoringExtraLines: true))
   }
+  
+  // MARK: Protocol conformance
+  
+  func test__render_selectionSet__givenTypeWithKeyFieldID_rendersIdentifiableConformance() async throws {
+    // given
+    schemaSDL = """
+    type Query {
+      allAnimals: [Animal!]
+    }
+
+    type Animal @typePolicy(keyFields: "id") {
+      id: ID!
+    }
+    """
+
+    document = """
+    query TestOperation {
+      allAnimals {
+        id
+      }
+    }
+    """
+
+    let expected = """
+    public struct AllAnimal: TestSchema.SelectionSet, Identifiable {
+    """
+
+    // when
+    try await buildSubjectAndOperation()
+    let allAnimals = try XCTUnwrap(
+      operation[field: "query"]?[field: "allAnimals"]?.selectionSet
+    )
+
+    let actual = subject.test_render(childEntity: allAnimals.computed)
+
+    // then
+    expect(actual).to(equalLineByLine(expected, atLine: 2, ignoringExtraLines: true))
+  }
+  
+  func test__render_selectionSet__givenInterfaceWithKeyFieldID_rendersIdentifiableConformance() async throws {
+    // given
+    schemaSDL = """
+    type Query {
+      allAnimals: [Animal!]
+    }
+
+    interface Animal @typePolicy(keyFields: "id") {
+      id: ID!
+      species: String!
+    }
+    """
+
+    document = """
+    query TestOperation {
+      allAnimals {
+        id
+      }
+    }
+    """
+
+    let expected = """
+    public struct AllAnimal: TestSchema.SelectionSet, Identifiable {
+    """
+
+    // when
+    try await buildSubjectAndOperation()
+    let allAnimals = try XCTUnwrap(
+      operation[field: "query"]?[field: "allAnimals"]?.selectionSet
+    )
+
+    let actual = subject.test_render(childEntity: allAnimals.computed)
+
+    // then
+    expect(actual).to(equalLineByLine(expected, atLine: 2, ignoringExtraLines: true))
+  }
+  
+  func test__render_selectionSet__givenTypeWithOtherKeyField_doesNotRenderIdentifiableConformance() async throws {
+    // given
+    schemaSDL = """
+    type Query {
+      allAnimals: [Animal!]
+    }
+
+    type Animal @typePolicy(keyFields: "species") {
+      species: String!
+    }
+    """
+
+    document = """
+    query TestOperation {
+      allAnimals {
+        species
+      }
+    }
+    """
+
+    let expected = """
+    public struct AllAnimal: TestSchema.SelectionSet {
+    """
+
+    // when
+    try await buildSubjectAndOperation()
+    let allAnimals = try XCTUnwrap(
+      operation[field: "query"]?[field: "allAnimals"]?.selectionSet
+    )
+
+    let actual = subject.test_render(childEntity: allAnimals.computed)
+
+    // then
+    expect(actual).to(equalLineByLine(expected, atLine: 2, ignoringExtraLines: true))
+  }
 
   // MARK: Parent Type
 
diff --git a/apollo-ios-codegen/Sources/ApolloCodegenLib/Templates/SelectionSetTemplate.swift b/apollo-ios-codegen/Sources/ApolloCodegenLib/Templates/SelectionSetTemplate.swift
index 98448edbb..7bafaf63c 100644
--- a/apollo-ios-codegen/Sources/ApolloCodegenLib/Templates/SelectionSetTemplate.swift
+++ b/apollo-ios-codegen/Sources/ApolloCodegenLib/Templates/SelectionSetTemplate.swift
@@ -102,7 +102,9 @@ struct SelectionSetTemplate {
       """
       \(SelectionSetNameDocumentation(selectionSet))
       \(renderAccessControl())\
-      struct \(fieldSelectionSetName): \(SelectionSetType()) {
+      struct \(fieldSelectionSetName): \(SelectionSetType())\
+      \(if: selectionSet.isIdentifiable, ", Identifiable")\
+       {
         \(BodyTemplate(context))
       }
       """
@@ -118,6 +120,7 @@ struct SelectionSetTemplate {
       \(renderAccessControl())\
       struct \(inlineFragment.renderedTypeName): \(SelectionSetType(asInlineFragment: true))\
       \(if: inlineFragment.isCompositeInlineFragment, ", \(config.ApolloAPITargetName).CompositeInlineFragment")\
+      \(if: inlineFragment.isIdentifiable, ", Identifiable")\
        {
         \(BodyTemplate(context))
       }
diff --git a/apollo-ios-codegen/Sources/IR/IR+ComputedSelectionSet.swift b/apollo-ios-codegen/Sources/IR/IR+ComputedSelectionSet.swift
index d9cf5e7f0..bd2f7c0fe 100644
--- a/apollo-ios-codegen/Sources/IR/IR+ComputedSelectionSet.swift
+++ b/apollo-ios-codegen/Sources/IR/IR+ComputedSelectionSet.swift
@@ -1,4 +1,5 @@
 import Foundation
+import GraphQLCompiler
 import OrderedCollections
 import Utilities
 
@@ -18,6 +19,21 @@ public struct ComputedSelectionSet {
 
   /// The `TypeInfo` for the selection set of the computed selections
   public let typeInfo: IR.SelectionSet.TypeInfo
+  
+  /// Indicates if the parent type has a single keyField named `id`.
+  public var isIdentifiable: Bool {
+    if let type = typeInfo.parentType as? GraphQLObjectType,
+       type.keyFields == ["id"] {
+      return true
+    }
+    
+    if let type = typeInfo.parentType as? GraphQLInterfaceType,
+       type.keyFields == ["id"] {
+      return true
+    }
+    
+    return false
+  }
 
   // MARK: Dynamic Member Subscript