Refactor unconstraining to use deserializer interface

.gitignore

@@ -32,3 +32,6 @@ test/*.log
 
 # Mac OS X hidden files
 *.DS_Store
+
+# .hpp files in test folder
+test/**/*.hpp

src/analysis_and_optimization/Mir_utils.ml

@@ -40,7 +40,8 @@ type bound_values =
   { lower: [`None | `Nonlit | `Lit of float]
   ; upper: [`None | `Nonlit | `Lit of float] }
 
-let trans_bounds_values (trans : Expr.Typed.t transformation) : bound_values =
+let trans_bounds_values (trans : Expr.Typed.t Transformation.t) : bound_values
+    =
   let bound_value e =
     match num_expr_value e with None -> `Nonlit | Some (f, _) -> `Lit f
   in
@@ -132,8 +133,7 @@ let map_rec_expr_state f state e =
 
 let rec map_rec_stmt_loc f stmt =
   let recurse = map_rec_stmt_loc f in
-  Stmt.Fixed.
-    {stmt with pattern= f (Pattern.map (fun x -> x) recurse stmt.pattern)}
+  Stmt.Fixed.{stmt with pattern= f (Pattern.map Fn.id recurse stmt.pattern)}
 
 let rec top_down_map_rec_stmt_loc f stmt =
   let recurse = top_down_map_rec_stmt_loc f in
@@ -250,7 +250,7 @@ let rec expr_var_set Expr.Fixed.({pattern; meta}) =
   match pattern with
   | Var s -> Set.Poly.singleton (VVar s, meta)
   | Lit _ -> Set.Poly.empty
-  | FunApp (_, exprs) -> union_recur exprs
+  | FunApp (kind, exprs) -> union_recur (exprs @ Fun_kind.collect_exprs kind)
   | TernaryIf (expr1, expr2, expr3) -> union_recur [expr1; expr2; expr3]
   | Indexed (expr, ix) ->
       Set.Poly.union_list (expr_var_set expr :: List.map ix ~f:index_var_set)
@@ -268,7 +268,8 @@ and index_var_set ix =
 let stmt_rhs stmt =
   match stmt with
   | Stmt.Fixed.Pattern.For vars -> Set.Poly.of_list [vars.lower; vars.upper]
-  | NRFunApp (_, exprs) -> Set.Poly.of_list exprs
+  | NRFunApp (kind, exprs) ->
+      Set.Poly.of_list (exprs @ Fun_kind.collect_exprs kind)
   | IfElse (rhs, _, _)
    |While (rhs, _)
    |Assignment (_, rhs)
@@ -354,7 +355,8 @@ let expr_subst_stmt m = map_rec_stmt_loc (expr_subst_stmt_base m)
 let rec expr_depth Expr.Fixed.({pattern; _}) =
   match pattern with
   | Var _ | Lit (_, _) -> 0
-  | FunApp (_, l) ->
+  | FunApp (kind, args) ->
+      let l = args @ Fun_kind.collect_exprs kind in
       1
       + Option.value ~default:0
           (List.max_elt ~compare:compare_int (List.map ~f:expr_depth l))
@@ -393,8 +395,11 @@ let rec update_expr_ad_levels autodiffable_variables
       else {e with meta= {e.meta with adlevel= DataOnly}}
   | Lit (_, _) -> {e with meta= {e.meta with adlevel= DataOnly}}
   | FunApp (kind, l) ->
+      let kind' =
+        Fun_kind.map (update_expr_ad_levels autodiffable_variables) kind
+      in
       let l = List.map ~f:(update_expr_ad_levels autodiffable_variables) l in
-      {pattern= FunApp (kind, l); meta= {e.meta with adlevel= ad_level_sup l}}
+      {pattern= FunApp (kind', l); meta= {e.meta with adlevel= ad_level_sup l}}
   | TernaryIf (e1, e2, e3) ->
       let e1 = update_expr_ad_levels autodiffable_variables e1 in
       let e2 = update_expr_ad_levels autodiffable_variables e2 in

src/analysis_and_optimization/Mir_utils.mli

@@ -9,7 +9,7 @@ type bound_values =
   { lower: [`None | `Nonlit | `Lit of float]
   ; upper: [`None | `Nonlit | `Lit of float] }
 
-val trans_bounds_values : Expr.Typed.t Program.transformation -> bound_values
+val trans_bounds_values : Expr.Typed.t Transformation.t -> bound_values
 val chop_dist_name : string -> string Option.t
 val top_var_declarations : Stmt.Located.t -> string Set.Poly.t
 
@@ -22,7 +22,7 @@ val data_set :
 val parameter_set :
      ?include_transformed:bool
   -> Program.Typed.t
-  -> (string * Expr.Typed.t Program.transformation) Set.Poly.t
+  -> (string * Expr.Typed.t Transformation.t) Set.Poly.t
 
 val parameter_names_set :
   ?include_transformed:bool -> Program.Typed.t -> string Set.Poly.t

src/analysis_and_optimization/Monotone_framework.ml

@@ -45,7 +45,9 @@ and free_vars_idx (i : Expr.Typed.t Index.t) =
   | Between (e1, e2) -> Set.Poly.union (free_vars_expr e1) (free_vars_expr e2)
 
 and free_vars_fnapp kind l =
-  let arg_vars = List.map ~f:free_vars_expr l in
+  let arg_vars =
+    List.map ~f:free_vars_expr (l @ Fun_kind.collect_exprs kind)
+  in
   match kind with
   | Fun_kind.UserDefined (f, _) ->
       Set.Poly.union_list (Set.Poly.singleton f :: List.map ~f:free_vars_expr l)
@@ -547,8 +549,9 @@ let rec used_subexpressions_expr (e : Expr.Typed.t) =
     (Expr.Typed.Set.singleton e)
     ( match e.pattern with
     | Var _ | Lit (_, _) -> Expr.Typed.Set.empty
-    | FunApp (_, l) ->
-        Expr.Typed.Set.union_list (List.map ~f:used_subexpressions_expr l)
+    | FunApp (k, l) ->
+        Expr.Typed.Set.union_list
+          (List.map ~f:used_subexpressions_expr (l @ Fun_kind.collect_exprs k))
     | TernaryIf (e1, e2, e3) ->
         Expr.Typed.Set.union_list
           [ used_subexpressions_expr e1
@@ -585,7 +588,8 @@ let rec used_expressions_stmt_help f
         [ f e
         ; used_expressions_stmt_help f b1.pattern
         ; used_expressions_stmt_help f b2.pattern ]
-  | NRFunApp (_, l) -> Expr.Typed.Set.union_list (List.map ~f l)
+  | NRFunApp (k, l) ->
+      Expr.Typed.Set.union_list (List.map ~f (l @ Fun_kind.collect_exprs k))
   | Decl _ | Return None | Break | Continue | Skip -> Expr.Typed.Set.empty
   | IfElse (e, b, None) | While (e, b) ->
       Expr.Typed.Set.union (f e) (used_expressions_stmt_help f b.pattern)
@@ -619,7 +623,8 @@ let top_used_expressions_stmt_help f
         (Expr.Typed.Set.union_list
            (List.map ~f:(used_expressions_idx_help f) l))
   | While (e, _) | IfElse (e, _, _) -> f e
-  | NRFunApp (_, l) -> Expr.Typed.Set.union_list (List.map ~f l)
+  | NRFunApp (k, l) ->
+      Expr.Typed.Set.union_list (List.map ~f (l @ Fun_kind.collect_exprs k))
   | Profile _ | Block _ | SList _ | Decl _
    |Return None
    |Break | Continue | Skip ->

src/analysis_and_optimization/Optimize.ml

@@ -690,13 +690,8 @@ and accum_any pred b e = b || expr_any pred e
 let can_side_effect_top_expr (e : Expr.Typed.t) =
   match e.pattern with
   | FunApp ((UserDefined (_, FnTarget) | StanLib (_, FnTarget)), _) -> true
-  | FunApp
-      ( CompilerInternal
-          ( FnReadParam _ | FnReadData | FnWriteParam | FnConstrain _
-          | FnValidateSize | FnValidateSizeSimplex | FnValidateSizeUnitVector
-          | FnUnconstrain _ )
-      , _ ) ->
-      true
+  | FunApp (CompilerInternal internal_fn, _) ->
+      Internal_fun.can_side_effect internal_fn
   | _ -> false
 
 let cannot_duplicate_expr (e : Expr.Typed.t) =

src/analysis_and_optimization/Pedantic_analysis.ml

@@ -238,7 +238,7 @@ let list_non_one_priors (fg : factor_graph) (mir : Program.Typed.t) :
 (* Collect useful information about an expression that's available at
    compile-time into a convenient form. *)
 let compiletime_value_of_expr
-    (params : (string * Expr.Typed.t Program.transformation) Set.Poly.t)
+    (params : (string * Expr.Typed.t Transformation.t) Set.Poly.t)
     (data : string Set.Poly.t) (expr : Expr.Typed.t) :
     compiletime_val * Expr.Typed.Meta.t =
   let v =

src/analysis_and_optimization/Pedantic_dist_warnings.ml

@@ -10,7 +10,7 @@ open Mir_utils
 type compiletime_val =
   | Opaque
   | Number of (float * string)
-  | Param of (string * Expr.Typed.t Program.transformation)
+  | Param of (string * Expr.Typed.t Transformation.t)
   | Data of string
 
 (* Info about a distribution occurrences that's useful for checking that
@@ -79,18 +79,17 @@ let bounds_out_of_range (range : range) (bounds : bound_values) : bool =
 (* Check for inconsistency between a distribution argument's constraint and the
    constraint transformation of a variable *)
 let transform_mismatch_constraint (constr : var_constraint)
-    (trans : Expr.Typed.t Program.transformation) : bool =
+    (trans : Expr.Typed.t Transformation.t) : bool =
   match constr with
   | Range range -> bounds_out_of_range range (trans_bounds_values trans)
-  | Ordered -> trans <> Program.Ordered
-  | PositiveOrdered -> trans <> Program.PositiveOrdered
-  | Simplex -> trans <> Program.Simplex
-  | UnitVector -> trans <> Program.UnitVector
-  | CholeskyCorr -> trans <> Program.CholeskyCorr
-  | CholeskyCov ->
-      trans <> Program.CholeskyCov && trans <> Program.CholeskyCorr
-  | Correlation -> trans <> Program.Correlation
-  | Covariance -> trans <> Program.Covariance && trans <> Program.Correlation
+  | Ordered -> trans <> Transformation.Ordered
+  | PositiveOrdered -> trans <> PositiveOrdered
+  | Simplex -> trans <> Simplex
+  | UnitVector -> trans <> UnitVector
+  | CholeskyCorr -> trans <> CholeskyCorr
+  | CholeskyCov -> trans <> CholeskyCov && trans <> CholeskyCorr
+  | Correlation -> trans <> Correlation
+  | Covariance -> trans <> Covariance && trans <> Correlation
 
 (* Check for inconsistency between a distribution argument's range and
    a literal value *)

src/frontend/Ast.ml

@@ -150,7 +150,7 @@ type ('e, 's, 'l, 'f) statement =
   | Block of 's list
   | VarDecl of
       { decl_type: 'e Middle.Type.t
-      ; transformation: 'e Middle.Program.transformation
+      ; transformation: 'e Transformation.t
       ; identifier: identifier
       ; initial_value: 'e option
       ; is_global: bool }