Source file patricia_storage.ml

open Utils
open Lwt.Syntax

(** Particia tree of commits *)
module Make(P : sig
    type key
    type value

    val parse : Storage.t -> Index.t -> [`Leaf of value | `Internal of Segment.t * Index.t * Segment.t * Index.t ]
    val write_leaf : Storage.t -> value -> (Index.t, Error.t) result
    val write_internal : Storage.t -> Segment.t * Index.t * Segment.t * Index.t -> (Index.t, Error.t) result

    val config : Storage.config

    val equal_value : value -> value -> bool
    val pp_value : Format.formatter -> value -> unit
    val segment_of_key : key -> Segment.segment
    val key_of_segment : Segment.segment -> key option
  end) = struct

  module Internal = struct

    include P

    type storage = Storage.t
    type node = desc ref
    and desc =
      | View of view
      | Disk of Index.t
    and view =
      | Leaf of
          Index.t option ref (* where saved *)
          * value
      | Internal of
          Index.t option ref (* where saved *)
          * Segment.t * node (* left *)
          * Segment.t * node (* right *)

    let index n = match !n with
      | Disk i -> Some i
      | View (Leaf (iopt, _)) -> !iopt
      | View (Internal (iopt, _, _, _, _)) -> !iopt

    (* Equality without loading *)
    let rec equal_node n1 n2 = equal_desc !n1 !n2

    and equal_desc n1 n2 =
      match n1, n2 with
      | View v1, View v2 -> equal_view v1 v2
      | Disk i1, Disk i2 -> i1 = i2
      | _ -> false

    and equal_view v1 v2 =
      match v1, v2 with
      | Leaf (io1, leaf_val1), Leaf (io2, leaf_val2) ->
          io1 = io2 && equal_value leaf_val1 leaf_val2
      | Internal (io1, sL1, nL1, sR1, nR1),
        Internal (io2, sL2, nL2, sR2, nR2) ->
          io1 = io2
          && Segment.equal sL1 sL2
          && equal_node nL1 nL2
          && Segment.equal sR1 sR2
          && equal_node nR1 nR2
      | _ -> false

    let rec pp_node ppf n = pp_desc ppf !n

    and pp_desc ppf = function
      | Disk i -> Format.fprintf ppf "Disk %a" Index.pp i
      | View v -> pp_view ppf v

    and pp_view ppf = function
      | Leaf (iopt, leaf_val) ->
          Format.fprintf ppf "@[Leaf (%a, %a)@]"
            (Format.option Index.pp) !iopt
            pp_value leaf_val
      | Internal (iopt, sL, nL, sR, nR) ->
          Format.fprintf ppf "@[<2>Internal@ (@[%a,@ %a, @[%a@],@ %a, @[%a@]@])@]"
            (Format.option Index.pp) !iopt
            Segment.pp sL
            pp_node nL
            Segment.pp sR
            pp_node nR

    module View = struct
      type t = view
      let pp = pp_view
      let equal = equal_view

      let parse storage i =
        match parse storage i with
        | `Leaf l -> Leaf (ref (Some i), l)
        | `Internal (l, li, r, ri) ->
            Internal (ref (Some i), l, ref (Disk li), r, ref (Disk ri))
    end

    module Node = struct
      open Storage

      type t = node

      let view storage n =
        match !n with
        | View v -> v
        | Disk i when i = Index.zero -> assert false
        | Disk i ->
            let v = View.parse storage i in
            n := View v;
            v

      let pp = pp_node

      let equal = equal_node

      let write storage n =
        let open Result.Syntax in
        let rec aux storage n =
          match !n with
          | Disk i -> Ok i
          | View v ->
              match v with
              | Leaf ({contents= Some i}, _) -> Ok i
              | Leaf (r, l) ->
                  let+ i = write_leaf storage l in
                  r := Some i;
                  i
              | Internal ({contents= Some i}, _,_,_,_) -> Ok i
              | Internal (r,sL,nL,sR,nR) ->
                  let* iL = aux storage nL in
                  let* iR = aux storage nR in
                  let+ i = write_internal storage (sL,iL,sR,iR) in
                  r := Some i;
                  i
        in
        let+ i = aux storage n in
        (set_last_root_index storage (Some i))

      let may_forget n =
        match !n with
        | View (Leaf ({contents= Some i}, _)) -> Some (ref (Disk i))
        | View (Internal ({contents= Some i}, _, _, _, _)) -> Some (ref (Disk i))
        | Disk i -> Some (ref (Disk i))
        | _ -> None

      let empty = ref (Disk Index.zero)

      let is_empty n =
        match !n with
        | Disk i -> Index.zero = i
        | _ -> false

      let mk_leaf leaf_val = ref (View (Leaf (ref None, leaf_val)))
    end

    let rec get storage node seg =
      if Node.is_empty node then None
      else
        match Segment.cut seg with
        | None -> Some node
        | Some (side, seg) ->
            match Node.view storage node with
            | Leaf _ -> None
            | Internal (_io, sL, nL, sR, nR) ->
                match side with
                | Left ->
                    begin match Segment.common_prefix seg sL with
                      | _prefix, seg', sL' when Segment.is_empty sL' ->
                          get storage nL seg'
                      | _ -> None
                    end
                | Right ->
                    begin match Segment.common_prefix seg sR with
                      | _prefix, seg', sR' when Segment.is_empty sR' ->
                          get storage nR seg'
                      | _ -> None
                    end

    let rec set storage node seg n =
      if Node.is_empty node then seg, n
      else
        let seg0 = seg in
        match Segment.cut seg with
        | None -> seg, n
        | Some (side, seg) ->
            match Node.view storage node with
            | Leaf _ -> seg0, n
            | Internal (_io, sL, nL, sR, nR) ->
                match side with
                | Left ->
                    begin match Segment.common_prefix seg sL with
                      | _prefix, seg', sL' when Segment.is_empty sL' ->
                          let seg'', nL' = set storage nL seg' n in
                          Segment.empty,
                          ref (View (Internal (ref None, Segment.append sL seg'', nL', sR, nR)))
                      | _prefix, seg', _sL' when Segment.is_empty seg' ->
                          Segment.empty,
                          ref (View (Internal (ref None, seg, n, sR, nR)))
                      | prefix, seg', sL' ->
                          let i =
                            match Segment.cut seg', Segment.cut sL' with
                            | Some (Left, seg'), Some (Right, sL') ->
                                ref (View (Internal (ref None, seg', n, sL', nL)))
                            | Some (Right, seg'), Some (Left, sL') ->
                                ref (View (Internal (ref None, sL', nL, seg', n)))
                            | _ -> assert false
                          in
                          Segment.empty, ref (View (Internal (ref None, prefix, i, sR, nR)))
                    end
                | Right ->
                    begin match Segment.common_prefix seg sR with
                      | _prefix, seg', sR' when Segment.is_empty sR' ->
                          let seg'', nR' = set storage nR seg' n in
                          Segment.empty,
                          ref (View (Internal (ref None, sL, nL, Segment.append sR seg'', nR')))
                      | _prefix, seg', _sR' when Segment.is_empty seg' ->
                          Segment.empty,
                          ref (View (Internal (ref None, sL, nL, seg, n)))
                      | prefix, seg', sR' ->
                          let i =
                            match Segment.cut seg', Segment.cut sR' with
                            | Some (Left, seg'), Some (Right, sR') ->
                                ref (View (Internal (ref None, seg', n, sR', nR)))
                            | Some (Right, seg'), Some (Left, sR') ->
                                ref (View (Internal (ref None, sR', nR, seg', n)))
                            | _ -> assert false
                          in
                          Segment.empty, ref (View (Internal (ref None, sL, nL, prefix, i)))
                    end

    let set storage node seg n =
      let seg, n = set storage node seg n in
      match Segment.cut seg with
      | None -> n
      | Some (Left, seg') ->
          ref (View (Internal (ref None, seg', n, Segment.empty, Node.empty)))
      | Some (Right, seg') ->
          ref (View (Internal (ref None, Segment.empty, Node.empty, seg', n)))

    type t = storage * node
    let get_storage = fst
    let get_node = snd
  end

  open Internal

  type value = P.value
  type key = P.key
  type t = Internal.t

  let config = P.config

  let empty storage = storage, Node.empty

  let write (storage, n) = Node.write storage n

  let commit (storage, _) = Storage.commit storage

  let flush (storage, _) = Storage.flush storage

  let update_reader (storage, _) =
    let* () = Storage.update_reader storage in
    match Storage.get_last_root_index storage with
    | Some idx -> Lwt.return (storage, ref (Disk idx))
    | None -> Lwt.return (storage, Node.empty)

  let create ?length ?resize_step_bytes ~key fn =
    let* storage = Storage.create ?length ?resize_step_bytes ~config ~key fn in
    Lwt.return @@ empty storage

  let open_existing_for_read fn =
    let open Result_lwt.Syntax in
    let*=? (conf, storage) = Storage.open_existing_for_read fn in
    if conf <> config then
      Lwt.return_error
        (Storage.Config_mismatch { actual= conf; expected= config })
    else
      let* t = update_reader (empty storage) in
      Lwt.return_ok t

  let open_for_write ?resize_step_bytes ~config ~key fn =
    let+ storage = Storage.open_for_write ?resize_step_bytes ~config ~key fn in
    (empty storage)

  let mem (storage, n) key =
    let seg = P.segment_of_key key in
    match get storage n seg with
    | None -> false
    | Some _ -> true

  let find (storage, n) key =
    let seg = P.segment_of_key key in
    match get storage n seg with
    | None -> None
    | Some n ->
        match Node.view storage n with
        | Leaf (_, t) -> Some t
        | _ -> assert false

  let add (storage, n) key leaf_val =
    let seg = P.segment_of_key key in
    let n = set storage n seg (Node.mk_leaf leaf_val) in
    (storage, n)

  let fold' f (storage, node) s =
    let rec aux seg node s =
      if Node.is_empty node then Lwt.return s
      else
        match Node.view storage node with
        | Internal (_, sL, nL, sR, nR) ->
            let* s' = aux (seg @ [`Left; `Segment sL]) nL s in
            aux (seg @ [`Right; `Segment sR]) nR s'
        | Leaf (_, leaf_val) -> f (Segment.unfat seg) leaf_val s
    in
    if Node.is_empty node then Lwt.return s
    else aux [] node s

  let fold f t s =
    fold' (fun seg leaf_val s ->
        match key_of_segment seg with
        | Some k -> f k leaf_val s
        | None -> assert false) t s

  let iter f t = fold (fun k v () -> f k v) t ()

  let mode (storage, _) = Storage.mode storage

  let close (storage,_ as t) =
    let* res = if mode t = Writer then
        let res = write t in
        let+ () = flush t in
        res
      else Lwt.return_ok ()
    in
    let+ () = Storage.close storage in
    res

  let may_forget (storage, n) =
    Option.map (fun n -> (storage, n)) @@ Node.may_forget n

  let copy_the_latest src dst =
    (* make sure it knows the latest *)
    let* src = update_reader src in
    (* copy the latest tree *)
    fold' (fun ch ent (storage, n) ->
        let n = set storage n ch (Node.mk_leaf ent) in
        Lwt.return (storage, n)) src dst

  let offline_gc ?resize_step_bytes fn =
    let open Result_lwt.Syntax in
    let fn' = fn ^ ".tmp" in
    let*=? src = open_existing_for_read fn in
    let* key = Storage.make_new_writer_key [] in
    let* dst = create ?resize_step_bytes ~key fn' in
    let* dst = copy_the_latest src dst in
    let* res =
      let open Result_lwt.Syntax in
      let*=? () = Lwt.return @@ write dst in
      let* () = flush dst in
      let*=? () = close src in
      close dst
    in
    match res with
    | Ok () ->
        let* () = Lwt_unix.rename fn (fn ^ ".old") in
        let+ () = Lwt_unix.rename fn' fn in
        Ok ()
    | Error _ as e ->
        let+ () = Lwt_unix.unlink fn' in
        e
end