Delimited Continuations in OCaml

delimcc.ml

(* Boilerplate *)
(* Identity: id :: a -> a *)
let id = fun x -> x
(* Function composition: (.) :: (b -> c) -> (a -> b) -> a -> c *)
let (<.>) (f : 'b -> 'c) (g : 'a -> 'b) : 'a -> 'c =
    fun x -> f (g x)


(* Cont type *)
type ('w, 'a) cont = Cont of (('a -> 'w) -> 'w)
(* Extraction: runCont :: Cont w a -> (a -> w) -> w *)
let runCont (c : ('w, 'a) cont) : ('a -> 'w) -> 'w =
    let (Cont c') = c in
    c'


(* Monad instance *)
(* return :: Monad m => a -> m a *)
let return (x : 'a) : ('w, 'a) cont =
    Cont (fun k -> k x)
(* (>>=) :: Monad m => m a -> (a -> m b) -> m b *)
let (>>=) (m : ('w, 'a) cont) (f : 'a -> ('w, 'b) cont) =
    Cont (fun k -> runCont m (fun v -> runCont (f v) k))

(* Functor instance *)
(* fmap :: Functor f => (a -> b) -> f a -> f b *)
let fmap (f : 'a -> 'b) (m : ('a, 'w) cont) : ('w, 'b) cont =
    Cont (fun k -> runCont m (k <.> f))

(* Applicative instance *)
(* pure :: Applicative f => a -> f a *)
let pure (x : 'a) : ('w, 'a) cont = return x
(* (<*>) :: Applicative f => f (a -> b) -> f a -> f b *)
let (<*>) (m : ('a, 'a -> 'b) cont) (a : ('a, 'a) cont) : ('a, 'b) cont = m >>= fun h -> fmap h a


(* Utilities *)
(* Run with `id` as final continuation *)
let runC (m : ('a, 'a) cont) : 'a = runCont m id
(* shift/reset, the 'magic' *)
let reset (x : ('a, 'a) cont) : ('w, 'a) cont = return (runC x)
let shift (f : ('a -> 'w) -> ('w, 'w) cont) : ('w, 'a) cont = Cont (runC <.> f)


(* Operator lifting *)
(* Not adding type info here since these should be generic *)
(* liftM2 :: (a -> b -> c) -> m a -> m b -> m c *)
let liftM2 f m1 m2 =
    m1 >>= fun x1 -> m2 >>= fun x2 -> return (f x1 x2)
(* These become :: (int, int) cont -> (int, int) cont -> (int, int) cont *)
let (-!) = liftM2 (-)
let (+!) = liftM2 (+)
let ( *! ) = liftM2 ( * )


(* Simple demo *)
let demo (f : (int -> int) -> (int, int) cont) : ('w, int) cont =
    reset (return 3 +! (shift f) -! return 1)

(* Short-cutting computation *)
let mul_list (l : int list) : ('w, int) cont =
    let rec loop = function
        | [] -> Printf.printf "loop: empty list, 1\n"; return 1
        | (0 :: _) -> Printf.printf "loop: found 0\n"; shift (fun _ -> return 0)
        | (x :: xs) -> Printf.printf "loop: (*) %d\n" x; return x *! loop xs
    in
    reset (loop l)


let main () =
    Printf.printf "demo drop = %d\n" (runC (demo (fun _ -> return (5 * 2))));
    Printf.printf "demo call = %d\n" (runC (demo (fun k -> return (k (5 * 2)))));
    Printf.printf "mul_list [1; 2; 3; 4; 5] %d\n" (runC (mul_list [1; 2; 3; 4; 5]));
    Printf.printf "mul_list [1; 2; 3; 0; 4; 5] = %d\n" (runC (mul_list [1; 2; 3; 0; 4; 5]))
;;

main ()

output.txt

$ ocaml delimcc.ml
demo drop = 10
demo call = 12
loop: (*) 1
loop: (*) 2
loop: (*) 3
loop: (*) 4
loop: (*) 5
loop: empty list, 1
mul_list [1; 2; 3; 4; 5] 120
loop: (*) 1
loop: (*) 2
loop: (*) 3
loop: found 0
mul_list [1; 2; 3; 0; 4; 5] = 0