Merge pull request #1498 from stan-dev/pre-generate-signatures

Generate signature map at build time

Author: WardBrian

Jenkinsfile

@@ -260,8 +260,8 @@ pipeline {
 
                             sh '''
                                 eval $(opam env)
-                                bisect-ppx-report summary --expect src/ --do-not-expect src/stancjs/
-                                bisect-ppx-report coveralls coverage.json --service-name jenkins --service-job-id $BUILD_ID --expect src/ --do-not-expect src/stancjs/
+                                bisect-ppx-report summary --expect src/ --do-not-expect src/stancjs/ --do-not-expect src/stan_math_signatures/Generate.ml
+                                bisect-ppx-report coveralls coverage.json --service-name jenkins --service-job-id $BUILD_ID --expect src/ --do-not-expect src/stancjs/ --do-not-expect src/stan_math_signatures/Generate.ml
                             '''
 
                             withCredentials([usernamePassword(credentialsId: 'stan-stanc3-codecov-token', usernameVariable: 'DUMMY_USERNAME', passwordVariable: 'CODECOV_TOKEN')]) {

Makefile

@@ -10,8 +10,8 @@ testcoverage:
 	@find . -name '*.coverage' | xargs rm -f
 	dune clean
 	BISECT_FILE=`pwd`/bisect dune runtest --instrument-with bisect_ppx --force
-	bisect-ppx-report html --expect src/ --do-not-expect src/stancjs/
-	bisect-ppx-report summary --expect src/ --do-not-expect src/stancjs/
+	bisect-ppx-report html --expect src/ --do-not-expect src/stancjs/ --do-not-expect src/stan_math_signatures/Generate.ml
+	bisect-ppx-report summary --expect src/ --do-not-expect src/stancjs/ --do-not-expect src/stan_math_signatures/Generate.ml
 	@rm *.coverage
 
 format:

docs/exposing_new_functions.mld

@@ -7,8 +7,10 @@
 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].
-The compiler uses the signatures defined there to do type checking.
+To do the latter, we have to add a corresponding line in [src/stan_math_signatures/Generate.ml].
+The compiler uses the signatures defined there to do type checking. (Aside: this file generates a
+binary representation of the signatures, which is then stored in the executable. This is faster
+than re-generating the list of signatures every time the compiler is run.)
 
 
 {1 Adding a distribution function }

docs/index.mld

@@ -54,6 +54,7 @@ Frontend
 Middle
 Common
 Analysis_and_optimization
+Stan_math_signatures
 Stan_math_backend
 }
 

src/analysis_and_optimization/Memory_patterns.ml

@@ -97,13 +97,12 @@ let query_stan_math_mem_pattern_support (name : string)
       let namematches = lookup_stan_math_function name in
       let filteredmatches =
         List.filter
-          ~f:(fun x ->
-            Frontend.SignatureMismatch.check_compatible_arguments_mod_conv
-              (snd3 x) args
+          ~f:(fun (x, _, _, _) ->
+            Frontend.SignatureMismatch.check_compatible_arguments_mod_conv x
+              args
             |> Result.is_ok)
           namematches in
-      let is_soa ((_ : UnsizedType.returntype), _, mem) =
-        mem = Mem_pattern.SoA in
+      let is_soa = function _, _, _, Mem_pattern.SoA -> true | _ -> false in
       List.exists ~f:is_soa filteredmatches
 
 (*Validate whether a function can support SoA matrices*)

src/analysis_and_optimization/dune

@@ -1,11 +1,11 @@
 (library
  (name analysis_and_optimization)
  (public_name stanc.analysis)
- (libraries core middle yojson frontend)
+ (libraries core middle yojson frontend stan_math_signatures)
  (instrumentation
   (backend bisect_ppx))
  (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)
  (preprocess
-  (pps ppx_jane ppx_deriving.map ppx_deriving.fold)))
+  (pps ppx_jane)))

src/common/dune

@@ -6,4 +6,4 @@
   (backend bisect_ppx))
  (inline_tests)
  (preprocess
-  (pps ppx_jane ppx_deriving.map ppx_deriving.fold ppx_deriving.create)))
+  (pps ppx_jane ppx_deriving.map ppx_deriving.fold)))

src/frontend/Environment.ml

@@ -29,11 +29,8 @@ let stan_math_environment =
     Stan_math_signatures.get_stan_math_signatures_alist ()
     |> List.map ~f:(fun (key, values) ->
            ( key
-           , List.map values ~f:(fun (rt, args, mem) ->
-                 let type_ =
-                   UnsizedType.UFun
-                     (args, rt, Fun_kind.suffix_from_name key, mem) in
-                 {type_; kind= `StanMath}) ))
+           , List.map values ~f:(fun s ->
+                 {type_= UnsizedType.UFun s; kind= `StanMath}) ))
     |> String.Map.of_alist_exn in
   functions
 

src/frontend/SignatureMismatch.ml

@@ -424,7 +424,7 @@ let pp_math_lib_assignmentoperator_sigs ppf (lt, op) =
       Stan_math_signatures.make_assignmentoperator_stan_math_signatures op in
     let errors =
       List.filter
-        ~f:(fun (_, args, _) ->
+        ~f:(fun (args, _, _, _) ->
           Result.is_ok (check_same_type 0 lt (snd (List.hd_exn args))))
         errors in
     match List.split_n errors max_n_errors with
@@ -433,7 +433,7 @@ let pp_math_lib_assignmentoperator_sigs ppf (lt, op) =
     | errors, _ -> Some (errors, true) in
   let pp_sigs ppf (signatures, omitted) =
     Fmt.pf ppf "@[<v>%a%a@]"
-      (Fmt.list ~sep:Fmt.cut Stan_math_signatures.pp_math_sig)
+      (Fmt.list ~sep:Fmt.cut UnsizedType.pp_math_sig)
       signatures
       (if omitted then Fmt.pf else Fmt.nop)
       "@ (Additional signatures omitted)" in

src/frontend/Typechecker.ml

@@ -501,12 +501,9 @@ let check_normal_fn ~is_cond_dist loc tenv id es =
   | [] ->
       (match Utils.split_distribution_suffix id.name with
       | Some (prefix, suffix) -> (
-          let known_families =
-            List.map
-              ~f:(fun (_, y, _, _) -> y)
-              Stan_math_signatures.distributions in
           let is_known_family s =
-            List.mem known_families s ~equal:String.equal in
+            List.Assoc.mem Stan_math_signatures.distributions s
+              ~equal:String.equal in
           match suffix with
           | ("lpmf" | "lupmf") when Env.mem tenv (prefix ^ "_lpdf") ->
               Semantic_error.returning_fn_expected_wrong_dist_suffix_found loc
@@ -663,8 +660,8 @@ and check_reduce_sum ~is_cond_dist loc cf tenv id tes =
       |> error
 
 and check_variadic ~is_cond_dist loc cf tenv id tes =
-  let Stan_math_signatures.
-        {control_args; required_fn_args; required_fn_rt; return_type} =
+  let UnsizedType.{control_args; required_fn_args; required_fn_rt; return_type}
+      =
     Stan_math_signatures.lookup_stan_math_variadic_function id.name
     |> Option.value_exn in
   let matching remaining_es Env.{type_= ftype; _} =