1. Sébastien Ferré
  2. ocaml-lib

Source

ocaml-lib / setrie.ml


module type PATH =
  sig
    type t

    val is_empty : t -> bool
    val compare_first : t -> t -> int
    val append : t -> t -> t

    val prefix_zip : t -> t -> t * t * t
    val prefix_inter : t -> t -> t * t
  end


let compare x y = Pervasives.compare y x (* for compatibility with LSet (Oops!) *)

type ('a,'b) tree = Node of 'a LSet.t * 'b option * ('a,'b) tree * ('a,'b) tree | Nil
  (* Node (label, value at label, elder child, next brother), to be specialized (?) when child and/or brother are Nil *)

type ('a,'b) t = 'b option * ('a,'b) tree (* (vopt, tree) ~ Node ([], vopt, tree, Nil) *)
      (* tree, and value associated to the empty set *)

let make_branch ys vopt child =
  if ys = []
  then vopt, child
  else None, Node (ys, vopt, child, Nil)

let rec first_vopt = function
  | [] -> None
  | None::l -> first_vopt l
  | Some v::l -> Some v

(* utilitary functions on LSets *)
(* ---------------------------- *)

let rec prefix_zip xs ys =
  prefix_zip_aux [] (xs,ys)
and prefix_zip_aux p = function
  | x::xs1, y::ys1 when x=y ->
      prefix_zip_aux (x::p) (xs1,ys1)
  | xs, ys ->
      List.rev p, xs, ys

let rec prefix_inter xs ys =
  prefix_inter_aux [] (xs,ys)
and prefix_inter_aux s (xs0,ys0) =
  match xs0, ys0 with
  | x::xs, y::ys ->
      let c = compare x y in
      if c = 0 then prefix_inter_aux (y::s) (xs,ys)
      else if c < 0 then prefix_inter_aux s (xs,ys0)
      else (* c > 0 *) prefix_inter_aux s (xs0,ys)
  | [], _ -> [], List.rev s
  | _, [] -> xs0, List.rev s

(* main interface *)
(* -------------- *)

let rec fold : ('a LSet.t -> 'b -> 'c -> 'c) -> ('a,'b) t -> 'c -> 'c =
  fun f (vopt,t) e ->
    let e1 = match vopt with None -> e | Some v -> f [] v e in
    fold2 f t e1 []
and fold2 f t e ys0 =
  match t with
  | Nil -> e
  | Node (ys, vopt, child, brother) ->
      let ys' = ys0 @ ys in
      let e1 = match vopt with None -> e | Some v -> f ys' v e in
      let e2 = fold2 f child e1 ys' in
      let e3 = fold2 f brother e2 ys0 in
      e3


let empty : ('a,'b) t = None, Nil

let rec cardinal : ('a,'b) t -> int =
  fun (vopt,tree) ->
    cardinal_vopt vopt + cardinal2 tree
and cardinal2 = function
  | Nil -> 0
  | Node (_,vopt,child,brother) ->
      cardinal_vopt vopt + cardinal2 child + cardinal2 brother
and cardinal_vopt = function
  | None -> 0
  | Some _ -> 1

let singleton : 'a LSet.t -> 'b -> ('a,'b) t =
  fun xs v ->
    if LSet.is_empty xs
    then Some v, Nil
    else None, Node (xs, Some v, Nil, Nil)

let rec add : ('b -> 'b -> 'b) -> 'a LSet.t -> 'b -> ('a,'b) t -> ('a,'b) t =
  fun u xs v t ->
    match xs, t with
    | [], (vopt,tree) -> (add_vopt u vopt v, tree)
    | x::xs1, (vopt,Nil) -> (vopt, Node (xs, Some v, Nil, Nil))
    | x::xs1, (vopt, (Node (y::ys2 as ys, vopt2, c2, b2) as tree)) ->
	let c = compare x y in
	if c = 0 then
	  let prefix, xs1', ys2' = prefix_zip xs1 ys2 in
	  let t2' = make_branch ys2' vopt2 c2 in
	  let vopt', child' = add u xs1' v t2' in
	  vopt, Node (y::prefix, vopt', child', b2)
	else if c < 0 then
	  vopt, Node (xs, Some v, Nil, tree)
	else (* c > 0 *)
	  let _, brother' = add u xs v (None, b2) in
	  vopt, Node (ys, vopt2, c2, brother')
    | _ -> assert false
and add_vopt u vopt v =
  match vopt with
  | None -> Some v
  | Some v0 -> Some (u v0 v)

let rec remove : 'a LSet.t -> ('a,'b) t -> ('a,'b) t =
  fun xs (vopt,tree) ->
    if xs = []
    then None, tree
    else vopt, remove2 xs tree
and remove2 xs t =
  match xs, t with
  | [], _ -> t
  | x::xs1, Nil -> Nil
  | x::xs1, Node (y::ys1 as ys, vopt, child, brother) ->
      let c = compare x y in
      if c = 0 then
	let prefix, xs1', ys1' = prefix_zip xs1 ys1 in
	if ys1' <> [] then (* xs is not present *)
	  t
	else if xs1' = [] then (* vopt must be set to None *)
	  if child = Nil then brother else Node (ys, None, child, brother)
	else (* element to be removed is in b *)
	  let child' = remove2 xs1' child in
	  if vopt = None && child' = Nil then brother else Node (ys, vopt, child', brother)
      else if c < 0 then (* xs is not present *)
	t
      else (* c > 0 *) (* xs is in b *)
	let brother' = remove2 xs brother in
	Node (ys, vopt, child, brother')
  | _ -> assert false

let rec union : ('b -> 'b -> 'b) -> ('a,'b) t -> ('a,'b) t -> ('a,'b) t =
  fun u (vopt1,tree1) (vopt2,tree2) ->
    union_vopt u vopt1 vopt2, union2 u tree1 tree2
and union2 u t1 t2 =
    match t1, t2 with
    | Nil, _ -> t2
    | _, Nil -> t1
    | Node (y1::ys1, vopt1, c1, b1), Node (y2::ys2, vopt2, c2, b2) ->
	let c = compare y1 y2 in
	if c = 0 then
	  let prefix, ys1', ys2' = prefix_zip ys1 ys2 in
	  let t1' = make_branch ys1' vopt1 c1 in
	  let t2' = make_branch ys2' vopt2 c2 in
	  let vopt', child' = union u t1' t2' in
	  let brother' = union2 u b1 b2 in
	  Node (y1::prefix, vopt', child', brother')
	else if c < 0 then
	  Node (y1::ys1, vopt1, c1, union2 u b1 t2)
	else (* c > 0 *)
	  Node (y2::ys2, vopt2, c2, union2 u t1 b2)
    | _ -> assert false
and union_vopt u vopt1 vopt2 =
  match vopt1, vopt2 with
  | None, _ -> vopt2
  | _, None -> vopt1
  | Some v1, Some v2 -> Some (u v1 v2)


let rec exists_contains : ('a,'b) t -> 'a LSet.t -> ('a LSet.t -> 'b -> bool) -> bool =
  fun (vopt,t) xs p ->
    exists_contains_vopt vopt p [] ||
    exists_contains2 t xs p []
and exists_contains2 t xs p path =
  match xs, t with
  | _, Nil -> false
  | [], Node (ys, vopt, child, brother) ->
      let path1 = path @ ys in
      exists_contains_vopt vopt p path1 ||
      exists_contains2 child [] p path1 ||
      exists_contains2 brother [] p path
  | x::xs1, Node (y::ys1 as ys, vopt, child, brother) ->
      let c = compare x y in
      if c = 0 then
	let xs1', ys1' = prefix_inter xs1 ys1 in
	if ys1' = ys1
	then
	  let path1 = path @ ys in
	  exists_contains_vopt vopt p path1 ||
	  exists_contains2 child xs1' p path1
	else false
      else if c < 0 then false  (* x is missing *)
      else (* c > 0 *)
	let xs1', ys1' = prefix_inter xs ys1 in
	if ys1' = ys1
	then
	  let path1 = path @ ys in
	  exists_contains_vopt vopt p path1 ||
	  exists_contains2 child xs1' p path1 ||
	  exists_contains2 brother xs p path
	else
	  exists_contains2 brother xs p path
  | _ -> assert false
and exists_contains_vopt vopt p path =
  match vopt with
  | None -> false
  | Some v -> p path v

let rec map_inter : ('a LSet.t -> 'b -> 'a LSet.t -> 'b option) -> ('b -> 'b -> 'b) -> ('a,'b) t -> 'a LSet.t -> ('a,'b) t =
  fun f u (vopt,t) xs ->
    let vopt' = map_inter_vopt f vopt [] [] in
    let vopt'', t' = map_inter2 f u t xs [] [] in
    union_vopt u vopt' vopt'', t'
and map_inter2 f u t xs path path' =
  match xs, t with
  | _, Nil -> None, Nil
  | [], Node (ys, vopt, child, brother) ->
      let vopt' = map_inter_vopt f vopt (path@ys) path' in
      let child_vopt', _ (* Nil *) = map_inter2 f u child [] (path@ys) path' in
      let brother_vopt', _ (* Nil *) = map_inter2 f u brother [] path path' in
      union_vopt u vopt' (union_vopt u child_vopt' brother_vopt'), Nil
  | x::xs1, Node (y::ys1, vopt, child, brother) ->
      let c = compare x y in
      if c = 0 then
	let xs1', ys1' = prefix_inter xs1 ys1 in
	let path1, path1' = path @ y::ys1, path' @ y::ys1' in
	let vopt' = map_inter_vopt f vopt path1 path1' in
	let child_vopt', child' = map_inter2 f u child xs1' path1 path1' in
	let brother_vopt', brother' = map_inter2 f u brother xs1 path path' in
	brother_vopt', Node (y::ys1', union_vopt u child_vopt' vopt', child', brother')
      else if c < 0 then
	map_inter2 f u t xs1 path path'
      else (* c > 0 *)
	let xs1', ys1' = prefix_inter xs ys1 in
	let path1, path1' = path @ y::ys1, path' @ ys1' in
	let vopt' = map_inter_vopt f vopt path1 path1' in
	let child_vopt', child' = map_inter2 f u child xs1' path1 path1' in
	let brother_vopt', brother' = map_inter2 f u brother xs path path' in
	let t1 = make_branch ys1' (union_vopt u vopt' child_vopt') child' in
	let t2 = brother_vopt', brother' in
	union u t1 t2
  | _ -> assert false
and map_inter_vopt f vopt ys ys' =
  match vopt with
  | None -> None
  | Some v -> f ys v ys'


let rec mapmin_inter : ('a LSet.t -> 'b -> 'b option) -> ('a,'b) t -> 'a LSet.t -> ('a,'b) t =
  fun f (vopt,t) xs ->
    let child_vopt, t' = mapmin_inter2 f t xs [] in
    let vopt' = mapmin_inter_vopt f (first_vopt [vopt; child_vopt]) [] in
    vopt', t'
and mapmin_inter2 f t xs path' =
  match xs, t with
  | _, Nil -> None, Nil
  | [], Node (ys, vopt, child, brother) ->
      let return_vopt =
	let brother_vopt, _ (* Nil *) = mapmin_inter2 f brother [] path' in
	if brother_vopt <> None then brother_vopt
	else if vopt <> None then vopt
	else
	  let child_vopt, _ (* Nil *) = mapmin_inter2 f child [] path' in
	  child_vopt in
      return_vopt, Nil
  | x::xs1, Node (y::ys1, vopt, child, brother) ->
      let c = compare x y in
      if c = 0 then
	let xs1', ys1' = prefix_inter xs1 ys1 in
	let path1' = path' @ y::ys1' in
	let child_vopt, child' = mapmin_inter2 f child xs1' path1' in
	let brother_vopt, brother' = mapmin_inter2 f brother xs1 path' in
	let vopt' = mapmin_inter_vopt f (first_vopt [vopt; child_vopt]) path1' in
	brother_vopt, Node (y::ys1', vopt', child', brother')
      else if c < 0 then
	mapmin_inter2 f t xs1 path'
      else (* c > 0 *)
	let xs1', ys1' = prefix_inter xs ys1 in
	let path1' = path' @ ys1' in
	let child_vopt, child' = mapmin_inter2 f child xs1' path1' in
	let brother_vopt, brother' = mapmin_inter2 f brother xs path' in
	let t1 =
	  let vopt1 = first_vopt [vopt; child_vopt] in
	  if ys1' = []
	  then
	    let return_vopt = if brother_vopt = None then vopt1 else None in
	    return_vopt, child'
	  else
	    let vopt' = mapmin_inter_vopt f vopt1 path1' in
	    None, Node (ys1', vopt', child', Nil) in
	let t2 = brother_vopt, brother' in
	union (fun _ v2 -> v2) t1 t2
  | _ -> assert false
and mapmin_inter_vopt f vopt ys' =
  match vopt with
  | None -> None
  | Some v -> f ys' v