From 66124eedb530350ed5557708224b6e008fd19f13 Mon Sep 17 00:00:00 2001
From: chuenchen309 <48723787+chuenchen309@users.noreply.github.com>
Date: Sun, 19 Jul 2026 11:51:12 +0800
Subject: [PATCH] Fix Transformer_NonRecursive corrupting the tree on a
 non-last Discard

When a callback returned Discard, Transformer_NonRecursive did not push a
result for that node, but a parent still sliced `stack[-len(children):]`.
For a discarded node that is not the last child processed, that slice
over-consumed the stack and pulled an unrelated sibling into the wrong
parent -- e.g. `start[keep, wrap[drop]]` came back as
`start[wrap[keep]]`, silently dropping and re-parenting nodes. The base
Transformer (documented to give the same result) was unaffected.

Push discarded results as placeholders so every node consumes exactly
`len(children)` stack slots, and filter the placeholders when building a
parent's args. A wholly-discarded tree now returns None (this also covers
the root-discard crash) instead of mis-unpacking an empty stack.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: chuenchen309 <48723787+chuenchen309@users.noreply.github.com>
---
 lark/visitors.py    | 13 ++++++++-----
 tests/test_trees.py | 17 +++++++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/lark/visitors.py b/lark/visitors.py
index 376732e9f..a3fabfcd0 100644
--- a/lark/visitors.py
+++ b/lark/visitors.py
@@ -310,23 +310,26 @@ def transform(self, tree: Tree[_Leaf_T]) -> _Return_T:
             if isinstance(x, Tree):
                 size = len(x.children)
                 if size:
-                    args = stack[-size:]
+                    args = [a for a in stack[-size:] if a is not Discard]
                     del stack[-size:]
                 else:
                     args = []
 
                 res = self._call_userfunc(x, args)
-                if res is not Discard:
-                    stack.append(res)
+                # Push discarded results too, as placeholders, so that a parent's
+                # ``stack[-len(children):]`` slice stays aligned with its children
+                # even when some of them were discarded (they are filtered above).
+                stack.append(res)
 
             elif self.__visit_tokens__ and isinstance(x, Token):
                 res = self._call_userfunc_token(x)
-                if res is not Discard:
-                    stack.append(res)
+                stack.append(res)
             else:
                 stack.append(x)
 
         result, = stack  # We should have only one tree remaining
+        if result is Discard:  # the whole tree was discarded
+            return None  # type: ignore[return-value]
         # There are no guarantees on the type of the value produced by calling a user func for a
         # child will produce. This means type system can't statically know that the final result is
         # _Return_T. As a result a cast is required.
diff --git a/tests/test_trees.py b/tests/test_trees.py
index 6a3b05ee1..f657af8d6 100644
--- a/tests/test_trees.py
+++ b/tests/test_trees.py
@@ -424,6 +424,23 @@ def IGNORE_TOKEN(self, token):
             result = T().transform(copied)
             self.assertEqual(result, Tree('start', [3, 7]))
 
+    def test_transformer_variants_discard_with_preceding_sibling(self):
+        # A discarded node whose parent is preceded by another subtree must not
+        # corrupt the tree. Transformer_NonRecursive used to mis-slice its stack
+        # here and steal `keep` into `wrap` (dropping it from `start`).
+        tree = Tree('start', [
+            Tree('keep', [Token('T', 'x')]),
+            Tree('wrap', [Tree('drop', [])]),
+            ])
+        expected = Tree('start', [Tree('keep', [Token('T', 'x')]), Tree('wrap', [])])
+        for base in (Transformer, Transformer_InPlace, Transformer_NonRecursive, Transformer_InPlaceRecursive):
+            class T(base):
+                def drop(self, children):
+                    return Discard
+
+            result = T().transform(copy.deepcopy(tree))
+            self.assertEqual(result, expected)
+
     def test_merge_transformers(self):
         tree = Tree('start', [
             Tree('main', [