From 76a9562aba9ad3f66a9ed8046b9afbe375959e15 Mon Sep 17 00:00:00 2001
From: Nat Quayle Nelson <natquaylenelson@gmail.com>
Date: Wed, 12 Apr 2023 10:34:12 -0600
Subject: [PATCH] Special case of (set). Close #187

---
 src/kiss/SpecialForms.hx          | 16 +++++++++++++++-
 src/test/cases/BasicTestCase.hx   |  4 ++++
 src/test/cases/BasicTestCase.kiss | 11 ++++++++++-
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/kiss/SpecialForms.hx b/src/kiss/SpecialForms.hx
index dd578c6..053b3c5 100644
--- a/src/kiss/SpecialForms.hx
+++ b/src/kiss/SpecialForms.hx
@@ -119,7 +119,21 @@ class SpecialForms {
 
         k.doc("set", 2, 2, "(set <variable> <value>)");
         map["set"] = (wholeExp:ReaderExp, args:Array<ReaderExp>, k:KissState) -> {
-            EBinop(OpAssign, k.convert(args[0]), k.convert(args[1])).withMacroPosOf(wholeExp);
+            // Special case: (set ~var value)
+            var printExp = null;
+            var setVar = switch (args[0].def) {
+                case CallExp({def:Symbol("print" | "Prelude.print" | "trace")}, printArgs):
+                    printExp = args[0];
+                    printArgs[0];
+                default:
+                    args[0];
+            }
+            var setExp = EBinop(OpAssign, k.convert(setVar), k.convert(args[1])).withMacroPosOf(wholeExp);
+            if (printExp != null) {
+                EBlock([k.convert(printExp), setExp]).withMacroPosOf(wholeExp);
+            } else {
+                setExp;
+            }
         };
 
         function varName(nameExp:ReaderExp) {
diff --git a/src/test/cases/BasicTestCase.hx b/src/test/cases/BasicTestCase.hx
index d7e742e..683bd68 100644
--- a/src/test/cases/BasicTestCase.hx
+++ b/src/test/cases/BasicTestCase.hx
@@ -439,6 +439,10 @@ class BasicTestCase extends Test {
         _testCaseMacroError();
     }
 
+    function testQuickPrintOnSetVar() {
+        _testQuickPrintOnSetVar();
+    }
+
     var aNullToPrint = null;
 }
 
diff --git a/src/test/cases/BasicTestCase.kiss b/src/test/cases/BasicTestCase.kiss
index bfa4d88..32f9bdc 100644
--- a/src/test/cases/BasicTestCase.kiss
+++ b/src/test/cases/BasicTestCase.kiss
@@ -960,4 +960,13 @@ From:[(assert false (+ \"false \" \"should \" \"have \" \"been \" \"true\"))]" m
                             e)
                         (Assert.fail))
                     e)
-                (Assert.fail)))))
\ No newline at end of file
+                (Assert.fail)))))
+        
+
+(function _testQuickPrintOnSetVar []
+    (let [&mut v 5]
+        (set (trace v) 7)
+        (set (print v) 8)
+        (set (Prelude.print v) 9)
+        (set ~v 6)
+        (Assert.equals v 6)))
\ No newline at end of file