Source file int_set.ml

open! Core
open! Import

module Range : sig
  (* A range represents a closed interval of integers *)

  type t = private
    { lo : int
    ; hi : int
    }

  (* Invariant: lo <= hi *)

  (** Create [t] from a range *)
  val make : int -> int -> t

  val to_string : t -> string

  (** [merge s t] merges mergeable ranges *)
  val merge : t -> t -> [ `Ok of t | `Lt_and_not_adjacent | `Gt_and_not_adjacent ]

  val contains : int -> t -> bool
end = struct
  type t =
    { lo : int
    ; hi : int
    }

  let make x y = if x <= y then { lo = x; hi = y } else { lo = y; hi = x }
  let to_string t = if t.lo = t.hi then Int.to_string t.lo else sprintf "%d-%d" t.lo t.hi

  (* on the number line, r1 is either on the left, on the right, or
     intersected with r2 *)
  let compare r1 r2 =
    if r1.hi < r2.lo - 1
    then `Lt_and_not_adjacent
    else if r1.lo > r2.hi + 1
    then `Gt_and_not_adjacent
    else `Mergeable
  ;;

  let merge r1 r2 =
    match compare r1 r2 with
    | `Lt_and_not_adjacent -> `Lt_and_not_adjacent
    | `Gt_and_not_adjacent -> `Gt_and_not_adjacent
    | `Mergeable -> `Ok { lo = Int.min r1.lo r2.lo; hi = Int.max r1.hi r2.hi }
  ;;

  let contains i r = r.lo <= i && i <= r.hi
end

type t = Range.t list

(* invariant : the elements of [t] must be pairwise discrete (not mergeable) and sorted
   in DECREASING order. *)

let empty = []
let to_string t = String.concat ~sep:"," (List.rev_map t ~f:Range.to_string)

let add_range t x y =
  (* note: not tail recursive *)
  let rec loop ranges to_add =
    match ranges with
    | r :: rest ->
      (* the following keeps the invariant: discrete + sorted *)
      (match Range.merge to_add r with
       | `Lt_and_not_adjacent -> r :: loop rest to_add
       | `Gt_and_not_adjacent -> to_add :: r :: rest
       | `Ok merged -> loop rest merged)
    | [] -> [ to_add ]
  in
  loop t (Range.make x y)
;;

let add t i = add_range t i i
let mem t i = List.exists t ~f:(fun x -> Range.contains i x)
let ranges t = List.map t ~f:(fun { Range.lo; hi } -> lo, hi)

let max t =
  match t with
  | [] -> None
  | { Range.hi; lo = _ } :: _ -> Some hi
;;

let min t =
  match List.last t with
  | None -> None
  | Some { Range.lo; hi = _ } -> Some lo
;;