Start AST-level promotion

Author: WardBrian

src/frontend/Ast.ml

@@ -44,6 +44,7 @@ type ('e, 'f) expression =
   | ImagNumeral of string
   | FunApp of 'f * identifier * 'e list
   | CondDistApp of 'f * identifier * 'e list
+  | Promotion of 'e * UnsizedType.t
   (* GetLP is deprecated *)
   | GetLP
   | GetTarget
@@ -304,7 +305,11 @@ let rec id_of_lvalue {lval; _} =
 
 let rec get_loc_expr (e : untyped_expression) =
   match e.expr with
-  | TernaryIf (e, _, _) | BinOp (e, _, _) | PostfixOp (e, _) | Indexed (e, _) ->
+  | TernaryIf (e, _, _)
+   |BinOp (e, _, _)
+   |PostfixOp (e, _)
+   |Indexed (e, _)
+   |Promotion (e, _) ->
       get_loc_expr e
   | PrefixOp (_, e) | ArrayExpr (e :: _) | RowVectorExpr (e :: _) | Paren e ->
       e.emeta.loc.begin_loc

src/frontend/Ast_to_Mir.ml

@@ -81,6 +81,7 @@ and trans_expr {Ast.expr; Ast.emeta} =
       FunApp (CompilerInternal FnMakeRowVec, trans_exprs eles) |> ewrap
   | Indexed (lhs, indices) ->
       Indexed (trans_expr lhs, List.map ~f:trans_idx indices) |> ewrap
+  | Promotion (e, ty) -> Promotion (trans_expr e, ty) |> ewrap
 
 and trans_idx = function
   | Ast.All -> All

src/frontend/Canonicalize.ml

@@ -142,7 +142,7 @@ let rec no_parens {expr; emeta} =
                   | i -> map_index keep_parens i )
                 l )
       ; emeta }
-  | ArrayExpr _ | RowVectorExpr _ | FunApp _ | CondDistApp _ ->
+  | ArrayExpr _ | RowVectorExpr _ | FunApp _ | CondDistApp _ | Promotion _ ->
       {expr= map_expression no_parens ident expr; emeta}
 
 and keep_parens {expr; emeta} =

src/frontend/Pretty_printing.ml

@@ -272,6 +272,7 @@ and pp_expression ppf ({expr= e_content; emeta= {loc; _}} : untyped_expression)
   | ArrayExpr es -> pf ppf "{@[%a}@]" pp_list_of_expression (es, loc)
   | RowVectorExpr es -> pf ppf "[@[%a]@]" pp_list_of_expression (es, loc)
   | Paren e -> pf ppf "(%a)" pp_expression e
+  | Promotion (e, _) -> pp_expression ppf e
   | Indexed (e, l) -> (
     match l with
     | [] -> pf ppf "%a" pp_expression e

src/frontend/SignatureMismatch.ml

@@ -111,38 +111,40 @@ let rec compare_errors e1 e2 =
       | InputMismatch _, _ | _, SuffixMismatch _ -> 1 ) )
 
 let rec check_same_type depth t1 t2 =
-  let wrap_func = Option.map ~f:(fun e -> TypeMismatch (t1, t2, Some e)) in
+  let wrap_func = Result.map_error ~f:(fun e -> TypeMismatch (t1, t2, Some e)) in
   match (t1, t2) with
-  | t1, t2 when t1 = t2 -> None
-  | UnsizedType.(UReal, UInt) when depth < 1 -> None
-  | UnsizedType.(UComplex, UInt) when depth < 1 -> None
-  | UnsizedType.(UComplex, UReal) when depth < 1 -> None
+  | t1, t2 when t1 = t2 -> Ok ()
+  | UnsizedType.(UReal, UInt) when depth < 1 -> Ok ()
+  | UnsizedType.(UComplex, UInt) when depth < 1 -> Ok ()
+  | UnsizedType.(UComplex, UReal) when depth < 1 -> Ok ()
   | UFun (_, _, s1, _), UFun (_, _, s2, _)
     when Fun_kind.without_propto s1 <> Fun_kind.without_propto s2 ->
-      Some
+      Error
         (SuffixMismatch (Fun_kind.without_propto s1, Fun_kind.without_propto s2))
       |> wrap_func
   | UFun (_, rt1, _, _), UFun (_, rt2, _, _) when rt1 <> rt2 ->
-      Some (ReturnTypeMismatch (rt1, rt2)) |> wrap_func
+      Error (ReturnTypeMismatch (rt1, rt2)) |> wrap_func
   | UFun (l1, _, _, _), UFun (l2, _, _, _) ->
       check_compatible_arguments (depth + 1) l2 l1
-      |> Option.map ~f:(fun e -> InputMismatch e)
+      |> Result.map_error ~f:(fun e -> InputMismatch e)
       |> wrap_func
-  | t1, t2 -> Some (TypeMismatch (t1, t2, None))
+  | t1, t2 -> Error (TypeMismatch (t1, t2, None))
 
 and check_compatible_arguments depth args1 args2 =
   match List.zip args1 args2 with
   | List.Or_unequal_lengths.Unequal_lengths ->
-      Some (ArgNumMismatch (List.length args1, List.length args2))
+      Error (ArgNumMismatch (List.length args1, List.length args2))
   | Ok l ->
       List.find_mapi l ~f:(fun i ((ad1, ut1), (ad2, ut2)) ->
           match check_same_type depth ut1 ut2 with
-          | Some e -> Some (ArgError (i + 1, e))
-          | None ->
+          | Error e -> Some (ArgError (i + 1, e))
+          | Ok _ ->
               if ad1 = ad2 then None
               else if depth < 2 && UnsizedType.autodifftype_can_convert ad1 ad2
               then None
               else Some (ArgError (i + 1, DataOnlyError)) )
+      |> Option.map ~f:Result.fail
+      |> Option.value ~default:(Ok ())
 
 let check_compatible_arguments_mod_conv = check_compatible_arguments 0
 let max_n_errors = 5
@@ -155,9 +157,13 @@ let extract_function_types f =
       Some (return, args, (fun x -> UserDefined x), mem)
   | _ -> None
 
-let returntype env name args =
+let arg_type x = Ast.(x.emeta.ad_level, x.emeta.type_)
+let get_arg_types = List.map ~f:arg_type
+
+let returntype env name arg_exprs =
   (* NB: Variadic arguments are special-cased in the typechecker and not handled here *)
   let name = Utils.stdlib_distribution_name name in
+  let args = get_arg_types arg_exprs in
   Environment.find env name
   |> List.filter_map ~f:extract_function_types
   |> List.sort ~compare:(fun (x, _, _, _) (y, _, _, _) ->
@@ -166,8 +172,9 @@ let returntype env name args =
   |> List.fold_until ~init:[]
        ~f:(fun errors (rt, tys, funkind_constructor, _) ->
          match check_compatible_arguments 0 tys args with
-         | None -> Stop (Ok (rt, funkind_constructor))
-         | Some e -> Continue (((rt, tys), e) :: errors) )
+         (* TODO instead of unit, return Ast.typed_expr list which could contain promotions*)
+         | Ok () -> Stop (Ok (rt, funkind_constructor))
+         | Error e -> Continue (((rt, tys), e) :: errors) )
        ~finish:(fun errors ->
          let errors =
            List.sort errors ~compare:(fun (_, e1) (_, e2) ->
@@ -182,7 +189,7 @@ let check_variadic_args allow_lpdf mandatory_arg_tys mandatory_fun_arg_tys
   in
   let minimal_args =
     (UnsizedType.AutoDiffable, minimal_func_type) :: mandatory_arg_tys in
-  let wrap_err x = Some (minimal_args, ArgError (1, x)) in
+  let wrap_err x = Error (minimal_args, ArgError (1, x)) in
   match args with
   | ( _
     , ( UnsizedType.UFun (fun_args, ReturnType return_type, suffix, _) as
@@ -195,22 +202,22 @@ let check_variadic_args allow_lpdf mandatory_arg_tys mandatory_fun_arg_tys
       let suffix = Fun_kind.without_propto suffix in
       if suffix = FnPlain || (allow_lpdf && suffix = FnLpdf ()) then
         match check_compatible_arguments 1 mandatory mandatory_fun_arg_tys with
-        | Some x -> wrap_func_error (InputMismatch x)
-        | None -> (
+        | Error x -> wrap_func_error (InputMismatch x)
+        | Ok () -> (
           match check_same_type 1 return_type fun_return with
-          | Some _ ->
+          | Error _ ->
               wrap_func_error
                 (ReturnTypeMismatch
                    (ReturnType fun_return, ReturnType return_type) )
-          | None ->
+          | Ok () ->
               let expected_args =
                 ((UnsizedType.AutoDiffable, func_type) :: mandatory_arg_tys)
                 @ variadic_arg_tys in
               check_compatible_arguments 0 expected_args args
-              |> Option.map ~f:(fun x -> (expected_args, x)) )
+              |> Result.map_error ~f:(fun x -> (expected_args, x)) )
       else wrap_func_error (SuffixMismatch (FnPlain, suffix))
   | (_, x) :: _ -> TypeMismatch (minimal_func_type, x, None) |> wrap_err
-  | [] -> Some ([], ArgNumMismatch (List.length mandatory_arg_tys, 0))
+  | [] -> Error ([], ArgNumMismatch (List.length mandatory_arg_tys, 0))
 
 let pp_signature_mismatch ppf (name, arg_tys, (sigs, omitted)) =
   let open Fmt in

src/frontend/SignatureMismatch.mli

@@ -22,12 +22,12 @@ type signature_error =
 val check_compatible_arguments_mod_conv :
      (UnsizedType.autodifftype * UnsizedType.t) list
   -> (UnsizedType.autodifftype * UnsizedType.t) list
-  -> function_mismatch option
+  -> (unit (* Ast.typed_expression list *), function_mismatch) result
 
 val returntype :
      Environment.t
   -> string
-  -> (UnsizedType.autodifftype * UnsizedType.t) list
+  -> Ast.typed_expression list
   -> ( UnsizedType.returntype * (bool Middle.Fun_kind.suffix -> Ast.fun_kind)
      , signature_error list * bool )
      result
@@ -38,8 +38,9 @@ val check_variadic_args :
   -> (UnsizedType.autodifftype * UnsizedType.t) list
   -> UnsizedType.t
   -> (UnsizedType.autodifftype * UnsizedType.t) list
-  -> ((UnsizedType.autodifftype * UnsizedType.t) list * function_mismatch)
-     option
+  -> ( unit
+     , (UnsizedType.autodifftype * UnsizedType.t) list * function_mismatch )
+     result
 
 val pp_signature_mismatch :
      Format.formatter

src/frontend/Typechecker.ml

@@ -424,7 +424,7 @@ let check_fn ~is_cond_dist loc tenv id es =
             (Env.nearest_ident tenv id.name) )
       |> error
   | _ (* a function *) -> (
-    match SignatureMismatch.returntype tenv id.name (get_arg_types es) with
+    match SignatureMismatch.returntype tenv id.name es with
     | Ok (Void, _) ->
         Semantic_error.returning_fn_expected_nonreturning_found loc id.name
         |> error
@@ -458,11 +458,11 @@ let check_reduce_sum ~is_cond_dist loc id es =
         SignatureMismatch.check_variadic_args true mandatory_args
           mandatory_fun_args UReal (get_arg_types es)
       with
-      | None ->
+      | Ok () ->
           mk_typed_expression
             ~expr:(mk_fun_app ~is_cond_dist (StanLib FnPlain, id, es))
             ~ad_level:(expr_ad_lub es) ~type_:UnsizedType.UReal ~loc
-      | Some (expected_args, err) ->
+      | Error (expected_args, err) ->
           Semantic_error.illtyped_reduce_sum loc id.name
             (List.map ~f:type_of_expr_typed es)
             expected_args err
@@ -477,7 +477,7 @@ let check_reduce_sum ~is_cond_dist loc id es =
       let expected_args, err =
         SignatureMismatch.check_variadic_args true mandatory_args
           mandatory_fun_args UReal (get_arg_types es)
-        |> Option.value_exn in
+        |> Result.error |> Option.value_exn in
       Semantic_error.illtyped_reduce_sum_generic loc id.name
         (List.map ~f:type_of_expr_typed es)
         expected_args err
@@ -498,12 +498,12 @@ let check_variadic_ode ~is_cond_dist loc id es =
       Stan_math_signatures.variadic_ode_mandatory_fun_args
       Stan_math_signatures.variadic_ode_fun_return_type (get_arg_types es)
   with
-  | None ->
+  | Ok () ->
       mk_typed_expression
         ~expr:(mk_fun_app ~is_cond_dist (StanLib FnPlain, id, es))
         ~ad_level:(expr_ad_lub es)
         ~type_:Stan_math_signatures.variadic_ode_return_type ~loc
-  | Some (expected_args, err) ->
+  | Error (expected_args, err) ->
       Semantic_error.illtyped_variadic_ode loc id.name
         (List.map ~f:type_of_expr_typed es)
         expected_args err
@@ -657,6 +657,10 @@ and check_expression cf tenv ({emeta; expr} : Ast.untyped_expression) :
       es |> List.map ~f:ce |> check_funapp loc cf tenv ~is_cond_dist:false id
   | CondDistApp ((), id, es) ->
       es |> List.map ~f:ce |> check_funapp loc cf tenv ~is_cond_dist:true id
+  | Promotion (e, _) ->
+      (* Should never happen: promotions are produced during typechecking *)
+      Common.FatalError.fatal_error_msg
+        [%message "Promotion in untyped AST" (e : Ast.untyped_expression)]
 
 and check_expression_of_int_type cf tenv e name =
   let te = check_expression cf tenv e in
@@ -698,7 +702,7 @@ let check_nrfn loc tenv id es =
         (Env.nearest_ident tenv id.name)
       |> error
   | _ (* a function *) -> (
-    match SignatureMismatch.returntype tenv id.name (get_arg_types es) with
+    match SignatureMismatch.returntype tenv id.name es with
     | Ok (Void, fnk) ->
         mk_typed_statement
           ~stmt:(NRFunApp (fnk (Fun_kind.suffix_from_name id.name), id, es))