Fix Transformer_NonRecursive corrupting the tree on a non-last Discard

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.

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', [