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.

110 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 Import0
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 =
with_return
@@ fun {return} ->
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 =
with_return
@@ fun {return} ->
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