Refactor "standard library" as a virtual module

docs/core_ideas.mld

@@ -65,6 +65,15 @@ This takes some getting used to, and also can lead to some unhelpful type signat
 VSCode, because abbreviations are not always used in hover-over text. For example, [Expr.Typed.t], the MIR's typed
 expression type, actually has a signature of [Expr.Typed.Meta.t Expr.Fixed.t].
 
+{1 The [Library] virtual module}
+
+The [Frontend] library is a {{:https://dune.readthedocs.io/en/latest/variants.html}virtual library}
+where the module [Library] is unimplemented. This allows the rest of the library to operate without
+making backend-specific assumptions about any one library.
+
+This must be supplied when the executable is built in the [dune] file.
+For Stanc, we supply the [stan_math_library] instantiation defined in [src/stan_math_backend/stan_math_library].
+
 {1 The [Fmt] library and pretty-printing}
 
 We extensively use the {{:https://erratique.ch/software/fmt}Fmt} library for our pretty-printing and code

docs/exposing_new_functions.mld

@@ -7,7 +7,7 @@
 For a function to be built into Stan, it has to be included in the Stan Math library and its signature has to be
 exposed to the compiler.
 
-To do the latter, we have to add a corresponding line in [src/middle/Stan_math_signatures.ml].
+To do the latter, we have to add a corresponding line in [src/stan_math_backend/Stan_math_signatures.ml].
 The compiler uses the signatures defined there to do type checking.
 
 
@@ -130,12 +130,13 @@ For example, the following line defines the signature [add(real, matrix) => matr
 Functions such as the ODE integrators or [reduce_sum], which take in user-functions and a variable-length
 list of arguments, are {b NOT} added to this list.
 
-"Nice" variadic functions are added to the hashtable [Stan_math_signatures.stan_math_variadic_signatures].
+"Nice" variadic functions are added to the hashtable [Library.variadic_signatures].
 This is probably sufficient for most variadic functions, e.g. all the ODE solvers and DAE solvers are done
 via this method.
 [reduce_sum] is not "nice", since it is both variadic and {e polymorphic}, requiring certain arguments to have the same
 (but {e not predetermined}) type. Therefore, [reduce_sum] is treated as special case in the [Typechecker]
-module in the frontend folder.
+module in the frontend folder. These are instead handled by special functions like [is_special_function_name]. They
+must also be given custom typechecking rules in the private sub-module [Special_typechecking].
 
 Note that higher-order functions also usually require changes to the C++ code generation to work properly.
 It is best to consult an existing example of how these are done before proceeding.

src/analysis_and_optimization/Debug_data_generation.ml

@@ -30,7 +30,7 @@ let rec vect_to_mat l m =
 
 let eval_expr m e =
   let e = Mir_utils.subst_expr m e in
-  let e = Partial_evaluator.eval_expr e in
+  let e = Partial_evaluator.try_eval_expr e in
   let rec strip_promotions (e : Middle.Expr.Typed.t) =
     match e.pattern with Promotion (e, _, _) -> strip_promotions e | _ -> e
   in

src/analysis_and_optimization/Dependence_analysis.ml

@@ -4,7 +4,7 @@ open Middle
 open Dataflow_types
 open Mir_utils
 open Dataflow_utils
-open Monotone_framework_sigs
+open Monotone_framework_intf
 open Monotone_framework
 
 (***********************************)
@@ -119,9 +119,8 @@ let mir_reaching_definitions (mir : Program.Typed.t) (stmt : Stmt.Located.t) :
   Map.Poly.map rd_map ~f:(fun {entry; exit} ->
       {entry= to_rd_set entry; exit= to_rd_set exit} )
 
-let all_labels
-    (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int)
-    : int Set.Poly.t =
+let all_labels (module Flowgraph : FLOWGRAPH with type labels = int) :
+    int Set.Poly.t =
   let step set =
     Set.Poly.union set
       (union_map set ~f:(fun l -> Map.Poly.find_exn Flowgraph.successors l))

src/analysis_and_optimization/Memory_patterns.ml

@@ -111,15 +111,14 @@ let rec is_uni_eigen_loop_indexing in_loop (ut : UnsizedType.t)
 
 let query_stan_math_mem_pattern_support (name : string)
     (args : (UnsizedType.autodifftype * UnsizedType.t) list) =
-  let open Stan_math_signatures in
   match name with
-  | x when is_stan_math_variadic_function_name x -> false
-  | x when is_reduce_sum_fn x -> false
+  | x when Frontend.Library.is_variadic_function_name x -> false
+  | x when Frontend.Library.is_special_function_name x -> false
   | _ ->
       let name =
-        string_operator_to_stan_math_fns (Utils.stdlib_distribution_name name)
-      in
-      let namematches = Hashtbl.find_multi stan_math_signatures name in
+        Frontend.Library.string_operator_to_function_name
+          (Utils.stdlib_distribution_name name) in
+      let namematches = Frontend.Library.get_signatures name in
       let filteredmatches =
         List.filter
           ~f:(fun x ->

src/analysis_and_optimization/Monotone_framework.ml

@@ -2,7 +2,7 @@
 
 open Core_kernel
 open Core_kernel.Poly
-open Monotone_framework_sigs
+open Monotone_framework_intf
 open Mir_utils
 open Middle
 
@@ -864,7 +864,7 @@ let rec declared_variables_stmt
         (List.map ~f:(fun x -> declared_variables_stmt x.pattern) l)
 
 let propagation_mfp (prog : Program.Typed.t)
-    (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int)
+    (module Flowgraph : FLOWGRAPH with type labels = int)
     (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t)
     (propagation_transfer :
          (int, Stmt.Located.Non_recursive.t) Map.Poly.t
@@ -897,7 +897,7 @@ let propagation_mfp (prog : Program.Typed.t)
   Mf.mfp ()
 
 let reaching_definitions_mfp (mir : Program.Typed.t)
-    (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int)
+    (module Flowgraph : FLOWGRAPH with type labels = int)
     (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) =
   let variables =
     ( module struct
@@ -918,7 +918,7 @@ let reaching_definitions_mfp (mir : Program.Typed.t)
   Mf.mfp ()
 
 let initialized_vars_mfp (total : string Set.Poly.t)
-    (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int)
+    (module Flowgraph : FLOWGRAPH with type labels = int)
     (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) =
   let (module Lattice) =
     dual_powerset_lattice_empty_initial
@@ -949,8 +949,7 @@ let globals (prog : Program.Typed.t) =
 (** Monotone framework instance for live_variables analysis. Expects reverse
     flowgraph. *)
 let live_variables_mfp (prog : Program.Typed.t)
-    (module Rev_Flowgraph : Monotone_framework_sigs.FLOWGRAPH
-      with type labels = int )
+    (module Rev_Flowgraph : FLOWGRAPH with type labels = int)
     (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) =
   let never_kill = globals prog in
   let variables =
@@ -970,10 +969,8 @@ let live_variables_mfp (prog : Program.Typed.t)
 
 (** Instantiate all four instances of the monotone framework for lazy
     code motion, reusing code between them *)
-let lazy_expressions_mfp
-    (module Flowgraph : Monotone_framework_sigs.FLOWGRAPH with type labels = int)
-    (module Rev_Flowgraph : Monotone_framework_sigs.FLOWGRAPH
-      with type labels = int )
+let lazy_expressions_mfp (module Flowgraph : FLOWGRAPH with type labels = int)
+    (module Rev_Flowgraph : FLOWGRAPH with type labels = int)
     (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t) =
   let all_expressions =
     used_subexpressions_stmt
@@ -1031,8 +1028,7 @@ let lazy_expressions_mfp
  *
  *)
 let minimal_variables_mfp
-    (module Circular_Fwd_Flowgraph : Monotone_framework_sigs.FLOWGRAPH
-      with type labels = int )
+    (module Circular_Fwd_Flowgraph : FLOWGRAPH with type labels = int)
     (flowgraph_to_mir : (int, Stmt.Located.Non_recursive.t) Map.Poly.t)
     (initial_variables : string Set.Poly.t)
     (gen_variable :

src/analysis_and_optimization/Optimize.ml

@@ -628,7 +628,7 @@ let list_collapsing (mir : Program.Typed.t) =
 let propagation
     (propagation_transfer :
          (int, Stmt.Located.Non_recursive.t) Map.Poly.t
-      -> (module Monotone_framework_sigs.TRANSFER_FUNCTION
+      -> (module Monotone_framework_intf.TRANSFER_FUNCTION
             with type labels = int
              and type properties = (string, Middle.Expr.Typed.t) Map.Poly.t
                                    option ) ) (mir : Program.Typed.t) =
@@ -711,7 +711,7 @@ let dead_code_elimination (mir : Program.Typed.t) =
     let dead_code_elim_stmt_base i stmt =
       (* NOTE: entry in the reverse flowgraph, so exit in the forward flowgraph *)
       let live_variables_s =
-        (Map.find_exn live_variables i).Monotone_framework_sigs.entry in
+        (Map.find_exn live_variables i).Monotone_framework_intf.entry in
       match stmt with
       | Stmt.Fixed.Pattern.Assignment ((x, _, []), rhs) ->
           if Set.Poly.mem live_variables_s x || cannot_remove_expr rhs then stmt
@@ -1183,11 +1183,6 @@ let optimize_soa (mir : Program.Typed.t) =
     List.fold ~init:Set.Poly.empty
       ~f:(Memory_patterns.query_initial_demotable_stmt false)
       mir.log_prob in
-  (*
-  let print_set s =
-    Set.Poly.iter ~f:print_endline s in
-  let () = print_set initial_variables in
-  *)
   let mod_exprs aos_exits mod_expr =
     Mir_utils.map_rec_expr
       (Memory_patterns.modify_expr_pattern aos_exits)
@@ -1207,8 +1202,6 @@ let optimize_soa (mir : Program.Typed.t) =
   in
   {mir with log_prob= transform' mir.log_prob}
 
-(* Apparently you need to completely copy/paste type definitions between
-   ml and mli files?*)
 type optimization_settings =
   { function_inlining: bool
   ; static_loop_unrolling: bool

src/analysis_and_optimization/Optimize.mli

@@ -1,4 +1,5 @@
 (* Code for optimization passes on the MIR *)
+
 open Middle
 
 val function_inlining : Program.Typed.t -> Program.Typed.t
@@ -59,7 +60,7 @@ val optimize_ad_levels : Program.Typed.t -> Program.Typed.t
 
 val allow_uninitialized_decls : Program.Typed.t -> Program.Typed.t
 (** Marks Decl types such that, if the first assignment after the decl
-    assigns to the full object, allow the object to be constructed but 
+    assigns to the full object, allow the object to be constructed but
     not uninitialized.  *)
 
 (** Interface for turning individual optimizations on/off. Useful for testing

src/analysis_and_optimization/Partial_evaluator.ml

@@ -106,11 +106,11 @@ let rec eval_expr ?(preserve_stability = false) (e : Expr.Typed.t) =
                 Operator.of_string_opt name
                 |> Option.value_map
                      ~f:(fun op ->
-                       Frontend.Typechecker.operator_stan_math_return_type op
+                       Frontend.Typechecker.operator_return_type op
                          argument_types
                        |> Option.map ~f:fst )
                      ~default:
-                       (Frontend.Typechecker.stan_math_return_type name
+                       (Frontend.Typechecker.library_function_return_type name
                           argument_types ) in
               let try_partially_evaluate_stanlib e =
                 Expr.Fixed.Pattern.(

src/analysis_and_optimization/dune

@@ -1,9 +1,8 @@
 (library
  (name analysis_and_optimization)
  (public_name stanc.analysis)
- (libraries core_kernel str fmt common middle frontend)
- (inline_tests)
- ;; TODO: Not sure what's going on but it's throwing an error that this module has no implementation
- (modules_without_implementation monotone_framework_sigs)
+ (libraries core_kernel str fmt common middle frontend stan_math_backend)
+ (inline_tests
+  (libraries stan_math_library))
  (preprocess
   (pps ppx_jane ppx_deriving.map ppx_deriving.fold)))

src/frontend/Ast_to_Mir.ml

@@ -107,8 +107,8 @@ let truncate_dist ud_dists (id : Ast.identifier) ast_obs ast_args t =
     | None ->
         ( Ast.StanLib FnPlain
         , Set.to_list possible_names |> List.hd_exn
-        , if Stan_math_signatures.is_stan_math_function_name (id.name ^ "_lpmf")
-          then UnsizedType.UInt
+        , if Library.is_stdlib_function_name (id.name ^ "_lpmf") then
+            UnsizedType.UInt
           else UnsizedType.UReal (* close enough *) ) in
   let trunc cond_op (x : Ast.typed_expression) y =
     let smeta = x.Ast.emeta.loc in
@@ -429,7 +429,8 @@ let rec trans_stmt ud_dists (declc : decl_context) (ts : Ast.typed_statement) =
   | Ast.IncrementLogProb e | Ast.TargetPE e -> TargetPE (trans_expr e) |> swrap
   | Ast.Tilde {arg; distribution; args; truncation} ->
       let suffix =
-        Stan_math_signatures.dist_name_suffix ud_dists distribution.name in
+        Std_library_utils.dist_name_suffix Library.is_stdlib_function_name
+          ud_dists distribution.name in
       let name = distribution.name ^ suffix in
       let kind =
         let possible_names =

src/frontend/Canonicalize.ml

@@ -43,7 +43,7 @@ let rec replace_deprecated_expr
         if is_deprecated_distribution name then
           CondDistApp
             ( StanLib suffix
-            , {name= rename_deprecated deprecated_distributions name; id_loc}
+            , {name= rename_deprecated_distribution name; id_loc}
             , List.map ~f:(replace_deprecated_expr deprecated_userdefined) e )
         else if String.is_suffix name ~suffix:"_cdf" then
           CondDistApp
@@ -53,7 +53,7 @@ let rec replace_deprecated_expr
         else
           FunApp
             ( StanLib suffix
-            , {name= rename_deprecated deprecated_functions name; id_loc}
+            , {name= rename_deprecated_function name; id_loc}
             , List.map ~f:(replace_deprecated_expr deprecated_userdefined) e )
     | FunApp (UserDefined suffix, {name; id_loc}, e) -> (
       match String.Map.find deprecated_userdefined name with

src/frontend/Deprecation_analysis.ml

@@ -2,48 +2,29 @@ open Core_kernel
 open Ast
 open Middle
 
-let deprecated_functions =
-  String.Map.of_alist_exn
-    [ ("multiply_log", ("lmultiply", "2.32.0"))
-    ; ("binomial_coefficient_log", ("lchoose", "2.32.0"))
-    ; ("cov_exp_quad", ("gp_exp_quad_cov", "2.32.0"))
-    ; ("fabs", ("abs", "2.33.0")) ]
-
-let deprecated_odes =
-  String.Map.of_alist_exn
-    [ ("integrate_ode", ("ode_rk45", "3.0"))
-    ; ("integrate_ode_rk45", ("ode_rk45", "3.0"))
-    ; ("integrate_ode_bdf", ("ode_bdf", "3.0"))
-    ; ("integrate_ode_adams", ("ode_adams", "3.0")) ]
-
-let deprecated_distributions =
-  String.Map.of_alist_exn
-    (List.map
-       ~f:(fun (x, y) -> (x, (y, "2.32.0")))
-       (List.concat_map Middle.Stan_math_signatures.distributions
-          ~f:(fun (fnkinds, name, _, _) ->
-            List.filter_map fnkinds ~f:(function
-              | Lpdf -> Some (name ^ "_log", name ^ "_lpdf")
-              | Lpmf -> Some (name ^ "_log", name ^ "_lpmf")
-              | Cdf -> Some (name ^ "_cdf_log", name ^ "_lcdf")
-              | Ccdf -> Some (name ^ "_ccdf_log", name ^ "_lccdf")
-              | Rng | Log | UnaryVectorized _ -> None ) ) ) )
-
 let stan_lib_deprecations =
-  Map.merge_skewed deprecated_distributions deprecated_functions
+  Map.merge_skewed Library.deprecated_distributions Library.deprecated_functions
     ~combine:(fun ~key x y ->
       Common.FatalError.fatal_error_msg
         [%message
           "Common key in deprecation map"
             (key : string)
-            (x : string * string)
-            (y : string * string)] )
+            (x : Std_library_utils.deprecation_info)
+            (y : Std_library_utils.deprecation_info)] )
 
 let is_deprecated_distribution name =
-  Option.is_some (Map.find deprecated_distributions name)
+  Map.mem Library.deprecated_distributions name
 
 let rename_deprecated map name =
-  Map.find map name |> Option.map ~f:fst |> Option.value ~default:name
+  Map.find map name
+  |> Option.map ~f:(fun Std_library_utils.{replacement; canonicalize_away; _} ->
+         if canonicalize_away then replacement else name )
+  |> Option.value ~default:name
+
+let rename_deprecated_distribution =
+  rename_deprecated Library.deprecated_distributions
+
+let rename_deprecated_function = rename_deprecated Library.deprecated_functions
 
 let distribution_suffix name =
   let open String in
@@ -101,41 +82,21 @@ let rec collect_deprecated_expr (acc : (Location_span.t * string) list)
     ({expr; emeta} : (typed_expr_meta, fun_kind) expr_with) :
     (Location_span.t * string) list =
   match expr with
-  | FunApp (StanLib FnPlain, {name= "if_else"; _}, l) ->
-      acc
-      @ [ ( emeta.loc
-          , "The function `if_else` is deprecated and will be removed in Stan \
-             2.32.0. Use the conditional operator (x ? y : z) instead; this \
-             can be automatically changed using the canonicalize flag for \
-             stanc" ) ]
-      @ List.concat_map l ~f:(fun e -> collect_deprecated_expr [] e)
   | FunApp ((StanLib _ | UserDefined _), {name; _}, l) ->
       let w =
         match Map.find stan_lib_deprecations name with
-        | Some (rename, version) ->
+        | Some {replacement; version; extra_message; _} ->
             [ ( emeta.loc
               , name ^ " is deprecated and will be removed in Stan " ^ version
-                ^ ". Use " ^ rename
-                ^ " instead. This can be automatically changed using the \
-                   canonicalize flag for stanc" ) ]
+                ^ ". Use " ^ replacement ^ " instead. " ^ extra_message ) ]
         | _ when String.is_suffix name ~suffix:"_cdf" ->
             [ ( emeta.loc
               , "Use of " ^ name
                 ^ " without a vertical bar (|) between the first two arguments \
                    of a CDF is deprecated and will be removed in Stan 2.32.0. \
                    This can be automatically changed using the canonicalize \
                    flag for stanc" ) ]
-        | _ -> (
-          match Map.find deprecated_odes name with
-          | Some (rename, version) ->
-              [ ( emeta.loc
-                , name ^ " is deprecated and will be removed in Stan " ^ version
-                  ^ ". Use " ^ rename
-                  ^ " instead. \n\
-                     The new interface is slightly different, see: \
-                     https://mc-stan.org/users/documentation/case-studies/convert_odes.html"
-                ) ]
-          | _ -> [] ) in
+        | _ -> [] in
       acc @ w @ List.concat_map l ~f:(fun e -> collect_deprecated_expr [] e)
   | PrefixOp (PNot, ({emeta= {type_= UReal; loc; _}; _} as e)) ->
       let acc =

src/frontend/Deprecation_analysis.mli

@@ -16,8 +16,7 @@ val collect_userdef_distributions :
 val distribution_suffix : string -> bool
 val without_suffix : string list -> string -> string
 val is_deprecated_distribution : string -> bool
-val deprecated_distributions : (string * string) String.Map.t
-val deprecated_functions : (string * string) String.Map.t
-val rename_deprecated : (string * string) String.Map.t -> string -> string
+val rename_deprecated_distribution : string -> string
+val rename_deprecated_function : string -> string
 val userdef_distributions : untyped_statement block option -> string list
 val collect_warnings : typed_program -> Warnings.t list