migrate Debug_data_generation to MIR-only

Author: nhuurre

src/analysis_and_optimization/Debug_data_generation.ml

@@ -1,7 +1,5 @@
 open Core_kernel
 open Middle
-open Frontend
-open Ast
 
 let rec transpose = function
   | [] :: _ -> []
@@ -30,15 +28,16 @@ let rec vect_to_mat l m =
     let hd, tl = List.split_n l m in
     hd :: vect_to_mat tl m
 
-let unwrap_num_exn m e =
-  let e = Ast_to_Mir.trans_expr e in
-  let m = Map.Poly.map m ~f:Ast_to_Mir.trans_expr in
+let eval_expr m e =
   let e = Mir_utils.subst_expr m e in
   let e = Partial_evaluator.eval_expr e in
   let rec strip_promotions (e : Middle.Expr.Typed.t) =
     match e.pattern with Promotion (e, _, _) -> strip_promotions e | _ -> e
   in
-  let e = strip_promotions e in
+  strip_promotions e
+
+let unwrap_num_exn m e =
+  let e = eval_expr m e in
   match e.pattern with
   | Lit (_, s) -> Float.of_string s
   | _ ->
@@ -74,92 +73,56 @@ let rec repeat n e =
 let rec repeat_th n f =
   match n with n when n <= 0 -> [] | m -> f () :: repeat_th (m - 1) f
 
-let wrap_int n =
-  { expr= IntNumeral (Int.to_string n)
-  ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= UInt} }
-
-let int_two = wrap_int 2
-
-let wrap_real r =
-  { expr= RealNumeral (Float.to_string r)
-  ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= UReal} }
-
-let wrap_row_vector l =
-  { expr= RowVectorExpr l
-  ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= URowVector} }
-
-let wrap_vector l =
-  { expr= PostfixOp (wrap_row_vector l, Transpose)
-  ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= UVector} }
-
-let gen_int m t = wrap_int (gen_num_int m t)
-let gen_real m t = wrap_real (gen_num_real m t)
-
 let gen_row_vector m n t =
-  let extract_var e =
-    match e with {expr= Variable x; _} -> Map.find_exn m x.name | _ -> e in
   let gen_bounded t e =
-    match e with
-    | {expr= RowVectorExpr unpacked_e; _}
-     |{expr= ArrayExpr unpacked_e; _}
-     |{expr= PostfixOp ({expr= RowVectorExpr unpacked_e; _}, Transpose); _} ->
-        wrap_row_vector (List.map ~f:(fun x -> gen_real m (t x)) unpacked_e)
-    | _ ->
+    match Expr.Helpers.try_unpack e with
+    | Some unpacked_e ->
+        Expr.Helpers.row_vector
+          (List.map ~f:(fun x -> gen_num_real m (t x)) unpacked_e)
+    | None ->
         Common.FatalError.fatal_error_msg
           [%message
             "Bad bounded (upper OR lower) expr: "
-              (e : (typed_expr_meta, fun_kind) expr_with)] in
+              (e : Expr.Typed.Meta.t Expr.Fixed.t)] in
   let gen_ul_bounded e1 e2 =
     let create_bounds l u =
-      wrap_row_vector
+      Expr.Helpers.row_vector
         (List.map2_exn
-           ~f:(fun x y -> gen_real m (Transformation.LowerUpper (x, y)))
+           ~f:(fun x y -> gen_num_real m (Transformation.LowerUpper (x, y)))
            l u ) in
-    match (e1, e2) with
-    | ( ( {expr= RowVectorExpr unpacked_e1 | ArrayExpr unpacked_e1; _}
-        | {expr= PostfixOp ({expr= RowVectorExpr unpacked_e1; _}, Transpose); _}
-          )
-      , ( {expr= RowVectorExpr unpacked_e2 | ArrayExpr unpacked_e2; _}
-        | {expr= PostfixOp ({expr= RowVectorExpr unpacked_e2; _}, Transpose); _}
-          ) ) ->
-        (* | {expr= ArrayExpr unpacked_e1; _}, {expr= ArrayExpr unpacked_e2; _} -> *)
+    match Expr.Helpers.(try_unpack e1, try_unpack e2) with
+    | Some unpacked_e1, Some unpacked_e2 ->
         create_bounds unpacked_e1 unpacked_e2
-    | ( ({expr= RealNumeral _; _} | {expr= IntNumeral _; _})
-      , ( {expr= RowVectorExpr unpacked_e2; _}
-        | {expr= ArrayExpr unpacked_e2; _}
-        | {expr= PostfixOp ({expr= RowVectorExpr unpacked_e2; _}, Transpose); _}
-          ) ) ->
+    | None, Some unpacked_e2 ->
         create_bounds
           (List.init (List.length unpacked_e2) ~f:(fun _ -> e1))
           unpacked_e2
-    | ( ( {expr= RowVectorExpr unpacked_e1; _}
-        | {expr= PostfixOp ({expr= RowVectorExpr unpacked_e1; _}, Transpose); _}
-        | {expr= ArrayExpr unpacked_e1; _} )
-      , ({expr= RealNumeral _; _} | {expr= IntNumeral _; _}) ) ->
+    | Some unpacked_e1, None ->
         create_bounds unpacked_e1
           (List.init (List.length unpacked_e1) ~f:(fun _ -> e2))
     | _ ->
         Common.FatalError.fatal_error_msg
           [%message
             "Bad bounded upper and lower expr: "
-              (e1 : (typed_expr_meta, fun_kind) expr_with)
+              (e1 : Expr.Typed.t)
               " and "
-              (e2 : (typed_expr_meta, fun_kind) expr_with)] in
-  match t with
-  | Transformation.Lower ({emeta= {type_= UVector | URowVector; _}; _} as e) ->
-      gen_bounded (fun x -> Transformation.Lower x) (extract_var e)
-  | Transformation.Upper ({emeta= {type_= UVector | URowVector; _}; _} as e) ->
-      gen_bounded (fun x -> Transformation.Upper x) (extract_var e)
+              (e2 : Expr.Typed.t)] in
+  match (t : Expr.Typed.t Transformation.t) with
+  | Transformation.Lower ({meta= {type_= UVector | URowVector; _}; _} as e) ->
+      gen_bounded (fun x -> Transformation.Lower x) (eval_expr m e)
+  | Transformation.Upper ({meta= {type_= UVector | URowVector; _}; _} as e) ->
+      gen_bounded (fun x -> Transformation.Upper x) (eval_expr m e)
   | Transformation.LowerUpper
-      ( ({emeta= {type_= UVector | URowVector | UReal | UInt; _}; _} as e1)
-      , ({emeta= {type_= UVector | URowVector; _}; _} as e2) )
+      ( ({meta= {type_= UVector | URowVector | UReal | UInt; _}; _} as e1)
+      , ({meta= {type_= UVector | URowVector; _}; _} as e2) )
    |Transformation.LowerUpper
-      ( ({emeta= {type_= UVector | URowVector; _}; _} as e1)
-      , ({emeta= {type_= UReal | UInt; _}; _} as e2) ) ->
-      gen_ul_bounded (extract_var e1) (extract_var e2)
+      ( ({meta= {type_= UVector | URowVector; _}; _} as e1)
+      , ({meta= {type_= UReal | UInt; _}; _} as e2) ) ->
+      gen_ul_bounded (eval_expr m e1) (eval_expr m e2)
   | _ ->
-      { expr= RowVectorExpr (repeat_th n (fun _ -> gen_real m t))
-      ; emeta= {loc= Location_span.empty; ad_level= DataOnly; type_= UMatrix} }
+      let e =
+        Expr.Helpers.row_vector (repeat_th n (fun _ -> gen_num_real m t)) in
+      {e with meta= {e.meta with type_= UMatrix}}
 
 let gen_vector m n t =
   let gen_ordered n =
@@ -173,38 +136,36 @@ let gen_vector m n t =
       let l = repeat_th n (fun _ -> Random.float 1.) in
       let sum = List.fold l ~init:0. ~f:(fun accum elt -> accum +. elt) in
       let l = List.map l ~f:(fun x -> x /. sum) in
-      wrap_vector (List.map ~f:wrap_real l)
+      Expr.Helpers.vector l
   | Ordered ->
       let l = gen_ordered n in
       let halfmax =
         Option.value_exn (List.max_elt l ~compare:compare_float) /. 2. in
       let l = List.map l ~f:(fun x -> (x -. halfmax) /. halfmax) in
-      wrap_vector (List.map ~f:wrap_real l)
+      Expr.Helpers.vector l
   | PositiveOrdered ->
       let l = gen_ordered n in
       let max = Option.value_exn (List.max_elt l ~compare:compare_float) in
       let l = List.map l ~f:(fun x -> x /. max) in
-      wrap_vector (List.map ~f:wrap_real l)
+      Expr.Helpers.vector l
   | UnitVector ->
       let l = repeat_th n (fun _ -> Random.float 1.) in
       let sum =
         Float.sqrt
           (List.fold l ~init:0. ~f:(fun accum elt -> accum +. (elt ** 2.)))
       in
       let l = List.map l ~f:(fun x -> x /. sum) in
-      wrap_vector (List.map ~f:wrap_real l)
-  | _ -> {int_two with expr= PostfixOp (gen_row_vector m n t, Transpose)}
+      Expr.Helpers.vector l
+  | _ ->
+      let v = Expr.Helpers.unary_op Transpose (gen_row_vector m n t) in
+      {v with meta= {v.meta with type_= UVector}}
 
 let gen_cov_unwrapped n =
   let l = repeat_th (n * n) (fun _ -> Random.float 2.) in
   let l_mat = vect_to_mat l n in
   matprod l_mat (transpose l_mat)
 
-let wrap_real_mat m =
-  let mat_wrapped =
-    List.map ~f:wrap_row_vector
-      (List.map ~f:(fun x -> List.map ~f:wrap_real x) m) in
-  {int_two with expr= RowVectorExpr mat_wrapped}
+let wrap_real_mat m = Expr.Helpers.matrix m
 
 let gen_diag_mat l =
   let n = List.length l in
@@ -263,20 +224,18 @@ let gen_matrix mm m n t =
   | CholeskyCov -> gen_cov_cholesky m n
   | CholeskyCorr -> gen_corr_cholesky m
   | _ ->
-      { int_two with
-        expr= RowVectorExpr (repeat_th m (fun () -> gen_row_vector mm n t)) }
-
-(* TODO: do some proper random generation of these special matrices *)
+      Expr.Helpers.matrix_from_rows
+        (repeat_th m (fun () -> gen_row_vector mm n t))
 
-let gen_array elt n _ = {int_two with expr= ArrayExpr (repeat_th n elt)}
+let gen_array elt n _ = Expr.Helpers.array_expr (repeat_th n elt)
 
 let rec generate_value m st t =
   match st with
-  | SizedType.SInt -> gen_int m t
-  | SReal -> gen_real m t
+  | SizedType.SInt -> Expr.Helpers.int (gen_num_int m t)
+  | SReal -> Expr.Helpers.float (gen_num_real m t)
   | SComplex ->
       (* when serialzied, a complex number looks just like a 2-array of reals *)
-      generate_value m (SArray (SReal, wrap_int 2)) t
+      generate_value m (SArray (SReal, Expr.Helpers.int 2)) t
   | SVector (_, e) -> gen_vector m (unwrap_int_exn m e) t
   | SRowVector (_, e) -> gen_row_vector m (unwrap_int_exn m e) t
   | SMatrix (_, e1, e2) ->
@@ -286,28 +245,23 @@ let rec generate_value m st t =
       gen_array element (unwrap_int_exn m e) t
 
 let rec pp_value_json ppf e =
-  match e.expr with
-  | PostfixOp (e, Transpose) -> pp_value_json ppf e
-  | IntNumeral s | RealNumeral s -> Fmt.string ppf s
-  | ArrayExpr l | RowVectorExpr l ->
+  match e.Expr.Fixed.pattern with
+  | Lit ((Int | Real), s) -> Fmt.string ppf s
+  | FunApp (CompilerInternal (FnMakeRowVec | FnMakeArray), l) ->
       Fmt.(pf ppf "[@[<hov 1>%a@]]" (list ~sep:comma pp_value_json) l)
+  | FunApp (StanLib (transpose, _, _), [e])
+    when String.equal transpose (Operator.to_string Transpose) ->
+      pp_value_json ppf e
   | _ ->
       Common.FatalError.fatal_error_msg
-        [%message "Could not evaluate expression " (e : typed_expression)]
+        [%message "Could not evaluate expression " (e : Expr.Typed.t)]
 
 let print_data_prog s =
-  let data = Ast.get_stmts s.datablock in
   let l, _ =
-    List.fold data ~init:([], Map.Poly.empty) ~f:(fun (l, m) decl ->
-        match decl.stmt with
-        | VarDecl
-            { decl_type= Sized sizedtype
-            ; transformation
-            ; identifier= {name; _}
-            ; _ } ->
-            let value = generate_value m sizedtype transformation in
-            ((name, value) :: l, Map.set m ~key:name ~data:value)
-        | _ -> (l, m) ) in
+    List.fold s ~init:([], Map.Poly.empty)
+      ~f:(fun (l, m) (sizedtype, transformation, name) ->
+        let value = generate_value m sizedtype transformation in
+        ((name, value) :: l, Map.set m ~key:name ~data:value) ) in
   let pp ppf (id, value) =
     Fmt.pf ppf {|@[<hov 2>"%s":@ %a@]|} id pp_value_json value in
   Fmt.(str "{@ @[<hov>%a@]@ }" (list ~sep:comma pp) (List.rev l))

src/analysis_and_optimization/Debug_data_generation.mli

@@ -1 +1,5 @@
-val print_data_prog : Frontend.Ast.typed_program -> string
+open Middle
+
+val print_data_prog :
+     (Expr.Typed.t SizedType.t * Expr.Typed.t Transformation.t * string) list
+  -> string

src/frontend/Ast_to_Mir.ml

@@ -679,6 +679,22 @@ let migrate_checks_to_end_of_block stmts =
   let checks, not_checks = List.partition_tf ~f:stmt_contains_check stmts in
   not_checks @ checks
 
+let trans_data (p : Ast.typed_program) =
+  let data = Ast.get_stmts p.datablock in
+  List.filter_map data ~f:(function
+    | { stmt=
+          VarDecl
+            { decl_type= Sized sizedtype
+            ; transformation
+            ; identifier= {name; _}
+            ; _ }
+      ; _ } ->
+        Some
+          ( SizedType.map trans_expr sizedtype
+          , Transformation.map trans_expr transformation
+          , name )
+    | _ -> None )
+
 let trans_prog filename (p : Ast.typed_program) : Program.Typed.t =
   let {Ast.functionblock; datablock; transformeddatablock; modelblock; _} = p in
   let map f list_op =

src/frontend/Ast_to_Mir.mli

@@ -1,4 +1,11 @@
 (** Translate from the AST to the MIR *)
+open Middle
 
-val trans_prog : string -> Ast.typed_program -> Middle.Program.Typed.t
-val trans_expr : Ast.typed_expression -> Middle.Expr.Typed.t
+val trans_data :
+     Ast.typed_program
+  -> ( Expr.Typed.Meta.t Expr.Fixed.t SizedType.t
+     * Expr.Typed.Meta.t Expr.Fixed.t Transformation.t
+     * string )
+     list
+
+val trans_prog : string -> Ast.typed_program -> Program.Typed.t

src/middle/Expr.ml

@@ -1,5 +1,4 @@
 open Core_kernel
-open Core_kernel.Poly
 open Common
 open Helpers
 
@@ -163,13 +162,18 @@ module Helpers = struct
   let int i = {Fixed.meta= Typed.Meta.empty; pattern= Lit (Int, string_of_int i)}
 
   let float i =
-    {Fixed.meta= Typed.Meta.empty; pattern= Lit (Real, string_of_float i)}
+    { Fixed.meta= {Typed.Meta.empty with type_= UReal}
+    ; pattern= Lit (Real, Float.to_string i) }
 
   let str i = {Fixed.meta= Typed.Meta.empty; pattern= Lit (Str, i)}
   let variable v = {Fixed.meta= Typed.Meta.empty; pattern= Var v}
   let zero = int 0
   let one = int 1
 
+  let unary_op op e =
+    { Fixed.meta= Typed.Meta.empty
+    ; pattern= FunApp (StanLib (Operator.to_string op, FnPlain, AoS), [e]) }
+
   let binop e1 op e2 =
     { Fixed.meta= Typed.Meta.empty
     ; pattern= FunApp (StanLib (Operator.to_string op, FnPlain, AoS), [e1; e2])
@@ -181,6 +185,41 @@ module Helpers = struct
     | head :: rest ->
         List.fold ~init:head ~f:(fun accum next -> binop accum op next) rest
 
+  let row_vector l =
+    { Fixed.meta= {Typed.Meta.empty with type_= URowVector}
+    ; pattern= FunApp (CompilerInternal FnMakeRowVec, List.map ~f:float l) }
+
+  let vector l =
+    let v = unary_op Transpose (row_vector l) in
+    {v with meta= {Typed.Meta.empty with type_= UVector}}
+
+  let matrix l =
+    { Fixed.meta= {Typed.Meta.empty with type_= UMatrix}
+    ; pattern= FunApp (CompilerInternal FnMakeRowVec, List.map ~f:row_vector l)
+    }
+
+  let matrix_from_rows l =
+    { Fixed.meta= {Typed.Meta.empty with type_= UMatrix}
+    ; pattern= FunApp (CompilerInternal FnMakeRowVec, l) }
+
+  let array_expr l =
+    let type_ =
+      List.hd l |> Option.value_map ~f:Typed.type_of ~default:UnsizedType.UReal
+    in
+    { Fixed.meta= {Typed.Meta.empty with type_= UArray type_}
+    ; pattern= FunApp (CompilerInternal FnMakeArray, l) }
+
+  let try_unpack e =
+    (* FIXME: what about matrices? *)
+    match e.Fixed.pattern with
+    | FunApp (CompilerInternal (FnMakeRowVec | FnMakeArray), l) -> Some l
+    | FunApp
+        ( StanLib (transpose, FnPlain, _)
+        , [{pattern= FunApp (CompilerInternal FnMakeRowVec, l); _}] )
+      when String.equal transpose (Operator.to_string Transpose) ->
+        Some l
+    | _ -> None
+
   let loop_bottom = one
 
   let internal_funapp fn args meta =
@@ -197,7 +236,9 @@ module Helpers = struct
 
   let%test "expr contains fn" =
     internal_funapp FnReadData [] ()
-    |> contains_fn_kind (fun kind -> kind = CompilerInternal FnReadData)
+    |> contains_fn_kind (function
+         | CompilerInternal FnReadData -> true
+         | _ -> false )
 
   let rec infer_type_of_indexed ut indices =
     match (ut, indices) with