Source file fs_nameenc_impl.ml

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

(* Separated from fs.mli so that some internal functions can be used
   in fs_tree.ml *)

open Fs_types

type t = (module NAMEENC)

module Bits8 = struct
  type name = Name.t

  let to_segment s =
    Segment.unsafe_of_bits (String.length s * 8 + 8)
      (s ^ "\000")

  let of_segment seg =
    let len, s = Segment.to_bits seg in
    let slen = String.length s in
    if slen = 0 then begin
      Format.eprintf "seg? <%a>@." Segment.pp seg;
      assert (slen <> 0);
    end;
    if slen * 8 = len && s.[slen - 1] = '\000' then
      Some (String.sub s 0 (slen - 1))
    else None

  let test s =
    let seg = to_segment s in
    Format.eprintf "%s => %a@." s Segment.pp seg;
    match of_segment seg with
    | None ->
        Format.eprintf "%s => None@." s;
        assert false
    | Some s' ->
        Format.eprintf "%s => %s@." s s';
        assert (s = s')
end

module Bits6 = struct
  type name = Name.t

  let encode =
    let c0 = Char.code '0' in
    let ca = Char.code 'a' in
    fun c ->
      let c' = Char.code c in
      match c with
      | '0'..'9' -> c' - c0 + 1 (* 1 .. 10 *)
      | 'a'..'z' -> c' - ca + 11 (* 11 .. 36 *)
      | '_' -> 37
      | '-' -> 38 (* big_map ID < 0 *)
      | c -> Format.eprintf "Unknown char %c@." c; assert false

  let decode =
    let c0 = Char.code '0' in
    let ca = Char.code 'a' in
    fun c ->
      if c = 0 then '\000'
      else if c <= 10 then Char.chr (c - 1 + c0)
      else if c <= 36 then Char.chr (c - 11 + ca)
      else if c = 37 then '_'
      else if c = 38 then '-'
      else assert false

  let to_segment s =
    let len = String.length s + 1 in
    let bits = len * 6 in
    let bytes = (bits + 7) / 8 in
    (* 543210 54|3210 5432|10 543210| .. *)
    let a = Array.init len (fun i ->
        if i = len - 1 then 0
        else encode @@ String.unsafe_get s i)
    in
    (*
    Array.iteri (fun _i n -> Format.eprintf "%d;" n) a;
    Format.eprintf "@.";
    *)
    let buf = String.init bytes (fun i ->                (*  i = 0, 1, 2, 3 *)
        let jl = i * 8 / 6 in                            (* jl = 0, 1, 2, 4 *)
        let jr = jl + 1 in                               (* jr = 1, 2, 3, 5 *)
        let shiftl = 8 - (jr * 6) mod 8 in           (* shiftl = 2, 4, 6, 2 *)
        let shiftr = 6 - shiftl in                   (* shiftr = 4, 2, 0, 4 *)
        let cl = Array.unsafe_get a jl in
        let cr = if jr >= len then 0 else Array.unsafe_get a jr in
        let c = ((cl lsl shiftl) + (cr lsr shiftr)) land 0xff in
        (* Format.eprintf "%d: %d <- %d %d -> %d ==> %d@." i cl shiftl cr shiftr c; *)
        Char.chr c)
    in
    Segment.unsafe_of_bits bits buf

  let of_segment seg =
    let len, s = Segment.to_bits seg in
    let slen = String.length s in
    if len mod 6 <> 0 then None
    else
      let chars = len / 6 - 1 in
      let f i =
        (* 543210 54|3210 5432|10 543210| .. *)
        let b = i * 6 / 8 in
        let shiftr = 10 - (i * 6) mod 8 in
        let c1 = String.unsafe_get s b in
        let c2 = if b + 1 < slen then String.unsafe_get s (b+1) else '\000' in
        decode (((Char.code c1 lsl 8 + Char.code c2) lsr shiftr) land 0b111111)
      in
      if f (len / 6) = '\000' then
        Some (String.init chars f)
      else None

  let test s =
    let seg = to_segment s in
    Format.eprintf "%s => %a@." s Segment.pp seg;
    match of_segment seg with
    | None ->
        Format.eprintf "%s => None@." s;
        assert false
    | Some s' ->
        Format.eprintf "%s => %s@." s s';
        assert (s = s')
end

let is_all_hex s len =
  let rec f i =
    match String.unsafe_get s i with
    | '0'..'9' | 'a'..'f' ->
        if i = 0 then true
        else f (i-1)
    | _ -> false
  in
  f (len-1)

let c0 = Char.code '0'
let ca = Char.code 'a'

let encode5 neg s =
  let len = String.length s in
  let bits = len * 5 + 1 in
  let bytes = (bits + 7) / 8 in
  let enc i =
    if i > len then 0
    else if i = len then 0b10000
    else
      match String.unsafe_get s i with
      | '0'..'9' as c -> Char.code c - c0
      | 'a'..'f' as c -> Char.code c - ca + 10
      | _ -> assert false
  in
  (* f3210 f32|10 f3210 f|3210 f321|0 .. *)
  let buf = String.init bytes (fun i ->
      let j1 = i * 8 / 5 in
      let c1 = if j1 = 0 && neg then 0b10000 + enc j1 else enc j1 in
      let shift1 = 16 - (5 - (i * 8) mod 5) in
      let x1 = c1 lsl shift1 in

      let c2 = enc (j1+1) in
      let shift2 = shift1 - 5 in
      let x2 = if shift2 >= 0 then c2 lsl shift2 else c2 lsr (-shift2) in

      let c3 = enc (j1+2) in
      let shift3 = shift2 - 5 in
      let x3 = if shift3 >= 0 then c3 lsl shift3 else c3 lsr (-shift3) in

      let n = ((x1 + x2 + x3) lsr 8) land 255 in
      Char.chr n)
  in
  Segment.unsafe_of_bits bits buf

let decode5 seg =
  let neg =
    match Segment.get_side seg 0 with
    | Some Right -> true
    | _ -> false
  in
  let len, s = Segment.to_bits seg in
  let slen = String.length s in
  if len mod 5 <> 1 then None
  else
    let chars = len / 5 in
    let f i =
      (* f3210 f32|10 f3210 f|3210 f321|0.. *)
      let b = i * 5 / 8 in
      let shiftr = 11 - ((i * 5) mod 8) in
      let c1 = Char.code @@ String.unsafe_get s b in
      let c2 = if b + 1 < slen then Char.code @@ String.unsafe_get s (b+1) else 0 in
      let x = ((c1 lsl 8 + c2) lsr shiftr) land 0b01111 in
      if x < 10 then Char.chr (c0 + x)
      else Char.chr (ca + x - 10)
    in
    let s = String.init chars f in
    if neg then Some ("-" ^ s) else Some s

module Compressed(Arg : sig val path : string end) = struct
  (*
     Hex [0-9a-f]+, 2n chars, 2n >= 16
     :  4bits/char + 1 bit L flag, no terminator.
        Assuming all the file names at the same directory
        have the same length.

     Other Hex [0-9a-f]+
     :  5bits/char + 2 bit RL flag, 1 bit terminator

     Others:
     :  32bit hash + 2 bit RR flag
  *)

  type name = Name.t

  module H = Stdlib.Hashtbl
  let keyword_bytes = 2
  module B = Hashfunc.Blake2B(struct let bytes = keyword_bytes end)

  let tbl = H.create 100
  let tbl' = H.create 100

  let gen_hash s =
    let b4 = B.hash_string s in
    (* XXX can be optimized *)
    let seg0 = Segment.unsafe_of_bits (keyword_bytes * 8) b4 in
    let seg = Segment.unfat [`Right; `Right; `Segment seg0] in
    let nseg0 = Segment.to_bits seg0 in
    (seg, nseg0)

  (** Recording known names

     Without the persistent record of the known names, Plebeia cannot
     decode segments properly to names which are not yet known to it.
  *)
  module Record = struct
    module Lock = struct
      let path = Arg.path ^ ".lock"

      let with_lock f =
        let fd = Unix.openfile path [O_CREAT; O_RDWR] 0o644 in
        Unix.lockf fd F_LOCK 0;
        Utils.Exn.protect f (fun () -> Unix.close fd);
    end

    let tmp_path = Arg.path ^ ".tmp"
    let new_path = Arg.path ^ ".new"

    let load () =
      let file_exists =
        if Sys.file_exists Arg.path then true
        else if Sys.file_exists new_path then begin
          (* System might have crashed during [Unix.rename new_path Arg.path]. *)
          Unix.rename new_path Arg.path;
          true
        end else false
      in
      if not file_exists then []
      else
        let ic = open_in Arg.path in
        let rec loop rev_acc =
          match input_line ic with
          | exception End_of_file -> List.rev rev_acc
          | s -> loop (s::rev_acc)
        in
        let names = loop [] in
        close_in ic;
        names

    (* This is crash recoverable:
       If the system crashes during writing the name file:

       - The original name file is there, untouched.
       - The tree with new names is not committed yet.
    *)
    let save names =
      let names' = load () in
      let names = List.sort_uniq compare (names @ names') in
      if names = names' then ()
      else
        let len = List.length names in
        if len > 1000 then
          invalid_arg "Fs_nameenc_impl.Compressed.save: too many names, likely a bug of encoding scheme"
        else
          if Sys.file_exists new_path then Unix.unlink new_path;
          let oc = open_out tmp_path in
          List.iter (fun n -> output_string oc n; output_char oc '\n') names;
          close_out oc;
          Unix.rename tmp_path new_path; (* If [new_] exists, it is consistent *)
          Unix.rename new_path Arg.path (* If [Arg.path] does not exist, [new_path] should carry the latest data *)

    let load () = Lock.with_lock load
    let save names = Lock.with_lock @@ fun () -> save names
  end

  let init dirs =
    let f s =
      let seg, nseg0 = gen_hash s in
      Hashtbl.replace tbl s seg;
      Hashtbl.replace tbl' nseg0 s
    in
    Hashtbl.reset tbl;
    Hashtbl.reset tbl';
    List.iter (fun dir ->
        let len = String.length dir in
        if is_all_hex dir len then ()
        else if String.unsafe_get dir 0 = '-' then begin
          let ss = String.sub dir 1 (len-1) in
          if is_all_hex ss (len-1) then ()
          else f dir end
        else f dir) dirs

  let to_segment_aux ~may_record s =
    let len = String.length s in
    let default s =
      match H.find_opt tbl s with
      | Some seg -> seg
      | None ->
          let seg, nseg0 = gen_hash s in
          H.replace tbl s seg;
          if H.mem tbl' nseg0 then assert false; (* collision *)
          H.replace tbl' nseg0 s;
          (* When a new name is found, update the name file *)
          if may_record then begin
            Log.notice "Fs_nameenc: new name: %s" s;
            Record.save @@ List.of_seq @@ H.to_seq_keys tbl;
          end;
          seg
    in
    if is_all_hex s len then
      if len >= 16 && len mod 2 = 0 then
        Segment.unfat [`Left; `Segment (Segment.unsafe_of_bits (len * 4) (Hex.to_string (`Hex s)))]
      else
        Segment.unfat [`Right; `Left; `Segment (encode5 false s)]
    else if String.unsafe_get s 0 = '-' then
      let ss = String.sub s 1 (len-1) in
      if is_all_hex ss (len-1) then
        Segment.unfat [`Right; `Left; `Segment (encode5 true ss)]
      else
        default s
    else
      default s

  let to_segment = to_segment_aux ~may_record:true

  let of_segment seg =
    match Segment.get_side seg 0 with
    | None -> None
    | Some Left ->
        let seg = Segment.drop 1 seg in
        let nsides, s = Segment.to_bits seg in
        assert (String.length s = nsides / 8);
        let `Hex h = Hex.of_string s in
        Some h
    | Some Right ->
        match Segment.get_side seg 1 with
        | None -> None
        | Some Left ->
            decode5 @@ Segment.drop 2 seg
        | Some Right ->
            let seg = Segment.drop 2 seg in
            let bits = Segment.to_bits seg in
            match H.find_opt tbl' bits  with
            | Some n -> Some n
            | None ->
                (* When the name for the segment is not found,
                   retry the resolution after reloading the names *)
                Log.notice "Fs_nameenc: reloading";
                let names = Record.load () in
                init names;
                match H.find_opt tbl' bits  with
                | Some n -> Some n
                | None ->
                    Format.eprintf "seg RR %a@." Segment.pp seg; assert false

  module Test = struct
    let to_segment_no_register = to_segment_aux ~may_record:false

    let test s =
      Format.eprintf "%s => ...@." s;
      let seg = to_segment s in
      Format.eprintf "=> %a@." Segment.pp seg;
      match of_segment seg with
      | None ->
          Format.eprintf "%s => None@." s;
          assert false
      | Some s' ->
          Format.eprintf "%s => %s@." s s';
          assert (s = s');
          Format.eprintf "ok@."
  end

end

module Compressed' = struct
  (*
     Hex [0-9a-f]+, 2n chars, 2n >= 16
     :  4bits/char + 1 bit L flag, no terminator.
        Assuming all the file names at the same directory
        have the same length.

     Other
     :  8bit/char
  *)

  type name = Name.t

  let to_segment s =
    let len = String.length s in
    if is_all_hex s len && len >= 16 && len mod 2 = 0 then
      Segment.unfat [`Left; `Segment (Segment.unsafe_of_bits (len * 4) (Hex.to_string (`Hex s)))]
    else
      Segment.unfat [`Right; `Segment (Bits8.to_segment s)]

  let of_segment seg =
    match Segment.get_side seg 0 with
    | None -> None
    | Some Left ->
        let seg = Segment.drop 1 seg in
        let nsides, s = Segment.to_bits seg in
        assert (String.length s = nsides / 8);
        let `Hex h = Hex.of_string s in
        Some h
    | Some Right ->
        Bits8.of_segment @@ Segment.drop 1 seg

  let test s =
    Format.eprintf "%s => ...@." s;
    let seg = to_segment s in
    Format.eprintf "=> %a@." Segment.pp seg;
    match of_segment seg with
    | None ->
        Format.eprintf "%s => None@." s;
        assert false
    | Some s' ->
        Format.eprintf "%s => %s@." s s';
        assert (s = s');
        Format.eprintf "ok@."
end

module Compressed'' = struct
  (*
     Hex [0-9a-f]+, 2n chars, 2n >= 16
     :  4bits/char + 1 bit L flag, no terminator.
        Assuming all the file names at the same directory
        have the same length.

     Other Hex [0-9a-f]+
     :  5bits/char + 2 bit RL flag, 1 bit terminator

     Other
     :  8bit/char
  *)

  type name = Name.t

  let default s =
    Segment.unfat [`Right; `Right; `Segment (Bits8.to_segment s)]

  let to_segment s =
    let len = String.length s in
    if is_all_hex s len then
      if len >= 16 && len mod 2 = 0 then
        Segment.unfat [`Left; `Segment (Segment.unsafe_of_bits (len * 4) (Hex.to_string (`Hex s)))]
      else
        Segment.unfat [`Right; `Left; `Segment (encode5 false s)]
    else if String.unsafe_get s 0 = '-' then
      let ss = String.sub s 1 (len-1) in
      if is_all_hex ss (len-1) then
        Segment.unfat [`Right; `Left; `Segment (encode5 true ss)]
      else
        default s
    else
      default s

  let of_segment seg =
    match Segment.get_side seg 0 with
    | None -> None
    | Some Left ->
        let seg = Segment.drop 1 seg in
        let nsides, s = Segment.to_bits seg in
        assert (String.length s = nsides / 8);
        let `Hex h = Hex.of_string s in
        Some h
    | Some Right ->
        match Segment.get_side seg 1 with
        | None -> None
        | Some Left ->
            decode5 @@ Segment.drop 2 seg
        | Some Right ->
            let seg = Segment.drop 2 seg in
            Bits8.of_segment seg

  let test s =
    Format.eprintf "%s => ...@." s;
    let seg = to_segment s in
    Format.eprintf "=> %a@." Segment.pp seg;
    match of_segment seg with
    | None ->
        Format.eprintf "%s => None@." s;
        assert false
    | Some s' ->
        Format.eprintf "%s => %s@." s s';
        assert (s = s');
        Format.eprintf "ok@."
end

let bits8 : t = (module Bits8)
let bits6 : t = (module Bits6)
let compressed path : t = (module Compressed(struct let path = path end))
let compressed' : t = (module Compressed')
let compressed'' : t = (module Compressed'')

let to_segment (module Enc : NAMEENC) = Enc.to_segment
let of_segment (module Enc : NAMEENC) = Enc.of_segment
let to_segments conv = List.map (to_segment conv)
let of_segments conv = Option.mapM (of_segment conv)