Add support for @annotations

src/analysis_and_optimization/Memory_patterns.ml

@@ -563,12 +563,14 @@ let rec modify_stmt_pattern
       ; initialize=
           Assign
             ({ pattern= FunApp (CompilerInternal (FnReadParam read_param), args)
-             ; _ } as assigner) } ->
+             ; _ } as assigner)
+      ; decl_annotations } ->
       let name = decl_id in
       if Set.mem modifiable_set name then
         Stmt.Fixed.Pattern.Decl
           { decl_id
           ; decl_adtype
+          ; decl_annotations
           ; decl_type=
               Type.Sized (SizedType.modify_sizedtype_mem AoS sized_type)
           ; initialize=
@@ -583,6 +585,7 @@ let rec modify_stmt_pattern
         Stmt.Fixed.Pattern.Decl
           { decl_id
           ; decl_adtype
+          ; decl_annotations
           ; decl_type=
               Type.Sized (SizedType.modify_sizedtype_mem SoA sized_type)
           ; initialize=

src/analysis_and_optimization/Optimize.ml

@@ -84,13 +84,18 @@ let gen_inline_var (name : string) (id_var : string) =
 
 let replace_fresh_local_vars (fname : string) stmt =
   let f (m : (string, string) Core.Map.Poly.t) = function
-    | Stmt.Fixed.Pattern.Decl {decl_adtype; decl_type; decl_id; initialize} ->
+    | Stmt.Fixed.Pattern.Decl
+        {decl_adtype; decl_type; decl_id; decl_annotations; initialize} ->
         let new_name =
           match Map.Poly.find m decl_id with
           | Some existing -> existing
           | None -> gen_inline_var fname decl_id in
         ( Stmt.Fixed.Pattern.Decl
-            {decl_adtype; decl_id= new_name; decl_type; initialize}
+            { decl_adtype
+            ; decl_id= new_name
+            ; decl_type
+            ; decl_annotations
+            ; initialize }
         , Map.Poly.set m ~key:decl_id ~data:new_name )
     | Stmt.Fixed.Pattern.For {loopvar; lower; upper; body} ->
         let new_name =
@@ -201,6 +206,7 @@ let handle_early_returns (fname : string) opt_var stmt =
                 { decl_adtype= DataOnly
                 ; decl_id= returned
                 ; decl_type= Sized SInt
+                ; decl_annotations= []
                 ; initialize= Default }
           ; meta= Location_span.empty }
       ; Stmt.Fixed.
@@ -294,6 +300,7 @@ let rec inline_function_expression propto adt fim (Expr.Fixed.{pattern; _} as e)
                             (Type.to_unsized decl_type)
                       ; decl_id= inline_return_name
                       ; decl_type
+                      ; decl_annotations= []
                       ; initialize= Uninit } ]
                   (* We should minimize the code that's having its variables
                      replaced to avoid conflict with the (two) new dummy
@@ -464,10 +471,20 @@ let rec inline_function_statement propto adt fim Stmt.Fixed.{pattern; meta} =
             Block (List.map l ~f:(inline_function_statement propto adt fim))
         | SList l ->
             SList (List.map l ~f:(inline_function_statement propto adt fim))
-        | Decl {decl_adtype; decl_id; decl_type; initialize= Assign expr} ->
+        | Decl
+            { decl_adtype
+            ; decl_id
+            ; decl_type
+            ; decl_annotations
+            ; initialize= Assign expr } ->
             let d, s, e = inline_function_expression propto adt fim expr in
             slist_concat_no_loc (d @ s)
-              (Decl {decl_adtype; decl_id; decl_type; initialize= Assign e})
+              (Decl
+                 { decl_adtype
+                 ; decl_id
+                 ; decl_type
+                 ; decl_annotations
+                 ; initialize= Assign e })
         | Decl r -> Decl r
         | Skip -> Skip
         | Break -> Break
@@ -981,6 +998,7 @@ let lazy_code_motion ?(preserve_stability = false) (mir : Program.Typed.t) =
                   { decl_adtype= Expr.Typed.adlevel_of key
                   ; decl_id= data
                   ; decl_type= Type.Unsized (Expr.Typed.type_of key)
+                  ; decl_annotations= [] (* TODO annotations: correct? *)
                   ; initialize= Default }
             ; meta= Location_span.empty }
           :: accum) in

src/frontend/Annotations.ml

@@ -0,0 +1,21 @@
+open Ast
+open Core
+
+let message a = "Unknown annotation '" ^ a ^ "' will be ignored by the compiler"
+
+let unused_list pred id anns =
+  List.filter_map
+    ~f:(fun a -> if pred a then None else Some (id.id_loc, message a))
+    anns
+
+let rec collect_stmt pred (acc : Warnings.t list) {stmt; _} =
+  match stmt with
+  | FunDef {annotations; funname; _} ->
+      acc @ (unused_list pred funname) annotations
+  | VarDecl {annotations; variables; _} ->
+      acc @ (unused_list pred (List.hd_exn variables).identifier) annotations
+  | _ -> fold_statement Fn.const (collect_stmt pred) Fn.const Fn.const acc stmt
+
+let find_unrecognized (pred : string -> bool) (prog : Ast.untyped_program) :
+    Warnings.t list =
+  fold_program (collect_stmt pred) [] prog

src/frontend/Annotations.mli

@@ -0,0 +1,2 @@
+val find_unrecognized :
+  (string -> bool) -> Ast.untyped_program -> Warnings.t list

src/frontend/Ast.ml

@@ -171,13 +171,15 @@ type ('e, 's, 'l, 'f) statement =
       { decl_type: 'e SizedType.t
       ; transformation: 'e Transformation.t
       ; is_global: bool
+      ; annotations: string list
       ; variables: 'e variable list }
   | FunDef of
       { returntype: UnsizedType.returntype
       ; funname: identifier
       ; arguments:
           (Middle.UnsizedType.autodifftype * Middle.UnsizedType.t * identifier)
           list
+      ; annotations: string list
       ; body: 's }
 [@@deriving sexp, hash, compare, map, fold]
 

src/frontend/Ast_to_Mir.ml

@@ -458,8 +458,8 @@ let var_constrain_check_stmts dconstrain loc adlevel decl_id decl_var trans
       @ check_decl decl_var st trans loc adlevel
   | _ -> []
 
-let create_decl_with_assign decl_id declc decl_type initial_value transform
-    smeta =
+let create_decl_with_assign decl_id declc decl_type initial_value
+    decl_annotations transform smeta =
   let rhs = Option.map ~f:trans_expr initial_value in
   let decl_adtype =
     UnsizedType.fill_adtype_for_type declc.dadlevel (Type.to_unsized decl_type)
@@ -474,7 +474,12 @@ let create_decl_with_assign decl_id declc decl_type initial_value transform
   let decl =
     Stmt.
       { Fixed.pattern=
-          Decl {decl_adtype; decl_id; decl_type; initialize= Default}
+          Decl
+            { decl_adtype
+            ; decl_id
+            ; decl_type
+            ; decl_annotations
+            ; initialize= Default }
       ; meta= smeta } in
   let rhs_assignment =
     Option.map
@@ -599,6 +604,7 @@ let rec trans_stmt ud_dists (declc : decl_context) (ts : Ast.typed_statement) =
                 { decl_adtype= Expr.Typed.adlevel_of iteratee'
                 ; decl_id= loopvar.name
                 ; decl_type= Unsized decl_type
+                ; decl_annotations= []
                 ; initialize= Default } } in
       let assignment var =
         Stmt.Fixed.
@@ -614,15 +620,16 @@ let rec trans_stmt ud_dists (declc : decl_context) (ts : Ast.typed_statement) =
       Common.ICE.internal_compiler_error
         [%message
           "Found function definition statement outside of function block"]
-  | Ast.VarDecl {decl_type; transformation; variables; is_global= _} ->
+  | Ast.VarDecl {decl_type; transformation; variables; annotations; is_global= _}
+    ->
       List.concat_map
         ~f:(fun {identifier; initial_value} ->
           let transform = Transformation.map trans_expr transformation in
           let decl_id = identifier.Ast.name in
           let size_checks, dt = check_sizedtype decl_id decl_type in
           size_checks
-          @ create_decl_with_assign decl_id declc dt initial_value transform
-              smeta)
+          @ create_decl_with_assign decl_id declc dt initial_value annotations
+              transform smeta)
         variables
   | Ast.Block stmts -> Block (List.concat_map ~f:trans_stmt stmts) |> swrap
   | Ast.Profile (name, stmts) ->
@@ -645,6 +652,7 @@ and trans_packed_assign loc trans_stmt lvals rhs assign_op =
           { decl_adtype= rhs.emeta.ad_level
           ; decl_id= sym
           ; decl_type= Unsized rhs_type
+          ; decl_annotations= []
           ; initialize= Uninit }
     ; meta= rhs.emeta.loc } in
   let assign =
@@ -712,12 +720,13 @@ and trans_single_assignment smeta assign_lhs assign_rhs assign_op =
 
 let trans_fun_def ud_dists (ts : Ast.typed_statement) =
   match ts.stmt with
-  | Ast.FunDef {returntype; funname; arguments; body} ->
+  | Ast.FunDef {returntype; funname; arguments; annotations; body} ->
       [ Program.
           { fdrt= returntype
           ; fdname= funname.name
           ; fdsuffix= Fun_kind.(suffix_from_name funname.name |> without_propto)
           ; fdargs= List.map ~f:trans_arg arguments
+          ; fdannotations= annotations
           ; fdbody=
               trans_stmt ud_dists
                 {transform_action= IgnoreTransform; dadlevel= AutoDiffable}
@@ -759,6 +768,7 @@ let rec trans_sizedtype_decl declc tr name st =
                 { decl_type= Sized SInt
                 ; decl_id
                 ; decl_adtype= DataOnly
+                ; decl_annotations= []
                 ; initialize= Default }
           ; meta= e.meta.loc } in
         let assign =
@@ -837,7 +847,12 @@ let trans_block ud_dists declc block prog =
   let f stmt (accum1, accum2, accum3) =
     match stmt with
     | { Ast.stmt=
-          VarDecl {decl_type= type_; variables; transformation; is_global= true}
+          VarDecl
+            { decl_type= type_
+            ; variables
+            ; transformation
+            ; annotations
+            ; is_global= true }
       ; smeta } ->
         let outvars, sizes, stmts =
           List.unzip3
@@ -855,10 +870,11 @@ let trans_block ud_dists declc block prog =
                        { out_constrained_st= type_
                        ; out_unconstrained_st= param_size transform type_
                        ; out_block= block
+                       ; out_annotations= annotations
                        ; out_trans= transform } ) in
                  let stmts =
                    create_decl_with_assign decl_id declc (Sized type_)
-                     initial_value transform smeta.loc in
+                     initial_value annotations transform smeta.loc in
                  (outvar, size, stmts))
                variables in
         ( outvars @ accum1

src/frontend/Canonicalize.ml

@@ -65,11 +65,13 @@ let parens_lval = map_lval_with no_parens Fn.id
 let rec parens_stmt ({stmt; smeta} : typed_statement) : typed_statement =
   let stmt =
     match stmt with
-    | VarDecl {decl_type= d; transformation= t; variables; is_global} ->
+    | VarDecl
+        {decl_type= d; transformation= t; variables; annotations; is_global} ->
         VarDecl
           { decl_type= Middle.SizedType.map no_parens d
           ; transformation= Middle.Transformation.map keep_parens t
           ; variables= List.map ~f:(map_variable no_parens) variables
+          ; annotations
           ; is_global }
     | For {loop_variable; lower_bound; upper_bound; loop_body} ->
         For

src/frontend/Environment.ml

@@ -21,6 +21,7 @@
   ; kind:
       [ `Variable of varinfo
       | `UserDeclared of Location_span.t
+      | `UserExtern of Location_span.t
       | `StanMath
       | `UserDefined ] }
 [@@deriving sexp]

src/frontend/Environment.mli

@@ -21,6 +21,7 @@ type info =
   ; kind:
       [ `Variable of varinfo
       | `UserDeclared of Location_span.t
+      | `UserExtern of Location_span.t
       | `StanMath
       | `UserDefined ] }
 
@@ -37,6 +38,7 @@ val add :
   -> string
   -> Middle.UnsizedType.t
   -> [ `UserDeclared of Location_span.t
+     | `UserExtern of Location_span.t
      | `StanMath
      | `UserDefined
      | `Variable of varinfo ]

src/frontend/Pretty_printing.ml

@@ -387,6 +387,13 @@ let rec pp_transformed_type ppf (st, trans) =
         pp_possibly_transformed_type (ty, trans)
   | _ -> pf ppf "%a" pp_possibly_transformed_type (st, trans)
 
+let pp_annotations ppf ann =
+  (* TODO better comment handling? *)
+  if List.is_empty ann then ()
+  else
+    let pp ppf s = pf ppf "%@%s" s in
+    pf ppf "%a@ " (list ~sep:sp pp) ann
+
 let rec pp_indent_unless_block ppf ((s : untyped_statement), loc) =
   match s.stmt with
   | Block _ -> pp_statement ppf s
@@ -453,16 +460,23 @@ and pp_statement ppf ({stmt= s_content; smeta= {loc}} as ss : untyped_statement)
       pf ppf "profile(%s) {@,%a@,}" name
         (indented_box pp_list_of_statements)
         (vdsl, loc)
-  | VarDecl {decl_type= pst; transformation= trans; variables; is_global= _} ->
+  | VarDecl
+      { decl_type= pst
+      ; transformation= trans
+      ; variables
+      ; annotations
+      ; is_global= _ } ->
       let pp_var ppf {identifier; initial_value} =
         pf ppf "%a%a" pp_identifier identifier
           (option (fun ppf e -> pf ppf " = %a" pp_expression e))
           initial_value in
-      pf ppf "@[<h>%a %a;@]" pp_transformed_type (pst, trans)
-        (list ~sep:comma pp_var) variables
-  | FunDef {returntype= rt; funname= id; arguments= args; body= b} -> (
+      pf ppf "@[<hv>%a@[<h>%a %a;@]@]" pp_annotations annotations
+        pp_transformed_type (pst, trans) (list ~sep:comma pp_var) variables
+  | FunDef {returntype= rt; funname= id; arguments= args; annotations; body= b}
+    -> (
       let loc_of (_, _, id) = id.id_loc in
-      pf ppf "%a %a(%a" pp_returntype rt pp_identifier id
+      pf ppf "@[<hv>%a@[<h>%a %a(@]%a@]" pp_annotations annotations
+        pp_returntype rt pp_identifier id
         (box (pp_list_of pp_args loc_of))
         (args, {loc with end_loc= b.smeta.loc.begin_loc});
       match b with