You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
108 lines
3.0 KiB
108 lines
3.0 KiB
(*
|
|
* Copyright (c) Facebook, Inc. and its affiliates.
|
|
*
|
|
* This source code is licensed under the MIT license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*)
|
|
|
|
open NS0
|
|
include Map_intf
|
|
|
|
module Make (Key : sig
|
|
type t [@@deriving compare, sexp_of]
|
|
end) : S with type key = Key.t = struct
|
|
module KeyMap = Core.Map.Make_plain (Key)
|
|
module Key = KeyMap.Key
|
|
|
|
type key = Key.t
|
|
|
|
include KeyMap.Tree
|
|
|
|
let compare = compare_direct
|
|
|
|
let to_map t =
|
|
Core.Map.Using_comparator.of_tree ~comparator:Key.comparator t
|
|
|
|
let of_map m = Base.Map.Using_comparator.to_tree m
|
|
|
|
let merge_skewed x y ~combine =
|
|
of_map (Core.Map.merge_skewed (to_map x) (to_map y) ~combine)
|
|
|
|
let map_endo t ~f = map_endo map t ~f
|
|
|
|
let merge_endo t u ~f =
|
|
let change = ref false in
|
|
let t' =
|
|
merge t u ~f:(fun ~key side ->
|
|
let f_side = f ~key side in
|
|
( match (side, f_side) with
|
|
| (`Both (data, _) | `Left data), Some data' when data' == data ->
|
|
()
|
|
| _ -> change := true ) ;
|
|
f_side )
|
|
in
|
|
if !change then t' else t
|
|
|
|
let fold_until m ~init ~f ~finish =
|
|
let fold m ~init ~f =
|
|
let f ~key ~data s = f s (key, data) in
|
|
fold m ~init ~f
|
|
in
|
|
let f s (k, v) = f ~key:k ~data:v s in
|
|
Container.fold_until ~fold ~init ~f ~finish m
|
|
|
|
let root_key_exn m =
|
|
let@ {return} = with_return in
|
|
binary_search_segmented m `Last_on_left ~segment_of:(fun ~key ~data:_ ->
|
|
return key )
|
|
|> ignore ;
|
|
raise (Not_found_s (Atom __LOC__))
|
|
|
|
let choose_exn m =
|
|
let@ {return} = with_return in
|
|
binary_search_segmented m `Last_on_left ~segment_of:(fun ~key ~data ->
|
|
return (key, data) )
|
|
|> ignore ;
|
|
raise (Not_found_s (Atom __LOC__))
|
|
|
|
let choose m = try Some (choose_exn m) with Not_found_s _ -> None
|
|
let pop m = choose m |> Option.map ~f:(fun (k, v) -> (k, v, remove m k))
|
|
|
|
let pop_min_elt m =
|
|
min_elt m |> Option.map ~f:(fun (k, v) -> (k, v, remove m k))
|
|
|
|
let is_singleton m =
|
|
try
|
|
let l, _, r = split m (root_key_exn m) in
|
|
is_empty l && is_empty r
|
|
with Not_found_s _ -> false
|
|
|
|
let find_and_remove m k =
|
|
let found = ref None in
|
|
let m =
|
|
change m k ~f:(fun v ->
|
|
found := v ;
|
|
None )
|
|
in
|
|
Option.map ~f:(fun v -> (v, m)) !found
|
|
|
|
let pp pp_k pp_v fs m =
|
|
Format.fprintf fs "@[<1>[%a]@]"
|
|
(List.pp ",@ " (fun fs (k, v) ->
|
|
Format.fprintf fs "@[%a@ @<2>↦ %a@]" pp_k k pp_v v ))
|
|
(to_alist m)
|
|
|
|
let pp_diff ~data_equal pp_key pp_val pp_diff_val fs (x, y) =
|
|
let pp_diff_elt fs = function
|
|
| k, `Left v ->
|
|
Format.fprintf fs "-- [@[%a@ @<2>↦ %a@]]" pp_key k pp_val v
|
|
| k, `Right v ->
|
|
Format.fprintf fs "++ [@[%a@ @<2>↦ %a@]]" pp_key k pp_val v
|
|
| k, `Unequal vv ->
|
|
Format.fprintf fs "[@[%a@ @<2>↦ %a@]]" pp_key k pp_diff_val vv
|
|
in
|
|
let sd = Sequence.to_list (symmetric_diff ~data_equal x y) in
|
|
if not (List.is_empty sd) then
|
|
Format.fprintf fs "[@[<hv>%a@]];@ " (List.pp ";@ " pp_diff_elt) sd
|
|
end
|