fix(compiler): Fix array infix assign semantics (#2080)

compiler/src/formatting/fmt.re

@@ -2008,12 +2008,12 @@ let print_expression = (fmt, ~infix_wrap=d => group(indent(d)), expr) => {
            )
            ++ break,
          )
-    | PExpArraySet(arr, elem, new_value) =>
-      fmt.print_grouped_access_expression(fmt, arr)
-      ++ fmt.print_comment_range(fmt, arr.pexp_loc, elem.pexp_loc)
+    | PExpArraySet({array, index, value, infix_op: None}) =>
+      fmt.print_grouped_access_expression(fmt, array)
+      ++ fmt.print_comment_range(fmt, array.pexp_loc, index.pexp_loc)
       ++ list_brackets(
            indent(
-             break ++ fmt.print_expression(fmt, ~infix_wrap=Fun.id, elem),
+             break ++ fmt.print_expression(fmt, ~infix_wrap=Fun.id, index),
            )
            ++ break,
          )
@@ -2023,10 +2023,38 @@ let print_expression = (fmt, ~infix_wrap=d => group(indent(d)), expr) => {
            ~none=space,
            ~lead=space,
            ~trail=space,
-           elem.pexp_loc,
-           new_value.pexp_loc,
+           index.pexp_loc,
+           value.pexp_loc,
          )
-      ++ fmt.print_expression(fmt, new_value)
+      ++ fmt.print_expression(fmt, value)
+    | PExpArraySet({array, index, value, infix_op: Some(infix)}) =>
+      fmt.print_grouped_access_expression(fmt, array)
+      ++ fmt.print_comment_range(fmt, array.pexp_loc, index.pexp_loc)
+      ++ list_brackets(
+           indent(
+             break ++ fmt.print_expression(fmt, ~infix_wrap=Fun.id, index),
+           )
+           ++ break,
+         )
+      ++ fmt.print_comment_range(
+           fmt,
+           ~none=space,
+           ~lead=space,
+           ~trail=space,
+           index.pexp_loc,
+           infix.pexp_loc,
+         )
+      ++ fmt.print_infix_prefix_op(fmt, infix)
+      ++ string("=")
+      ++ fmt.print_comment_range(
+           fmt,
+           ~none=space,
+           ~lead=space,
+           ~trail=space,
+           infix.pexp_loc,
+           value.pexp_loc,
+         )
+      ++ fmt.print_expression(fmt, value)
     | PExpRecord(base, labels) =>
       braces(
         indent(

compiler/src/middle_end/linearize.re

@@ -790,16 +790,49 @@ let rec transl_imm =
         ),
       ],
     );
-  | TExpArraySet(arr, idx, arg) =>
+  | TExpArraySet({array, index, value, infix_op}) =>
     let tmp = gensym("array_access");
-    let (arr_var, arr_setup) = transl_imm(arr);
-    let (idx_var, idx_setup) = transl_imm(idx);
-    let (arg_var, arg_setup) = transl_imm(arg);
+    let (arr_var, arr_setup) = transl_imm(array);
+    let (idx_var, idx_setup) = transl_imm(index);
+    let (arg_var, arg_setup) = transl_imm(value);
+    let (infix_app_var, infix_app_setup) =
+      switch (infix_op) {
+      | Some(infix_op) =>
+        let infix_arg1 = gensym("infix_arg1");
+        let infix_app = gensym("infix_app");
+        let (infix_func, infix_func_setup) = transl_imm(infix_op);
+        let (infix_alloc_args, infix_alloc_ret) =
+          get_fn_allocation_type(env, infix_op.exp_type);
+        (
+          Imm.id(~loc, ~env, infix_app),
+          infix_func_setup
+          @ [
+            BLet(
+              infix_arg1,
+              Comp.array_get(~loc, ~env, ~allocation_type, idx_var, arr_var),
+              Nonglobal,
+            ),
+            BLet(
+              infix_app,
+              Comp.app(
+                ~loc=infix_op.exp_loc,
+                ~env,
+                ~allocation_type=infix_alloc_ret,
+                (infix_func, (infix_alloc_args, infix_alloc_ret)),
+                [Imm.id(~loc, ~env, infix_arg1), arg_var],
+              ),
+              Nonglobal,
+            ),
+          ],
+        );
+      | None => (arg_var, [])
+      };
     (
       Imm.id(~loc, ~env, tmp),
       arr_setup
       @ idx_setup
       @ arg_setup
+      @ infix_app_setup
       @ [
         BLet(
           tmp,
@@ -809,7 +842,7 @@ let rec transl_imm =
             ~allocation_type,
             idx_var,
             arr_var,
-            arg_var,
+            infix_app_var,
           ),
           Nonglobal,
         ),

compiler/src/parsing/ast_helper.re

@@ -266,8 +266,23 @@ module Expression = {
     mk(~loc, ~core_loc, ~attributes?, PExpArray(a));
   let array_get = (~loc, ~core_loc, ~attributes=?, a, b) =>
     mk(~loc, ~core_loc, ~attributes?, PExpArrayGet(a, b));
-  let array_set = (~loc, ~core_loc, ~attributes=?, a, b, c) =>
-    mk(~loc, ~core_loc, ~attributes?, PExpArraySet(a, b, c));
+  let array_set =
+      (
+        ~loc,
+        ~core_loc,
+        ~attributes=?,
+        ~infix_op=?,
+        ~lhs_loc,
+        array,
+        index,
+        value,
+      ) =>
+    mk(
+      ~loc,
+      ~core_loc,
+      ~attributes?,
+      PExpArraySet({lhs_loc, array, index, value, infix_op}),
+    );
   let let_ = (~loc, ~core_loc, ~attributes=?, a, b, c) =>
     mk(~loc, ~core_loc, ~attributes?, PExpLet(a, b, c));
   let match = (~loc, ~core_loc, ~attributes=?, a, b) =>

compiler/src/parsing/ast_helper.rei

@@ -228,6 +228,8 @@ module Expression: {
       ~loc: loc,
       ~core_loc: loc,
       ~attributes: attributes=?,
+      ~infix_op: expression=?,
+      ~lhs_loc: loc,
       expression,
       expression,
       expression

compiler/src/parsing/ast_mapper.re

@@ -102,14 +102,16 @@ module E = {
         sub.expr(sub, a),
         sub.expr(sub, i),
       )
-    | PExpArraySet(a, i, arg) =>
+    | PExpArraySet({lhs_loc, array, index, value, infix_op}) =>
       array_set(
         ~loc,
         ~core_loc,
         ~attributes,
-        sub.expr(sub, a),
-        sub.expr(sub, i),
-        sub.expr(sub, arg),
+        ~infix_op=?Option.map(sub.expr(sub), infix_op),
+        ~lhs_loc=sub.location(sub, lhs_loc),
+        sub.expr(sub, array),
+        sub.expr(sub, index),
+        sub.expr(sub, value),
       )
     | PExpRecord(b, es) =>
       record(

compiler/src/parsing/parser.mly

@@ -656,8 +656,8 @@ array_get:
   | left_accessor_expr lbrack expr rbrack { Expression.array_get ~loc:(to_loc $loc) ~core_loc:(to_loc $loc) $1 $3 }
 
 array_set:
-  | left_accessor_expr lbrack expr rbrack equal expr { Expression.array_set ~loc:(to_loc $loc) ~core_loc:(to_loc $loc) $1 $3 $6 }
-  | left_accessor_expr lbrack expr rbrack assign_binop_op expr { Expression.array_set ~loc:(to_loc $loc) ~core_loc:(to_loc $loc) $1 $3 (Expression.apply ~loc:(to_loc $loc) ~core_loc:(to_loc $loc) (mkid_expr $loc($5) [$5]) [{paa_label=Unlabeled; paa_expr=Expression.array_get ~loc:(to_loc $loc) ~core_loc:(to_loc $loc) $1 $3; paa_loc=(to_loc $loc($6))}; {paa_label=Unlabeled; paa_expr=$6; paa_loc=(to_loc $loc($6))}]) }
+  | left_accessor_expr lbrack expr rbrack equal expr { Expression.array_set ~loc:(to_loc $loc) ~core_loc:(to_loc $loc) ~lhs_loc:(to_loc (fst $loc($1), snd $loc($4))) $1 $3 $6 }
+  | left_accessor_expr lbrack expr rbrack assign_binop_op expr { Expression.array_set ~loc:(to_loc $loc) ~core_loc:(to_loc $loc) ~infix_op:(mkid_expr $loc($5) [$5]) ~lhs_loc:(to_loc (fst $loc($1), snd $loc($4))) $1 $3 $6 }
 
 record_get:
   | left_accessor_expr dot lid { Expression.record_get ~loc:(to_loc $loc) ~core_loc:(to_loc $loc) $1 $3 }

compiler/src/parsing/parsetree.re

@@ -521,7 +521,13 @@ and expression_desc =
   | PExpList(list(list_item(expression)))
   | PExpArray(list(expression))
   | PExpArrayGet(expression, expression)
-  | PExpArraySet(expression, expression, expression)
+  | PExpArraySet({
+      lhs_loc: Location.t,
+      array: expression,
+      index: expression,
+      value: expression,
+      infix_op: option(expression),
+    })
   | PExpRecord(option(expression), list((loc(Identifier.t), expression)))
   | PExpRecordGet(expression, loc(Identifier.t))
   | PExpRecordSet(expression, loc(Identifier.t), expression)

compiler/src/parsing/parsetree_iter.re

@@ -276,10 +276,11 @@ and iter_expression =
   | PExpArrayGet(a, i) =>
     iter_expression(hooks, a);
     iter_expression(hooks, i);
-  | PExpArraySet(a, i, arg) =>
-    iter_expression(hooks, a);
-    iter_expression(hooks, i);
-    iter_expression(hooks, arg);
+  | PExpArraySet({array, index, value, infix_op}) =>
+    iter_expression(hooks, array);
+    iter_expression(hooks, index);
+    iter_expression(hooks, value);
+    Option.iter(iter_expression(hooks), infix_op);
   | PExpRecord(b, es) =>
     Option.iter(iter_expression(hooks), b);
     iter_record_fields(hooks, es);

compiler/src/parsing/well_formedness.re

@@ -843,11 +843,9 @@ let array_index_non_integer = (errs, super) => {
   let enter_expression = ({pexp_desc: desc, pexp_loc: loc} as e) => {
     switch (desc) {
     | PExpArrayGet(_, {pexp_desc: PExpConstant(PConstNumber(number_type))})
-    | PExpArraySet(
-        _,
-        {pexp_desc: PExpConstant(PConstNumber(number_type))},
-        _,
-      ) =>
+    | PExpArraySet({
+        index: {pexp_desc: PExpConstant(PConstNumber(number_type))},
+      }) =>
       switch (number_type) {
       | PConstNumberFloat({txt}) =>
         let warning = Warnings.ArrayIndexNonInteger(txt);

compiler/src/typed/typecore.re

@@ -941,7 +941,7 @@ and type_expect_ =
       exp_type: instance(env, array_type),
       exp_env: env,
     });
-  | PExpArraySet(sarrexp, sidx, se) =>
+  | PExpArraySet({array: sarrexp, index: sidx, value: se, infix_op: None}) =>
     let array_type = newvar(~name="a", ());
     let arrexp =
       type_expect(
@@ -963,7 +963,88 @@ and type_expect_ =
       );
     let e = type_expect(env, se, mk_expected(array_type));
     rue({
-      exp_desc: TExpArraySet(arrexp, idx, e),
+      exp_desc:
+        TExpArraySet({array: arrexp, index: idx, value: e, infix_op: None}),
+      exp_loc: loc,
+      exp_extra: [],
+      exp_attributes: attributes,
+      exp_type: Builtin_types.type_void,
+      exp_env: env,
+    });
+  | PExpArraySet({
+      lhs_loc,
+      array: sarrexp,
+      index: sidx,
+      value: se,
+      infix_op: Some(infix),
+    }) =>
+    let array_type = newvar(~name="a", ());
+    let arrexp =
+      type_expect(
+        env,
+        sarrexp,
+        mk_expected(
+          ~explanation=Assign_not_array,
+          Builtin_types.type_array(array_type),
+        ),
+      );
+    let idx =
+      type_expect(
+        env,
+        sidx,
+        mk_expected(
+          ~explanation=Assign_not_array_index,
+          Builtin_types.type_number,
+        ),
+      );
+    let infix = type_exp(env, infix);
+    let ty_fun = expand_head(env, infix.exp_type);
+    let (ty_args, ty_ret) =
+      switch (ty_fun.desc) {
+      | TTyVar(_) =>
+        let t_args = [(Unlabeled, newvar()), (Unlabeled, newvar())]
+        and t_ret = newvar();
+        unify(
+          env,
+          ty_fun,
+          newty(TTyArrow(t_args, t_ret, TComLink(ref(TComUnknown)))),
+        );
+        (t_args, t_ret);
+      | TTyArrow(t_args, t_ret, _) => (t_args, t_ret)
+      | _ =>
+        raise(
+          Error(
+            infix.exp_loc,
+            env,
+            Apply_non_function(expand_head(env, infix.exp_type)),
+          ),
+        )
+      };
+    let (ty_arg1, ty_arg2) =
+      switch (ty_args) {
+      | [(_, arg1), (_, arg2)] => (arg1, arg2)
+      | _ =>
+        raise(
+          Error(
+            infix.exp_loc,
+            env,
+            Arity_mismatch(expand_head(env, infix.exp_type), None),
+          ),
+        )
+      };
+
+    unify_exp_types(lhs_loc, env, array_type, ty_arg1);
+    let e = type_expect(env, se, mk_expected(ty_arg2));
+    let assignment_loc = {...infix.exp_loc, loc_end: se.pexp_loc.loc_end};
+    unify_exp_types(assignment_loc, env, ty_ret, array_type);
+    rue({
+      exp_desc:
+        TExpArraySet({
+          array: arrexp,
+          index: idx,
+          value: e,
+          infix_op: Some(infix),
+        }),
       exp_loc: loc,
       exp_extra: [],
       exp_attributes: attributes,

compiler/src/typed/typedtree.re

@@ -467,7 +467,12 @@ and expression_desc =
   | TExpTuple(list(expression))
   | TExpArray(list(expression))
   | TExpArrayGet(expression, expression)
-  | TExpArraySet(expression, expression, expression)
+  | TExpArraySet({
+      array: expression,
+      index: expression,
+      value: expression,
+      infix_op: option(expression),
+    })
   | TExpRecord(
       option(expression),
       array((Types.label_description, record_label_definition)),

compiler/src/typed/typedtree.rei

@@ -434,7 +434,12 @@ and expression_desc =
   | TExpTuple(list(expression))
   | TExpArray(list(expression))
   | TExpArrayGet(expression, expression)
-  | TExpArraySet(expression, expression, expression)
+  | TExpArraySet({
+      array: expression,
+      index: expression,
+      value: expression,
+      infix_op: option(expression),
+    })
   | TExpRecord(
       option(expression),
       array((Types.label_description, record_label_definition)),

compiler/src/typed/typedtreeIter.re

@@ -243,10 +243,11 @@ module MakeIterator =
     | TExpArrayGet(a1, a2) =>
       iter_expression(a1);
       iter_expression(a2);
-    | TExpArraySet(a1, a2, a3) =>
+    | TExpArraySet({array: a1, index: a2, value: a3, infix_op: a4}) =>
       iter_expression(a1);
       iter_expression(a2);
       iter_expression(a3);
+      Option.iter(iter_expression, a4);
     | TExpIf(c, t, f) =>
       iter_expression(c);
       iter_expression(t);

compiler/src/typed/typedtreeMap.re

@@ -232,12 +232,13 @@ module MakeMap =
       | TExpArray(args) => TExpArray(List.map(map_expression, args))
       | TExpArrayGet(a1, a2) =>
         TExpArrayGet(map_expression(a1), map_expression(a2))
-      | TExpArraySet(a1, a2, a3) =>
-        TExpArraySet(
-          map_expression(a1),
-          map_expression(a2),
-          map_expression(a3),
-        )
+      | TExpArraySet({array, index, value, infix_op}) =>
+        TExpArraySet({
+          array: map_expression(array),
+          index: map_expression(index),
+          value: map_expression(value),
+          infix_op: Option.map(map_expression, infix_op),
+        })
       | TExpRecord(b, args) =>
         TExpRecord(Option.map(map_expression, b), map_record_fields(args))
       | TExpRecordGet(record, field, ld) =>