Initial annotation parsing

Author: WardBrian

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= true }
           ; 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= false } ]
                   (* We should minimize the code that's having its variables
                      replaced to avoid conflict with the (two) new dummy
@@ -972,6 +979,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= true }
             ; meta= Location_span.empty }
           :: accum) in

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

@@ -443,8 +443,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)
@@ -458,7 +458,9 @@ let create_decl_with_assign decl_id declc decl_type initial_value transform
             () } in
   let decl =
     Stmt.
-      { Fixed.pattern= Decl {decl_adtype; decl_id; decl_type; initialize= true}
+      { Fixed.pattern=
+          Decl
+            {decl_adtype; decl_id; decl_type; decl_annotations; initialize= true}
       ; meta= smeta } in
   let rhs_assignment =
     Option.map
@@ -583,6 +585,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= true } } in
       let assignment var =
         Stmt.Fixed.
@@ -598,15 +601,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) ->
@@ -629,6 +633,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= false }
     ; meta= rhs.emeta.loc } in
   let assign =
@@ -696,12 +701,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}
@@ -743,6 +749,7 @@ let rec trans_sizedtype_decl declc tr name st =
                 { decl_type= Sized SInt
                 ; decl_id
                 ; decl_adtype= DataOnly
+                ; decl_annotations= []
                 ; initialize= true }
           ; meta= e.meta.loc } in
         let assign =
@@ -821,7 +828,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
@@ -839,10 +851,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/Pretty_printing.ml

@@ -382,6 +382,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
@@ -448,16 +455,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

src/frontend/Typechecker.ml

@@ -1659,7 +1659,7 @@ and check_transformation cf tenv ut trans =
       TupleTransformation tes
 
 and check_var_decl loc cf tenv sized_ty trans
-    (variables : untyped_expression Ast.variable list) is_global =
+    (variables : untyped_expression Ast.variable list) annotations is_global =
   let checked_type =
     check_sizedtype {cf with in_toplevel_decl= is_global} tenv sized_ty in
   let unsized_type = SizedType.to_unsized checked_type in
@@ -1684,6 +1684,7 @@ and check_var_decl loc cf tenv sized_ty trans
       { decl_type= checked_type
       ; transformation= checked_trans
       ; variables= tvariables
+      ; annotations
       ; is_global } in
   (tenv, mk_typed_statement ~stmt ~loc ~return_type:Incomplete)
 
@@ -1786,7 +1787,7 @@ and add_function tenv name type_ defined =
     Env.set_raw tenv name (new_fn :: defns)
   else Env.add tenv name type_ defined
 
-and check_fundef loc cf tenv return_ty id args body =
+and check_fundef loc cf tenv return_ty id args annotations body =
   List.iter args ~f:(fun (_, _, id) -> verify_identifier id);
   verify_identifier id;
   let arg_types = List.map ~f:(fun (w, y, _) -> (w, y)) args in
@@ -1838,8 +1839,11 @@ and check_fundef loc cf tenv return_ty id args body =
   verify_fundef_return_tys loc return_ty checked_body;
   let stmt =
     FunDef
-      {returntype= return_ty; funname= id; arguments= args; body= checked_body}
-  in
+      { returntype= return_ty
+      ; funname= id
+      ; arguments= args
+      ; annotations
+      ; body= checked_body } in
   (* NB: **not** tenv_body, so args don't leak out *)
   (tenv, mk_typed_statement ~return_type:Incomplete ~loc ~stmt)
 
@@ -1872,10 +1876,11 @@ and check_statement (cf : context_flags_record) (tenv : Env.t)
   | Block stmts -> (tenv, check_block loc cf tenv stmts)
   | Profile (name, vdsl) -> (tenv, check_profile loc cf tenv name vdsl)
   (* these two are special in that they're allowed to change the type environment *)
-  | VarDecl {decl_type; transformation; variables; is_global} ->
-      check_var_decl loc cf tenv decl_type transformation variables is_global
-  | FunDef {returntype; funname; arguments; body} ->
-      check_fundef loc cf tenv returntype funname arguments body
+  | VarDecl {decl_type; transformation; variables; annotations; is_global} ->
+      check_var_decl loc cf tenv decl_type transformation variables annotations
+        is_global
+  | FunDef {returntype; funname; arguments; annotations; body} ->
+      check_fundef loc cf tenv returntype funname arguments annotations body
 
 let verify_fun_def_body_in_block = function
   | {stmt= FunDef {body= {stmt= Block _; _}; _}; _}
@@ -1904,7 +1909,8 @@ let add_userdefined_functions tenv stmts_opt =
   | Some {stmts; _} ->
       let f tenv (s : Ast.untyped_statement) =
         match s with
-        | {stmt= FunDef {returntype; funname; arguments; body}; smeta= {loc}} ->
+        | { stmt= FunDef {returntype; funname; arguments; body; annotations= _}
+          ; smeta= {loc} } ->
             let arg_types = Ast.type_of_arguments arguments in
             verify_fundef_overloaded loc tenv funname arg_types returntype;
             let defined =

src/frontend/lexer.mll

@@ -197,6 +197,12 @@ rule token = parse
   | identifier as id          { lexer_logger ("identifier " ^ id) ;
                                 lexer_pos_logger (lexeme_start_p lexbuf);
                                 Parser.IDENTIFIER (lexeme lexbuf) }
+  (* TODO annotation proper paren lexxing *)
+  | "@" (non_space_or_newline+ as ann)
+                              { lexer_logger ("annotation " ^ ann) ;
+                                lexer_pos_logger (lexeme_start_p lexbuf);
+                                add_separator lexbuf ;
+                                Parser.ANNOTATION (lexeme lexbuf |> Core.String.chop_prefix_exn ~prefix:"@") }
 (* End of file *)
   | eof                       { lexer_logger "eof" ;
                                 if Preprocessor.size () = 1