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 Transfer
  (Collection: Ringo.COLLECTION)
  (H: Hashtbl.HashedType)
: Sigs.TRANSFER with type key = H.t
= struct

  module Table = Hashtbl.Make(H)

  type key = H.t

  type 'resource t = {
     table: (key * 'resource) Collection.node Table.t;
     collection: (key * 'resource) Collection.t;
     destroy: (key -> 'resource -> unit);
  }

  let create destroy n = {
    table = Table.create n;
    collection = Collection.create n;
    destroy;
  }

  let put {collection; table; destroy} k v =
    begin match Table.find_opt table k with
       | Some node ->
           Collection.remove collection node;
           let (key, data) = Collection.data node in
           assert (H.equal key k);
           destroy k data
       | None -> ()
    end ;
    match Collection.add_and_return_erased collection (k, v) with
    | node, Some (kerased, rerased) ->
        Table.remove table kerased;
        destroy kerased rerased;
        Table.replace table k node
    | node, None ->
        Table.replace table k node

  let take {table; collection; destroy=_ } k =
    match Table.find_opt table k with
    | None -> None
    | Some node ->
       let (key, data) = Collection.data node in
       assert (H.equal key k);
       Collection.remove collection node;
       Table.remove table k;
       Some data

  let take_all {table; collection; destroy=_ } =
    let vs = Collection.elements_data collection in
    Table.clear table;
    Collection.clear collection;
    vs

  let take_some {table; collection; destroy=_ } f =
    let f ((kvts, nodes) as acc) node =
      let (k, v) as kvt = Collection.data node in
      if f k v then ((kvt :: kvts), (node :: nodes)) else acc
    in
    let kvts, nodes = Collection.fold_oldest_first collection ~init:([], []) ~f in
    List.iter2
      (fun (k, _) node ->
         Collection.remove collection node;
         Table.remove table k)
      kvts
      nodes;
    kvts

  let borrow {table; collection; destroy=_ } k f =
    match Table.find_opt table k with
    | None -> None
    | Some node ->
       Collection.promote_read collection node;
       let (key, data) = Collection.data node in
       assert (H.equal key k);
       Some (f data)

  let fold f {collection; _} init =
    let f acc kv =
      let (k, v) = Collection.data kv in
      f k v acc in
    Collection.fold collection ~init ~f

  let fold_oldest_first f {collection; _} init =
    let f acc kv =
      let (k, v) = Collection.data kv in
      f k v acc in
    Collection.fold_oldest_first collection ~init ~f

  let remove {table; collection; destroy} k =
    match Table.find_opt table k with
    | None -> ()
    | Some node ->
       let (key, data) = Collection.data node in
       assert (H.equal key k);
       Collection.remove collection node;
       destroy key data;
       Table.remove table k

  let filter {table; collection; destroy} f =
    Table.iter
      (fun key node ->
         let (k, v) = Collection.data node in
         assert (H.equal k key);
         if not (f k v) then begin
           Collection.remove collection node;
           Table.remove table k;
           destroy key v
         end)
      table

  let clear {table; collection; destroy} =
    Table.iter (fun k n -> let (key, v) = Collection.data n in assert (H.equal key k); destroy k v) table;
    Table.clear table;
    Collection.clear collection

  let length {table; _} = Table.length table

  let capacity {collection; _} = Collection.capacity collection

  module H : Hashtbl.HashedType with type t = H.t = H

end

module Borrow
  (Collection: Ringo.COLLECTION)
  (H: Hashtbl.HashedType)
: Sigs.BORROW with type key = H.t
= struct

  module Table = Hashtbl.Make(H)

  type key = H.t

  type 'resource t = {
     table: (key * 'resource) Collection.node Table.t;
     collection: (key * 'resource) Collection.t;
     destroy: (key -> 'resource -> unit);
  }

  let create destroy n = {
    table = Table.create n;
    collection = Collection.create n;
    destroy;
  }

  let borrow_or_make {collection; table; destroy} k make f =
    match Table.find_opt table k with
       | Some node ->
           Collection.promote_read collection node;
           let (key, data) = Collection.data node in
           assert (H.equal key k);
           f data
       | None ->
          let v = make k in
          let (node, discard) = Collection.add_and_return_erased collection (k, v) in
          begin match discard with
            | Some (kerased, rerased) ->
                Table.remove table kerased;
                destroy kerased rerased
            | None ->
                ()
          end;
          Table.replace table k node;
          f v

  let borrow {table; collection; _} k f =
    match Table.find_opt table k with
    | None -> None
    | Some node ->
       Collection.promote_read collection node;
       let (key, data) = Collection.data node in
       assert (H.equal key k);
       Some (f data)

  let fold f {collection; _} init =
    let f acc kv =
      let (k, v) = Collection.data kv in
      f k v acc in
    Collection.fold collection ~init ~f

  let fold_oldest_first f {collection; _} init =
    let f acc kv =
      let (k, v) = Collection.data kv in
      f k v acc in
    Collection.fold_oldest_first collection ~init ~f

  let remove {table; collection; destroy} k =
    match Table.find_opt table k with
    | None -> ()
    | Some node ->
       let (key, data) = Collection.data node in
       assert (H.equal key k);
       Collection.remove collection node;
       destroy key data;
       Table.remove table k

  let filter {table; collection; destroy} f =
    Table.iter
      (fun key node ->
         let (k, v) = Collection.data node in
         assert (H.equal k key);
         if not (f k v) then begin
           Collection.remove collection node;
           Table.remove table k;
           destroy key v
         end)
      table

  let clear {table; collection; destroy} =
    Table.iter (fun k n -> let (key, v) = Collection.data n in assert (H.equal key k); destroy k v) table;
    Table.clear table;
    Collection.clear collection

  let length {table; _} = Table.length table

  let capacity {collection; _} = Collection.capacity collection

  module H : Hashtbl.HashedType with type t = H.t = H

end