New ast to represent index access

Author: shulhi

compiler/ml/ast_helper.ml

@@ -174,6 +174,7 @@ module Exp = struct
   let record ?loc ?attrs a b = mk ?loc ?attrs (Pexp_record (a, b))
   let field ?loc ?attrs a b = mk ?loc ?attrs (Pexp_field (a, b))
   let setfield ?loc ?attrs a b c = mk ?loc ?attrs (Pexp_setfield (a, b, c))
+  let index ?loc ?attrs a b c = mk ?loc ?attrs (Pexp_index (a, b, c))
   let array ?loc ?attrs a = mk ?loc ?attrs (Pexp_array a)
   let ifthenelse ?loc ?attrs a b c = mk ?loc ?attrs (Pexp_ifthenelse (a, b, c))
   let sequence ?loc ?attrs a b = mk ?loc ?attrs (Pexp_sequence (a, b))

compiler/ml/ast_helper.mli

@@ -163,6 +163,13 @@ module Exp : sig
   val field : ?loc:loc -> ?attrs:attrs -> expression -> lid -> expression
   val setfield :
     ?loc:loc -> ?attrs:attrs -> expression -> lid -> expression -> expression
+  val index :
+    ?loc:loc ->
+    ?attrs:attrs ->
+    expression ->
+    expression ->
+    expression option ->
+    expression
   val array : ?loc:loc -> ?attrs:attrs -> expression list -> expression
   val ifthenelse :
     ?loc:loc ->

compiler/ml/ast_iterator.ml

@@ -321,6 +321,10 @@ module E = struct
       sub.expr sub e1;
       iter_loc sub lid;
       sub.expr sub e2
+    | Pexp_index (e1, e2, e3) ->
+      sub.expr sub e1;
+      sub.expr sub e2;
+      iter_opt (sub.expr sub) e3
     | Pexp_array el -> List.iter (sub.expr sub) el
     | Pexp_ifthenelse (e1, e2, e3) ->
       sub.expr sub e1;

compiler/ml/ast_mapper.ml

@@ -313,6 +313,9 @@ module E = struct
       field ~loc ~attrs (sub.expr sub e) (map_loc sub lid)
     | Pexp_setfield (e1, lid, e2) ->
       setfield ~loc ~attrs (sub.expr sub e1) (map_loc sub lid) (sub.expr sub e2)
+    | Pexp_index (e1, e2, e3) ->
+      index ~loc ~attrs (sub.expr sub e1) (sub.expr sub e2)
+        (map_opt (sub.expr sub) e3)
     | Pexp_array el -> array ~loc ~attrs (List.map (sub.expr sub) el)
     | Pexp_ifthenelse (e1, e2, e3) ->
       ifthenelse ~loc ~attrs (sub.expr sub e1) (sub.expr sub e2)

compiler/ml/parsetree.ml

@@ -279,6 +279,9 @@ and expression_desc =
     *)
   | Pexp_field of expression * Longident.t loc (* E.l *)
   | Pexp_setfield of expression * Longident.t loc * expression (* E1.l <- E2 *)
+  | Pexp_index of expression * expression * expression option
+    (* E1[E2]        (None)     - read access
+       E1[E2] = E3   (Some E3)  - write access *)
   | Pexp_array of expression list (* [| E1; ...; En |] *)
   | Pexp_ifthenelse of expression * expression * expression option
     (* if E1 then E2 else E3 *)

compiler/ml/pprintast.ml

@@ -691,6 +691,11 @@ and expression ctxt f x =
     | Pexp_setfield (e1, li, e2) ->
       pp f "@[<2>%a.%a@ <-@ %a@]" (simple_expr ctxt) e1 longident_loc li
         (simple_expr ctxt) e2
+    | Pexp_index (e1, e2, None) ->
+      pp f "%a.(%a)" (expression ctxt) e1 (expression ctxt) e2
+    | Pexp_index (e1, e2, Some e3) ->
+      pp f "%a.(%a)@ <-@ %a" (expression ctxt) e1 (expression ctxt) e2
+        (expression ctxt) e3
     | Pexp_ifthenelse (e1, e2, eo) ->
       (* @;@[<2>else@ %a@]@] *)
       let fmt : (_, _, _) format =

compiler/ml/printast.ml

@@ -294,6 +294,16 @@ and expression i ppf x =
     expression i ppf e1;
     longident_loc i ppf li;
     expression i ppf e2
+  | Pexp_index (e1, e2, e3) -> (
+    line i ppf "Pexp_index\n";
+    expression i ppf e1;
+    line i ppf "index:\n";
+    expression i ppf e2;
+    match e3 with
+    | None -> line i ppf "read access\n"
+    | Some e ->
+      line i ppf "write access:\n";
+      expression i ppf e)
   | Pexp_array l ->
     line i ppf "Pexp_array\n";
     list i expression ppf l

compiler/ml/tast_iterator.ml

@@ -176,6 +176,10 @@ let expr sub {exp_extra; exp_desc; exp_env; _} =
   | Texp_setfield (exp1, _, _, exp2) ->
     sub.expr sub exp1;
     sub.expr sub exp2
+  | Texp_index (exp1, exp2, expo) ->
+    sub.expr sub exp1;
+    sub.expr sub exp2;
+    Option.iter (sub.expr sub) expo
   | Texp_array list -> List.iter (sub.expr sub) list
   | Texp_ifthenelse (exp1, exp2, expo) ->
     sub.expr sub exp1;

compiler/ml/tast_mapper.ml

@@ -232,6 +232,8 @@ let expr sub x =
     | Texp_field (exp, lid, ld) -> Texp_field (sub.expr sub exp, lid, ld)
     | Texp_setfield (exp1, lid, ld, exp2) ->
       Texp_setfield (sub.expr sub exp1, lid, ld, sub.expr sub exp2)
+    | Texp_index (exp1, exp2, expo) ->
+      Texp_index (sub.expr sub exp1, sub.expr sub exp2, opt (sub.expr sub) expo)
     | Texp_array list -> Texp_array (List.map (sub.expr sub) list)
     | Texp_ifthenelse (exp1, exp2, expo) ->
       Texp_ifthenelse

compiler/ml/translcore.ml

@@ -891,6 +891,18 @@ and transl_exp0 (e : Typedtree.expression) : Lambda.lambda =
         Psetfield (lbl.lbl_pos + 1, Lambda.fld_record_extension_set lbl)
     in
     Lprim (access, [transl_exp arg; transl_exp newval], e.exp_loc)
+  | Texp_index (container, index, value_opt) -> (
+    let container_lambda = transl_exp container in
+    let index_lambda = transl_exp index in
+    match value_opt with
+    | None ->
+      (* Read: translate to Parrayrefu primitive (unsafe array get) *)
+      Lprim (Parrayrefu, [container_lambda; index_lambda], e.exp_loc)
+    | Some value ->
+      (* Write: translate to Parraysetu primitive (unsafe array set) *)
+      let value_lambda = transl_exp value in
+      Lprim
+        (Parraysetu, [container_lambda; index_lambda; value_lambda], e.exp_loc))
   | Texp_array expr_list ->
     let ll = transl_list expr_list in
     Lprim (Pmakearray Mutable, ll, e.exp_loc)