Source file weighted_dll.ml

(*****************************************************************************)
(*                                                                           *)
(* Open Source License                                                       *)
(* Copyright (c) 2024 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.                                                 *)
(*                                                                           *)
(*****************************************************************************)

type 'a data = {
  data : 'a;
  weight : int;
  mutable newer : 'a data option;
  mutable older : 'a data option;
}

let data {data; _} = data

let weight {weight; _} = weight

type 'a node = 'a data

type 'a tt =
  | Empty of {capacity : int}
  | Inited of {
      capacity : int;
      mutable length : int;
      mutable total_weight : int;
      mutable newest : 'a data;
      mutable oldest : 'a data;
    }

type 'a t = 'a tt ref

let create capacity =
  if capacity <= 0 then
    raise
      (Invalid_argument "Ringo.Weight_dll.create: negative or null capacity")
  else ref (Empty {capacity})

let capacity dll =
  match !dll with Empty {capacity} | Inited {capacity; _} -> capacity

let length dll = match !dll with Empty _ -> 0 | Inited {length; _} -> length

let total_weight dll =
  match !dll with Empty _ -> 0 | Inited {total_weight; _} -> total_weight

let pop_until_it_fits dll data_weight =
  assert (
    match dll with
    | Empty {capacity} | Inited {capacity; _} -> capacity >= data_weight) ;
  let rec loop erased = function
    | Empty _ -> (erased, dll)
    | Inited {capacity; total_weight; length; newest; oldest} as dll -> (
        assert (newest.newer = None) ;
        assert (oldest.older = None) ;
        if capacity - total_weight >= data_weight then (erased, dll)
        else
          let erased_elt = oldest in
          match oldest.newer with
          | None -> (erased_elt :: erased, Empty {capacity})
          | Some new_oldest ->
              new_oldest.older <- None ;
              erased_elt.newer <- None ;
              let new_dll =
                Inited
                  {
                    capacity;
                    total_weight = total_weight - erased_elt.weight;
                    length = pred length;
                    newest;
                    oldest = new_oldest;
                  }
              in
              loop (erased_elt :: erased) new_dll)
  in
  (* NOTE: this could return an empty Inited dll, but it is only used to prepare
     for an inserting a new element anyway so it'll be fixed immediately. *)
  loop [] dll

let add_and_return_erased dll (data, data_weight) : 'a node option * 'a list =
  if data_weight > capacity dll then (None, [])
  else
    let erased_nodes, new_dll = pop_until_it_fits !dll data_weight in
    let erased_data = List.rev_map (fun {data; _} -> data) erased_nodes in
    let new_node = {data; newer = None; older = None; weight = data_weight} in
    (match new_dll with
    | Empty {capacity} ->
        dll :=
          Inited
            {
              capacity;
              total_weight = data_weight;
              length = 1;
              newest = new_node;
              oldest = new_node;
            }
    | Inited {capacity; total_weight; length; newest; oldest} ->
        assert (newest.newer = None) ;
        new_node.older <- Some newest ;
        newest.newer <- Some new_node ;
        dll :=
          Inited
            {
              capacity;
              total_weight = total_weight + data_weight;
              length = succ length;
              newest = new_node;
              oldest;
            }) ;
    (Some new_node, erased_data)

let add dll data = fst @@ add_and_return_erased dll data

let add_list dll l =
  let capacity = capacity dll in
  let rec suffix_that_fits suffix_weight acc = function
    | [] -> acc
    | ((_, w) as h) :: t ->
        let new_suffix_weight = w + suffix_weight in
        if new_suffix_weight > capacity then
          suffix_that_fits suffix_weight acc t
        else suffix_that_fits new_suffix_weight (h :: acc) t
  in
  (* TODO? avoid calling [add] multiple times which calls [pop_until_it_fits]
     multiple times; instead call [pop_until_it_fits] once and then add in bulk
     possibly by-hand. *)
  List.map
    (fun e -> match add dll e with None -> assert false | Some e -> e)
    (suffix_that_fits 0 [] (List.rev l))

let clear dll =
  match !dll with
  | Empty _ -> ()
  | Inited {capacity; _} -> dll := Empty {capacity}

let rec fold_node f acc node =
  let acc = f acc node in
  match node.older with None -> acc | Some older -> fold_node f acc older

let fold dll ~init ~f =
  match !dll with
  | Empty _ -> init
  | Inited {newest; _} -> fold_node f init newest

let rec fold_node_oldest_first f acc node =
  let acc = f acc node in
  match node.newer with
  | None -> acc
  | Some newer -> fold_node_oldest_first f acc newer

let fold_oldest_first dll ~init ~f =
  match !dll with
  | Empty _ -> init
  | Inited {oldest; _} -> fold_node_oldest_first f init oldest

let oldest_element dll =
  match !dll with
  | Empty _ -> None
  | Inited {oldest; _} -> Some oldest

let newest_element dll =
  match !dll with
  | Empty _ -> None
  | Inited {newest; _} -> Some newest

let elements t = fold t ~init:[] ~f:(fun acc elt -> elt :: acc)

let elements_data t = fold t ~init:[] ~f:(fun acc elt -> elt.data :: acc)

let remove c node =
  match !c with
  | Empty _ -> assert false
  | Inited dll ->
      (match (node.newer, node.older) with
      | None, None ->
          assert (node == dll.newest) ;
          assert (node == dll.oldest) ;
          assert (dll.length = 1) ;
          c := Empty {capacity = dll.capacity}
      | None, Some older ->
          assert (dll.length >= 2) ;
          older.newer <- None ;
          dll.newest <- older
      | Some newer, None ->
          assert (dll.length >= 2) ;
          newer.older <- None ;
          dll.oldest <- newer
      | Some newer, Some older ->
          assert (dll.length >= 3) ;
          newer.older <- node.older ;
          older.newer <- node.newer) ;
      node.newer <- None ;
      node.older <- None ;
      dll.total_weight <- dll.total_weight - node.weight ;
      dll.length <- pred dll.length

let remove_newest c =
  match !c with
  | Empty _ -> None
  | Inited {capacity; length; total_weight; newest; oldest} ->
      assert (length > 0);
      assert (newest.newer = None) ;
      (match newest.older with
      | None ->
          (* [newest] is the only element *)
          assert (newest == oldest);
          assert (length = 1);
          c := Empty {capacity}
      | Some older ->
          assert (length >= 2);
          newest.older <- None ;
          older.newer <- None ;
          c := Inited {capacity; length = length - 1; total_weight =
            total_weight - newest.weight; newest = older; oldest}) ;
      Some newest

let remove_oldest c =
  match !c with
  | Empty _ -> None
  | Inited {capacity; length; total_weight; newest; oldest} ->
      assert (length > 0);
      assert (oldest.older = None) ;
      (match oldest.newer with
      | None ->
          (* [oldest] is the only element *)
          assert (oldest == newest);
          assert (length = 1);
          c := Empty {capacity}
      | Some newer ->
          assert (length >= 2);
          oldest.newer <- None ;
          newer.older <- None ;
          c := Inited {capacity; length = length - 1; total_weight =
            total_weight - oldest.weight; newest; oldest = newer}) ;
      Some oldest

let promote dll node =
  match !dll with
  | Empty _ -> assert false
  | Inited dll ->
      if dll.newest == node then
        ()
      else begin
        (* newest, promote neighbors *)
        begin match (node.newer, node.older) with
        | (None, None) -> assert false
        | (None, Some _next) -> assert false
        | (Some newer, None) ->
            newer.older <- None;
            dll.oldest <- newer
        | (Some newer, Some older) ->
            newer.older <- node.older;
            older.newer <- node.newer
        end;
        (* promote node to newest *)
        dll.newest.newer <- Some node ;
        node.newer <- None;
        node.older <- Some dll.newest;
        dll.newest <- node
     end