Source file node_type.ml

(*****************************************************************************)
(*                                                                           *)
(* Open Source License                                                       *)
(* Copyright (c) 2019 Arthur Breitman                                        *)
(* Copyright (c) 2019,2020 DaiLambda, Inc. <contact@dailambda.jp>            *)
(*                                                                           *)
(* 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.                                                 *)
(*                                                                           *)
(*****************************************************************************)

open Utils
open Result_lwt.Infix

type hashed =
  | Hashed of Hash.Prefix.t
  | Not_Hashed
  (** Type used to prove that if a node is hashed then so are its children.
      The type also provides the hash prefix as a witness.
  *)

type indexed =
  | Indexed of Index.t
  | Not_Indexed
  (** This rule expresses the following invariant : if a node is indexed, then
      its children are necessarily indexed. Less trivially, if an internal node is not
      indexed then at least one of its children is not yet indexed. The reason
      is that we never construct new nodes that just point to only existing nodes.
      This property guarantees that when we write internal nodes on
      disk, at least one of the child can be written adjacent to its parent. *)

type extender_witness =
  | Maybe_Extender
  | Not_Extender
  | Is_Extender

type node =
  | Disk of Index.t * extender_witness
  (* Represents a node stored on the disk at a given index, the node hasn't
     been loaded yet. Although it's considered hashed for simplicity's sake,
     reading the hash requires a disk access and is expensive. *)

  | View of view
  (* A view node is the in-memory structure manipulated by programs in order
     to compute edits to the context. New view nodes can be committed to disk
     once the computations are done. *)

  | Hash of Hash.t

and view =
  | Internal of node * node
               * indexed
               * hashed
  (* An internal node, left and right children *)

  | Bud of node option
          * indexed
          * hashed
  (* Buds represent the end of a segment and the beginning of a new tree. They
     are used whenever there is a natural hierarchical separation in the key
     or, in general, when one wants to be able to grab sub-trees. For instance
     the big_map storage of a contract in Tezos would start from a bud. *)

  | Leaf of Value.t
          * indexed
          * hashed
  (* Leaf of a tree, the end of a path, contains or points to a value.

     XXX Value can be huge.  It would be nice if we can load the hash without
     loading the value *)

  | Extender of Segment.t
                * node
                * indexed
                * hashed
  (* Extender node, contains a path to the next node. Represents implicitely
     a collection of internal nodes where one child is Null. *)
  (* XXX Q: [Extender] should only take encoded segments?
         A: Sounds reasonable once we stop using temporary _Extender creations
            in diff.ml *)

(* A trail represents the content of the memory stack when recursively exploring a tree.
   Constructing these trails from closure would be easier, but it would make it harder
   to port the code to C. The type parameters of the trail keep track of the type of each
   element on the "stack" using a product type. *)

type t = node

let indexed = function
  | Disk _ -> true
  | View ( Bud (_, Indexed _, _)
         | Leaf (_, Indexed _, _)
         | Internal (_, _, Indexed _, _)
         | Extender (_, _, Indexed _, _) ) -> true
  | View _ -> false
  | Hash _ -> false

let index_of_view = function
  | Bud (_, Indexed i, _) -> Some i
  | Leaf (_, Indexed i, _) -> Some i
  | Internal (_, _, Indexed i, _) -> Some i
  | Extender (_, _, Indexed i, _) -> Some i
  | Bud _ | Leaf _ | Internal _ | Extender _ -> None

let index = function
  | Disk (i,_) -> Some i
  | View v -> index_of_view v
  | Hash _ -> None

let hashed = function
  | Disk _ -> true
  | View (Bud (_, _, Hashed _)) -> true
  | View (Bud (_, _, Not_Hashed)) -> false
  | View (Leaf (_, _, Hashed _)) -> true
  | View (Leaf (_, _, Not_Hashed)) -> false
  | View (Internal (_, _, _, Hashed _)) -> true
  | View (Internal (_, _, _, Not_Hashed)) -> false
  | View (Extender (_, _, _, Hashed _)) -> true
  | View (Extender (_, _, _, Not_Hashed)) -> false
  | Hash _ -> true

(* Even if not available, quickly obtained from the disk *)
let hash_available = function
  | Disk _ -> true
  | View (Bud (_, _, Hashed _)) -> true
  | View (Bud (_, Indexed _, Not_Hashed)) -> true
  | View (Bud (_, _, Not_Hashed)) -> false
  | View (Leaf (_, _, Hashed _)) -> true
  | View (Leaf (_, Indexed _, Not_Hashed)) -> true
  | View (Leaf (_, _, Not_Hashed)) -> false
  | View (Internal (_, _, _, Hashed _)) -> true
  | View (Internal (_, _, Indexed _, Not_Hashed)) -> true
  | View (Internal (_, _, _, Not_Hashed)) -> false
  | View (Extender (_, _, _, Hashed _)) -> true
  | View (Extender (_, _, Indexed _, Not_Hashed)) -> true
  | View (Extender (_, _, _, Not_Hashed)) -> false
  | Hash _ -> true

let hash_prefix_of_view v =
  match v with
  | Bud (_, _, Hashed h)
  | Leaf (_, _, Hashed h)
  | Internal (_, _, _, Hashed h)
  | Extender (_, _, _, Hashed h) -> Some h
  | Bud (_, _, Not_Hashed)
  | Leaf (_, _, Not_Hashed)
  | Internal (_, _, _, Not_Hashed)
  | Extender (_, _, _, Not_Hashed) -> None

(* XXX compile time switch the invariant checks *)
module Invariant_view : sig
  val check : view -> view
end = struct
  (* XXX Do we use error_monad ? *)
  type Error.t += Node_invariant of string

  let () = Error.register_printer (function
      | Node_invariant s -> Some ("Node invariant: " ^ s)
      | _ -> None)

  let error_node_invariant s = Error (Node_invariant s)

  let view_shape_invariant : view -> (unit, Error.t) Result.t = function
    | Bud (None, _, _) -> Ok ()
    | Bud (Some (Disk _), _, _) ->
        (* We do not check whether Disk points non Bud node here *)
        Ok ()
    | Bud (Some (View (Bud _)), _, _) -> error_node_invariant "Bud: cannot have Bud"
    | Bud (Some (View (Leaf _)), _, _) -> error_node_invariant "Bud: cannot have Leaf"
    | Bud (Some (View (Internal _)), _, _) -> Ok ()
    | Bud (Some (View (Extender _)), _, _) -> Ok ()
    | Bud (Some (Hash _), _, _) -> Ok ()
    | Extender (_, Disk (_, Not_Extender), _, _) -> Ok ()
    | Extender (_, Disk (_, Is_Extender), _, _) -> error_node_invariant "Extender: cannot have Disk with Is_Extender"
    | Extender (_, Disk (_, Maybe_Extender), _, _) -> error_node_invariant "Extender: cannot have Disk with Maybe_Extender"
    | Extender (_, View (Extender _), _, _) -> error_node_invariant "Extender: cannot have Extender"
    | Extender (_, View _, _, _) -> Ok ()
    | Extender (_, Hash (_,""), _, _) -> Ok ()
    | Extender (_, Hash (_,_), _, _) ->
        error_node_invariant "Extender: cannot have Hash with postfix"
    | Leaf _ -> Ok ()
    | Internal _ -> Ok ()

  let view_indexed_invariant : view -> (unit, Error.t) Result.t = function
    | Bud (None, Indexed _, _) -> Ok ()
    | Bud (Some n, Indexed _, _) when indexed n -> Ok ()
    | Bud (_, Not_Indexed, _) -> Ok ()
    | Leaf (_, Indexed _, _) -> Ok ()
    | Leaf (_, Not_Indexed, _) -> Ok ()
    | Internal (l, r, Indexed _i, _) ->
        begin match index l, index r with
          | None, _ -> error_node_invariant "Internal: invalid Indexed"
          | _, None -> error_node_invariant "Internal: invalid Indeced"
  (* We abandoned this strict invariant on internals.
          | Some li, Some ri ->
              if Index.(i - li = one || i - ri = one) then Ok ()
              else error_node_invariant "Internal: invalid indices"
  *)
          | _ -> Ok () (* we now use fat internals *)
        end
    | Internal (_l, _r, Not_Indexed, _) -> Ok ()
    | Extender (_, n, Indexed _, _) when indexed n -> Ok ()
    | Extender (_, _, Not_Indexed, _) -> Ok ()
    | Bud (_, Indexed _, _)
    | Extender (_, _, Indexed _, _)  -> error_node_invariant "Invalid Indexed"

  let view_hashed_invariant : view -> (unit, Error.t) Result.t = function
    | Leaf _ -> Ok ()
    | Bud (None, _, _) -> Ok ()
    | Bud (_, _, Not_Hashed) -> Ok ()
    | Bud (Some n, _, Hashed _) when hash_available n -> Ok ()
    | Internal (l, r, _, Hashed _) when hash_available l && hash_available r -> Ok ()
    | Internal (_, _, _, Not_Hashed) -> Ok ()
    | Extender (_, n, _, Hashed _) when hash_available n -> Ok ()
    | Extender (_, _, _, Not_Hashed) -> Ok ()
    | _ -> error_node_invariant "Invalid Hashed"

  let extender_segment_length_invariant = function
    | Extender (seg, _, _, _) ->
        let len = Segment.length seg in
        if 0 < len && len <= Segment.max_length then Ok ()
        else error_node_invariant (Printf.sprintf "invalid segment length %d" len)
    | _ -> Ok ()

  let view_invariant : view -> (unit, Error.t) Result.t = fun v ->
    view_shape_invariant v >>? fun () ->
    view_indexed_invariant v >>? fun () ->
    view_hashed_invariant v >>? fun () ->
    extender_segment_length_invariant v

  let _check v =
    match view_invariant v with
    | Ok _ -> v
    | Error e -> failwith (Error.show e)

  (* Currently the check is off *)
  let check x = x
end

let _Internal (n1, n2, ir, hit) =
  Invariant_view.check @@ Internal (n1, n2, ir, hit)

let _Bud (nopt, ir, hit) =
  Invariant_view.check @@ Bud (nopt, ir, hit)

let _Leaf (v, ir, hit) =
  Invariant_view.check @@ Leaf (v, ir, hit)

let _Extender (seg, n, ir, hit) =
  Invariant_view.check @@ Extender (seg, n, ir, hit)

let new_leaf v = View (_Leaf (v, Not_Indexed, Not_Hashed))

(* XXX Why only new_extender is so "kind"?
   Sato: it was used in Cursor.remove_up
*)
let new_extender : Segment.segment -> node -> node = fun segment node ->
  if Segment.is_empty segment then node (* XXX We should simply reject it *)
  else
    (* This connects segments *)
    match node with
    | View (Extender (seg, n, _, _)) ->
        View (_Extender (Segment.append segment seg, n, Not_Indexed, Not_Hashed))
    | _ ->
        (* XXXX If the subnode is Disk, we have to make sure merging *)
        View (_Extender (segment, node, Not_Indexed, Not_Hashed))

let new_bud no = View (_Bud (no, Not_Indexed, Not_Hashed))

let new_internal n1 n2 = View (_Internal (n1, n2, Not_Indexed, Not_Hashed))

let may_forget = function
  | Disk _ as n -> Some n
  | View (Internal (_, _, Indexed i, _)) -> Some (Disk (i, Not_Extender))
  | View (Bud (_, Indexed i, _)) -> Some (Disk (i, Not_Extender))
  | View (Leaf (_, Indexed i, _)) -> Some (Disk (i, Not_Extender))
  | View (Extender (_, _, Indexed i, _)) -> Some (Disk (i, Is_Extender))
  | _ -> None

let rec pp ppf = function
  | Hash h ->
      Format.fprintf ppf "Hash %a" Hash.pp h
  | Disk (i, ew) ->
      Format.fprintf ppf "Disk (%a, %s)" Index.pp i
        (match ew with
         | Maybe_Extender -> "Maybe_Ex"
         | Not_Extender -> "Not_Ex"
         | Is_Extender -> "Is_Ex")
  | View v ->
      let f = function
        | Internal (n1, n2, _i, _h) ->
            Format.fprintf ppf "@[<v2>Internal@ L %a@ R %a@]"
              pp n1
              pp n2
        | Bud (None, _, _) ->
            Format.fprintf ppf "Bud None"
        | Bud (Some n, _, _) ->
            Format.fprintf ppf "@[<v2>Bud@ %a@]"
              pp n
        | Leaf (v, _, _) ->
            Format.fprintf ppf "Leaf %s"
              (let s = Value.to_hex_string v in
               if String.length s > 16 then String.sub s 0 16 ^ "..." else s)
        | Extender (seg, n, _, _) ->
            Format.fprintf ppf "@[<v2>Extender %s@ %a@]"
              (Segment.to_string seg)
              pp n
      in
      f v

module Mapper = struct
  type job =
    | BudSome of indexed * hashed
    | Do of t
    | Extender of Segment.segment * indexed * hashed
    | Internal of indexed * hashed

  type 'a state = job list * (t * 'a) list

  type 'a mkview =
    { bud : (node * 'a) option -> indexed -> hashed -> node * 'a
    ; extender : Segment.t -> node * 'a -> indexed -> hashed -> node * 'a
    ; internal : node * 'a -> node * 'a -> indexed -> hashed -> node * 'a
    ; leaf : Value.t -> indexed -> hashed -> node * 'a
    }

  let default_mkview =
    { bud= (function
          | None -> fun i h -> View (_Bud (None, i, h)), ()
          | Some (n,()) -> fun i h -> View (_Bud (Some n, i, h)), ())
    ; extender= (fun seg (n,()) i h -> View (_Extender (seg, n, i, h)), ())
    ; internal= (fun (nl,()) (nr,()) i h -> View (_Internal (nl, nr, i, h)), ())
    ; leaf= (fun v i h -> View (_Leaf (v, i, h)), ())
    }

  let step ~node ~view:mkview (js, ss) = match js, ss with
    | [], [s] -> `Right s
    | [], _ -> assert false
    | BudSome (i,h)::js, s::ss ->
        let s = mkview.bud (Some s) i h in
        `Left (js, (s::ss))
    | Extender (seg,i,h)::js, s::ss ->
        let s = mkview.extender seg s i h  in
        `Left (js, (s::ss))
    | Internal (i,h)::js, sr::sl::ss ->
        let s = mkview.internal sl sr i h in
        `Left (js, (s::ss))
    | (BudSome _ | Extender _ | Internal _)::_, _ -> assert false
    | Do n::js, ss ->
        match node n with
        | `Return s -> `Left (js, (s::ss))
        | `Continue v ->
            match v with
            | Leaf (v, i, h) ->
                let s = mkview.leaf v i h in
                `Left (js, s::ss)
            | Bud (None, i, h) ->
                let s = mkview.bud None i h in
                `Left (js, s::ss)
            | Bud (Some n, i, h) ->
                `Left ((Do n :: BudSome (i,h) :: js), ss)
            | Extender (seg, n, i, h) ->
                `Left ((Do n :: Extender (seg, i,h) :: js), ss)
            | Internal (nl, nr, i, h) ->
                `Left ((Do nl :: Do nr :: Internal (i,h) :: js), ss)

  let interleaved ~node ~view n =
    let js_ss = [Do n], [] in
    let stepper x = step ~node ~view x in
    stepper, js_ss

  let map ~node ~view n =
    let rec loop js_ss = match step ~node ~view js_ss with
      | `Right res -> res
      | `Left js_ss -> loop js_ss
    in
    loop ([Do n], [])
end

module Fold = struct
  type job =
    | BudSome of indexed * hashed
    | Do of t
    | Extender of Segment.segment * indexed * hashed
    | Internal of indexed * hashed

  type 'a state = job list * 'a list

  type 'a leave =
    { bud : 'a option -> indexed -> hashed -> 'a
    ; extender : Segment.t -> 'a -> indexed -> hashed -> 'a
    ; internal : 'a -> 'a -> indexed -> hashed -> 'a
    ; leaf : Value.t -> indexed -> hashed -> 'a
    }

  let default_leave_rebuild =
    { bud= (function
          | None -> fun i h -> View (_Bud (None, i, h))
          | Some n -> fun i h -> View (_Bud (Some n, i, h)))
    ; extender= (fun seg n i h -> View (_Extender (seg, n, i, h)))
    ; internal= (fun nl nr i h -> View (_Internal (nl, nr, i, h)))
    ; leaf= (fun v i h -> View (_Leaf (v, i, h)))
    }

  let step ~enter ~leave (js, ss) = match js, ss with
    | [], [s] -> `Right s
    | [], _ -> assert false
    | BudSome (i,h)::js, s::ss ->
        let s = leave.bud (Some s) i h in
        `Left (js, (s::ss))
    | Extender (seg,i,h)::js, s::ss ->
        let s = leave.extender seg s i h  in
        `Left (js, (s::ss))
    | Internal (i,h)::js, sr::sl::ss ->
        let s = leave.internal sl sr i h in
        `Left (js, (s::ss))
    | (BudSome _ | Extender _ | Internal _)::_, _ -> assert false
    | Do n::js, ss ->
        match enter n with
        | `Return s -> `Left (js, (s::ss))
        | `Continue v ->
            match v with
            | Leaf (v, i, h) ->
                let s = leave.leaf v i h in
                `Left (js, s::ss)
            | Bud (None, i, h) ->
                let s = leave.bud None i h in
                `Left (js, s::ss)
            | Bud (Some n, i, h) ->
                `Left ((Do n :: BudSome (i,h) :: js), ss)
            | Extender (seg, n, i, h) ->
                `Left ((Do n :: Extender (seg, i,h) :: js), ss)
            | Internal (nl, nr, i, h) ->
                `Left ((Do nl :: Do nr :: Internal (i,h) :: js), ss)

  let interleaved ~enter ~leave n =
    let js_ss = [Do n], [] in
    let stepper x = step ~enter ~leave x in
    stepper, js_ss

  let fold ~enter ~leave n =
    let rec loop js_ss = match step ~enter ~leave js_ss with
      | `Right res -> res
      | `Left js_ss -> loop js_ss
    in
    loop ([Do n], [])
end

module Fold' = struct
  type 'b job =
    | BudSome of indexed * hashed
    | Do of 'b * t
    | Extender of Segment.segment * indexed * hashed
    | Internal of indexed * hashed

  type ('b, 'a) state = 'b job list * 'a list

  type 'a leave =
    { bud : 'a option -> indexed -> hashed -> 'a
    ; extender : Segment.t -> 'a -> indexed -> hashed -> 'a
    ; internal : 'a -> 'a -> indexed -> hashed -> 'a
    ; leaf : Value.t -> indexed -> hashed -> 'a
    }

  let default_leave_rebuild =
    { bud= (function
          | None -> fun i h -> View (_Bud (None, i, h))
          | Some n -> fun i h -> View (_Bud (Some n, i, h)))
    ; extender= (fun seg n i h -> View (_Extender (seg, n, i, h)))
    ; internal= (fun nl nr i h -> View (_Internal (nl, nr, i, h)))
    ; leaf= (fun v i h -> View (_Leaf (v, i, h)))
    }

  let step ~enter ~leave (js, ss) = match js, ss with
    | [], [s] -> `Right s
    | [], _ -> assert false
    | BudSome (i,h)::js, s::ss ->
        let s = leave.bud (Some s) i h in
        `Left (js, (s::ss))
    | Extender (seg,i,h)::js, s::ss ->
        let s = leave.extender seg s i h  in
        `Left (js, (s::ss))
    | Internal (i,h)::js, sr::sl::ss ->
        let s = leave.internal sl sr i h in
        `Left (js, (s::ss))
    | (BudSome _ | Extender _ | Internal _)::_, _ -> assert false
    | Do (x, n)::js, ss ->
        match enter x n with
        | `Return s -> `Left (js, (s::ss))
        | `Continue (xs, v) ->
            match xs, v with
            | [], Leaf (v, i, h) ->
                let s = leave.leaf v i h in
                `Left (js, s::ss)
            | [], Bud (None, i, h) ->
                let s = leave.bud None i h in
                `Left (js, s::ss)
            | [x], Bud (Some n, i, h) ->
                `Left ((Do (x, n) :: BudSome (i,h) :: js), ss)
            | [x], Extender (seg, n, i, h) ->
                `Left ((Do (x, n) :: Extender (seg, i,h) :: js), ss)
            | [xl;xr], Internal (nl, nr, i, h) ->
                `Left ((Do (xl,nl) :: Do (xr,nr) :: Internal (i,h) :: js), ss)
            | _ -> assert false

  let interleaved ~enter ~leave x n =
    let js_ss = [Do (x, n)], [] in
    let stepper x = step ~enter ~leave x in
    stepper, js_ss

  let fold ~enter ~leave x n =
    let rec loop js_ss = match step ~enter ~leave js_ss with
      | `Right res -> res
      | `Left js_ss -> loop js_ss
    in
    loop ([Do (x, n)], [])
end

let hash_of_view =
  function
  | ( Bud (_, _, Hashed hp)
    | Leaf (_, _, Hashed hp))
    | (Internal (_, _, _, Hashed hp)) -> Some (hp,"")
  | Extender (seg, _, _, Hashed hp) -> Some (hp, Segment.Serialization.encode seg)
  | _ -> None

let hash_of_node node =
  match node with
  | Disk _ -> (None, None)
  | Hash h -> (None, Some h)
  | View v -> (Some v, hash_of_view v)

open Gen
open Gen.Infix

(* Tricky.  In one Bud level, we cannot have more than [Segment.max_length]
     level of segments, otherwise removing nodes from it might create an Extender
     with very long segment over this limit.
  *)
let rec gen_leaf = value >>| new_leaf

and gen_bud_none = return (new_bud None)

and gen_bud depth =
  let depth0 = depth in
  if depth <= 1 then return (new_bud None)
  else
    int depth >>= fun n ->
    if n = 0 then return (new_bud None)
    else begin
      let depth = depth - 1 in
      one_of [gen_extender depth; gen_internal depth] >>= fun n ->
      begin match n with
        | View (Leaf _) -> gen_bud depth0
        | View (Bud _) -> gen_bud depth0
        | _ -> return (new_bud (Some n))
      end
    end

and gen_internal depth =
  let f =
    if depth <= 1 then
      (* we cannot create an internal*)
      one_of [gen_bud_none; gen_leaf]
    else
      let depth = depth - 1 in
      int depth >>= fun n ->
      if n = 0 then one_of [gen_bud_none; gen_leaf]
      else one_of [gen_internal depth;
                   gen_extender depth;
                   gen_bud depth]
  in
  f >>= fun n1 ->
  f >>| fun n2 ->
  let n = new_internal n1 n2 in
  match n with
  | View (Internal _) -> n
  | _ -> assert false

and gen_extender depth =
  if depth <= 1 then
    (* we cannot create an extender *)
    one_of [gen_bud_none; gen_leaf]
  else
    int depth >>= fun n ->
    if n = 0 then one_of [gen_bud_none; gen_leaf]
    else
      let limit = Int.max 1 (Int.min Segment.max_length depth) in
      segment (int_range (1, limit)) >>= fun seg ->

      let depth = depth - Segment.length seg in
      one_of [gen_internal depth; gen_bud depth] >>| fun n ->
      new_extender seg n