Fix tagged key subtrees not working correctly

1 files changed, 8 insertions, 2 deletions

diff --git a/abdl/_vm.py b/abdl/_vm.py
index f8b8c80..41f28eb 100644
--- a/abdl/_vm.py
+++ b/abdl/_vm.py
@@ -88,18 +88,24 @@ class KeySubtree(PatternElement):
         self.key = toks[0]
         self.skippable = toks[1] == '?'
 
+    def on_in_key(self, frame, path, defs):
+        return self.on_not_in_key(frame, path, defs)
+
     def on_not_in_key(self, frame, path, defs):
         path[-1].subtree = True
-        filtered_iterator = self.filter(path[-1].iterator, defs)
+        filtered_iterator = self.filter(path[-1].iterator, defs, name=path[-1].name)
         del path[-1].iterator
         path[-1].iterator = filtered_iterator
         del filtered_iterator
         path[-1].empty = False
         return True
 
-    def filter(self, iter_, defs):
+    def filter(self, iter_, defs, name):
         for x in iter_:
             for y in match_helper(self.key, defs, x[0]):
+                if name:
+                    # FIXME this "name" thing is a bit suboptimal
+                    y.setdefault(name, x)
                 yield (y, x[1])
 
     def collect_params(self, res: list):