One recent arrival in Core is with_return, a function that lets you return early from a computation. Here’s a trivial example:

let sum_until_first_negative list =
  with_return (fun r ->
    List.fold list ~init:0 ~f:(fun acc x ->
      if x >= 0 then acc + x else r.return acc))

One thing that might not be obvious in the above example is what the type of r is. It turns out it’s this:

type 'a return = { return : 'b . 'a -> 'b }

The reason for this is to have the return function be truly polymorphic in its return value. That way, like raise, it can be used in any context because its return value will unify with any type.

Note that the sum_until_first_negative example will work even without the record trick, because r.return is used in only one place (and more importantly, with only one type.) But if you want this to work in the general case, you need the record.

So, how does one go about implementing this? Here’s the implementation that’s currently in Core. Note that we use a locally defined exception, to make sure that only this exception handler can catch the exceptions in question. Also, we use a ref to store the value being returned.

let with_return f =
  let module M =
    struct exception Return end
  in
  let r = ref None in
  let return = {
    return = (fun x ->
      r := Some x;
      raise M.Return);
  }
  in
  try f return
  with M.Return ->
    match !r with
    | None -> assert false
    | Some x -> x

In OCaml 3.12, we can make this code, simpler, safer and more efficient:

let with_return (type t) (f : _ -> t) =
  let module M =
    struct exception Return of t end
  in
  let return = { return = (fun x -> raise (M.Return x)); } in
  try f return with M.Return x -> x

I’d like to be able to improve this further by getting rid of the overhead of the record, but I suspect it’s not possible.

It’s worth noting that with_return has its pitfalls. For example, there’s no guarantee that the return always terminates the full computation. For example, the following code:

let foo = with_return
  (fun r -> try r.return 0 with _ -> 1)

will return 1, not 0. Another interesting behavior is to see what happens when r.return is called outside the scope of the closure passed into with_return. Consider the following convoluted function, which returns an optional function which, when called, calls r.return.

let foo = with_return
  (fun r -> Some (fun () -> r.return None))

If you call the function contained in foo, you’ll get the following response:

# let f = Option.value_exn foo;;
val f : unit -> 'a = <fun>
# f ();;
Exception: Return 0.

This isn’t particularly surprising, once you understand the implementation, and it’s semantically fairly reasonable. The only thing you could really ask for beyond this is a warning that the return has escaped its scope, but I think this is beyond the powers of the type-checker.