From e88646b586a9f8ea796a75461c5c9d7f47788c93 Mon Sep 17 00:00:00 2001
From: Chanic Panic <chanicpanic@gmail.com>
Date: Mon, 28 Dec 2020 15:54:20 -0800
Subject: [PATCH] Fix an SPPF cycle edge case (Issue #820)

---
 lark/parsers/earley_forest.py |  6 +++---
 tests/test_parser.py          | 13 +++++++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/lark/parsers/earley_forest.py b/lark/parsers/earley_forest.py
index 03c4573..2317d73 100644
--- a/lark/parsers/earley_forest.py
+++ b/lark/parsers/earley_forest.py
@@ -507,7 +507,7 @@ class ForestToParseTree(ForestTransformer):
 
     def _check_cycle(self, node):
         if self._on_cycle_retreat:
-            if id(node) == id(self._cycle_node):
+            if id(node) == id(self._cycle_node) or id(node) in self._successful_visits:
                 self._cycle_node = None
                 self._on_cycle_retreat = False
                 return
@@ -541,16 +541,16 @@ class ForestToParseTree(ForestTransformer):
     def transform_symbol_node(self, node, data):
         if id(node) not in self._successful_visits:
             raise Discard()
-        self._successful_visits.remove(id(node))
         self._check_cycle(node)
+        self._successful_visits.remove(id(node))
         data = self._collapse_ambig(data)
         return self._call_ambig_func(node, data)
 
     def transform_intermediate_node(self, node, data):
         if id(node) not in self._successful_visits:
             raise Discard()
-        self._successful_visits.remove(id(node))
         self._check_cycle(node)
+        self._successful_visits.remove(id(node))
         if len(data) > 1:
             children = [self.tree_class('_inter', c) for c in data]
             return self.tree_class('_iambig', children)
diff --git a/tests/test_parser.py b/tests/test_parser.py
index daa7299..8500b01 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -775,6 +775,19 @@ def _make_full_earley_test(LEXER):
             tree = l.parse('')
             self.assertEqual(tree, Tree('start', []))
 
+        def test_cycle2(self):
+            grammar = """
+            start: _operation
+            _operation:  value
+            value: "b"
+                  | "a" value
+                  | _operation
+            """
+
+            l = Lark(grammar, ambiguity="explicit", lexer=LEXER)
+            tree = l.parse("ab")
+            self.assertEqual(tree, Tree('start', [Tree('value', [Tree('value', [])])]))
+
         def test_cycles(self):
             grammar = """
             a: b