Codegen EntitySelectionTree + MergedSelections Refactor + Perf Improvments

Tests/ApolloCodegenInternalTestHelpers/IR+Mocking.swift → Tests/ApolloCodegenInternalTestHelpers/IRBuilder+Mocking.swift

@@ -42,10 +42,10 @@ extension IRBuilder {
     return .mock(compilationResult: compilationResult)
   }
 
-  public static func mock(    
+  public static func mock(
     compilationResult: CompilationResult
   ) -> IRBuilder {
-    return IRBuilder(compilationResult: compilationResult)
+    return IRBuilder(compilationResult: compilationResult)    
   }
 
 }
@@ -63,24 +63,31 @@ extension IR.NamedFragment {
       location: .init(source: .namedFragment(definition), fieldPath: nil),
       rootTypePath: LinkedList(type)
     )
+    let typeInfo = SelectionSet.TypeInfo(
+      entity: rootEntity,
+      scopePath: [
+        .descriptor(
+          forType: type,
+          inclusionConditions: nil,
+          givenAllTypesInSchema: .init([type], schemaRootTypes: .mock())
+        )
+      ]
+    )
     let rootEntityField = IR.EntityField.init(
       rootField,
       inclusionConditions: nil,
       selectionSet: .init(
-        entity: rootEntity,
-        scopePath: LinkedList(
-          .descriptor(
-            forType: type,
-            inclusionConditions: nil,
-            givenAllTypesInSchema: .init([type], schemaRootTypes: .mock())
-          ))))
+        typeInfo: typeInfo,
+        selections: DirectSelections()
+      )
+    )
 
     return IR.NamedFragment(
       definition: definition,
       rootField: rootEntityField,
       referencedFragments: [],
-      containsDeferredFragment: false,
-      entities: [rootEntity.location: rootEntity]
+      entityStorage: .init(rootEntity: rootEntity),
+      containsDeferredFragment: false
     )
   }
 }
@@ -93,23 +100,30 @@ extension IR.Operation {
     containsDeferredFragment: Bool = false
   ) -> IR.Operation {
     let definition = definition ?? .mock()
+    let typeInfo = SelectionSet.TypeInfo(
+      entity: .init(
+        location: .init(source: .operation(definition), fieldPath: nil),
+        rootTypePath: [.mock()]
+      ),
+      scopePath: [.descriptor(
+        forType: .mock(),
+        inclusionConditions: nil,
+        givenAllTypesInSchema: .init([], schemaRootTypes: .mock()))
+      ]
+    )
+    let rootField = IR.EntityField(
+      .mock(),
+      inclusionConditions: nil,
+      selectionSet: .init(
+        typeInfo: typeInfo,
+        selections: DirectSelections()
+      )
+    )
     return IR.Operation.init(
       definition: definition,
-      rootField: .init(
-        .mock(),
-        inclusionConditions: nil,
-        selectionSet: .init(
-          entity: .init(
-            location: .init(source: .operation(definition), fieldPath: nil),
-            rootTypePath: [.mock()]
-          ),
-          scopePath: [.descriptor(
-            forType: .mock(),
-            inclusionConditions: nil,
-            givenAllTypesInSchema: .init([], schemaRootTypes: .mock()))
-          ])
-      ),
+      rootField: rootField,
       referencedFragments: referencedFragments,
+      entityStorage: .init(rootEntity: rootField.entity),
       containsDeferredFragment: containsDeferredFragment
     )
   }

Tests/ApolloCodegenInternalTestHelpers/IRBuilderTestWrapper.swift

@@ -0,0 +1,65 @@
+@testable import ApolloCodegenLib
+@testable import IR
+import GraphQLCompiler
+import OrderedCollections
+import Utilities
+
+/// This class wraps an `IRBuilder` for so that the built definitions are wrapped in
+/// `IRTestWrapper` instances.
+///
+/// These wrappers are used to make testing `ComputedSelectionSet` easier.
+@dynamicMemberLookup
+public class IRBuilderTestWrapper {
+  public let irBuilder: IRBuilder
+
+  public private(set) lazy var builtFragmentStorage : BuiltFragmentStorage = {
+    BuiltFragmentStorage(self.irBuilder.builtFragmentStorage)
+  }()
+
+  public init(_ irBuilder: IRBuilder) {
+    self.irBuilder = irBuilder
+  }
+
+  public func build(
+    operation operationDefinition: CompilationResult.OperationDefinition
+  ) async -> IRTestWrapper<IR.Operation> {
+    let operation = await irBuilder.build(operation: operationDefinition)
+    return IRTestWrapper(
+      irObject: operation,
+      computedSelectionSetCache: .init(entityStorage: operation.entityStorage)
+    )
+  }
+
+  public func build(
+    fragment fragmentDefinition: CompilationResult.FragmentDefinition
+  ) async -> IRTestWrapper<IR.NamedFragment> {
+    let fragment = await irBuilder.build(fragment: fragmentDefinition)
+    return IRTestWrapper(
+      irObject: fragment,
+      computedSelectionSetCache: .init(entityStorage: fragment.entityStorage)
+    )
+  }
+
+  public subscript<T>(dynamicMember keyPath: KeyPath<IRBuilder, T>) -> T {
+    irBuilder[keyPath: keyPath]
+  }
+
+  public class BuiltFragmentStorage {
+    private var wrappedStorage: IRBuilder.BuiltFragmentStorage    
+
+    fileprivate init(_ wrappedStorage: IRBuilder.BuiltFragmentStorage) {
+      self.wrappedStorage = wrappedStorage
+    }
+
+    public func getFragmentIfBuilt(named name: String) async -> IRTestWrapper<NamedFragment>? {
+      guard let fragment = await wrappedStorage.getFragmentIfBuilt(named: name) else {
+        return nil
+      }
+
+      return IRTestWrapper(
+        irObject: fragment,
+        computedSelectionSetCache: .init(entityStorage: fragment.entityStorage)
+      )
+    }
+  }
+}