Add support for visitors to skip sub-trees

digiaonline · Mar 4, 2018 · 50847c4 · 50847c4
1 parent 5b5c3b9
commit 50847c4
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 9 deletions.
diff --git a/src/Language/AST/Visitor/AcceptVisitorTrait.php b/src/Language/AST/Visitor/AcceptVisitorTrait.php
@@ -41,8 +41,8 @@ public function accept(VisitorInterface $visitor, ?string $key = null, array $pa
             return null;
         }
 
-        // If the node has been edited, we have to return early, because otherwise
-        // the edited AST won't be what we'd expect.
+        // If the node was edited, we want to return early
+        // to avoid visiting its sub-tree completely.
         if ($newNode->determineIsEdited($this)) {
             return $newNode;
         }
@@ -67,13 +67,7 @@ public function accept(VisitorInterface $visitor, ?string $key = null, array $pa
             }
         }
 
-        if (null === ($newNode = $visitor->leaveNode($newNode, $key, $this->path))) {
-            return null;
-        }
-
-        $newNode->determineIsEdited($this);
-
-        return $newNode;
+        return $visitor->leaveNode($newNode, $key, $this->path);
     }
 
     /**

diff --git a/tests/Functional/Language/VisitorTest.php b/tests/Functional/Language/VisitorTest.php
@@ -8,6 +8,7 @@
 use Digia\GraphQL\Language\AST\Node\FieldNode;
 use Digia\GraphQL\Language\AST\Node\FieldsTrait;
 use Digia\GraphQL\Language\AST\Node\NameNode;
+use Digia\GraphQL\Language\AST\Node\NameTrait;
 use Digia\GraphQL\Language\AST\Node\NodeInterface;
 use Digia\GraphQL\Language\AST\Node\OperationDefinitionNode;
 use Digia\GraphQL\Language\AST\NodeKindEnum;
@@ -219,6 +220,55 @@ function (NodeInterface $node, ?string $key, array $path = []) use (&$didVisitEd
 
         $this->assertTrue($didVisitEditedNode);
     }
+
+    /**
+     * @throws \Digia\GraphQL\Error\GraphQLError
+     * @throws \Exception
+     */
+    public function testAllowsSkippingSubTree()
+    {
+        $visited = [];
+
+        /** @var DocumentNode $ast */
+        $ast = parse('{ a, b { x }, c }', ['noLocation' => true]);
+
+        $visitor = new Visitor(
+            function (NodeInterface $node, ?string $key, array $path = []) use (&$visited): ?NodeInterface {
+                $visited[] = ['enter', $node->getKind(), $node instanceof NameNode ? $node->getValue() : null];
+
+                if ($node instanceof FieldNode && $node->getNameValue() === 'b') {
+                    return null;
+                }
+
+                return $node;
+            },
+            function (NodeInterface $node, ?string $key, array $path = []) use (&$visited): ?NodeInterface {
+                $visited[] = ['leave', $node->getKind(), $node instanceof NameNode ? $node->getValue() : null];
+
+                return $node;
+            }
+        );
+
+        $ast->accept($visitor);
+
+        $this->assertEquals([
+            ['enter', 'Document', null],
+            ['enter', 'OperationDefinition', null],
+            ['enter', 'SelectionSet', null],
+            ['enter', 'Field', null],
+            ['enter', 'Name', 'a'],
+            ['leave', 'Name', 'a'],
+            ['leave', 'Field', null],
+            ['enter', 'Field', null],
+            ['enter', 'Field', null],
+            ['enter', 'Name', 'c'],
+            ['leave', 'Name', 'c'],
+            ['leave', 'Field', null],
+            ['leave', 'SelectionSet', null],
+            ['leave', 'OperationDefinition', null],
+            ['leave', 'Document', null],
+        ], $visited);
+    }
 }
 
 class Visitor extends AbstractVisitor