Merge branch 'dminkovsky-visitor-edit-root-node'

Author: leebyron

src/language/__tests__/visitor.js

@@ -20,6 +20,80 @@ import { getNamedType, isCompositeType } from '../../type';
 
 
 describe('Visitor', () => {
+
+  it('allows editing a node both on enter and on leave', () => {
+
+    const ast = parse('{ a, b, c { a, b, c } }', { noLocation: true });
+
+    let selectionSet;
+
+    const editedAst = visit(ast, {
+      OperationDefinition: {
+        enter(node) {
+          selectionSet = node.selectionSet;
+          return {
+            ...node,
+            selectionSet: {
+              kind: 'SelectionSet',
+              selections: []
+            },
+            didEnter: true,
+          };
+        },
+        leave(node) {
+          return {
+            ...node,
+            selectionSet,
+            didLeave: true,
+          };
+        }
+      }
+    });
+
+    expect(editedAst).to.deep.equal({
+      ...ast,
+      definitions: [
+        {
+          ...ast.definitions[0],
+          didEnter: true,
+          didLeave: true
+        }
+      ]
+    });
+  });
+
+  it('allows editing the root node on enter and on leave', () => {
+
+    const ast = parse('{ a, b, c { a, b, c } }', { noLocation: true });
+
+    const { definitions } = ast;
+
+    const editedAst = visit(ast, {
+      Document: {
+        enter(node) {
+          return {
+            ...node,
+            definitions: [],
+            didEnter: true,
+          };
+        },
+        leave(node) {
+          return {
+            ...node,
+            definitions,
+            didLeave: true,
+          };
+        }
+      }
+    });
+
+    expect(editedAst).to.deep.equal({
+      ...ast,
+      didEnter: true,
+      didLeave: true
+    });
+  });
+
   it('allows for editing on enter', () => {
 
     const ast = parse('{ a, b, c { a, b, c } }', { noLocation: true });

src/language/visitor.js

@@ -254,7 +254,7 @@ export function visit(root, visitor, keyMap) {
   } while (stack !== undefined);
 
   if (edits.length !== 0) {
-    newRoot = edits[0][1];
+    newRoot = edits[edits.length - 1][1];
   }
 
   return newRoot;