Skip to content

Instantly share code, notes, and snippets.

@polytypic
Last active November 22, 2024 08:21
Show Gist options
  • Save polytypic/334c6958f16ca536576bb12185c8ae57 to your computer and use it in GitHub Desktop.
Save polytypic/334c6958f16ca536576bb12185c8ae57 to your computer and use it in GitHub Desktop.
Fix OCaml

Every time I look at ocaml/ocaml code for more than a few seconds, I spot issues. This collects some of them. Some day I'll just learn to shrug and move on.

The Stdlib.at_exit

let exit_function = atomic_make flush_all

let rec at_exit f =
  (* MPR#7253, MPR#7796: make sure "f" is executed only once *)
  let f_yet_to_run = atomic_make true in
  let old_exit = atomic_get exit_function in
  let new_exit () =
    if atomic_compare_and_set f_yet_to_run true false then f () ;
    old_exit ()
  in
  let success = atomic_compare_and_set exit_function old_exit new_exit in
  if not success then at_exit f

let do_domain_local_at_exit = ref (fun () -> ())

let do_at_exit () =
  (!do_domain_local_at_exit) ();
  (atomic_get exit_function) ()

let exit retcode =
  do_at_exit ();
  sys_exit retcode

leaves a lot to be desired.

First of all, it doesn't guarantee that actions registered through at_exit are called. This is because after an action raises an exception, the remaining actions are not called:

  let new_exit () =
    if atomic_compare_and_set f_yet_to_run true false then f () ;
    old_exit ()
  in

Did I already mention that Domain.at_exit has the same problem?

At minimum this behaviour should be documented. It currently is not.

Also, at_exit unnecessarily allocates an atomic per registered action. The idea there is to prevent an action from being called multiple times. However, there is a much simpler and less expensive way to achieve that by atomically exchanging the exit_function atomic:

 let do_at_exit () =
   (!do_domain_local_at_exit) ();
-  (atomic_get exit_function) ()
+  (atomic_exchange exit_function (fun x -> x)) ()

The Stdlib.Fun.protect function allocates for no good reason. Here is how to fix it:

modified   stdlib/fun.ml
@@ -25,15 +25,15 @@ let () = Printexc.register_printer @@ function
 | Finally_raised exn -> Some ("Fun.Finally_raised: " ^ Printexc.to_string exn)
 | _ -> None

+let finally_no_exn finally =
+  try finally () with e ->
+    let bt = Printexc.get_raw_backtrace () in
+    Printexc.raise_with_backtrace (Finally_raised e) bt
+
 let protect ~(finally : unit -> unit) work =
-  let finally_no_exn () =
-    try finally () with e ->
-      let bt = Printexc.get_raw_backtrace () in
-      Printexc.raise_with_backtrace (Finally_raised e) bt
-  in
   match work () with
-  | result -> finally_no_exn () ; result
+  | result -> finally_no_exn finally ; result
   | exception work_exn ->
       let work_bt = Printexc.get_raw_backtrace () in
-      finally_no_exn () ;
+      finally_no_exn finally ;
       Printexc.raise_with_backtrace work_exn work_bt

The Stdlib.Effect.Deep.match_with function allocates a wrapper for the effc function:

let match_with comp arg handler =
  let effc eff k last_fiber =
    match handler.effc eff with
    | Some f -> f k
    | None -> reperform eff k last_fiber
  in
  let s = alloc_stack handler.retc handler.exnc effc in
  runstack s comp arg

This means that every fiber requires an extra closure. Furthermore, the wrapper also refers to the handler record. This means that every effc call also incurs an extra indirection to access handler.effc. This also means that the handler record is also required to be in memory.

There are many ways in which the per fiber allocation (and keeping the handler record also in memory) could be avoided:

  • The reperform logic could be moved to the runtime — no wrapper would be needed at all.
  • The reperform operation could be exposed and the pattern to implement handlers could be changed such that the responsibility to call reperform would be left to the user. This could also be used to avoid the allocation of Some f.
  • There could be a separate function, say wrap, that performs the wrapping and an abstract type that hides the result of said wrapping. User would then need to wrap effc before calling match_with. This would allow one to avoid wrapping effc separately for every fiber.

Consider the following function:

let space_leak fn gn ab =
  fn (fst ab);
  gn (snd ab)

That function has an obvious space leak: the value (fst ab) is not needed after fn (fst ab) is entered, but since the value ab is needed for the (snd ab), the value (fst ab) also remains reachable.

Now consider the following:

let still_space_leak fn gn (a, b) =
  fn a;
  gn b

One might think that OCaml would translate the above essentially as:

let no_space_leak fn gn ab =
  let a = fst ab in
  let b = snd ab in
  fn a;
  gn b

But it does not. The "wide" match (a, b) is translated like the first space_leak version.

The issue with this is that the value (a, b) (and notably a) is reachable even after fn a has been entered. This is essentially a space leak. If you care for space safety, you should avoid "wide" matches in OCaml.

Here is a godbolt link with the above code snippets.

Proper tail recursion and space efficiency

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment