[Codegen]: Partial fragments

apollo-ios-codegen/Sources/IR/IR+DefinitionEntityStorage.swift

@@ -11,6 +11,14 @@ public class DefinitionEntityStorage {
     self.entitiesForFields[rootEntity.location] = rootEntity
   }
 
+  private init(
+    sourceDefinition: Entity.Location.SourceDefinition,
+    entitiesForFields: [Entity.Location: Entity]
+  ) {
+    self.sourceDefinition = sourceDefinition
+    self.entitiesForFields = entitiesForFields
+  }
+
   func entity(
     for field: CompilationResult.Field,
     on enclosingEntity: Entity
@@ -66,5 +74,37 @@ public class DefinitionEntityStorage {
     entitiesForFields[location] = entity
     return entity
   }
+}
+
+extension DefinitionEntityStorage {
+  func partial() -> DefinitionEntityStorage {
+    return DefinitionEntityStorage(
+      sourceDefinition: sourceDefinition,
+      entitiesForFields: makePartialEntitiesForFields()
+    )
+  }
+
+  private func makePartialEntitiesForFields() -> [Entity.Location: Entity] {
+    var partialEntitiesForFields: [Entity.Location: Entity] = [:]
+    for (location, entity) in entitiesForFields {
+      /*
+       Based on my understanding:
+       Fragment rendering points to these entities.
+       Whenever we try to render a top-level fragment spread, it resolves the fragment from the entity storage.
+       In theory, the generation flow needs the types and entities to spread, fragments, and their underlying entities,
+       hence the hardcoded 'count < 4'.
 
+       Possibly, could iterate through the field path and exclude irrelevant entities rather than a hardcoded value?
+
+       */
+      if let count = location.fieldPath?.count, count < 4 {
+        partialEntitiesForFields[location] = entity
+      }
+      if location.fieldPath == nil {
+        partialEntitiesForFields[location] = entity
+      }
+    }
+    return partialEntitiesForFields
+  }
 }
+

apollo-ios-codegen/Sources/IR/IR+NamedFragment.swift

@@ -51,3 +51,15 @@ public class NamedFragment: Definition, Hashable, CustomDebugStringConvertible {
     definition.debugDescription
   }
 }
+
+public extension NamedFragment {
+  func partialFragment() -> NamedFragment {
+    NamedFragment(
+      definition: definition,
+      rootField: rootField,
+      referencedFragments: referencedFragments,
+      entityStorage: entityStorage.partial(),
+      containsDeferredFragment: containsDeferredFragment
+    )
+  }
+}

apollo-ios-codegen/Sources/IR/IR+RootFieldBuilder.swift

@@ -398,7 +398,7 @@ class RootFieldBuilder {
     spreadIntoParentWithTypePath parentTypeInfo: SelectionSet.TypeInfo,
     deferCondition: CompilationResult.DeferCondition? = nil
   ) async -> NamedFragmentSpread {
-    let fragment = await ir.build(fragment: fragmentSpread.fragment)
+    let fragment = await ir.build(fragment: fragmentSpread.fragment, shouldGeneratePartialFragment: true)
     referencedFragments.append(fragment)
     referencedFragments.append(contentsOf: fragment.referencedFragments)
 
@@ -410,7 +410,7 @@ class RootFieldBuilder {
 
     self.containsDeferredFragment = fragment.containsDeferredFragment || scope.deferCondition != nil
 
-    let scopePath = scope.isEmpty 
+    let scopePath = scope.isEmpty
       ? parentTypeInfo.scopePath
       : parentTypeInfo.scopePath.mutatingLast { $0.appending(scope) }
 

apollo-ios-codegen/Sources/IR/IRBuilder.swift

@@ -23,7 +23,8 @@ public class IRBuilder {
   }
 
   public func build(
-    operation operationDefinition: CompilationResult.OperationDefinition
+    operation operationDefinition: CompilationResult.OperationDefinition,
+    usePartialEntities: Bool = false
   ) async -> Operation {
     let rootField = CompilationResult.Field(
       name: operationDefinition.operationType.rawValue,
@@ -86,9 +87,11 @@ public class IRBuilder {
   }
 
   public func build(
-    fragment fragmentDefinition: CompilationResult.FragmentDefinition
+    fragment fragmentDefinition: CompilationResult.FragmentDefinition,
+    shouldGeneratePartialFragment: Bool = false
   ) async -> NamedFragment {
-    await builtFragmentStorage.getFragment(named: fragmentDefinition.name) {
+
+    let namedFragment = await builtFragmentStorage.getFragment(named: fragmentDefinition.name) {
       let rootField = CompilationResult.Field(
         name: fragmentDefinition.name,
         type: .nonNull(.entity(fragmentDefinition.type)),
@@ -113,6 +116,10 @@ public class IRBuilder {
         containsDeferredFragment: result.containsDeferredFragment
       )
     }
-  }
 
+    if shouldGeneratePartialFragment {
+      return namedFragment.partialFragment()
+    }
+    return namedFragment
+  }
 }