Source file functors.ml

(*****************************************************************************)
(*                                                                           *)
(* Open Source License                                                       *)
(* Copyright (c) 2022 Nomadic Labs, <contact@nomadic-labs.com>               *)
(*                                                                           *)
(* Permission is hereby granted, free of charge, to any person obtaining a   *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
(* and/or sell copies of the Software, and to permit persons to whom the     *)
(* Software is furnished to do so, subject to the following conditions:      *)
(*                                                                           *)
(* The above copyright notice and this permission notice shall be included   *)
(* in all copies or substantial portions of the Software.                    *)
(*                                                                           *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
(* DEALINGS IN THE SOFTWARE.                                                 *)
(*                                                                           *)
(*****************************************************************************)

(* Main functor that is exported to the outside *)
module Make
  (R: Rache.TRANSFER)
: Sigs.MAP with type key = R.key
= struct

  type key = R.key
  type 'a promise = {
    p : 'a Lwt.t;
    mutable binded : bool;
  }
  let cancel { p; binded } =
    if binded then
      ()
    else
      Lwt.cancel p
  let of_lwt p = { p; binded = false }
  let to_lwt { p; _ } = p
  type 'a t = 'a promise R.t

  let create n =
    (* Promises are resources which are destroyed by cancelation *)
    R.create (fun _k p -> cancel p) n
  let put c k r = R.put c k (of_lwt r)
  let take c k = Option.map to_lwt (R.take c k)
  let take_all c = List.map (fun (k, r) -> (k, to_lwt r)) (R.take_all c)
  let take_some c f =
    List.map
      (fun (k, r) -> (k, to_lwt r))
      (* You can't meaningfully test on the promise so we only offer the key *)
      (R.take_some c (fun k _p -> f k))

  let bind c k f =
    R.borrow c k
      (fun r ->
        r.binded <- true;
        let p = to_lwt r in
        let q = Lwt.bind p f in
        (* NOTE: you might be thinking that the code is missing
           [r.binded <- false] to signify the promise being returned. This
           could happen when [q] resolves (which in some sense is when the
           function [f] finally "returns"). (Or actually when [f] is called.)
           However, note that in this (these) cases, the promise [p] is already
           resolved and [Lwt.cancel] is a no-op. *)
        q
      )

  let bind_or_put c k make f =
    match bind c k f with
    | Some q -> q
    | None ->
        let p = make k in
        let r = of_lwt p in
        R.put c k r;
        r.binded <- true;
        let q = Lwt.bind p f in
        q

  let fold f c acc =
    R.fold
      (fun k r acc ->
        r.binded <- true;
        let open Lwt.Syntax in
        let* acc = acc in
        let* x = to_lwt r in
        f k x acc
      )
      c
      (Lwt.return acc)

  let fold_oldest_first f c acc =
    R.fold_oldest_first
      (fun k r acc ->
        r.binded <- true;
        let open Lwt.Syntax in
        let* acc = acc in
        let* x = to_lwt r in
        f k x acc
      )
      c
      (Lwt.return acc)

  let remove c k = R.remove c k
  let clear c = R.clear c
  let filter c f = R.filter c (fun k _r -> f k)
  let length c = R.length c
  let capacity c = R.capacity c

end

module Make_option
  (R: Rache.TRANSFER)
: Sigs.MAP_OPTION with type key = R.key
= struct

  type key = R.key
  type 'a promise = {
    p : 'a option Lwt.t;
    mutable binded : bool;
    on_take : unit Lwt.t;
    resolve_on_take : unit Lwt.u;
  }
  let cancel { p; binded; _ } =
    if binded then
      ()
    else
      Lwt.cancel p
  let of_lwt p =
    let (on_take, resolve_on_take) = Lwt.wait () in
    { p; binded = false; on_take; resolve_on_take; }
  let to_lwt { p; _ } = p
  let monitor_none c k { p; on_take; _ } =
    Lwt.on_failure
      (Lwt.choose [
        on_take; (* on_take is successfully resolved when the promise is taken
                    out of the cache so this callback is called and can be
                    garbage collected. *)
        (Lwt.try_bind
          (fun () -> p)
          (function None -> raise Exit | Some _ -> Lwt.return_unit)
          (fun (_ : exn) -> Lwt.return_unit))
      ])
      (fun (_ : exn) -> R.remove c k)
  type 'a t = 'a promise R.t

  let create n =
    (* Promises are resources which are destroyed by cancelation *)
    R.create (fun _k p -> cancel p) n
  let put c k p =
    match Lwt.state p with
    | Lwt.Return None -> ()
    | _ ->
      let r = of_lwt p in
      monitor_none c k r;
      R.put c k r
  let take c k =
    match R.take c k with
    | None -> None
    | Some r ->
        Lwt.wakeup r.resolve_on_take ();
        Some (to_lwt r)
  let take_all c =
    List.map
      (fun (k, r) ->
        Lwt.wakeup r.resolve_on_take ();
        (k, to_lwt r))
      (R.take_all c)
  let take_some c f =
    List.map
      (fun (k, r) ->
        Lwt.wakeup r.resolve_on_take ();
        (k, to_lwt r))
      (* You can't meaningfully test on the promise so we only offer the key *)
      (R.take_some c (fun k _p -> f k))

  let bind c k f =
    R.borrow c k
      (fun r ->
        r.binded <- true;
        let p = to_lwt r in
        let q = Lwt.bind p f in
        (* NOTE: you might be thinking that the code is missing
           [r.binded <- false] to signify the promise being returned. This
           could happen when [q] resolves (which in some sense is when the
           function [f] finally "returns"). (Or actually when [f] is called.)
           However, note that in this (these) cases, the promise [p] is already
           resolved and [Lwt.cancel] is a no-op. *)
        q
      )

  let bind_or_put c k make f =
    match bind c k f with
    | Some q -> q
    | None ->
        let p = make k in
        match Lwt.state p with
        | Lwt.Return None -> f None
        | _ ->
          let r = of_lwt p in
          monitor_none c k r;
          R.put c k r;
          r.binded <- true;
          let q = Lwt.bind p f in
          q

  let fold f c acc =
    R.fold
      (fun k r acc ->
        r.binded <- true;
        let open Lwt.Syntax in
        let* acc = acc in
        let* x = to_lwt r in
        match x with
        | None -> Lwt.return acc
        | Some x -> f k x acc
      )
      c
      (Lwt.return acc)

  let fold_oldest_first f c acc =
    R.fold_oldest_first
      (fun k r acc ->
        r.binded <- true;
        let open Lwt.Syntax in
        let* acc = acc in
        let* x = to_lwt r in
        match x with
        | None -> Lwt.return acc
        | Some x -> f k x acc
      )
      c
      (Lwt.return acc)

  let remove c k = R.remove c k
  let clear c = R.clear c
  let filter c f = R.filter c (fun k _r -> f k)
  let length c = R.length c
  let capacity c = R.capacity c

end

module Make_result
  (R: Rache.TRANSFER)
: Sigs.MAP_RESULT with type key = R.key
= struct

  type key = R.key
  type ('a, 'err) promise = {
    p : ('a, 'err) result Lwt.t;
    mutable binded : bool;
    on_take : unit Lwt.t;
    resolve_on_take : unit Lwt.u;
  }
  let cancel { p; binded; _ } =
    if binded then
      ()
    else
      Lwt.cancel p
  let of_lwt p =
    let (on_take, resolve_on_take) = Lwt.wait () in
    { p; binded = false; on_take; resolve_on_take; }
  let to_lwt { p; _ } = p
  let monitor_error c k { p; on_take; _ } =
    Lwt.on_failure
      (Lwt.choose [
        on_take; (* on_take is successfully resolved when the promise is taken
                    out of the cache so this callback is called and can be
                    garbage collected. *)
        (Lwt.try_bind
          (fun () -> p)
          (function Error _ -> raise Exit | Ok _ -> Lwt.return_unit)
          (fun (_ : exn) -> Lwt.return_unit))
      ])
      (fun (_ : exn) -> R.remove c k)
  type ('a, 'err) t = ('a, 'err) promise R.t

  let create n =
    (* Promises are resources which are destroyed by cancelation *)
    R.create (fun _k p -> cancel p) n
  let put c k p =
    match Lwt.state p with
    | Lwt.Return (Error _) -> ()
    | _ ->
      let r = of_lwt p in
      monitor_error c k r;
      R.put c k r
  let take c k =
    match R.take c k with
    | None -> None
    | Some r ->
        Lwt.wakeup r.resolve_on_take ();
        Some (to_lwt r)
  let take_all c =
    List.map
      (fun (k, r) ->
        Lwt.wakeup r.resolve_on_take ();
        (k, to_lwt r))
      (R.take_all c)
  let take_some c f =
    List.map
      (fun (k, r) ->
        Lwt.wakeup r.resolve_on_take ();
        (k, to_lwt r))
      (* You can't meaningfully test on the promise so we only offer the key *)
      (R.take_some c (fun k _p -> f k))

  let bind c k f =
    R.borrow c k
      (fun r ->
        r.binded <- true;
        let p = to_lwt r in
        let q = Lwt.bind p f in
        (* NOTE: you might be thinking that the code is missing
           [r.binded <- false] to signify the promise being returned. This
           could happen when [q] resolves (which in some sense is when the
           function [f] finally "returns"). (Or actually when [f] is called.)
           However, note that in this (these) cases, the promise [p] is already
           resolved and [Lwt.cancel] is a no-op. *)
        q
      )

  let bind_or_put c k make f =
    match bind c k f with
    | Some q -> q
    | None ->
        let p = make k in
        match Lwt.state p with
        | Lwt.Return (Error _ as error) -> f error
        | _ ->
          let r = of_lwt p in
          monitor_error c k r;
          R.put c k r;
          r.binded <- true;
          let q = Lwt.bind p f in
          q

  let fold f c acc =
    R.fold
      (fun k r acc ->
        r.binded <- true;
        let open Lwt.Syntax in
        let* acc = acc in
        let* x = to_lwt r in
        match x with
        | Error _ -> Lwt.return acc
        | Ok x -> f k x acc
      )
      c
      (Lwt.return acc)

  let fold_oldest_first f c acc =
    R.fold_oldest_first
      (fun k r acc ->
        r.binded <- true;
        let open Lwt.Syntax in
        let* acc = acc in
        let* x = to_lwt r in
        match x with
        | Error _ -> Lwt.return acc
        | Ok x -> f k x acc
      )
      c
      (Lwt.return acc)

  let remove c k = R.remove c k
  let clear c = R.clear c
  let filter c f = R.filter c (fun k _r -> f k)
  let length c = R.length c
  let capacity c = R.capacity c

end