Added mapi iterator

The last argument of the iterated function is the
index of the element in the array.

compiler/heptagon/analysis/typing.ml

@@ -48,6 +48,7 @@ type error =
   | Eempty_record
   | Eempty_array
   | Efoldi_bad_args of ty
+  | Emapi_bad_args of ty
   | Emerge_missing_constrs of QualSet.t
   | Emerge_uniq of qualname
   | Emerge_mix of qualname
@@ -161,6 +162,12 @@ let message loc kind =
                as the last but one argument (found: %a).@."
           print_location loc
           print_type ty
+    | Emapi_bad_args  ty ->
+        eprintf
+          "%aThe function given to mapi should expect an integer \
+               as the last argument (found: %a).@."
+          print_location loc
+          print_type ty
   end;
   raise Errors.Error
 
@@ -757,6 +764,18 @@ and typing_iterator const_env h
         args_ty_list e_list in
       prod result_ty_list, typed_e_list
 
+  | Imapi ->
+      let args_ty_list, idx_ty = split_last args_ty_list in
+      let args_ty_list = List.map (fun ty -> Tarray(ty, n)) args_ty_list in
+      let result_ty_list =
+        List.map (fun ty -> Tarray(ty, n)) result_ty_list in
+      (* Last but one arg of the function should be integer *)
+        ( try unify idx_ty (Tid Initial.pint)
+          with TypingError _ -> raise (TypingError (Emapi_bad_args idx_ty)));
+      let typed_e_list = typing_args const_env h
+        args_ty_list e_list in
+      prod result_ty_list, typed_e_list
+
   | Ifold ->
       let args_ty_list =
         map_butlast (fun ty -> Tarray (ty, n)) args_ty_list in

compiler/heptagon/hept_printer.ml

@@ -24,6 +24,7 @@ open Heptagon
 let iterator_to_string i =
   match i with
     | Imap -> "map"
+    | Imapi -> "mapi"
     | Ifold -> "fold"
     | Ifoldi -> "foldi"
     | Imapfold -> "mapfold"

compiler/heptagon/heptagon.ml

@@ -21,6 +21,7 @@ type state_name = name
 
 type iterator_type =
   | Imap
+  | Imapi
   | Ifold
   | Ifoldi
   | Imapfold

compiler/heptagon/parsing/hept_lexer.mll

@@ -56,6 +56,7 @@ List.iter (fun (str,tok) -> Hashtbl.add keyword_table str tok) [
  "when", WHEN;
  "merge", MERGE;
  "map", MAP;
+ "mapi", MAPI;
  "fold", FOLD;
  "foldi", FOLDI;
  "mapfold", MAPFOLD;

compiler/heptagon/parsing/hept_parser.mly

@@ -46,7 +46,7 @@ open Hept_parsetree
 %token DOUBLE_DOT
 %token AROBASE
 %token DOUBLE_LESS DOUBLE_GREATER
-%token MAP FOLD FOLDI MAPFOLD
+%token MAP MAPI FOLD FOLDI MAPFOLD
 %token <string> PREFIX
 %token <string> INFIX0
 %token <string> INFIX1
@@ -526,6 +526,7 @@ call_params:
 
 iterator:
   | MAP { Imap }
+  | MAPI { Imapi }
   | FOLD { Ifold }
   | FOLDI { Ifoldi }
   | MAPFOLD { Imapfold }

compiler/heptagon/parsing/hept_parsetree.ml

@@ -50,6 +50,7 @@ and static_exp_desc =
 
 type iterator_type =
   | Imap
+  | Imapi
   | Ifold
   | Ifoldi
   | Imapfold

compiler/heptagon/parsing/hept_scoping.ml

@@ -164,6 +164,7 @@ let build_const loc vd_list =
 (** { 3 Translate the AST into Heptagon. } *)
 let translate_iterator_type = function
   | Imap -> Heptagon.Imap
+  | Imapi -> Heptagon.Imapi
   | Ifold -> Heptagon.Ifold
   | Ifoldi -> Heptagon.Ifoldi
   | Imapfold -> Heptagon.Imapfold

compiler/main/hept2mls.ml

@@ -190,6 +190,7 @@ let translate_reset = function
 
 let translate_iterator_type = function
   | Heptagon.Imap -> Imap
+  | Heptagon.Imapi -> Imapi
   | Heptagon.Ifold -> Ifold
   | Heptagon.Ifoldi -> Ifoldi
   | Heptagon.Imapfold -> Imapfold

compiler/main/mls2obc.ml

@@ -384,6 +384,19 @@ and translate_iterator map call_context it name_list
           [Afor (xd, mk_static_int 0, n, bi)], j,
            [Afor (xd, mk_static_int 0, n, b)]
 
+    | Minils.Imapi ->
+        let c_list = array_of_input c_list in
+        let ty_list = List.map unarray (Types.unprod ty) in
+        let name_list = array_of_output name_list ty_list in
+        let node_out_ty = Types.prod ty_list in
+        let v, si, j, action = mk_node_call map call_context
+          app loc name_list (p_list@c_list@[mk_evar_int x]) node_out_ty in
+        let v = translate_var_dec v in
+        let b = mk_block ~locals:v action in
+        let bi = mk_block si in
+          [Afor (xd, mk_static_int 0, n, bi)], j,
+           [Afor (xd, mk_static_int 0, n, b)]
+
     | Minils.Imapfold ->
         let (c_list, acc_in) = split_last c_list in
         let c_list = array_of_input c_list in

compiler/minils/minils.ml

@@ -23,6 +23,7 @@ let minils_format_version = "1"
 
 type iterator_type =
   | Imap
+  | Imapi
   | Ifold
   | Ifoldi
   | Imapfold

compiler/minils/mls_printer.ml

@@ -18,6 +18,7 @@ open Minils
 let iterator_to_string i =
   match i with
     | Imap -> "map"
+    | Imapi -> "mapi"
     | Ifold -> "fold"
     | Ifoldi -> "foldi"
     | Imapfold -> "mapfold"

test/good/array_iterators.ept

@@ -42,4 +42,14 @@ let
   reset
    o = map<<n>> plusone (a);
   every (r & r)
+tel
+
+node m(t:int^n; a:int; i:int) returns (o:int)
+let
+  o = a * i + t[0];
+tel
+
+node itmapi(a:int^n) returns (o:int^n)
+let
+  o = mapi <<n>> m(<a>)(a);
 tel