From 3ab04fb7d831a8d960611095f0135a03e3941af9 Mon Sep 17 00:00:00 2001
From: Benjie Gillam <benjie@jemjie.com>
Date: Thu, 21 May 2020 15:26:37 +0100
Subject: [PATCH] feat(merge): make fragment merging type-aware

---
 packages/graphiql/src/utility/mergeAst.ts | 148 +++++++++++++++++-----
 1 file changed, 114 insertions(+), 34 deletions(-)

diff --git a/packages/graphiql/src/utility/mergeAst.ts b/packages/graphiql/src/utility/mergeAst.ts
index 42c19a5763b..4f272ea730e 100644
--- a/packages/graphiql/src/utility/mergeAst.ts
+++ b/packages/graphiql/src/utility/mergeAst.ts
@@ -1,5 +1,5 @@
 /**
- *  Copyright (c) 2019 GraphQL Contributors.
+ *  Copyright (c) 2020 GraphQL Contributors.
  *
  *  This source code is licensed under the MIT license found in the
  *  LICENSE file in the root directory of this source tree.
@@ -7,59 +7,139 @@
 
 import {
   DocumentNode,
+  FieldNode,
   FragmentDefinitionNode,
-  visit,
+  GraphQLOutputType,
+  GraphQLSchema,
   SelectionNode,
-} from 'graphql';
+  TypeInfo,
+  getNamedType,
+  visit,
+  visitWithTypeInfo,
+} from "graphql";
 
-export function uniqueBy(
+export function uniqueBy<T>(
   array: readonly SelectionNode[],
-  iteratee: (item: any) => any,
+  iteratee: (item: FieldNode) => T
 ) {
-  const FilteredMap = new Map();
-  const result = [];
+  const FilteredMap = new Map<T, FieldNode>();
+  const result: SelectionNode[] = [];
   for (const item of array) {
-    const uniqeValue = iteratee(item);
-    if (!FilteredMap.has(uniqeValue)) {
-      FilteredMap.set(uniqeValue, true);
+    if (item.kind === "Field") {
+      const uniqueValue = iteratee(item);
+      const existing = FilteredMap.get(uniqueValue);
+      if (item.directives && item.directives.length) {
+        // Cannot inline fields with directives (yet)
+        result.push(itemClone);
+      } else if (existing && existing.selectionSet && item.selectionSet) {
+        // Merge the selection sets
+        existing.selectionSet.selections = [
+          ...existing.selectionSet.selections,
+          ...item.selectionSet.selections,
+        ];
+      } else if (!existing) {
+        const itemClone = { ...item };
+        FilteredMap.set(uniqueValue, itemClone);
+        result.push(itemClone);
+      }
+    } else {
       result.push(item);
     }
   }
   return result;
 }
 
+export function inlineRelevantFragmentSpreads(
+  fragmentDefinitions: {
+    [key: string]: FragmentDefinitionNode | undefined;
+  },
+  selectionSetType: GraphQLOutputType,
+  selections: readonly SelectionNode[]
+): readonly SelectionNode[] {
+  const selectionSetTypeName = getNamedType(selectionSetType).name;
+  const outputSelections = [];
+  for (let selection of selections) {
+    if (selection.kind === "FragmentSpread") {
+      const fragmentDefinition = fragmentDefinitions[selection.name.value];
+      if (fragmentDefinition) {
+        const { typeCondition, directives, selectionSet } = fragmentDefinition;
+        selection = {
+          kind: "InlineFragment",
+          typeCondition,
+          directives,
+          selectionSet,
+        };
+      }
+    }
+    if (
+      selection.kind === "InlineFragment" &&
+      // Cannot inline if there are directives
+      (!selection.directives || selection.directives?.length === 0)
+    ) {
+      const fragmentTypeName = selection.typeCondition
+        ? selection.typeCondition.name.value
+        : null;
+      if (!fragmentTypeName || fragmentTypeName === selectionSetTypeName) {
+        outputSelections.push(
+          ...inlineRelevantFragmentSpreads(
+            fragmentDefinitions,
+            selectionSetType,
+            selection.selectionSet.selections
+          )
+        );
+        continue;
+      }
+    }
+    outputSelections.push(selection);
+  }
+  return outputSelections;
+}
+
 /**
- * Given a document AST, inline all named fragment definitions
+ * Given a document AST, inline all named fragment definitions.
  */
-export default function mergeAST(documentAST: DocumentNode): DocumentNode {
+export default function mergeAST(
+  documentAST: DocumentNode,
+  schema: GraphQLSchema
+): DocumentNode {
+  const typeInfo = new TypeInfo(schema);
   const fragmentDefinitions: {
-    [key: string]: FragmentDefinitionNode;
+    [key: string]: FragmentDefinitionNode | undefined;
   } = Object.create(null);
 
   for (const definition of documentAST.definitions) {
-    if (definition.kind === 'FragmentDefinition') {
+    if (definition.kind === "FragmentDefinition") {
       fragmentDefinitions[definition.name.value] = definition;
     }
   }
 
-  return visit(documentAST, {
-    FragmentSpread(node) {
-      return {
-        ...fragmentDefinitions[node.name.value],
-        kind: 'InlineFragment',
-      };
-    },
-    SelectionSet(node) {
-      return {
-        ...node,
-        selections: uniqueBy(
-          node.selections,
-          selection => selection.name.value,
-        ),
-      };
-    },
-    FragmentDefinition() {
-      return null;
-    },
-  });
+  return visit(
+    documentAST,
+    visitWithTypeInfo(typeInfo, {
+      SelectionSet(node) {
+        const selectionSetType = typeInfo.getParentType();
+        let { selections } = node;
+
+        if (selectionSetType) {
+          selections = inlineRelevantFragmentSpreads(
+            fragmentDefinitions,
+            selectionSetType,
+            selections
+          );
+        }
+
+        selections = uniqueBy(selections, selection =>
+          selection.alias ? selection.alias.value : selection.name.value
+        );
+
+        return {
+          ...node,
+          selections,
+        };
+      },
+      FragmentDefinition() {
+        return null;
+      },
+    })
+  );
 }