[Validation] Perf improvements for fragment cycle detection

Related GraphQL-js commit graphql/graphql-js@4cf2190

Author: syrusakbary

graphql/core/validation/rules/no_fragment_cycles.py

@@ -5,71 +5,70 @@
 
 
 class NoFragmentCycles(ValidationRule):
-    __slots__ = 'spreads_in_fragment', 'known_to_lead_to_cycle'
+    __slots__ = 'errors', 'visited_frags', 'spread_path', 'spread_path_index_by_name'
 
     def __init__(self, context):
         super(NoFragmentCycles, self).__init__(context)
-        self.spreads_in_fragment = {
-            node.name.value: self.gather_spreads(node)
-            for node in context.get_ast().definitions
-            if isinstance(node, ast.FragmentDefinition)
-            }
-        self.known_to_lead_to_cycle = set()
+        self.errors = []
+        self.visited_frags = set()
+        self.spread_path = []
+        self.spread_path_index_by_name = {}
 
-    def enter_FragmentDefinition(self, node, key, parent, path, ancestors):
-        errors = []
-        initial_name = node.name.value
-        spread_path = []
-
-        def detect_cycle_recursive(fragment_name):
-            spread_nodes = self.spreads_in_fragment.get(fragment_name)
-            if not spread_nodes:
-                return
-
-            for spread_node in spread_nodes:
-                if spread_node in self.known_to_lead_to_cycle:
-                    continue
-
-                if spread_node.name.value == initial_name:
-                    cycle_path = spread_path + [spread_node]
-                    self.known_to_lead_to_cycle |= set(cycle_path)
-
-                    errors.append(GraphQLError(
-                        self.cycle_error_message(initial_name, [s.name.value for s in spread_path]),
-                        cycle_path
-                    ))
-                    continue
-
-                if any(spread is spread_node for spread in spread_path):
-                    continue
+    def leave_Document(self, node, key, parent, path, ancestors):
+        if self.errors:
+            return self.errors
 
-                spread_path.append(spread_node)
-                detect_cycle_recursive(spread_node.name.value)
-                spread_path.pop()
+    def enter_OperationDefinition(self, node, key, parent, path, ancestors):
+        return False
 
-        detect_cycle_recursive(initial_name)
-        if errors:
-            return errors
+    def enter_FragmentDefinition(self, node, key, parent, path, ancestors):
+        if node.name.value not in self.visited_frags:
+            self.detect_cycle_recursive(node)
+        return False
+
+    def detect_cycle_recursive(self, fragment):
+        fragment_name = fragment.name.value
+        self.visited_frags.add(fragment_name)
+
+        spread_nodes = []
+        self.gather_spreads(spread_nodes, fragment.selection_set)
+        if not spread_nodes:
+            return
+
+        self.spread_path_index_by_name[fragment_name] = len(self.spread_path)
+
+        for spread_node in spread_nodes:
+            spread_name = spread_node.name.value
+            cycle_index = self.spread_path_index_by_name.get(spread_name)
+
+            if cycle_index is None:
+                self.spread_path.append(spread_node)
+                if spread_name not in self.visited_frags:
+                    spread_fragment = self.context.get_fragment(spread_name)
+                    if spread_fragment:
+                        self.detect_cycle_recursive(spread_fragment)
+                self.spread_path.pop()
+            else:
+                cycle_path = self.spread_path[cycle_index:]
+                self.errors.append(GraphQLError(
+                    self.cycle_error_message(
+                        spread_name,
+                        [s.name.value for s in cycle_path]
+                    ),
+                    cycle_path+[spread_node]
+                ))
+
+        self.spread_path_index_by_name[fragment_name] = None
 
     @staticmethod
     def cycle_error_message(fragment_name, spread_names):
         via = ' via {}'.format(', '.join(spread_names)) if spread_names else ''
         return 'Cannot spread fragment "{}" within itself{}.'.format(fragment_name, via)
 
     @classmethod
-    def gather_spreads(cls, node):
-        visitor = cls.CollectFragmentSpreadNodesVisitor()
-        visit(node, visitor)
-        return visitor.collect_fragment_spread_nodes()
-
-    class CollectFragmentSpreadNodesVisitor(Visitor):
-        __slots__ = 'spread_nodes',
-
-        def __init__(self):
-            self.spread_nodes = []
-
-        def enter_FragmentSpread(self, node, key, parent, path, ancestors):
-            self.spread_nodes.append(node)
-
-        def collect_fragment_spread_nodes(self):
-            return self.spread_nodes
+    def gather_spreads(cls, spreads, node):
+        for selection in node.selections:
+            if isinstance(selection, ast.FragmentSpread):
+                spreads.append(selection)
+            elif selection.selection_set:
+                cls.gather_spreads(spreads, selection.selection_set)

tests/core_validation/test_no_fragment_cycles.py

@@ -124,14 +124,15 @@ def test_no_spreading_itself_deeply():
     fragment fragX on Dog { ...fragY }
     fragment fragY on Dog { ...fragZ }
     fragment fragZ on Dog { ...fragO }
-    fragment fragO on Dog { ...fragA, ...fragX }
+    fragment fragO on Dog { ...fragP }
+    fragment fragP on Dog { ...fragA, ...fragX }
     ''', [
-        cycle_error_message('fragA', ['fragB', 'fragC', 'fragO'], L(2, 29), L(3, 29), L(4, 29), L(8, 29)),
-        cycle_error_message('fragX', ['fragY', 'fragZ', 'fragO'], L(5, 29), L(6, 29), L(7, 29), L(8, 39))
+        cycle_error_message('fragA', ['fragB', 'fragC', 'fragO', 'fragP'], L(2, 29), L(3, 29), L(4, 29), L(8, 29), L(9, 29)),
+        cycle_error_message('fragO', ['fragP', 'fragX', 'fragY', 'fragZ'], L(8, 29), L(9, 39), L(5, 29), L(6, 29), L(7, 29))
     ])
 
 
-def test_no_spreading_itself_deeply_two_paths():  # -- new rule
+def test_no_spreading_itself_deeply_two_paths():
     expect_fails_rule(NoFragmentCycles, '''
     fragment fragA on Dog { ...fragB, ...fragC }
     fragment fragB on Dog { ...fragA }
@@ -140,3 +141,26 @@ def test_no_spreading_itself_deeply_two_paths():  # -- new rule
         cycle_error_message('fragA', ['fragB'], L(2, 29), L(3, 29)),
         cycle_error_message('fragA', ['fragC'], L(2, 39), L(4, 29))
     ])
+
+
+def test_no_spreading_itself_deeply_two_paths_alt_reverse_order():
+    expect_fails_rule(NoFragmentCycles, '''
+    fragment fragA on Dog { ...fragC }
+    fragment fragB on Dog { ...fragC }
+    fragment fragC on Dog { ...fragA, ...fragB }
+    ''', [
+        cycle_error_message('fragA', ['fragC'], L(2, 29), L(4, 29)),
+        cycle_error_message('fragC', ['fragB'], L(4, 39), L(3, 29))
+    ])
+
+
+def test_no_spreading_itself_deeply_and_immediately():
+    expect_fails_rule(NoFragmentCycles, '''
+    fragment fragA on Dog { ...fragB }
+    fragment fragB on Dog { ...fragB, ...fragC }
+    fragment fragC on Dog { ...fragA, ...fragB }
+    ''', [
+        cycle_error_message('fragB', [], L(3, 29)),
+        cycle_error_message('fragA', ['fragB', 'fragC'], L(2, 29), L(3, 39), L(4, 29)),
+        cycle_error_message('fragB', ['fragC'], L(3, 39), L(4, 39))
+    ])