Refactor array bracket access to dedicated AST node

analysis/src/Utils.ml

@@ -95,6 +95,7 @@ let identifyPexp pexp =
   | Pexp_record _ -> "Pexp_record"
   | Pexp_field _ -> "Pexp_field"
   | Pexp_setfield _ -> "Pexp_setfield"
+  | Pexp_index _ -> "Pexp_index"
   | Pexp_array _ -> "Pexp_array"
   | Pexp_ifthenelse _ -> "Pexp_ifthenelse"
   | Pexp_sequence _ -> "Pexp_sequence"

compiler/frontend/bs_ast_mapper.ml

@@ -350,6 +350,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/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/ast_mapper_to0.ml

@@ -439,6 +439,27 @@ 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) -> (
+      (* Map back to Array.get/Array.set for parsetree0 compatibility *)
+      let container = sub.expr sub e1 in
+      let index = sub.expr sub e2 in
+      match e3 with
+      | None ->
+        (* Read: Array.get(container, index) *)
+        let array_get =
+          ident ~loc
+            (mknoloc (Longident.Ldot (Longident.Lident "Array", "get")))
+        in
+        apply ~loc ~attrs array_get [(Nolabel, container); (Nolabel, index)]
+      | Some value ->
+        (* Write: Array.set(container, index, value) *)
+        let array_set =
+          ident ~loc
+            (mknoloc (Longident.Ldot (Longident.Lident "Array", "set")))
+        in
+        let value_expr = sub.expr sub value in
+        apply ~loc ~attrs array_set
+          [(Nolabel, container); (Nolabel, index); (Nolabel, value_expr)])
     | 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/depend.ml

@@ -248,6 +248,10 @@ let rec add_expr bv exp =
     add_expr bv e1;
     add bv fld;
     add_expr bv e2
+  | Pexp_index (e1, e2, e3) ->
+    add_expr bv e1;
+    add_expr bv e2;
+    add_opt add_expr bv e3
   | Pexp_array el -> List.iter (add_expr bv) el
   | Pexp_ifthenelse (e1, e2, opte3) ->
     add_expr bv e1;

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/printtyped.ml

@@ -329,6 +329,16 @@ and expression i ppf x =
     expression i ppf e1;
     longident i ppf li;
     expression i ppf e2
+  | Texp_index (e1, e2, e3) -> (
+    line i ppf "Texp_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)
   | Texp_array l ->
     line i ppf "Texp_array\n";
     list i expression ppf l

compiler/ml/rec_check.ml

@@ -196,8 +196,8 @@ let rec classify_expression : Typedtree.expression -> sd =
     classify_expression e
   | Texp_ident _ | Texp_for _ | Texp_constant _ | Texp_tuple _ | Texp_array _
   | Texp_construct _ | Texp_variant _ | Texp_record _ | Texp_setfield _
-  | Texp_while _ | Texp_pack _ | Texp_function _ | Texp_extension_constructor _
-    ->
+  | Texp_index _ | Texp_while _ | Texp_pack _ | Texp_function _
+  | Texp_extension_constructor _ ->
     Static
   | Texp_apply {funct = {exp_desc = Texp_ident (_, _, vd)}} when is_ref vd ->
     Static
@@ -273,6 +273,11 @@ let rec expression : Env.env -> Typedtree.expression -> Use.t =
         (join (expression env ifso) (option expression env ifnot)))
   | Texp_setfield (e1, _, _, e2) ->
     Use.(join (inspect (expression env e1)) (inspect (expression env e2)))
+  | Texp_index (e1, e2, e3) ->
+    Use.(
+      join
+        (join (inspect (expression env e1)) (inspect (expression env e2)))
+        (inspect (option expression env e3)))
   | Texp_sequence (e1, e2) ->
     Use.(join (discard (expression env e1)) (expression env e2))
   | Texp_while (e1, e2) ->

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)

compiler/ml/typecore.ml

@@ -182,6 +182,10 @@ let iter_expression f e =
       ->
       expr e1;
       expr e2
+    | Pexp_index (e1, e2, eo) ->
+      expr e1;
+      expr e2;
+      may expr eo
     | Pexp_ifthenelse (e1, e2, eo) ->
       expr e1;
       expr e2;
@@ -2834,6 +2838,42 @@ and type_expect_ ?deprecated_context ~context ?in_function ?(recarg = Rejected)
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
       }
+  | Pexp_index (scontainer, sindex, svalue_opt) -> (
+    (* Type check as array access (same as current Array.get/set behavior) *)
+    let container = type_exp ~context:None env scontainer in
+    let index =
+      type_expect ~context:None env sindex (instance_def Predef.type_int)
+    in
+    match svalue_opt with
+    | None ->
+      (* Read access: arr[i] -> array<'a> -> int -> 'a *)
+      let element_type = newgenvar () in
+      let array_type = instance_def (Predef.type_array element_type) in
+      unify_exp ~context:None env container array_type;
+      rue
+        {
+          exp_desc = Texp_index (container, index, None);
+          exp_loc = loc;
+          exp_extra = [];
+          exp_type = instance env element_type;
+          exp_attributes = sexp.pexp_attributes;
+          exp_env = env;
+        }
+    | Some svalue ->
+      (* Write access: arr[i] = v -> array<'a> -> int -> 'a -> unit *)
+      let element_type = newgenvar () in
+      let array_type = instance_def (Predef.type_array element_type) in
+      unify_exp ~context:None env container array_type;
+      let value = type_expect ~context:None env svalue element_type in
+      rue
+        {
+          exp_desc = Texp_index (container, index, Some value);
+          exp_loc = loc;
+          exp_extra = [];
+          exp_type = instance_def Predef.type_unit;
+          exp_attributes = sexp.pexp_attributes;
+          exp_env = env;
+        })
   | Pexp_array sargl ->
     let ty = newgenvar () in
     let to_unify = Predef.type_array ty in

compiler/ml/typedtree.ml

@@ -106,6 +106,7 @@ and expression_desc =
   | Texp_field of expression * Longident.t loc * label_description
   | Texp_setfield of
       expression * Longident.t loc * label_description * expression
+  | Texp_index of expression * expression * expression option
   | Texp_array of expression list
   | Texp_ifthenelse of expression * expression * expression option
   | Texp_sequence of expression * expression

compiler/ml/typedtree.mli

@@ -207,6 +207,7 @@ and expression_desc =
   | Texp_field of expression * Longident.t loc * label_description
   | Texp_setfield of
       expression * Longident.t loc * label_description * expression
+  | Texp_index of expression * expression * expression option
   | Texp_array of expression list
   | Texp_ifthenelse of expression * expression * expression option
   | Texp_sequence of expression * expression

compiler/ml/typedtreeIter.ml

@@ -261,6 +261,12 @@ end = struct
     | Texp_setfield (exp1, _, _label, exp2) ->
       iter_expression exp1;
       iter_expression exp2
+    | Texp_index (exp1, exp2, expo) -> (
+      iter_expression exp1;
+      iter_expression exp2;
+      match expo with
+      | None -> ()
+      | Some exp -> iter_expression exp)
     | Texp_array list -> List.iter iter_expression list
     | Texp_ifthenelse (exp1, exp2, expo) -> (
       iter_expression exp1;

compiler/syntax/src/res_ast_debugger.ml

@@ -641,6 +641,13 @@ module SexpAst = struct
             longident longident_loc.Asttypes.txt;
             expression expr2;
           ]
+      | Pexp_index (e1, e2, e3) ->
+        Sexp.list
+          ([Sexp.atom "Pexp_index"; expression e1; expression e2]
+          @
+          match e3 with
+          | None -> []
+          | Some e -> [expression e])
       | Pexp_array exprs ->
         Sexp.list
           [Sexp.atom "Pexp_array"; Sexp.list (map_empty ~f:expression exprs)]

compiler/syntax/src/res_comments_table.ml

@@ -476,6 +476,7 @@ let rec is_block_expr expr =
   | Pexp_constraint (expr, _) when is_block_expr expr -> true
   | Pexp_field (expr, _) when is_block_expr expr -> true
   | Pexp_setfield (expr, _, _) when is_block_expr expr -> true
+  | Pexp_index (expr, _, _) when is_block_expr expr -> true
   | _ -> false
 
 let is_if_then_else_expr expr =
@@ -1313,6 +1314,12 @@ and walk_expression expr t comments =
       attach t.leading expr2.pexp_loc leading;
       walk_expression expr2 t inside;
       attach t.trailing expr2.pexp_loc trailing
+  | Pexp_index (container, index, value_opt) -> (
+    walk_expression container t comments;
+    walk_expression index t comments;
+    match value_opt with
+    | None -> ()
+    | Some value -> walk_expression value t comments)
   | Pexp_ifthenelse (if_expr, then_expr, else_expr) -> (
     let leading, rest = partition_leading_trailing comments expr.pexp_loc in
     attach t.leading expr.pexp_loc leading;

compiler/syntax/src/res_printer.ml

@@ -3374,6 +3374,31 @@ and print_expression ~state (e : Parsetree.expression) cmt_tbl =
     | Pexp_setfield (expr1, longident_loc, expr2) ->
       print_set_field_expr ~state e.pexp_attributes expr1 longident_loc expr2
         e.pexp_loc cmt_tbl
+    | Pexp_index (container, index, value_opt) -> (
+      let container_doc =
+        let doc = print_expression_with_comments ~state container cmt_tbl in
+        match Parens.field_expr container with
+        | Parens.Parenthesized -> add_parens doc
+        | Braced braces -> print_braces doc container braces
+        | Nothing -> doc
+      in
+      let index_doc = print_expression_with_comments ~state index cmt_tbl in
+      match value_opt with
+      | None ->
+        (* Read: container[index] *)
+        Doc.concat [container_doc; Doc.lbracket; index_doc; Doc.rbracket]
+      | Some value ->
+        (* Write: container[index] = value *)
+        let value_doc = print_expression_with_comments ~state value cmt_tbl in
+        Doc.concat
+          [
+            container_doc;
+            Doc.lbracket;
+            index_doc;
+            Doc.rbracket;
+            Doc.text " = ";
+            value_doc;
+          ])
     | Pexp_ifthenelse (_ifExpr, _thenExpr, _elseExpr)
       when ParsetreeViewer.is_ternary_expr e ->
       let parts, alternate = ParsetreeViewer.collect_ternary_parts e in