[@poll error] attribute

backend/cfg/cfgize.ml

@@ -761,7 +761,12 @@ let fundecl :
         fun_codegen_options;
         fun_dbg;
         fun_num_stack_slots;
-        fun_contains_calls
+        fun_contains_calls;
+        (* CR-someday mshinwell: [fun_poll] will need to be propagated in the
+           future, e.g. when writing a [Polling] equivalent on [Cfg]. We don't
+           do this at present since there is no need, and because
+           [Linear_to_cfg] doesn't have [fun_poll] available. *)
+        fun_poll = _
       } =
     fundecl
   in

backend/cmm.ml

@@ -260,6 +260,7 @@ type fundecl =
     fun_args: (Backend_var.With_provenance.t * machtype) list;
     fun_body: expression;
     fun_codegen_options : codegen_option list;
+    fun_poll: Lambda.poll_attribute;
     fun_dbg : Debuginfo.t;
   }
 

backend/cmm.mli

@@ -269,6 +269,7 @@ type fundecl =
     fun_args: (Backend_var.With_provenance.t * machtype) list;
     fun_body: expression;
     fun_codegen_options : codegen_option list;
+    fun_poll: Lambda.poll_attribute;
     fun_dbg : Debuginfo.t;
   }
 

backend/cmm_helpers.ml

@@ -2385,7 +2385,8 @@ let send_function (arity, mode) =
       fun_args = List.map (fun (arg, ty) -> VP.create arg, ty) fun_args;
       fun_body = body;
       fun_codegen_options = [];
-      fun_dbg
+      fun_dbg;
+      fun_poll = Default_poll
     }
 
 let apply_function arity =
@@ -2398,7 +2399,8 @@ let apply_function arity =
       fun_args = List.map (fun arg -> VP.create arg, typ_val) all_args;
       fun_body = body;
       fun_codegen_options = [];
-      fun_dbg
+      fun_dbg;
+      fun_poll = Default_poll
     }
 
 (* Generate tuplifying functions:
@@ -2430,7 +2432,8 @@ let tuplify_function arity =
             @ [Cvar clos],
             dbg () );
       fun_codegen_options = [];
-      fun_dbg
+      fun_dbg;
+      fun_poll = Default_poll
     }
 
 (* Generate currying functions:
@@ -2505,7 +2508,8 @@ let final_curry_function ~nlocal ~arity =
       fun_args = [VP.create last_arg, typ_val; VP.create last_clos, typ_val];
       fun_body = curry_fun [] last_clos (arity - 1);
       fun_codegen_options = [];
-      fun_dbg
+      fun_dbg;
+      fun_poll = Default_poll
     }
 
 let rec intermediate_curry_functions ~nlocal ~arity num =
@@ -2552,7 +2556,8 @@ let rec intermediate_curry_functions ~nlocal ~arity num =
                   Cvar clos ],
                 dbg () ));
         fun_codegen_options = [];
-        fun_dbg
+        fun_dbg;
+        fun_poll = Default_poll
       }
     ::
     (if arity <= max_arity_optimized && arity - num > 2
@@ -2598,7 +2603,8 @@ let rec intermediate_curry_functions ~nlocal ~arity num =
                 (List.map (fun (arg, _) -> Cvar arg) direct_args)
                 clos;
             fun_codegen_options = [];
-            fun_dbg
+            fun_dbg;
+            fun_poll = Default_poll
           }
       in
       cf :: intermediate_curry_functions ~nlocal ~arity (num + 1)
@@ -3639,7 +3645,8 @@ let entry_point namelist =
       fun_args = [];
       fun_body = body;
       fun_codegen_options = [Reduce_code_size];
-      fun_dbg
+      fun_dbg;
+      fun_poll = Default_poll
     }
 
 (* Generate the table of globals *)
@@ -4155,8 +4162,8 @@ let cfunction decl = Cmm.Cfunction decl
 
 let cdata d = Cmm.Cdata d
 
-let fundecl fun_name fun_args fun_body fun_codegen_options fun_dbg =
-  { Cmm.fun_name; fun_args; fun_body; fun_codegen_options; fun_dbg }
+let fundecl fun_name fun_args fun_body fun_codegen_options fun_dbg fun_poll =
+  { Cmm.fun_name; fun_args; fun_body; fun_codegen_options; fun_dbg; fun_poll }
 
 (* Gc root table *)
 

backend/cmm_helpers.mli

@@ -1188,6 +1188,7 @@ val fundecl :
   expression ->
   codegen_option list ->
   Debuginfo.t ->
+  Lambda.poll_attribute ->
   fundecl
 
 (** Create a cmm phrase for a function declaration. *)

backend/cmmgen.ml

@@ -1481,6 +1481,7 @@ let transl_function f =
              fun_args = List.map (fun (id, _) -> (id, typ_val)) f.params;
              fun_body = cmm_body;
              fun_codegen_options;
+             fun_poll = f.poll;
              fun_dbg  = f.dbg}
 
 (* Translate all function definitions *)
@@ -1583,7 +1584,8 @@ let compunit (ulam, preallocated_blocks, constants) =
                            Use_linscan_regalloc;
                          ]
                          else [ Reduce_code_size; Use_linscan_regalloc ];
-                       fun_dbg  = Debuginfo.none }] in
+                       fun_dbg  = Debuginfo.none;
+                       fun_poll = Default_poll }] in
   let c2 = transl_clambda_constants constants c1 in
   let c3 = transl_all_functions c2 in
   Cmmgen_state.set_structured_constants [];

backend/mach.ml

@@ -108,6 +108,7 @@ type fundecl =
     fun_body: instruction;
     fun_codegen_options : Cmm.codegen_option list;
     fun_dbg : Debuginfo.t;
+    fun_poll: Lambda.poll_attribute;
     fun_num_stack_slots: int array;
     fun_contains_calls: bool;
   }

backend/mach.mli

@@ -118,6 +118,7 @@ type fundecl =
     fun_body: instruction;
     fun_codegen_options : Cmm.codegen_option list;
     fun_dbg : Debuginfo.t;
+    fun_poll: Lambda.poll_attribute;
     fun_num_stack_slots: int array;
     fun_contains_calls: bool;
   }

backend/polling.ml

@@ -18,6 +18,7 @@
 (**************************************************************************)
 
 open Mach
+open Format
 
 module Int = Numbers.Int
 module String = Misc.Stdlib.String
@@ -26,6 +27,12 @@ let function_is_assumed_to_never_poll func =
   String.begins_with ~prefix:"caml_apply" func
   || String.begins_with ~prefix:"caml_send" func
 
+(* These are used for the poll error annotation later on*)
+type polling_point = Alloc | Poll | Function_call | External_call
+type error = Poll_error of (polling_point * Debuginfo.t) list
+
+exception Error of error
+
 (* Detection of recursive handlers that are not guaranteed to poll
    at every loop iteration. *)
 
@@ -169,7 +176,7 @@ let contains_polls = ref false
 
 let add_poll i =
   contains_polls := true;
-  Mach.instr_cons (Iop (Ipoll { return_label = None })) [||] [||] i
+  Mach.instr_cons_debug (Iop (Ipoll { return_label = None })) [||] [||] i.dbg i
 
 let instr_body handler_safe i =
   let add_unsafe_handler ube (k, _trap_stack, _) =
@@ -226,12 +233,46 @@ let instr_body handler_safe i =
   in
   instr Int.Set.empty i
 
+let find_poll_alloc_or_calls instr =
+  let f_match i =
+      match i.desc with
+      | Iop(Ipoll _) -> Some (Poll, i.dbg)
+      | Iop(Ialloc _) -> Some (Alloc, i.dbg)
+      | Iop(Icall_ind | Icall_imm _ |
+            Itailcall_ind | Itailcall_imm _ ) -> Some (Function_call, i.dbg)
+      | Iop(Iextcall { alloc = true }) -> Some (External_call, i.dbg)
+      | Iop(Imove | Ispill | Ireload | Iconst_int _ | Iconst_float _ |
+            Iconst_symbol _ | Iextcall { alloc = false } | Istackoffset _ |
+            Iload _ | Istore _ | Iintop _ | Iintop_imm _ | Ifloatofint |
+            Iintoffloat | Inegf | Iabsf | Iaddf | Isubf | Imulf | Idivf |
+            Iopaque | Ispecific _ | Ibeginregion | Iendregion |
+            Icompf _ | Iname_for_debugger _ | Iprobe _ |
+            Iprobe_is_enabled _ | Ivalueofint | Iintofvalue)-> None
+      | Iend | Ireturn _ | Iifthenelse _ | Iswitch _ | Icatch _ | Iexit _ |
+        Itrywith _ | Iraise _ -> None
+    in
+  let matches = ref [] in
+    Mach.instr_iter
+      (fun i ->
+        match f_match i with
+        | Some(x) -> matches := x :: !matches
+        | None -> ())
+      instr;
+  List.rev !matches
+
 let instrument_fundecl ~future_funcnames:_ (f : Mach.fundecl) : Mach.fundecl =
   if function_is_assumed_to_never_poll f.fun_name then f
   else begin
     let handler_needs_poll = polled_loops_analysis f.fun_body in
     contains_polls := false;
     let new_body = instr_body handler_needs_poll f.fun_body in
+    begin match f.fun_poll with
+    | Error_poll -> begin
+        match find_poll_alloc_or_calls new_body with
+        | [] -> ()
+        | poll_error_instrs -> raise (Error(Poll_error poll_error_instrs))
+      end
+    | Default_poll -> () end;
     let new_contains_calls = f.fun_contains_calls || !contains_polls in
     { f with fun_body = new_body; fun_contains_calls = new_contains_calls }
   end
@@ -242,3 +283,49 @@ let requires_prologue_poll ~future_funcnames ~fun_name i =
     match potentially_recursive_tailcall ~future_funcnames i with
     | Might_not_poll -> true
     | Always_polls -> false
+
+(* Error report *)
+
+let instr_type p =
+  match p with
+  | Poll -> "inserted poll"
+  | Alloc -> "allocation"
+  | Function_call -> "function call"
+  | External_call -> "external call that allocates"
+
+let report_error ppf = function
+| Poll_error instrs ->
+  begin
+    let num_inserted_polls =
+      List.fold_left
+      (fun s (p,_) -> s + match p with Poll -> 1
+                      | Alloc | Function_call | External_call -> 0
+      ) 0 instrs in
+      let num_user_polls = (List.length instrs) - num_inserted_polls in
+      if num_user_polls = 0 then
+        fprintf ppf "Function with poll-error attribute contains polling \
+        points (inserted by the compiler)\n"
+      else begin
+        fprintf ppf
+        "Function with poll-error attribute contains polling points:\n";
+        List.iter (fun (p,dbg) ->
+          begin match p with
+          | Poll -> ()
+          | Alloc | Function_call | External_call ->
+            fprintf ppf "\t%s at " (instr_type p);
+            Location.print_loc ppf (Debuginfo.to_location dbg);
+            fprintf ppf "\n"
+          end
+        ) instrs;
+        if num_inserted_polls > 0 then
+          fprintf ppf "\t(plus compiler-inserted polling point(s) in prologue \
+          and/or loop back edges)\n"
+      end
+  end
+
+let () =
+  Location.register_error_of_exn
+    (function
+      | Error err -> Some (Location.error_of_printer_file report_error err)
+      | _ -> None
+    )

backend/reloadgen.ml

@@ -152,6 +152,7 @@ method fundecl f num_stack_slots =
   ({fun_name = f.fun_name; fun_args = f.fun_args;
     fun_body = new_body; fun_codegen_options = f.fun_codegen_options;
     fun_dbg  = f.fun_dbg;
+    fun_poll = f.fun_poll;
     fun_contains_calls = f.fun_contains_calls;
     fun_num_stack_slots = Array.copy num_stack_slots;
    },

backend/selectgen.ml

@@ -1558,7 +1558,8 @@ method emit_fundecl ~future_funcnames f =
     if Polling.requires_prologue_poll ~future_funcnames
          ~fun_name:f.Cmm.fun_name body
       then
-      instr_cons (Iop(Ipoll { return_label = None })) [||] [||] body
+        instr_cons_debug
+          (Iop(Ipoll { return_label = None })) [||] [||] f.Cmm.fun_dbg body
     else
       body
     in
@@ -1569,6 +1570,7 @@ method emit_fundecl ~future_funcnames f =
     fun_body = body_with_prologue;
     fun_codegen_options = f.Cmm.fun_codegen_options;
     fun_dbg  = f.Cmm.fun_dbg;
+    fun_poll = f.Cmm.fun_poll;
     fun_num_stack_slots = Array.make Proc.num_register_classes 0;
     fun_contains_calls = !contains_calls;
   }

backend/spill.ml

@@ -633,6 +633,7 @@ let fundecl f =
     fun_args = f.fun_args;
     fun_body = new_body;
     fun_codegen_options = f.fun_codegen_options;
+    fun_poll = f.fun_poll;
     fun_dbg  = f.fun_dbg;
     fun_num_stack_slots = f.fun_num_stack_slots;
     fun_contains_calls = f.fun_contains_calls;

backend/split.ml

@@ -222,6 +222,7 @@ let fundecl f =
     fun_args = new_args;
     fun_body = new_body;
     fun_codegen_options = f.fun_codegen_options;
+    fun_poll = f.fun_poll;
     fun_dbg  = f.fun_dbg;
     fun_num_stack_slots = f.fun_num_stack_slots;
     fun_contains_calls = f.fun_contains_calls;

middle_end/clambda.ml

@@ -99,6 +99,7 @@ and ufunction = {
   body   : ulambda;
   dbg    : Debuginfo.t;
   env    : Backend_var.t option;
+  poll   : poll_attribute;
   mode   : Lambda.alloc_mode;
   check  : Lambda.check_attribute;
 }
@@ -117,6 +118,7 @@ type function_description =
     mutable fun_closed: bool;           (* True if environment not used *)
     mutable fun_inline: (Backend_var.With_provenance.t list * ulambda) option;
     mutable fun_float_const_prop: bool; (* Can propagate FP consts *)
+    fun_poll: poll_attribute;           (* Error on poll/alloc/call *)
     fun_region: bool;                   (* If false, may locally allocate
                                            in caller's region *)
   }

middle_end/clambda.mli

@@ -110,6 +110,7 @@ and ufunction = {
   body   : ulambda;
   dbg    : Debuginfo.t;
   env    : Backend_var.t option;
+  poll   : poll_attribute;
   mode   : Lambda.alloc_mode;
   check  : Lambda.check_attribute;
 }
@@ -128,6 +129,7 @@ type function_description =
     mutable fun_closed: bool;           (* True if environment not used *)
     mutable fun_inline: (Backend_var.With_provenance.t list * ulambda) option;
     mutable fun_float_const_prop: bool; (* Can propagate FP consts *)
+    fun_poll: poll_attribute;           (* Behaviour for polls *)
     fun_region: bool;                   (* If false, may locally allocate
                                            in caller's region *)
   }

middle_end/closure/closure.ml

@@ -1412,6 +1412,7 @@ and close_functions { backend; fenv; cenv; mutable_vars } fun_defs =
                fun_closed = initially_closed;
                fun_inline = None;
                fun_float_const_prop = !Clflags.float_const_prop;
+               fun_poll = attr.poll;
                fun_region = region} in
             let dbg = Debuginfo.from_location loc in
             (id, params, return, body, mode, attrib, fundesc, dbg)
@@ -1465,6 +1466,7 @@ and close_functions { backend; fenv; cenv; mutable_vars } fun_defs =
         body   = ubody;
         dbg;
         env = Some env_param;
+        poll = fundesc.fun_poll;
         mode;
         check;
       }

middle_end/flambda/augment_specialised_args.ml

@@ -549,6 +549,7 @@ module Make (T : S) = struct
         ~check:Default_check
         ~is_a_functor:false
         ~closure_origin:function_decl.closure_origin
+        ~poll:Default_poll (* don't propagate attribute to wrappers *)
     in
     new_fun_var, new_function_decl, rewritten_existing_specialised_args,
       benefit
@@ -643,6 +644,7 @@ module Make (T : S) = struct
           ~check:function_decl.check
           ~is_a_functor:function_decl.is_a_functor
           ~closure_origin
+          ~poll:function_decl.poll
       in
       let funs, direct_call_surrogates =
         if for_one_function.make_direct_call_surrogates then