effects.ml

main.ml

type _ io =
  | Sync : (unit -> 'a) -> 'a io
  | FlatMap : 'a io * ('a -> 'b io) -> 'b io
  | Async : ((('a -> unit) -> unit) -> unit) -> 'a io
  | All : 'a io list -> 'a list io
  | Sleep : float -> unit io

let flatMap f io = FlatMap (io, f)
let map f io = flatMap (fun a -> Sync (fun () -> f a)) io

let rec repeat n io =
  if n <= 0 then io else flatMap (fun _ -> repeat (n - 1) io) io

let sleep duration = Sleep duration

let all effects = All effects

let pipe io fns = List.fold_left (fun acc f -> f acc) io fns

module Gen = struct
  let ( let* ) m f = flatMap f m
  let ( let+ ) m f = map f m
  let return x = Sync (fun () -> x)
  let ( and* ) m1 m2 = flatMap (fun x1 -> map (fun x2 -> (x1, x2)) m2) m1
end

let rec run : type a. a io -> a = function
  | Sync eval -> eval ()
  | Sleep duration -> Unix.sleepf duration
  | All effects -> List.map run effects
  | Async _ -> failwith "Async operations need special runtime support"
  | FlatMap (io, cont) ->
      let rec loop : type b. b io -> b = function
        | Sync eval -> eval ()
        | Sleep duration -> Unix.sleepf duration
        | All effects -> List.map run effects
        | Async _ -> failwith "Async operations need special runtime support"
        | FlatMap (inner, inner_cont) ->
            let result = loop inner in
            loop (inner_cont result)
      in
      let result = loop io in
      run (cont result)

let io thunk = Sync thunk

let succeed value = Sync (fun () -> value)
let fail error = Sync (fun () -> failwith error)

let ( *> ) io1 io2 = flatMap (fun _ -> io2) io1
let ( <* ) io1 io2 = flatMap (fun x -> map (fun _ -> x) io2) io1

let rec race : type a. a io list -> a io = function
  | [] -> fail "No effects to race"
  | [ x ] -> x
  | _ -> failwith "Race requires async runtime support"

let tap f io = flatMap (fun x -> map (fun () -> x) (f x)) io

let rec forever io = flatMap (fun _ -> forever io) io

let rec times n io =
  if n <= 0 then succeed []
  else flatMap (fun x -> map (fun xs -> x :: xs) (times (n - 1) io)) io

let example () =
  let hello =
    io (fun () ->
        print_endline "Hello";
        42)
  in
  let result =
    hello
    |> flatMap (fun x ->
           io (fun () ->
               print_endline ("Got: " ^ string_of_int x);
               x + 1))
    |> map (fun x -> x * 2)
    |> repeat 2
  in
  let final = run result in
  print_endline ("Final result: " ^ string_of_int final);

  let gen_example =
    let open Gen in
    let* x =
      io (fun () ->
          print_endline "Gen: step 1";
          10)
    in
    let* y =
      io (fun () ->
          print_endline "Gen: step 2";
          20)
    in
    let+ z =
      io (fun () ->
          print_endline "Gen: step 3";
          x + y)
    in
    z * 2
  in
  let gen_result = run gen_example in
  print_endline ("Gen result: " ^ string_of_int gen_result);

  let sleep_example =
    let open Gen in
    let* _ = io (fun () -> print_endline "Before sleep") in
    let* _ = sleep 0.5 in
    let* _ = io (fun () -> print_endline "After 0.5 second sleep") in
    return "Sleep done"
  in
  let sleep_result = run sleep_example in
  print_endline ("Sleep result: " ^ sleep_result);

  let all_example =
    all
      [
        io (fun () ->
            print_endline "Effect 1";
            1);
        io (fun () ->
            print_endline "Effect 2";
            2);
        io (fun () ->
            print_endline "Effect 3";
            3);
      ]
  in
  let all_result = run all_example in
  print_endline
    ("All result: " ^ String.concat ", " (List.map string_of_int all_result));

  let pipe_example =
    pipe (succeed 5)
      [
        map (fun x -> x * 2);
        flatMap (fun x ->
            io (fun () ->
                print_endline ("Piped value: " ^ string_of_int x);
                x));
        map (fun x -> x + 10);
        tap (fun x ->
            io (fun () -> print_endline ("Tapped value: " ^ string_of_int x)));
      ]
  in
  let pipe_result = run pipe_example in
  print_endline ("Pipe result: " ^ string_of_int pipe_result);

  let combinator_example =
    let effect1 =
      io (fun () ->
          print_endline "First";
          100)
    in
    let effect2 =
      io (fun () ->
          print_endline "Second";
          200)
    in
    let combined = effect1 <* effect2 in
    run combined
  in
  print_endline ("Combinator result: " ^ string_of_int combinator_example);

  let times_example =
    times 3
      (io (fun () ->
           print_endline "Running effect";
           Random.int 100))
  in
  let times_result = run times_example in
  print_endline
    ("Times result: " ^ String.concat ", " (List.map string_of_int times_result));

  final

let () = ignore (example ())