experiment: simplify array iter optimization, don't cheat the type system

src/codegen/compile.ml

@@ -55,7 +55,7 @@ module TaggingScheme = struct
      Flags.sanity_check will check tags, but not further locate them.
   *)
 
-  let debug = false
+  let debug = false (* should never be true in master! *)
 
   type bit = I
            | O
@@ -1677,7 +1677,7 @@ module BitTagged = struct
   (* static *)
   let can_tag_const pty (n : int64) = Type.(
     match pty with
-    |  Nat | Int | Int64 | Int32 ->
+    | Nat | Int | Int64 | Int32 ->
       let sbits = sbits_of pty in
       let lower_bound = Int64.(neg (shift_left 1L sbits)) in
       let upper_bound = Int64.shift_left 1L sbits in
@@ -1832,7 +1832,7 @@ module BitTagged = struct
       G.nop
 
   let if_can_tag_i32 env pty retty is1 is2 = Type.(match pty with
-    |  Nat | Int | Int64 | Int32 ->
+    | Nat | Int | Int64 | Int32 ->
       Func.share_code1 Func.Never env
         (prim_fun_name pty "if_can_tag_i32") ("x", I32Type) [I32Type] (fun env get_x ->
           (* checks that all but the low sbits are both either 0 or 1 *)
@@ -1857,7 +1857,7 @@ module BitTagged = struct
 
   let if_can_tag_u32 env pty retty is1 is2 = Type.(
     match pty with
-    |  Nat | Int | Int64 | Int32 ->
+    | Nat | Int | Int64 | Int32 ->
       let sbits = sbits_of pty in
       compile_shrU_const (Int32.of_int sbits) ^^
       E.if_ env retty is2 is1 (* NB: swapped branches *)
@@ -1874,7 +1874,7 @@ module BitTagged = struct
     compile_bitor_const (TaggingScheme.tag_of_typ pty)
 
   let untag_i32 line env pty = Type.(match pty with
-    |  Nat | Int | Int64 | Int32 ->
+    | Nat | Int | Int64 | Int32 ->
       let ubits = ubits_of pty in
       (* check tag *)
       sanity_check_tag line env pty ^^
@@ -3061,11 +3061,11 @@ sig
   val truncate_to_word64 : E.t -> G.t
 
   (* unsigned word to SR.Vanilla *)
-  val from_word30 : E.t -> G.t
   val from_word32 : E.t -> G.t
   val from_word64 : E.t -> G.t
 
   (* signed word to SR.Vanilla *)
+  val from_signed_word_compact : E.t -> G.t
   val from_signed_word32 : E.t -> G.t
   val from_signed_word64 : E.t -> G.t
 
@@ -3394,7 +3394,7 @@ module MakeCompact (Num : BigNumType) : BigNumType = struct
           G.if1 I32Type
             (get_res ^^ compile_bitor_const (TaggingScheme.tag_of_typ Type.Int))
             (get_n ^^ compile_shrS_const (Int32.of_int (32 - BitTagged.ubits_of Type.Int)) ^^
-             Num.from_word30 env ^^ get_amount ^^ Num.compile_lsh env) (* TBR: from_word30? *)
+             Num.from_signed_word_compact env ^^ get_amount ^^ Num.compile_lsh env)
         )
         (get_n ^^ get_amount ^^ Num.compile_lsh env))
 
@@ -3736,11 +3736,20 @@ module MakeCompact (Num : BigNumType) : BigNumType = struct
     let set_a, get_a = new_local64 env "a" in
     set_a ^^
     get_a ^^ BitTagged.if_can_tag_i64 env Type.Int [I32Type]
-      (get_a ^^ BitTagged.tag env Type.Int) (*TBR*)
+      (get_a ^^ BitTagged.tag env Type.Int)
       (get_a ^^ Num.from_signed_word64 env)
 
-  let from_word30 env = (* TBR: don't think we've got 30 bits! *)
-    (*    compile_shl_const (Int32.sub 32l BitTagged.ubitsl) ^^ *) (* TBR *)
+  let from_signed_word_compact env =
+    begin
+      if TaggingScheme.debug || !(Flags.sanity)
+     then
+      let set_a, get_a = new_local env "a" in
+      set_a ^^
+      get_a ^^ BitTagged.if_can_tag_i32 env Type.Int [I32Type]
+        get_a
+        (E.trap_with env "from_signed_word_compact")
+      else G.nop
+    end ^^
     BitTagged.tag_i32 env Type.Int
 
   let from_word32 env =
@@ -3804,7 +3813,7 @@ module BigNumLibtommath : BigNumType = struct
   let truncate_to_word32 env = E.call_import env "rts" "bigint_to_word32_wrap"
   let truncate_to_word64 env = E.call_import env "rts" "bigint_to_word64_wrap"
 
-  let from_word30 env = E.call_import env "rts" "bigint_of_word32"
+  let from_signed_word_compact env = E.call_import env "rts" "bigint_of_int32"
   let from_word32 env = E.call_import env "rts" "bigint_of_word32"
   let from_word64 env = E.call_import env "rts" "bigint_of_word64"
   let from_signed_word32 env = E.call_import env "rts" "bigint_of_int32"
@@ -4536,6 +4545,8 @@ module Arr = struct
      No difference between mutable and immutable arrays.
   *)
 
+  let max_size env = Int32.shift_left 1l 29 (* inclusive *)
+
   let header_size env = Int32.add (Tagged.header_size env) 1l
   let element_size = 4l
   let len_field env = Int32.add (Tagged.header_size env) 0l
@@ -10630,21 +10641,23 @@ and compile_prim_invocation (env : E.t) ae p es at =
     SR.Vanilla,
     compile_array_index env ae e1 e2 ^^
     load_ptr
-  | NextArrayOffset spacing, [e] ->
-    let advance_by =
-      match spacing with
-      | ElementSize
-      | One -> Int32.shift_left 1l (32 - BitTagged.ubits_of Type.Int) (* 1 : Nat *) in
+  (* NB: all these operations assume a valid array offset fits in a compact bignum *)
+  | NextArrayOffset, [e] ->
+    let one_untagged = Int32.shift_left 1l (32 - BitTagged.ubits_of Type.Int) in
     SR.Vanilla,
     compile_exp_vanilla env ae e ^^ (* previous byte offset to array *)
-    compile_add_const advance_by
-  | ValidArrayOffset, [e1; e2] ->
+    compile_add_const one_untagged (* preserving the tag in low bits *)
+  | EqArrayOffset, [e1; e2] ->
     SR.bool,
-    compile_exp_vanilla env ae e1 ^^ BitTagged.untag_i32 __LINE__ env Type.Int ^^
-    compile_exp_vanilla env ae e2 ^^ BitTagged.untag_i32 __LINE__ env Type.Int ^^
-    G.i (Compare (Wasm.Values.I32 I32Op.LtU))
+    compile_exp_vanilla env ae e1 ^^
+    BitTagged.sanity_check_tag __LINE__ env Type.Int ^^
+    compile_exp_vanilla env ae e2 ^^
+    BitTagged.sanity_check_tag __LINE__ env Type.Int ^^
+    (* equate (without untagging) *)
+    G.i (Compare (Wasm.Values.I32 I32Op.Eq))
   | DerefArrayOffset, [e1; e2] ->
     SR.Vanilla,
+    (* NB: no bounds check on index *)
     compile_exp_vanilla env ae e1 ^^ (* skewed pointer to array *)
     Tagged.load_forwarding_pointer env ^^
     compile_exp_vanilla env ae e2 ^^ (* byte offset *)
@@ -10657,15 +10670,13 @@ and compile_prim_invocation (env : E.t) ae p es at =
     G.i (Binary (Wasm.Values.I32 I32Op.Add)) ^^
     (* Not using Tagged.load_field since it is not a proper pointer to the array start *)
     Heap.load_field (Arr.header_size env) (* loads the element at the byte offset *)
-  | GetPastArrayOffset spacing, [e] ->
-    let shift =
-      match spacing with
-      | ElementSize
-      | One -> BigNum.from_word30 env in    (* make it a compact bignum *) (* TBR: from_word30? *)
+  | GetLastArrayOffset, [e] ->
+    assert (BitTagged.can_tag_const Type.Int (Int64.of_int32 (Int32.sub (Arr.max_size env) 1l)));
     SR.Vanilla,
     compile_exp_vanilla env ae e ^^ (* array *)
     Arr.len env ^^
-    shift
+    compile_sub_const 1l ^^
+    BigNum.from_signed_word_compact env
 
   | BreakPrim name, [e] ->
     let d = VarEnv.get_label_depth ae name in

src/ir_def/arrange_ir.ml

@@ -69,10 +69,10 @@ and prim = function
   | ActorDotPrim n    -> "ActorDotPrim" $$ [Atom n]
   | ArrayPrim (m, t)  -> "ArrayPrim"  $$ [mut m; typ t]
   | IdxPrim           -> Atom "IdxPrim"
-  | NextArrayOffset _ -> Atom "NextArrayOffset"
-  | ValidArrayOffset  -> Atom "ValidArrayOffset"
+  | NextArrayOffset   -> Atom "NextArrayOffset"
+  | EqArrayOffset     -> Atom "EqArrayOffset"
   | DerefArrayOffset  -> Atom "DerefArrayOffset"
-  | GetPastArrayOffset _ -> Atom "GetPastArrayOffset"
+  | GetLastArrayOffset -> Atom "GetLastArrayOffset"
   | BreakPrim i       -> "BreakPrim"  $$ [id i]
   | RetPrim           -> Atom "RetPrim"
   | AwaitPrim Type.Fut -> Atom "AwaitPrim"

src/ir_def/check_ir.ml

@@ -482,20 +482,20 @@ let rec check_exp env (exp:Ir.exp) : unit =
       in
       typ exp2 <: T.nat;
       T.as_immut t2 <: t
-    | GetPastArrayOffset _, [exp1] ->
+    | GetLastArrayOffset, [exp1] ->
       let t1 = T.promote (typ exp1) in
       ignore
         (try T.as_array_sub t1 with
          | Invalid_argument _ ->
            error env exp1.at "expected array type, but expression produces type\n  %s"
              (T.string_of_typ_expand t1));
-      T.nat <: t
-    | NextArrayOffset _, [exp1] ->
+      T.int <: t
+    | NextArrayOffset, [exp1] ->
       typ exp1 <: T.nat;
       T.nat <: t
-    | ValidArrayOffset, [exp1; exp2] ->
-      typ exp1 <: T.nat;
-      typ exp2 <: T.nat;
+    | EqArrayOffset, [exp1; exp2] ->
+      typ exp1 <: T.int;
+      typ exp2 <: T.int;
       T.bool <: t
     | BreakPrim id, [exp1] ->
       begin

src/ir_def/construct.ml

@@ -95,9 +95,9 @@ let primE prim es =
     | ICStableSize _ -> T.nat64
     | IdxPrim
     | DerefArrayOffset -> T.(as_immut (as_array_sub (List.hd es).note.Note.typ))
-    | ValidArrayOffset -> T.bool
-    | NextArrayOffset _
-    | GetPastArrayOffset _ -> T.nat
+    | EqArrayOffset -> T.bool
+    | NextArrayOffset -> T.nat
+    | GetLastArrayOffset -> T.int
     | IcUrlOfBlob -> T.text
     | ActorOfIdBlob t -> t
     | BinPrim (t, _) -> t
@@ -265,6 +265,12 @@ let natE n =
     note = Note.{ def with typ = T.nat }
   }
 
+let intE n =
+  { it = LitE (IntLit n);
+    at = no_region;
+    note = Note.{ def with typ = T.int }
+  }
+
 let textE s =
   { it = LitE (TextLit s);
     at = no_region;

src/ir_def/construct.mli

@@ -64,6 +64,7 @@ val thenE : exp -> exp -> exp
 val blockE : dec list -> exp -> exp
 val let_else_switch : pat -> exp -> exp -> exp
 val natE : Mo_values.Numerics.Nat.t -> exp
+val intE : Mo_values.Numerics.Int.t -> exp
 val nat32E : Mo_values.Numerics.Nat32.t -> exp
 val textE : string -> exp
 val blobE : string -> exp

src/ir_def/ir.ml

@@ -146,10 +146,10 @@ and prim =
   | SelfRef of Type.typ               (* returns the self actor ref *)
   | SystemTimePrim
   (* Array field iteration/access *)
-  | NextArrayOffset of spacing        (* advance array offset *)
-  | ValidArrayOffset                  (* verify array offset *)
-  | DerefArrayOffset                  (* array offset indexing *)
-  | GetPastArrayOffset of spacing     (* array offset past the last element *)
+  | NextArrayOffset                   (* advance compact array offset,as Nat *)
+  | EqArrayOffset                     (* equate compact array offset at type Int *)
+  | DerefArrayOffset                  (* compact array offset indexing (unchecked) *)
+  | GetLastArrayOffset                (* compact array offset of the last element, or -1, as Int *)
   (* Funds *)
   | SystemCyclesAddPrim
   | SystemCyclesAcceptPrim
@@ -176,8 +176,6 @@ and prim =
   | ICStableRead of Type.typ           (* deserialize value of stable type from stable memory *)
   | ICStableSize of Type.typ
 
-and spacing = One | ElementSize        (* increment units when iterating over arrays *)
-
 (* Declarations *)
 
 and dec = dec' Source.phrase
@@ -271,10 +269,10 @@ let map_prim t_typ t_id p =
   | ActorDotPrim _ -> p
   | ArrayPrim (m, t) -> ArrayPrim (m, t_typ t)
   | IdxPrim
-  | NextArrayOffset _
-  | ValidArrayOffset
+  | NextArrayOffset
+  | EqArrayOffset
   | DerefArrayOffset
-  | GetPastArrayOffset _ -> p
+  | GetLastArrayOffset -> p
   | BreakPrim id -> BreakPrim (t_id id)
   | RetPrim
   | AwaitPrim _

src/ir_interpreter/interpret_ir.ml

@@ -354,12 +354,12 @@ and interpret_exp_mut env exp (k : V.value V.cont) =
       | (IdxPrim | DerefArrayOffset), [v1; v2] ->
         k (try (V.as_array v1).(Numerics.Int.to_int (V.as_int v2))
            with Invalid_argument s -> trap exp.at "%s" s)
-      | NextArrayOffset _, [v1] ->
+      | NextArrayOffset , [v1] ->
         k (V.Int Numerics.Nat.(of_int ((to_int (V.as_int v1)) + 1)))
-      | ValidArrayOffset, [v1; v2] ->
-        k (V.Bool Numerics.Nat.(to_int (V.as_int v1) < to_int (V.as_int v2)))
-      | GetPastArrayOffset _, [v1] ->
-        k (V.Int Numerics.Nat.(of_int (Array.length (V.as_array v1))))
+      | EqArrayOffset, [v1; v2] ->
+        k (V.Bool Numerics.Int.(to_int (V.as_int v1) = to_int (V.as_int v2)))
+      | GetLastArrayOffset,  [v1] ->
+        k (V.Int Numerics.Int.(of_int (Array.length (V.as_array v1) - 1)))
       | BreakPrim id, [v1] -> find id env.labs v1
       | RetPrim, [v1] -> Option.get env.rets v1
       | ThrowPrim, [v1] -> Option.get env.throws v1

src/lowering/desugar.ml

@@ -269,33 +269,46 @@ 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 last = arr.size(e1) : Int - 1 ;
+     var indx = 0;
+     if (last == -1) { }
+     else {
+       label l loop {
+         let p = arr[indx]; // sans bound check
+         e2;
+         if (indx == last)
+         else { break l }
+         then { indx += 1 }
+       }
+     }
+  *)
   let arr_typ = arr_exp.note.note_typ in
   let arrv = fresh_var "arr" arr_typ in
   let indx = fresh_var "indx" T.(Mut nat) in
-  let spacing, indexing_exp = match proj.it with
-    | "vals" -> I.ElementSize, primE I.DerefArrayOffset [varE arrv; varE indx]
-    | "keys" -> I.One, varE indx
+  let indexing_exp = match proj.it with
+    | "vals" -> primE I.DerefArrayOffset [varE arrv; varE indx]
+    | "keys" -> varE indx
     | _ -> assert false in
-  let size_exp = primE I.(GetPastArrayOffset spacing) [varE arrv] in
-  let size = fresh_var "size" T.nat in
+  let last = fresh_var "last" T.int in
+  let lab = fresh_id "done" () in
   blockE
     [ letD arrv (exp arr_exp)
     ; expD (exp e1)
-    ; letD size size_exp
+    ; letD last (primE I.GetLastArrayOffset [varE arrv]) (* -1 for empty array *)
     ; varD indx (natE Numerics.Nat.zero)]
-    (whileE (primE I.ValidArrayOffset
-               [varE indx; varE size])
-       (blockE [ letP (pat p) indexing_exp
-               ; expD (exp e2)]
-          (assignE indx
-             (primE I.(NextArrayOffset spacing) [varE indx]))))
+    (ifE (primE I.EqArrayOffset [varE last; intE (Numerics.Int.of_int (-1))])
+      (* empty array, do nothing *)
+      (unitE())
+      (labelE lab T.unit (
+        loopE (
+          (blockE
+            [ letP (pat p) indexing_exp
+            ; expD (exp e2)]
+           (ifE (primE I.EqArrayOffset [varE indx; varE last])
+             (* last, exit loop *)
+             (breakE lab (tupE []))
+             (* else increment and continue *)
+             (assignE indx (primE I.NextArrayOffset [varE indx]))))))))
 
 and mut m = match m.it with
   | S.Const -> Ir.Const

test/run/array-iter-max.mo

@@ -4,7 +4,7 @@ let max_size = 2**29; // maximum array size
 let a = Prim.Array_tabulate<Nat>(max_size,func i = i+1);
 var c = 0;
 for (i in a.vals()) {
- assert i == c+1; c += 1;
+  assert i == c+1; c += 1;
 }
 ;
 assert c == max_size;