feat: transform for-loops on arrays into while-loops (#2831)

## How it works

- we transform AST `ForE` loops by `Construct.whileE` _while_ lowering to IR
- the transformation kicks in when a `DotE` expression is applied to a value with an array type, by projecting out the `vals` accessor
- we get hold of the expression in the `DotE` node and use that (i.e. its constituents) as the starting point to obtain the array's size, and then index into it (for `keys` accessor the indexing is trivial)
- we build a classic `while` loop with an index running from zero below the size
- we don't distinguish between immutable/mutable arrays and also cater for potentially effectful loop bodies, as well as array expressions (and unit expressions passed to the call)
- we always CBV-bind the array expression to a variable (even if it is already a `VarE`, otherwise we end up in name capture hell)

## Pain points

We lose a bunch of source locations in the process of conversion to IR. This is not really a problem now, but will make the debugging experience less enjoyable some day.
 
## TODOs:
- [x] `FileCheck` on final IR
- [x] `async` (talk to @crusso)
- [x] `S.`-only transformation?
- [x] Are parens interfering? — No
- [x] test all combinations
- [ ] check why we have a perf regression

## Further optimisation opportunities
- eliminate unknown call following `call $@(im)mut_array_size`
- use `u32` arithmetic and indexing (currently _bignum_ calls)
- [x] we could optimise `DotE(..., "keys")` similarly — DONE

Changelog.md

@@ -1,5 +1,9 @@
 # Motoko compiler changelog
 
+* `for` loops over arrays are now converted to more efficient
+  index-based iteration (#2831). This can result in significant cycle
+  savings for tight loops, as well as slightly less memory usage.
+
 * Add type union and intersection. The type expression
 
   ```motoko

src/codegen/compile.ml

@@ -6044,7 +6044,7 @@ module Var = struct
 end (* Var *)
 
 (* Calling well-known prelude functions *)
-(* FIX ME: calling into the prelude will not work if we ever need to compile a program
+(* FIXME: calling into the prelude will not work if we ever need to compile a program
    that requires top-level cps conversion;
    use new prims instead *)
 module Internals = struct

src/ir_def/construct.ml

@@ -89,8 +89,10 @@ let primE prim es =
     | ICCallerPrim -> T.caller
     | ICStableRead t -> t
     | ICStableWrite _ -> T.unit
+    | IdxPrim -> T.(as_immut (as_array (List.hd es).note.Note.typ))
     | IcUrlOfBlob -> T.text
     | ActorOfIdBlob t -> t
+    | BinPrim (t, _) -> t
     | CastPrim (t1, t2) -> t2
     | RelPrim _ -> T.bool
     | SerializePrim _ -> T.blob
@@ -115,7 +117,7 @@ let selfRefE typ =
 let assertE e =
   { it = PrimE (AssertPrim, [e]);
     at = no_region;
-    note = Note.{ def with typ = T.unit; eff = eff e}
+    note = Note.{ def with typ = T.unit; eff = eff e }
   }
 
 
@@ -219,6 +221,12 @@ let blockE decs exp =
       note = Note.{ def with typ; eff }
     }
 
+let natE n =
+  { it = LitE (NatLit n);
+    at = no_region;
+    note = Note.{ def with typ = T.nat }
+  }
+
 let textE s =
   { it = LitE (TextLit s);
     at = no_region;
@@ -601,21 +609,21 @@ let loopWhileE exp1 exp2 =
     )
 
 let forE pat exp1 exp2 =
-  (* for p in e1 e2
+  (* for (p in e1) e2
      ~~>
      let nxt = e1.next ;
      label l loop {
        switch nxt () {
          case null { break l };
-         case p    { e2 };
+         case ?p    { e2 };
        }
      } *)
   let lab = fresh_id "done" () in
   let ty1 = exp1.note.Note.typ in
-  let _, tfs = T.as_obj_sub ["next"] ty1 in
-  let tnxt = T.lookup_val_field "next" tfs in
+  let _, tfs = T.as_obj_sub [nextN] ty1 in
+  let tnxt = T.lookup_val_field nextN tfs in
   let nxt = fresh_var "nxt" tnxt in
-  letE nxt (dotE exp1 (nameN "next") tnxt) (
+  letE nxt (dotE exp1 nextN tnxt) (
     labelE lab T.unit (
       loopE (
         switch_optE (callE (varE nxt) [] (tupE []))

src/ir_def/construct.mli

@@ -60,6 +60,7 @@ val projE : exp -> int -> exp
 val optE : exp -> exp
 val tagE : id -> exp -> exp
 val blockE : dec list -> exp -> exp
+val natE : Mo_values.Numerics.Nat.t -> exp
 val textE : string -> exp
 val blobE : string -> exp
 val letE : var -> exp -> exp -> exp

src/ir_def/ir_effect.mli

@@ -3,7 +3,7 @@ open Mo_types.Type
 
 val max_eff : eff -> eff -> eff
 
-(* (incremental)  effect inference on IR *)
+(* (incremental) effect inference on IR *)
 
 val typ : ('a, Note.t) annotated_phrase -> typ
 val eff : ('a, Note.t) annotated_phrase -> eff

src/lowering/desugar.ml

@@ -35,7 +35,7 @@ let apply_sign op l = Syntax.(match op, l with
 let phrase f x = { x with it = f x.it }
 
 let typ_note : S.typ_note -> Note.t =
-  fun {S.note_typ;S.note_eff} -> Note.{def with typ = note_typ; eff = note_eff}
+  fun S.{ note_typ; note_eff } -> Note.{ def with typ = note_typ; eff = note_eff }
 
 let phrase' f x =
   { x with it = f x.at x.note x.it }
@@ -198,6 +198,9 @@ and exp' at note = function
   | S.WhileE (e1, e2) -> (whileE (exp e1) (exp e2)).it
   | S.LoopE (e1, None) -> I.LoopE (exp e1)
   | S.LoopE (e1, Some e2) -> (loopWhileE (exp e1) (exp e2)).it
+  | S.ForE (p, {it=S.CallE ({it=S.DotE (arr, proj); _}, _, e1); _}, e2)
+      when T.is_array arr.note.S.note_typ && (proj.it = "vals" || proj.it = "keys")
+    -> (transform_for_to_while p arr proj e1 e2).it
   | S.ForE (p, e1, e2) -> (forE (pat p) (exp e1) (exp e2)).it
   | S.DebugE e -> if !Mo_config.Flags.release_mode then (unitE ()).it else (exp e).it
   | S.LabelE (l, t, e) -> I.LabelE (l.it, t.Source.note, exp e)
@@ -239,6 +242,39 @@ and lexp' = function
   | S.IdxE (e1, e2) -> I.IdxLE (exp e1, exp e2)
   | _ -> raise (Invalid_argument ("Unexpected expression as lvalue"))
 
+and transform_for_to_while p arr_exp proj e1 e2 =
+  (* for (p in (arr_exp : [_]).proj(e1)) e2 when proj in {"keys", "vals"}
+     ~~>
+     let arr = arr_exp ;
+     let size = arr.size(e1) ;
+     var indx = 0 ;
+     label l loop {
+       if indx < size
+       then { let p = arr[indx]; e2; indx += 1 }
+       else { break l }
+     } *)
+  let arr_typ = arr_exp.note.note_typ in
+  let arrv = fresh_var "arr" arr_typ in
+  let size_exp = array_dotE arr_typ "size" (varE arrv) -*- exp e1 in
+  let indx = fresh_var "indx" T.(Mut nat) in
+  let indexing_exp = match proj.it with
+    | "vals" -> primE I.IdxPrim [varE arrv; varE indx]
+    | "keys" -> varE indx
+    | _ -> assert false in
+  let size = fresh_var "size" T.nat in
+  blockE
+    [ letD arrv (exp arr_exp)
+    ; letD size size_exp
+    ; varD indx (natE Numerics.Nat.zero)]
+    (whileE (primE (I.RelPrim (T.nat, LtOp))
+               [varE indx; varE size])
+       (blockE [ letP (pat p) indexing_exp
+               ; expD (exp e2)]
+          (assignE indx
+             (primE (I.BinPrim (T.nat, AddOp))
+                [ varE indx
+                ; natE (Numerics.Nat.of_int 1)]))))
+
 and mut m = match m.it with
   | S.Const -> Ir.Const
   | S.Var -> Ir.Var
@@ -436,8 +472,8 @@ and array_dotE array_ty proj e =
     let f = var name (fun_ty [ty_param] [poly_array_ty] [fun_ty [] t1 t2]) in
     callE (varE f) [element_ty] e in
   match T.is_mut (T.as_array array_ty), proj with
-    | true,  "size"  -> call "@mut_array_size"    [] [T.nat]
-    | false, "size"  -> call "@immut_array_size"  [] [T.nat]
+    | true,  "size" -> call "@mut_array_size"   [] [T.nat]
+    | false, "size" -> call "@immut_array_size" [] [T.nat]
     | true,  "get"  -> call "@mut_array_get"    [T.nat] [varA]
     | false, "get"  -> call "@immut_array_get"  [T.nat] [varA]
     | true,  "put"  -> call "@mut_array_put"    [T.nat; varA] []

test/run-drun/ok/optimise-for-array-eff.drun-run.ok

@@ -0,0 +1,15 @@
+ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
+ingress Completed: Reply: 0x4449444c0000
+debug.print: effect
+debug.print: hello
+debug.print: world
+debug.print: hello
+debug.print: world
+debug.print: effect
+debug.print: hello
+debug.print: bound
+debug.print: world
+debug.print: hello
+debug.print: bound
+debug.print: world
+ingress Completed: Reply: 0x4449444c0000

test/run-drun/ok/optimise-for-array-eff.ic-ref-run.ok

@@ -0,0 +1,18 @@
+→ update create_canister(record {settings = null})
+← replied: (record {hymijyo = principal "cvccv-qqaaq-aaaaa-aaaaa-c"})
+→ update install_code(record {arg = blob ""; kca_xin = blob "\00asm\01\00\00\00\0…
+← replied: ()
+→ update go()
+debug.print: effect
+debug.print: hello
+debug.print: world
+debug.print: hello
+debug.print: world
+debug.print: effect
+debug.print: hello
+debug.print: bound
+debug.print: world
+debug.print: hello
+debug.print: bound
+debug.print: world
+← replied: ()

test/run-drun/ok/optimise-for-array-eff.run-ir.ok

@@ -0,0 +1,12 @@
+effect
+hello
+world
+hello
+world
+effect
+hello
+bound
+world
+hello
+bound
+world

test/run-drun/ok/optimise-for-array-eff.run-low.ok

@@ -0,0 +1,12 @@
+effect
+hello
+world
+hello
+world
+effect
+hello
+bound
+world
+hello
+bound
+world

test/run-drun/ok/optimise-for-array-eff.run.ok

@@ -0,0 +1,12 @@
+effect
+hello
+world
+hello
+world
+effect
+hello
+bound
+world
+hello
+bound
+world

test/run-drun/optimise-for-array-eff.mo

@@ -0,0 +1,18 @@
+import Prim "mo:⛔";
+
+actor a {
+  public func go() : async () {
+    for (check1 in (await async ["effect", "hello", "world"]).vals()) { Prim.debugPrint check1 };
+
+    for (check2 in ["hello", "world", "effect"].vals()) { await async { Prim.debugPrint check2 } };
+
+    let array = ["hello", "bound", "world"];
+
+    for (check3 in (await async array).vals()) { Prim.debugPrint check3 };
+
+    for (check4 in array.vals()) { await async { Prim.debugPrint check4 } };
+
+    for (_ in array.vals(await async ())) { }
+  }
+};
+a.go(); //OR-CALL ingress go "DIDL\x00\x00"

test/run/ok/iter-no-alloc.wasm-run.ok

@@ -2,7 +2,7 @@ Allocation per iteration (Nat): 0
 Allocation per iteration (Nat16): 0
 Allocation per iteration (Nat32): 0
 Allocation per iteration (?Nat, all null): 0
-Allocation per iteration (FixOpt, all ?null): 8
-Allocation per iteration (FixOpt, all ??null): 8
+Allocation per iteration (FixOpt, all ?null): 0
+Allocation per iteration (FixOpt, all ??null): 0
 Allocation per iteration (?Nat, all values): 0
 Allocation per iteration (record): 0

test/run/ok/optimise-for-array.run-ir.ok

@@ -0,0 +1,12 @@
+hello
+world
+hello
+mutable
+world
+hello
+mutable
+world
+hello
+immutable
+world
+want to see you

test/run/ok/optimise-for-array.run-low.ok

@@ -0,0 +1,12 @@
+hello
+world
+hello
+mutable
+world
+hello
+mutable
+world
+hello
+immutable
+world
+want to see you

test/run/ok/optimise-for-array.run.ok

@@ -0,0 +1,12 @@
+hello
+world
+hello
+mutable
+world
+hello
+mutable
+world
+hello
+immutable
+world
+want to see you

test/run/ok/optimise-for-array.wasm-run.ok

@@ -0,0 +1,12 @@
+hello
+world
+hello
+mutable
+world
+hello
+mutable
+world
+hello
+immutable
+world
+want to see you