Source file fQueue.ml

type 'a digit = Zero | One of 'a | Two of 'a * 'a | Three of 'a * 'a * 'a

type 'a t =
  | Shallow of 'a digit
  | Deep of {s: int; f: 'a digit; m: ('a * 'a) t Lazy.t; r: 'a digit}

let empty = Shallow Zero

exception Empty

let _one x = Shallow (One x)

let _two x y = Shallow (Two (x, y))

let _deep s f m r =
  assert (f <> Zero && r <> Zero) ;
  Deep {s; f; m; r}

let is_empty = function
  | Shallow Zero -> true
  | Shallow (One _ | Two _ | Three _) | Deep _ -> false

let _empty = Lazy.from_val empty

let rec push : 'a. 'a t -> 'a -> 'a t =
 fun q x ->
  match q with
  | Shallow Zero -> _one x
  | Shallow (One y) -> Shallow (Two (y, x))
  | Shallow (Two (y, z)) -> Shallow (Three (y, z, x))
  | Shallow (Three (y, z, z')) -> _deep 4 (Two (y, z)) _empty (Two (z', x))
  | Deep {r= Zero; _} -> assert false
  | Deep {s; f; m; r= One y} -> _deep (s + 1) f m (Two (y, x))
  | Deep {s; f; m; r= Two (y, z)} -> _deep (s + 1) f m (Three (y, z, x))
  | Deep {s; f; m= (lazy q'); r= Three (y, z, z')} ->
      _deep (s + 1) f (lazy (push q' (y, z))) (Two (z', x))

let map_last_digit f = function
  | Zero -> Zero
  | One x -> One (f x)
  | Two (x, y) -> Two (x, f y)
  | Three (x, y, z) -> Three (x, y, f z)

let map_last : 'a. ('a -> 'a) -> 'a t -> 'a t =
 fun f -> function
  | Shallow v -> Shallow (map_last_digit f v)
  | Deep ({r; _} as deep) -> Deep {deep with r= map_last_digit f r}

let rec shift : 'a. 'a t -> 'a * 'a t =
 fun q ->
  match q with
  | Shallow Zero -> raise Empty
  | Shallow (One x) -> (x, empty)
  | Shallow (Two (x, y)) -> (x, Shallow (One y))
  | Shallow (Three (x, y, z)) -> (x, Shallow (Two (y, z)))
  | Deep {f= Zero; _} -> assert false
  | Deep {s; f= One x; m= (lazy q'); r} ->
      if is_empty q' then (x, Shallow r)
      else
        let (y, z), q' = shift q' in
        (x, _deep (s - 1) (Two (y, z)) (Lazy.from_val q') r)
  | Deep {s; f= Two (x, y); m; r} -> (x, _deep (s - 1) (One y) m r)
  | Deep {s; f= Three (x, y, z); m; r} -> (x, _deep (s - 1) (Two (y, z)) m r)

let rec cons : 'a. 'a t -> 'a -> 'a t =
 fun q x ->
  match q with
  | Shallow Zero -> Shallow (One x)
  | Shallow (One y) -> Shallow (Two (x, y))
  | Shallow (Two (y, z)) -> Shallow (Three (x, y, z))
  | Shallow (Three (y, z, z')) -> _deep 4 (Two (x, y)) _empty (Two (z, z'))
  | Deep {f= Zero; _} -> assert false
  | Deep {s; f= One y; m; r} -> _deep (s + 1) (Two (x, y)) m r
  | Deep {s; f= Two (y, z); m; r} -> _deep (s + 1) (Three (x, y, z)) m r
  | Deep {s; f= Three (y, z, z'); m= (lazy q'); r} ->
      _deep (s + 1) (Two (x, y)) (lazy (cons q' (z, z'))) r

let _digit_to_seq d k =
  match d with
  | Zero -> ()
  | One x -> k x
  | Two (x, y) -> k x ; k y
  | Three (x, y, z) -> k x ; k y ; k z

type 'a sequence = ('a -> unit) -> unit

let rec to_seq : 'a. 'a t -> 'a sequence =
 fun q k ->
  match q with
  | Shallow d -> _digit_to_seq d k
  | Deep {f; m= (lazy q'); r; _} ->
      _digit_to_seq f k ;
      to_seq q' (fun (x, y) -> k x ; k y) ;
      _digit_to_seq r k

let iter f q = to_seq q f

let _fold_digit f acc d =
  match d with
  | Zero -> acc
  | One x -> f acc x
  | Two (x, y) -> f (f acc x) y
  | Three (x, y, z) -> f (f (f acc x) y) z

let rec fold : 'a 'b. ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b =
 fun func acc q ->
  match q with
  | Shallow d -> _fold_digit func acc d
  | Deep {f; m= (lazy q'); r; _} ->
      let acc = _fold_digit func acc f in
      let acc = fold (fun acc (x, y) -> func (func acc x) y) acc q' in
      _fold_digit func acc r

let to_list q =
  let l = ref [] in
  to_seq q (fun x -> l := x :: !l) ;
  List.rev !l

let of_list l = List.fold_left push empty l

let pp ppv ppf q =
  Fmt.pf ppf "[ %a ]"
    (Fmt.hvbox (Fmt.list ~sep:(Fmt.unit ";@ ") ppv))
    (to_list q)