From 981c93bfd432dfd5be852dd556c6ecb1d8210a60 Mon Sep 17 00:00:00 2001
From: Virgil Dupras <hsoft@hardcoded.net>
Date: Sun, 29 Dec 2019 16:15:48 -0500
Subject: [PATCH] lib/expr: fix stack imbalance on failure

---
 apps/lib/expr.asm              |  3 ++-
 tools/tests/unit/common.asm    | 20 ++++++++++++++++++++
 tools/tests/unit/test_expr.asm | 13 +++++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/apps/lib/expr.asm b/apps/lib/expr.asm
index ccfa4f0..20ac914 100644
--- a/apps/lib/expr.asm
+++ b/apps/lib/expr.asm
@@ -45,10 +45,11 @@ _parseApply:
 	; Here we do some stack kung fu. We have, in HL, a string pointer we
 	; want to keep. We have, in (SP), our left result we want to use.
 	ex	(sp), hl	; <-> lvl 1
-	ret	nz
+	jr	nz, .end
 	push	af	; --> lvl 2, save ending operator
 	call	callIX
 	pop	af	; <-- lvl 2, restore operator.
+.end:
 	pop	hl	; <-- lvl 1, restore str pointer
 	ret
 
diff --git a/tools/tests/unit/common.asm b/tools/tests/unit/common.asm
index 2139009..9d3c3a8 100644
--- a/tools/tests/unit/common.asm
+++ b/tools/tests/unit/common.asm
@@ -6,6 +6,12 @@
 ; lib/fmt
 
 testNum:	.db 1
+; Each time we call assertSP, we verify that our stack isn't imbalanced by
+; comparing SP to its saved value. Whenever your "base" SP value change,
+; generally at the beginning of a test routine, run "ld (testSP), sp" to have
+; proper value saved to heap.
+testSP:		.dw 0xffff
+
 
 STDIO_PUTC:
 	out	(0), a
@@ -69,6 +75,20 @@ testList:
 	pop	hl		; <-- lvl 1
 	ret
 
+; test that SP == testSP
+assertSP:
+	ld	hl, (testSP)
+	; offset the fact that we call assertSP
+	dec	hl \ dec hl
+	or	a		; reset carry
+	sbc	hl, sp
+	ret	z
+	ld	hl, .msg
+	call	printstr
+	jr	fail
+.msg:
+	.db	"Wrong SP", CR, LF, 0
+
 nexttest:
 	ld	a, (testNum)
 	inc	a
diff --git a/tools/tests/unit/test_expr.asm b/tools/tests/unit/test_expr.asm
index f2961c0..fb57049 100644
--- a/tools/tests/unit/test_expr.asm
+++ b/tools/tests/unit/test_expr.asm
@@ -53,6 +53,7 @@ test:
 	jp	nz, fail
 
 	call	testParseExpr
+	call	testSPOnFail
 
 	; success
 	xor	a
@@ -130,3 +131,15 @@ testParseExpr:
 .alltests:
 	.dw	.t1, .t2, .t3, .t4, .t5, .t6, .t7, .t8, .t9, .t10, .t11, .t12
 	.dw	.t13, .t14, .t15, 0
+
+; Ensure that stack is balanced on failure
+testSPOnFail:
+	ld	(testSP), sp
+	ld	hl, .sFail
+	call	parseExpr
+	call	assertNZ
+	call	assertSP
+	jp	nexttest
+
+.sFail:	.db "1+abc123", 0
+