From ab10209b2ef148ef9ce6ba13f421dbf57735e17b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=B4me=20Le=20Breton?= <come.le-breton@inria.fr>
Date: Wed, 20 Dec 2023 14:48:58 +0100
Subject: [PATCH] Pretyping: accept words as array index

A int-of-word cast is silently introduced where needed

(cherry picked from commit f943682a5df03acbd5731ec43ef87d0357057094)
---
 CHANGELOG.md                                  |  5 ++
 compiler/src/pretyping.ml                     | 12 +++--
 .../tests/success/arm-m4/array_access.jazz    | 48 +++++++++++++++++++
 .../tests/success/x86-64/array_access.jazz    | 41 ++++++++++++++++
 4 files changed, 103 insertions(+), 3 deletions(-)
 create mode 100644 compiler/tests/success/arm-m4/array_access.jazz
 create mode 100644 compiler/tests/success/x86-64/array_access.jazz

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 53e24aea6..26b2723ff 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,11 @@
   `SMULTT`, `SMLABB`, `SMLABT`, `SMLATB`, `SMLATT`, `SMULWB`, and `SMULWT`
   ([PR #644](https://github.com/jasmin-lang/jasmin/pull/644)).
 
+- Array indexing expressions are automatically and silently casted from word to
+  int during pretyping
+  ([PR #673](https://github.com/jasmin-lang/jasmin/pull/673);
+  fixes [#672](https://github.com/jasmin-lang/jasmin/issues/672)).
+
 ## Bug fixes
 
 - Fix printing to EasyCrypt of ARMv7 instruction `bic`
diff --git a/compiler/src/pretyping.ml b/compiler/src/pretyping.ml
index dc76f0ae1..bdaaa1d12 100644
--- a/compiler/src/pretyping.ml
+++ b/compiler/src/pretyping.ml
@@ -1070,7 +1070,13 @@ let combine_flags =
 
 let is_combine_flags id =
   List.mem_assoc (L.unloc id) combine_flags
-  
+
+let ensure_int loc i ty =
+  match ty with
+  | P.Bty Int -> i
+  | P.Bty (P.U ws) -> P.Papp1(E.Oint_of_word ws,i)
+  | _ -> rs_tyerror ~loc (TypeMismatch (ty, P.tint))
+
 (* -------------------------------------------------------------------- *)
 let rec tt_expr pd ?(mode=`AllVar) (env : 'asm Env.env) pe =
   match L.unloc pe with
@@ -1097,7 +1103,7 @@ let rec tt_expr pd ?(mode=`AllVar) (env : 'asm Env.env) pe =
     let ws = Option.map_default tt_ws (P.ws_of_ty ty) ws in
     let ty = P.tu ws in
     let i,ity  = tt_expr ~mode pd env pi in
-    check_ty_eq ~loc:(L.loc pi) ~from:ity ~to_:P.tint;
+    let i = ensure_int (L.loc pi) i ity in
     begin match olen with
     | None -> P.Pget (aa, ws, x, i), ty
     | Some plen ->
@@ -1274,7 +1280,7 @@ let tt_lvalue pd (env : 'asm Env.env) { L.pl_desc = pl; L.pl_loc = loc; } =
     let ws = Option.map_default tt_ws (P.ws_of_ty ty) ws in
     let ty = P.tu ws in
     let i,ity  = tt_expr ~mode:`AllVar pd env pi in
-    check_ty_eq ~loc:(L.loc pi) ~from:ity ~to_:P.tint;
+    let i = ensure_int (L.loc pi) i ity in
     begin match olen with
     | None -> 
       loc, (fun _ -> P.Laset (aa, ws, L.mk_loc xlc x, i)), Some ty
diff --git a/compiler/tests/success/arm-m4/array_access.jazz b/compiler/tests/success/arm-m4/array_access.jazz
new file mode 100644
index 000000000..ba9d7eaf6
--- /dev/null
+++ b/compiler/tests/success/arm-m4/array_access.jazz
@@ -0,0 +1,48 @@
+inline fn array_init () -> stack u32[5] {
+  stack u32[5] a;
+  reg u32 t;
+  inline int i;
+  for i = 0 to 5 {
+     t = i; a[i] = t;
+  }
+  return a;
+}
+
+export
+fn test_u16 () -> reg u32 {
+   stack u32[5] a;
+   reg ptr u32[5] pa;
+   a = array_init();
+   pa = a;
+   reg u32 i;
+   i = 0;
+
+   reg u32 t;
+   t = 3;
+   pa[i] = t;
+   reg u32 r;
+   r = (32u) pa[u16 i];
+
+   return r;
+}
+
+fn test_u32 () -> reg u32 {
+   stack u32[5] a;
+   reg ptr u32[5] pa;
+   a = array_init();
+   pa = a;
+
+   reg u32 r;
+   r = 0;
+
+   reg u32 i;
+   i = 0;
+   reg u32 t;
+   while (i < 5) {
+     t = pa[i];
+     r = r + t;
+     i += 1;
+   }
+
+   return r;
+}
diff --git a/compiler/tests/success/x86-64/array_access.jazz b/compiler/tests/success/x86-64/array_access.jazz
new file mode 100644
index 000000000..676a4bcea
--- /dev/null
+++ b/compiler/tests/success/x86-64/array_access.jazz
@@ -0,0 +1,41 @@
+inline fn array_init () -> stack u32[5] {
+  stack u32[5] a;
+  inline int i;
+  for i = 0 to 5 {
+     a[i] = i;
+  }
+  return a;
+}
+
+export
+fn test_u16 () -> reg u16 {
+   stack u32[5] a;
+   a = array_init();
+
+   reg u64 i;
+   i = 0;
+   a[i] = 3;
+   reg u16 r;
+   r = a[u16 i];
+
+   return r;
+}
+
+fn test_u32 () -> reg u32 {
+   stack u32[5] a;
+   a = array_init();
+
+   reg u32 r;
+   r = 0;
+
+   reg u64 i;
+   i = 0;
+   reg u32 t;
+   while (i < 5) {
+     t = a[i];
+     r = r + t;
+     i += 1;
+   }
+
+   return r;
+}