Merge pull request #1059 from nhuurre/lvalue-indexing-bug

nhuurre · web-flow · commit 92092ea951ed · 2022-01-10T19:36:05.000+02:00
Add warning for lvalue indexing inconsistency.
diff --git a/src/frontend/Canonicalize.ml b/src/frontend/Canonicalize.ml
@@ -91,8 +91,28 @@ let rec replace_deprecated_expr
           ident expr in
   {expr; emeta}
 
-let replace_deprecated_lval deprecated_userdefined =
-  map_lval_with (replace_deprecated_expr deprecated_userdefined) ident
+let replace_deprecated_lval deprecated_userdefined {lval; lmeta} =
+  let is_multiindex = function
+    | Single {emeta= {type_= Middle.UnsizedType.UInt; _}; _} -> false
+    | _ -> true in
+  let rec flatten_multi = function
+    | LVariable id -> (LVariable id, None)
+    | LIndexed ({lval; lmeta}, idcs) -> (
+        let outer =
+          List.map idcs
+            ~f:(map_index (replace_deprecated_expr deprecated_userdefined))
+        in
+        let unwrap = Option.value_map ~default:[] ~f:fst in
+        match flatten_multi lval with
+        | lval, inner when List.exists ~f:is_multiindex outer ->
+            (lval, Some (unwrap inner @ outer, lmeta))
+        | lval, None -> (LIndexed ({lval; lmeta}, outer), None)
+        | lval, Some (inner, _) -> (lval, Some (inner @ outer, lmeta)) ) in
+  let lval =
+    match flatten_multi lval with
+    | lval, None -> lval
+    | lval, Some (idcs, lmeta) -> LIndexed ({lval; lmeta}, idcs) in
+  {lval; lmeta}
 
 let rec replace_deprecated_stmt
     (deprecated_userdefined : Middle.UnsizedType.t Core_kernel.String.Map.t)
diff --git a/src/frontend/Semantic_error.ml b/src/frontend/Semantic_error.ml
@@ -320,6 +320,7 @@ module StatementError = struct
   type t =
     | CannotAssignToReadOnly of string
     | CannotAssignToGlobal of string
+    | LValueMultiIndexing
     | InvalidSamplingPDForPMF
     | InvalidSamplingCDForCCDF of string
     | InvalidSamplingNoSuchDistribution of string
@@ -352,6 +353,9 @@ module StatementError = struct
         Fmt.pf ppf
           "Cannot assign to global variable '%s' declared in previous blocks."
           name
+    | LValueMultiIndexing ->
+        Fmt.pf ppf
+          "Left hand side of an assignment cannot have nested multi-indexing."
     | TargetPlusEqualsOutsideModelOrLogProb ->
         Fmt.pf ppf
           "Target can only be accessed in the model block or in definitions of \
@@ -604,6 +608,9 @@ let cannot_assign_to_read_only loc name =
 let cannot_assign_to_global loc name =
   StatementError (loc, StatementError.CannotAssignToGlobal name)
 
+let cannot_assign_to_multiindex loc =
+  StatementError (loc, StatementError.LValueMultiIndexing)
+
 let invalid_sampling_pdf_or_pmf loc =
   StatementError (loc, StatementError.InvalidSamplingPDForPMF)
 
diff --git a/src/frontend/Semantic_error.mli b/src/frontend/Semantic_error.mli
@@ -103,6 +103,7 @@ val empty_array : Location_span.t -> t
 val bad_int_literal : Location_span.t -> t
 val cannot_assign_to_read_only : Location_span.t -> string -> t
 val cannot_assign_to_global : Location_span.t -> string -> t
+val cannot_assign_to_multiindex : Location_span.t -> t
 val invalid_sampling_pdf_or_pmf : Location_span.t -> t
 val invalid_sampling_cdf_or_ccdf : Location_span.t -> string -> t
 val invalid_sampling_no_such_dist : Location_span.t -> string -> t
diff --git a/src/frontend/Typechecker.ml b/src/frontend/Typechecker.ml
@@ -207,7 +207,7 @@ let check_postfixop loc op te =
         ~type_ ~loc
   | _ -> Semantic_error.illtyped_postfix_op loc op te.emeta.type_ |> error
 
-let check_variable cf loc tenv id =
+let check_id cf loc tenv id =
   match Env.find tenv (Utils.stdlib_distribution_name id.name) with
   | [] ->
       (* OCaml in these situations suggests similar names
@@ -217,9 +217,8 @@ let check_variable cf loc tenv id =
         (Env.nearest_ident tenv id.name)
       |> error
   | {kind= `StanMath; _} :: _ ->
-      mk_typed_expression ~expr:(Variable id)
-        ~ad_level:(calculate_autodifftype cf MathLibrary UMathLibraryFunction)
-        ~type_:UMathLibraryFunction ~loc
+      ( calculate_autodifftype cf MathLibrary UMathLibraryFunction
+      , UnsizedType.UMathLibraryFunction )
   | {kind= `Variable {origin= Param | TParam | GQuant; _}; _} :: _
     when cf.in_toplevel_decl ->
       Semantic_error.non_data_variable_size_decl loc |> error
@@ -230,9 +229,7 @@ let check_variable cf loc tenv id =
               || cf.current_block = Model ) ->
       Semantic_error.invalid_unnormalized_fn loc |> error
   | {kind= `Variable {origin; _}; type_} :: _ ->
-      mk_typed_expression ~expr:(Variable id)
-        ~ad_level:(calculate_autodifftype cf origin type_)
-        ~type_ ~loc
+      (calculate_autodifftype cf origin type_, type_)
   (* TODO - When it's time for overloading, will this need
      some kind of filter/match on arg types? *)
   | { kind= `UserDefined | `UserDeclared _
@@ -241,13 +238,13 @@ let check_variable cf loc tenv id =
       let type_ =
         UnsizedType.UFun
           (args, rt, Fun_kind.suffix_from_name id.name, mem_pattern) in
-      mk_typed_expression ~expr:(Variable id)
-        ~ad_level:(calculate_autodifftype cf Functions type_)
-        ~type_ ~loc
+      (calculate_autodifftype cf Functions type_, type_)
   | {kind= `UserDefined | `UserDeclared _; type_} :: _ ->
-      mk_typed_expression ~expr:(Variable id)
-        ~ad_level:(calculate_autodifftype cf Functions type_)
-        ~type_ ~loc
+      (calculate_autodifftype cf Functions type_, type_)
+
+let check_variable cf loc tenv id =
+  let ad_level, type_ = check_id cf loc tenv id in
+  mk_typed_expression ~expr:(Variable id) ~ad_level ~type_ ~loc
 
 let get_consistent_types ad_level type_ es =
   let f state e =
@@ -298,6 +295,9 @@ let indexing_type idx =
   | Single {emeta= {type_= UnsizedType.UInt; _}; _} -> `Single
   | _ -> `Multi
 
+let is_multiindex i =
+  match indexing_type i with `Single -> false | `Multi -> true
+
 let inferred_unsizedtype_of_indexed ~loc ut indices =
   let rec aux type_ idcs =
     match (type_, idcs) with
@@ -730,34 +730,75 @@ let verify_assignment_global loc cf block is_global id =
   if (not is_global) || block = cf.current_block then ()
   else Semantic_error.cannot_assign_to_global loc id.name |> error
 
-let mk_assignment_from_indexed_expr assop lhs rhs =
-  Assignment
-    {assign_lhs= Ast.lvalue_of_expr lhs; assign_op= assop; assign_rhs= rhs}
-
-let check_assignment_operator loc assop lhs rhs =
+let verify_assignment_operator loc assop lhs rhs =
   let err op =
-    Semantic_error.illtyped_assignment loc op lhs.emeta.type_ rhs.emeta.type_
+    Semantic_error.illtyped_assignment loc op lhs.lmeta.type_ rhs.emeta.type_
   in
-  let () =
-    match assop with
-    | Assign | ArrowAssign ->
-        if
-          UnsizedType.check_of_same_type_mod_array_conv "" lhs.emeta.type_
-            rhs.emeta.type_
-        then ()
-        else err Operator.Equals |> error
-    | OperatorAssign op -> (
-        let args = List.map ~f:arg_type [lhs; rhs] in
-        let return_type =
-          Stan_math_signatures.assignmentoperator_stan_math_return_type op args
-        in
-        match return_type with Some Void -> () | _ -> err op |> error ) in
-  mk_typed_statement ~return_type:NoReturnType ~loc
-    ~stmt:(mk_assignment_from_indexed_expr assop lhs rhs)
+  match assop with
+  | Assign | ArrowAssign ->
+      if
+        UnsizedType.check_of_same_type_mod_array_conv "" lhs.lmeta.type_
+          rhs.emeta.type_
+      then ()
+      else err Operator.Equals |> error
+  | OperatorAssign op -> (
+      let args = List.map ~f:arg_type [Ast.expr_of_lvalue lhs; rhs] in
+      let return_type =
+        Stan_math_signatures.assignmentoperator_stan_math_return_type op args
+      in
+      match return_type with Some Void -> () | _ -> err op |> error )
+
+let check_lvalue cf tenv = function
+  | {lval= LVariable id; lmeta= ({loc} : located_meta)} ->
+      verify_identifier id ;
+      let ad_level, type_ = check_id cf loc tenv id in
+      {lval= LVariable id; lmeta= {ad_level; type_; loc}}
+  | {lval= LIndexed (lval, idcs); lmeta= {loc}} ->
+      let rec check_inner = function
+        | {lval= LVariable id; lmeta= ({loc} : located_meta)} ->
+            verify_identifier id ;
+            let ad_level, type_ = check_id cf loc tenv id in
+            let var = {lval= LVariable id; lmeta= {ad_level; type_; loc}} in
+            (var, var, [])
+        | {lval= LIndexed (lval, idcs); lmeta= {loc}} ->
+            let lval, var, flat = check_inner lval in
+            let idcs = List.map ~f:(check_index cf tenv) idcs in
+            let ad_level =
+              inferred_ad_type_of_indexed lval.lmeta.ad_level idcs in
+            let type_ =
+              inferred_unsizedtype_of_indexed ~loc lval.lmeta.type_ idcs in
+            ( {lval= LIndexed (lval, idcs); lmeta= {ad_level; type_; loc}}
+            , var
+            , flat @ idcs ) in
+      let lval, var, flat = check_inner lval in
+      let idcs = List.map ~f:(check_index cf tenv) idcs in
+      let ad_level = inferred_ad_type_of_indexed lval.lmeta.ad_level idcs in
+      let type_ = inferred_unsizedtype_of_indexed ~loc lval.lmeta.type_ idcs in
+      if List.exists ~f:is_multiindex flat then (
+        add_warning loc
+          "Nested multi-indexing on the left hand side of assignment does not \
+           behave the same as nested indexing in expressions. This is \
+           considered a bug and will be disallowed in Stan 2.32.0. The \
+           indexing can be automatically fixed using the canonicalize flag for \
+           stanc." ;
+        let lvalue_rvalue_types_differ =
+          try
+            let flat_type =
+              inferred_unsizedtype_of_indexed ~loc var.lmeta.type_ (flat @ idcs)
+            in
+            let rec can_assign = function
+              | UnsizedType.(UArray t1, UArray t2) -> can_assign (t1, t2)
+              | UVector, URowVector | URowVector, UVector -> false
+              | t1, t2 -> UnsizedType.compare t1 t2 <> 0 in
+            can_assign (flat_type, type_)
+          with Errors.SemanticError _ -> true in
+        if lvalue_rvalue_types_differ then
+          Semantic_error.cannot_assign_to_multiindex loc |> error ) ;
+      {lval= LIndexed (lval, idcs); lmeta= {ad_level; type_; loc}}
 
 let check_assignment loc cf tenv assign_lhs assign_op assign_rhs =
   let assign_id = Ast.id_of_lvalue assign_lhs in
-  let lhs = assign_lhs |> expr_of_lvalue |> check_expression cf tenv in
+  let lhs = check_lvalue cf tenv assign_lhs in
   let rhs = check_expression cf tenv assign_rhs in
   let block, global, readonly =
     let var = Env.find tenv assign_id.name in
@@ -772,7 +813,9 @@ let check_assignment loc cf tenv assign_lhs assign_op assign_rhs =
         |> error in
   verify_assignment_global loc cf block global assign_id ;
   verify_assignment_read_only loc readonly assign_id ;
-  check_assignment_operator loc assign_op lhs rhs
+  verify_assignment_operator loc assign_op lhs rhs ;
+  mk_typed_statement ~return_type:NoReturnType ~loc
+    ~stmt:(Assignment {assign_lhs= lhs; assign_op; assign_rhs= rhs})
 
 (* target plus-equals / increment log-prob *)
 
@@ -1167,20 +1210,17 @@ and check_sizedtype cf tenv sizedty =
 
 and check_var_decl_initial_value loc cf tenv id init_val_opt =
   match init_val_opt with
-  | Some e -> (
-      let stmt =
-        Assignment
-          { assign_lhs= {lval= LVariable id; lmeta= {loc}}
-          ; assign_op= Assign
-          ; assign_rhs= e } in
-      mk_untyped_statement ~loc ~stmt
-      |> check_statement cf tenv |> snd
-      |> fun ts ->
-      match (ts.stmt, ts.smeta.return_type) with
-      | Assignment {assign_rhs= ue; _}, NoReturnType -> Some ue
-      | _ ->
-          Common.FatalError.fatal_error_msg
-            [%message " check_var_decl: `Assignment` expected."] )
+  | Some e ->
+      let lhs = check_lvalue cf tenv {lval= LVariable id; lmeta= {loc}} in
+      let rhs = check_expression cf tenv e in
+      if
+        UnsizedType.check_of_same_type_mod_array_conv "" lhs.lmeta.type_
+          rhs.emeta.type_
+      then Some rhs
+      else
+        Semantic_error.illtyped_assignment loc Equals lhs.lmeta.type_
+          rhs.emeta.type_
+        |> error
   | None -> None
 
 and check_transformation cf tenv ut trans =
diff --git a/test/integration/bad/lvalue_indexes1.stan b/test/integration/bad/lvalue_indexes1.stan
@@ -0,0 +1,7 @@
+data {
+    array[5] vector[5] x;
+}
+transformed data {
+    array[5] vector[5] y;
+    y[:][1] = x[:][1];
+}
diff --git a/test/integration/bad/lvalue_indexes2.stan b/test/integration/bad/lvalue_indexes2.stan
@@ -0,0 +1,7 @@
+data {
+    vector[5] x;
+}
+transformed data {
+    vector[5] y;
+    y[:][:] = x[:][:];
+}
diff --git a/test/integration/bad/stanc.expected b/test/integration/bad/stanc.expected
@@ -1358,6 +1358,28 @@ Semantic error in 'lp-error.stan', line 5, column 2 to column 6:
    -------------------------------------------------
 
 Identifier 'lp__' clashes with reserved keyword.
+  $ ../../../../install/default/bin/stanc lvalue_indexes1.stan
+Semantic error in 'lvalue_indexes1.stan', line 6, column 4 to column 11:
+   -------------------------------------------------
+     4:  transformed data {
+     5:      array[5] vector[5] y;
+     6:      y[:][1] = x[:][1];
+             ^
+     7:  }
+   -------------------------------------------------
+
+Left hand side of an assignment cannot have nested multi-indexing.
+  $ ../../../../install/default/bin/stanc lvalue_indexes2.stan
+Semantic error in 'lvalue_indexes2.stan', line 6, column 4 to column 11:
+   -------------------------------------------------
+     4:  transformed data {
+     5:      vector[5] y;
+     6:      y[:][:] = x[:][:];
+             ^
+     7:  }
+   -------------------------------------------------
+
+Left hand side of an assignment cannot have nested multi-indexing.
   $ ../../../../install/default/bin/stanc matrix_expr_bad1.stan
 Semantic error in 'matrix_expr_bad1.stan', line 2, column 2 to column 34:
    -------------------------------------------------
diff --git a/test/integration/cli-args/canonicalize/canonical.expected b/test/integration/cli-args/canonicalize/canonical.expected
@@ -183,6 +183,9 @@ transformed data {
   array[0] int x_i;
   array[0] real x_r;
   matrix[N, N] K = gp_exp_quad_cov(x_quad, 1.0, 1.0);
+  array[5, 5] real idxs;
+  idxs[1][ : ] = idxs[1][ : ];
+  idxs[ : , 1] = idxs[ : ][2];
 }
 parameters {
   real x;
@@ -232,6 +235,11 @@ generated quantities {
                                                      x_r, x_i);
 }
 
+Warning in 'deprecated.stan', line 36, column 2: Nested multi-indexing on the
+    left hand side of assignment does not behave the same as nested indexing
+    in expressions. This is considered a bug and will be disallowed in Stan
+    2.32.0. The indexing can be automatically fixed using the canonicalize
+    flag for stanc.
   $ ../../../../../install/default/bin/stanc --print-canonical funs.stanfunctions
 // comment test comment
 void test(int x, int y) {
diff --git a/test/integration/cli-args/canonicalize/deprecated.stan b/test/integration/cli-args/canonicalize/deprecated.stan
@@ -31,7 +31,10 @@ transformed data {
   array[0] int x_i;
   array[0] real x_r;
   matrix[N, N] K = cov_exp_quad(x_quad, 1.0, 1.0);
-}
+  real idxs[5,5];
+  idxs[1][:] = idxs[1][:];
+  idxs[:][1] = idxs[:][2];
+ }
 parameters {
   real x;
   array[3] real theta;
diff --git a/test/integration/cli-args/canonicalize/deprecations-only.expected b/test/integration/cli-args/canonicalize/deprecations-only.expected
@@ -163,6 +163,9 @@ transformed data {
   array[0] int x_i;
   array[0] real x_r;
   matrix[N, N] K = gp_exp_quad_cov(x_quad, 1.0, 1.0);
+  array[5, 5] real idxs;
+  idxs[1][ : ] = idxs[1][ : ];
+  idxs[ : , 1] = idxs[ : ][2];
 }
 parameters {
   real x;
@@ -212,6 +215,11 @@ generated quantities {
                                                      x_r, x_i);
 }
 
+Warning in 'deprecated.stan', line 36, column 2: Nested multi-indexing on the
+    left hand side of assignment does not behave the same as nested indexing
+    in expressions. This is considered a bug and will be disallowed in Stan
+    2.32.0. The indexing can be automatically fixed using the canonicalize
+    flag for stanc.
   $ ../../../../../install/default/bin/stanc --auto-format --canonicalize deprecations funs.stanfunctions
 // comment test comment
 void test(int x, int y) {
diff --git a/test/integration/cli-args/canonicalize/everything-but.expected b/test/integration/cli-args/canonicalize/everything-but.expected
diff --git a/test/integration/good/warning/lvalue_multiindex.stan b/test/integration/good/warning/lvalue_multiindex.stan
diff --git a/test/integration/good/warning/pretty.expected b/test/integration/good/warning/pretty.expected
