From 9fc088dcb0e26fd4ad6a29ef362557fe238023e3 Mon Sep 17 00:00:00 2001
From: Rangi <35663410+Rangi42@users.noreply.github.com>
Date: Tue, 7 Nov 2023 17:43:46 -0500
Subject: [PATCH] Fix the FOR loop count formula (#1222)

---
 src/asm/fstack.cpp          | 10 +++++++---
 test/asm/for-loop-count.asm | 19 +++++++++++++++++++
 test/asm/for-loop-count.err |  2 ++
 test/asm/for-loop-count.out | 18 ++++++++++++++++++
 4 files changed, 46 insertions(+), 3 deletions(-)
 create mode 100644 test/asm/for-loop-count.asm
 create mode 100644 test/asm/for-loop-count.err
 create mode 100644 test/asm/for-loop-count.out

diff --git a/src/asm/fstack.cpp b/src/asm/fstack.cpp
index 5ddb8e17..107f1e9e 100644
--- a/src/asm/fstack.cpp
+++ b/src/asm/fstack.cpp
@@ -240,7 +240,11 @@ bool yywrap(void)
 
 		// If this is a FOR, update the symbol value
 		if (contextStack->forName && fileInfo->iters[0] <= contextStack->nbReptIters) {
-			contextStack->forValue += contextStack->forStep;
+			// Avoid arithmetic overflow runtime error
+			uint32_t forValue = (uint32_t)contextStack->forValue +
+				(uint32_t)contextStack->forStep;
+			contextStack->forValue = forValue >= 0 ? (int32_t)forValue
+				: -(int32_t)~forValue - 1;
 			struct Symbol *sym = sym_AddVar(contextStack->forName,
 				contextStack->forValue);
 
@@ -513,9 +517,9 @@ void fstk_RunFor(char const *symName, int32_t start, int32_t stop, int32_t step,
 	uint32_t count = 0;
 
 	if (step > 0 && start < stop)
-		count = (stop - start - 1) / step + 1;
+		count = ((int64_t)stop - start - 1) / step + 1;
 	else if (step < 0 && stop < start)
-		count = (start - stop - 1) / -step + 1;
+		count = ((int64_t)start - stop - 1) / -(int64_t)step + 1;
 	else if (step == 0)
 		error("FOR cannot have a step value of 0\n");
 
diff --git a/test/asm/for-loop-count.asm b/test/asm/for-loop-count.asm
new file mode 100644
index 00000000..bcaa49fc
--- /dev/null
+++ b/test/asm/for-loop-count.asm
@@ -0,0 +1,19 @@
+; `FOR value, start, stop, step` calculates an initial count of loop repetitions as follows:
+; - if (step > 0 && start < stop) then count = (stop - start - 1) / step + 1
+; - if (step < 0 && stop < start) then count = (start - stop - 1) / -step + 1
+; - else count = 0
+; The absolute vaue |stop - start| needs to be 64-bit for this formula to work.
+; Results may be compared with Python range(start, stop, step).
+
+MACRO test
+  PRINTLN STRFMT("FOR x, %d, %d, %d", \1, \2, \3)
+  FOR x, \1, \2, \3
+    PRINTLN "  {08x:x} = {d:x}"
+  ENDR
+  PRINTLN "  done {08x:x} = {d:x}"
+ENDM
+
+  test $8000_0000, $6000_0000, $4000_0000 ; [-2147483648, -1073741824, 0, 1073741824]
+  test $8000_0000, $7000_0000, $4400_0000 ; [-2147483648, -1006632960, 134217728, 1275068416]
+  test $7fff_ffff, $8000_0000, $8000_0000 ; [2147483647, -1]
+  test $8000_0000, $ffff_ffff, $8000_0000 ; []
diff --git a/test/asm/for-loop-count.err b/test/asm/for-loop-count.err
new file mode 100644
index 00000000..0f8110b8
--- /dev/null
+++ b/test/asm/for-loop-count.err
@@ -0,0 +1,2 @@
+warning: for-loop-count.asm(19) -> for-loop-count.asm::test(12): [-Wbackwards-for]
+    FOR goes backwards from -2147483648 to -1 by -2147483648
diff --git a/test/asm/for-loop-count.out b/test/asm/for-loop-count.out
new file mode 100644
index 00000000..4787b8c9
--- /dev/null
+++ b/test/asm/for-loop-count.out
@@ -0,0 +1,18 @@
+FOR x, -2147483648, 1610612736, 1073741824
+  80000000 = -2147483648
+  c0000000 = -1073741824
+  00000000 = 0
+  40000000 = 1073741824
+  done 80000000 = -2147483648
+FOR x, -2147483648, 1879048192, 1140850688
+  80000000 = -2147483648
+  c4000000 = -1006632960
+  08000000 = 134217728
+  4c000000 = 1275068416
+  done 90000000 = -1879048192
+FOR x, 2147483647, -2147483648, -2147483648
+  7fffffff = 2147483647
+  ffffffff = -1
+  done 7fffffff = 2147483647
+FOR x, -2147483648, -1, -2147483648
+  done 80000000 = -2147483648