Changed the MIR to represent function declarations as Options

src/analysis_and_optimization/Dependence_analysis.ml

@@ -217,9 +217,10 @@ let mir_uninitialized_variables (mir : Program.Typed.t) :
                Set.Poly.of_list
                  (List.map fdargs ~f:(fun (_, arg_name, _) -> arg_name))
              in
-             stmt_uninitialized_variables
-               (Set.Poly.union arg_vars globals)
-               fdbody )) ]
+             Option.value_map fdbody ~default:Set.Poly.empty ~f:(fun fdbody ->
+                 stmt_uninitialized_variables
+                   (Set.Poly.union arg_vars globals)
+                   fdbody ) )) ]
 
 let build_dep_info_map (mir : Program.Typed.t)
     (stmt : (Expr.Typed.Meta.t, Stmt.Located.Meta.t) Stmt.Fixed.t) :

src/analysis_and_optimization/Optimize.ml

@@ -26,7 +26,7 @@ let transform_program (mir : Program.Typed.t)
   let transformed_prog_body = transform packed_prog_body in
   let transformed_functions =
     List.map mir.functions_block ~f:(fun fs ->
-        {fs with fdbody= transform fs.fdbody} )
+        {fs with fdbody= Option.map ~f:transform fs.fdbody} )
   in
   match transformed_prog_body with
   | { pattern=
@@ -59,15 +59,9 @@ let transform_program_blockwise (mir : Program.Typed.t)
         raise
           (Failure "Something went wrong with program transformation packing!")
   in
-  (* Right now, we have an implicit constraint where if fdbody = Skip, the
-     fun_def is a function declaration. When that's the case we don't want
-     to change it from Skip.*)
-  let non_decl_functions =
-    List.filter ~f:(fun def -> def.fdbody.pattern <> Skip) mir.functions_block
-  in
   let transformed_functions =
-    List.map non_decl_functions ~f:(fun fs ->
-        {fs with fdbody= transform (Some fs) fs.fdbody} )
+    List.map mir.functions_block ~f:(fun fs ->
+        {fs with fdbody= Option.map ~f:(transform (Some fs)) fs.fdbody} )
   in
   { mir with
     functions_block= transformed_functions
@@ -437,8 +431,8 @@ let create_function_inline_map adt l =
     else
       let accum' =
         match fdbody with
-        | Stmt.Fixed.({pattern= Stmt.Fixed.Pattern.Skip; _}) -> accum
-        | _ -> (
+        | None -> accum
+        | Some fdbody -> (
             let data =
               ( Option.map ~f:(fun x -> Type.Unsized x) fdrt
               , List.map ~f:(fun (_, name, _) -> name) fdargs

src/analysis_and_optimization/Pedantic_analysis.ml

@@ -123,19 +123,23 @@ let list_arg_dependant_fundef_cf (mir : Program.Typed.t)
     (fun_def : 'a Program.fun_def) :
     (Location_span.t * int * string) Set.Poly.t =
   let args = List.map ~f:(fun (_, name, _) -> name) fun_def.fdargs in
-  (* build dataflow data structure *)
-  let info_map = build_dep_info_map mir fun_def.fdbody in
-  let cf_deps = list_target_dependant_cf info_map (Set.Poly.of_list args) in
-  union_map cf_deps ~f:(fun (loc, names) ->
-      Set.Poly.map names ~f:(fun name ->
-          let ix, _ =
-            Option.value_exn
-              ~message:
-                "INTERNAL ERROR: Pedantic mode found CF dependent on an \
-                 arg,but the arg is mismatched. Please report a bug.\n"
-              (List.findi args ~f:(fun _ arg -> arg = name))
-          in
-          (loc, ix, name) ) )
+  (* Only look for control flow if this function definition has a body *)
+  Option.value_map fun_def.fdbody ~default:Set.Poly.empty ~f:(fun body ->
+      (* build dataflow data structure *)
+      let info_map = build_dep_info_map mir body in
+      let cf_deps =
+        list_target_dependant_cf info_map (Set.Poly.of_list args)
+      in
+      union_map cf_deps ~f:(fun (loc, names) ->
+          Set.Poly.map names ~f:(fun name ->
+              let ix, _ =
+                Option.value_exn
+                  ~message:
+                    "INTERNAL ERROR: Pedantic mode found CF dependent on an \
+                     arg,but the arg is mismatched. Please report a bug.\n"
+                  (List.findi args ~f:(fun _ arg -> arg = name))
+              in
+              (loc, ix, name) ) ) )
 
 let expr_collect_exprs (expr : Expr.Typed.t) ~f : 'a Set.Poly.t =
   let collect_expr s (expr : Expr.Typed.t) =
@@ -269,7 +273,7 @@ let list_distributions (mir : Program.Typed.t) : dist_info Set.Poly.t =
   in
   stmts_collect_exprs
     (List.append mir.log_prob
-       (List.map ~f:(fun f -> f.fdbody) mir.functions_block))
+       (List.filter_map ~f:(fun f -> f.fdbody) mir.functions_block))
     ~f:take_dist
 
 (* Our definition of 'unscaled' for constants used in distributions *)

src/frontend/Ast_to_Mir.ml

@@ -492,8 +492,8 @@ let trans_decl {dconstrain; dadlevel} smeta decl_type transform identifier
   else size_checks @ (decl :: rhs_assignment)
 
 let unwrap_block_or_skip = function
-  | [({Stmt.Fixed.pattern= Block _; _} as b)] | [({pattern= Skip; _} as b)] ->
-      b
+  | [({Stmt.Fixed.pattern= Block _; _} as b)] -> Some b
+  | [{pattern= Skip; _}] -> None
   | x ->
       raise_s
         [%message "Expecting a block or skip, not" (x : Stmt.Located.t list)]

src/middle/Program.ml

@@ -8,8 +8,10 @@ type fun_arg_decl = (UnsizedType.autodifftype * string * UnsizedType.t) list
 type 'a fun_def =
   { fdrt: UnsizedType.t option
   ; fdname: string
-  ; fdargs: (UnsizedType.autodifftype * string * UnsizedType.t) list
-  ; fdbody: 'a
+  ; fdargs:
+      (UnsizedType.autodifftype * string * UnsizedType.t) list
+      (* If fdbody is None, this is a function declaration without body. *)
+  ; fdbody: 'a option
   ; fdloc: Location_span.t sexp_opaque [@compare.ignore] }
 [@@deriving compare, hash, map, sexp, map, fold]
 
@@ -68,15 +70,19 @@ let pp_fun_arg_decl ppf (autodifftype, name, unsizedtype) =
 
 let pp_fun_def pp_s ppf = function
   | {fdrt; fdname; fdargs; fdbody; _} -> (
-    match fdrt with
-    | Some rt ->
-        Fmt.pf ppf {|@[<v2>%a %s%a {@ %a@]@ }|} UnsizedType.pp rt fdname
-          Fmt.(list pp_fun_arg_decl ~sep:comma |> parens)
-          fdargs pp_s fdbody
-    | None ->
-        Fmt.pf ppf {|@[<v2>%s %s%a {@ %a@]@ }|} "void" fdname
-          Fmt.(list pp_fun_arg_decl ~sep:comma |> parens)
-          fdargs pp_s fdbody )
+      let pp_body_opt ppf = function
+        | None -> Fmt.pf ppf ";"
+        | Some body -> pp_s ppf body
+      in
+      match fdrt with
+      | Some rt ->
+          Fmt.pf ppf {|@[<v2>%a %s%a {@ %a@]@ }|} UnsizedType.pp rt fdname
+            Fmt.(list pp_fun_arg_decl ~sep:comma |> parens)
+            fdargs pp_body_opt fdbody
+      | None ->
+          Fmt.pf ppf {|@[<v2>%s %s%a {@ %a@]@ }|} "void" fdname
+            Fmt.(list pp_fun_arg_decl ~sep:comma |> parens)
+            fdargs pp_body_opt fdbody )
 
 let pp_io_block ppf = function
   | Parameters -> Fmt.string ppf "parameters"
@@ -203,7 +209,9 @@ module Labelled = struct
       ~f:associate_outvar
 
   and associate_fun_def assocs {fdbody; _} =
-    Stmt.Labelled.associate ~init:assocs fdbody
+    match fdbody with
+    | None -> assocs
+    | Some fdbody -> Stmt.Labelled.associate ~init:assocs fdbody
 
   and associate_outvar assocs (_, {out_constrained_st; out_unconstrained_st; _})
       =

src/stan_math_backend/Stan_math_code_gen.ml

@@ -231,9 +231,9 @@ let pp_fun_def ppf Program.({fdrt; fdname; fdargs; fdbody; _})
     pf ppf "%s(@[<hov>%a@]) " name (list ~sep:comma string) arg_strs
   in
   pp_sig ppf fdname ;
-  match Stmt.Fixed.(fdbody.pattern) with
-  | Skip -> pf ppf ";@ "
-  | _ ->
+  match fdbody with
+  | None -> pf ppf ";@ "
+  | Some fdbody ->
       pp_block ppf (pp_body, fdbody) ;
       pf ppf "@,@,struct %s%s {@,%a const @,{@,return %a;@,}@,};@," fdname
         functor_suffix pp_sig "operator()" pp_call_str

src/tfp_backend/Code_gen.ml

@@ -256,10 +256,13 @@ let pp_methods ppf p =
   pf ppf "@ %a" pp_param_names p
 
 let pp_fundef ppf {Program.fdname; fdargs; fdbody; _} =
+  let no_body_default : Stmt.Located.t =
+    {pattern= Stmt.Fixed.Pattern.Skip; meta= Location_span.empty}
+  in
   pp_method ppf fdname
     (List.map ~f:(fun (_, name, _) -> name) fdargs)
     []
-    (fun ppf -> pp_stmt ppf fdbody)
+    (fun ppf -> pp_stmt ppf (Option.value ~default:no_body_default fdbody))
 
 let imports =
   {|