From 2053d2ec5ea1c313cec29920ba1802022f30974a Mon Sep 17 00:00:00 2001
From: Ico Doornekamp <ico@zevv.nl>
Date: Mon, 22 Apr 2024 10:42:54 +0200
Subject: [PATCH] Fixes issue #40

When stubbing and updating jump target addresses in `if`, `else` and
`fi`, variable-length encoding was used; this caused problems when the
stub address and final adress were encoded in different length,
corrupting the dictionary.

This change provides a fix by always using the maximum size
variable-length encoding for the initial stub and final target address.
---
 forth/core.zf       | 21 ++++++++++++++++-----
 src/zforth/zforth.c | 24 +++++++++++++-----------
 2 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/forth/core.zf b/forth/core.zf
index f7e3995..f81bf03 100644
--- a/forth/core.zf
+++ b/forth/core.zf
@@ -10,13 +10,21 @@
 : save    131 sys ;
 
 
-( dictionary access. These are shortcuts through the primitive operations are !!, @@ and ,, )
+( dictionary access for regular variable-length cells. These are shortcuts
+  through the primitive operations are !!, @@ and ,, )
 
 : !    0 !! ;
 : @    0 @@ ;
 : ,    0 ,, ;
 : #    0 ## ;
 
+( dictionary access for jmp instructions; these make sure to always use the
+  maximium cell size for the target address to allow safe stubbing and
+  updating of a jump address. `64` is the magic number for ZF_ACCESS_VAR_MAX,
+  see zforth.c for details )
+
+: !j	 64 !! ;
+: ,j	 64 ,, ;
 
 ( compiler state )
 
@@ -73,12 +81,15 @@
 : times ( XT n -- ) { >r dup >r exe r> r> dup x} drop drop ;
 
 
-( 'if' prepares conditional jump, address will be filled in by 'else' or 'fi' )
+( 'if' prepares conditional jump, the target address '0' will later be
+  overwritten by the 'else' or 'fi' words. Note that ,j and !j are used for
+  writing the jump target address to the dictionary, this makes sure that the
+  target address is always written with the same cell size)
 
-: if      ' jmp0 , here 999 , ; immediate
+: if      ' jmp0 , here 0 ,j ; immediate
 : unless  ' not , postpone if ; immediate
-: else    ' jmp , here 999 , swap here swap ! ; immediate
-: fi      here swap ! ; immediate
+: else    ' jmp , here 0 ,j swap here swap !j ; immediate
+: fi      here swap !j ; immediate
 
 
 ( forth style 'do' and 'loop', including loop iterators 'i' and 'j' )
diff --git a/src/zforth/zforth.c b/src/zforth/zforth.c
index 667a2b9..42e2b20 100644
--- a/src/zforth/zforth.c
+++ b/src/zforth/zforth.c
@@ -22,6 +22,17 @@
 #define CHECK(exp, abort)
 #endif
 
+typedef enum {
+	ZF_MEM_SIZE_VAR = 0,      /* Variable size encoding, 1, 2 or 1+sizeof(zf_cell) bytes */
+	ZF_MEM_SIZE_CELL = 1,     /* sizeof(zf_cell) bytes */
+	ZF_MEM_SIZE_U8 = 2,
+	ZF_MEM_SIZE_U16 = 3,
+	ZF_MEM_SIZE_U32 = 4,
+	ZF_MEM_SIZE_S8 = 5,
+	ZF_MEM_SIZE_S16 = 6,
+	ZF_MEM_SIZE_S32 = 7,
+	ZF_MEM_SIZE_VAR_MAX = 64, /* Variable size encoding, 1+sizeof(zf_cell) bytes */
+} zf_mem_size;
 
 /* Define all primitives, make sure the two tables below always match.  The
  * names are defined as a \0 separated list, terminated by double \0. This
@@ -31,17 +42,6 @@
 
 #define _(s) s "\0"
 
-typedef enum {
-	ZF_MEM_SIZE_VAR,
-	ZF_MEM_SIZE_CELL,
-	ZF_MEM_SIZE_U8,
-	ZF_MEM_SIZE_U16,
-	ZF_MEM_SIZE_U32,
-	ZF_MEM_SIZE_S8,
-	ZF_MEM_SIZE_S16,
-	ZF_MEM_SIZE_S32
-} zf_mem_size;
-
 typedef enum {
 	PRIM_EXIT,    PRIM_LIT,       PRIM_LTZ,  PRIM_COL,     PRIM_SEMICOL,  PRIM_ADD,
 	PRIM_SUB,     PRIM_MUL,       PRIM_DIV,  PRIM_MOD,     PRIM_DROP,     PRIM_DUP,
@@ -282,7 +282,9 @@ static zf_addr dict_put_cell_typed(zf_addr addr, zf_cell v, zf_mem_size size)
 				return dict_put_bytes(addr, t, sizeof(t));
 			}
 		}
+	}
 
+	if(size == ZF_MEM_SIZE_VAR || size == ZF_MEM_SIZE_VAR_MAX) {
 		trace(" ⁵");
 		t[0] = 0xff;
 		return dict_put_bytes(addr+0, t, 1) +