From b2eab150f205370ee201eaa57c5bdb342ed43f64 Mon Sep 17 00:00:00 2001
From: Nat Quayle Nelson <natquaylenelson@gmail.com>
Date: Mon, 20 Feb 2023 05:52:51 -0700
Subject: [PATCH] Add for and doFor capture vars to locals

---
 kiss/src/kiss/SpecialForms.hx          | 21 ++++++++++++--
 kiss/src/test/cases/BasicTestCase.kiss | 39 +++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/kiss/src/kiss/SpecialForms.hx b/kiss/src/kiss/SpecialForms.hx
index 06ddb53c..7c0977ca 100644
--- a/kiss/src/kiss/SpecialForms.hx
+++ b/kiss/src/kiss/SpecialForms.hx
@@ -265,8 +265,15 @@ class SpecialForms {
             var m = macro $i{uniqueVarName};
 
             var innerLet = false;
+            var varsInScope = [];
             var loopVarExpr:Expr = switch (namesExp.def) {
-                case KeyValueExp(_, _) | Symbol(_): k.convert(namesExp);
+                case KeyValueExp({pos: _, def: Symbol(s1)}, {pos: _, def: Symbol(s2)}): 
+                    varsInScope.push({name:s1});
+                    varsInScope.push({name:s2});
+                    k.convert(namesExp);
+                case Symbol(s): 
+                    varsInScope.push({name:s});
+                    k.convert(namesExp);
                 case ListExp(_) | TypedExp(_, {pos:_, def:Symbol(_)}): 
                     innerLet = true;
                     b.haxeExpr(m);
@@ -274,12 +281,22 @@ class SpecialForms {
                     throw KissError.fromExp(namesExp, 'invalid pattern in `$formName`');
             };
 
+
             var body = if (innerLet) {
                 b.let([namesExp, b.symbol(uniqueVarName)], bodyExps);
             } else {
                 b.begin(bodyExps);
             };
-            return EFor(EBinop(OpIn, loopVarExpr, k.convert(listExp)).withMacroPosOf(wholeExp), k.convert(body)).withMacroPosOf(wholeExp);
+
+            for (v in varsInScope) {
+                k.addVarInScope(v, true, false);
+            }
+            var body = k.convert(body);
+            for (v in varsInScope) {
+                k.removeVarInScope(v, true);
+            }
+
+            return EFor(EBinop(OpIn, loopVarExpr, k.convert(listExp)).withMacroPosOf(wholeExp), body).withMacroPosOf(wholeExp);
         }
 
         k.doc("doFor", 3, null, '(doFor <var> <iterable> <body...>)');
diff --git a/kiss/src/test/cases/BasicTestCase.kiss b/kiss/src/test/cases/BasicTestCase.kiss
index e163d0f6..ca46d710 100644
--- a/kiss/src/test/cases/BasicTestCase.kiss
+++ b/kiss/src/test/cases/BasicTestCase.kiss
@@ -815,7 +815,44 @@ From:[(assert false (+ \"false \" \"should \" \"have \" \"been \" \"true\"))]" m
             (set savedPrints [])
             (printLocalNulls)
             (Assert.isFalse (savedPrints.contains "anotherStaticNullToPrint: null"))
-            (Assert.isTrue (savedPrints.contains "u: null")))))
+            (Assert.isTrue (savedPrints.contains "u: null")))
+        
+        // Test for loop capture variables
+        (set savedPrints [])
+        (doFor a (for _ (range 5) null)
+            (printLocalNulls))
+        (Assert.isTrue (savedPrints.contains "a: null"))
+        (Assert.equals 5 savedPrints.length)
+        
+        (set savedPrints [])
+        (let [:Map<String,String> m (for a (range 5) =>"$a" "a")]
+            (doFor =>k v m
+                (set k null)
+                (printLocalNulls)))
+        (Assert.isTrue (savedPrints.contains "k: null"))
+        (Assert.equals 5 savedPrints.length)
+
+        (set savedPrints [])
+        (let [:Map<Int,String> m (for a (range 5) =>a null)]
+            (doFor =>k v m
+                (printLocalNulls)))
+        (Assert.isTrue (savedPrints.contains "v: null"))
+        (Assert.equals 5 savedPrints.length)
+        
+        (set savedPrints [])
+        (doFor [a b c] (for _ (range 5) [1 null 5])
+            (printLocalNulls))
+        (Assert.isTrue (savedPrints.contains "b: null"))
+        (Assert.equals 5 savedPrints.length)
+
+        (set savedPrints [])
+        (doFor :String s (for _ (range 5) null)
+            (printLocalNulls))
+        (Assert.isTrue (savedPrints.contains "s: null"))
+        (Assert.equals 5 savedPrints.length)
+
+        // TODO test case extraction locals:
+        ))
 
 (function :Void _testTypeCase []
     (typeCase ["a"]