Source

spotlib / lib / xlist.ml

The default branch has multiple heads

Full commit
open List

let iter_until f ls =
  let st = ref `Continue in
  List.iter (fun x ->
    match f x with
    | `Break b -> st := `Break b; raise Exit
    | `Continue -> ()) ls;
  match !st with
  | `Continue -> None
  | `Break v -> Some v

let iteri f l =
  let rec iter n = function
    | [] -> ()
    | x::xs -> f n x; iter (n+1) xs
  in
  iter 0 l

let mapi f l =
  let rec map n = function
    | [] -> []
    | x::xs -> f n x :: map (n+1) xs
  in
  map 0 l

let from_to f t =
  (* CR jfuruse: we should build from 'to' *)
  let rec from_to st f t =
    if f > t then rev st
    else from_to (f::st) (f+1) t
  in
  from_to [] f t

let (--) = from_to

let init_from_to f t fn =
  let rec loop st i =
    if i > t then rev st
    else loop (fn i :: st) (i+1)
  in
  loop [] f

let find_opt f ls = try Some (List.find f ls) with Not_found -> None

let rec find_map_opt f = function
  | [] -> None
  | x::xs ->
      match f x with
      | Some v -> Some v
      | None -> find_map_opt f xs

(** mapMaybe of Haskell *)
let rev_filter_map f lst =
  List.fold_left (fun st x ->
    match f x with
    | Some v -> v :: st
    | None -> st) [] lst

(** mapMaybe of Haskell *)
let filter_map f lst = List.rev (rev_filter_map f lst)

(** concatMap of Haskell *)
let concat_map f l = List.concat (List.map f l)

module TR = struct
  let map f l = List.rev (List.rev_map f l)

  TEST "TL.map" =
      map (fun x -> x) [1;2;3] = [1;2;3]

  let rev_concat_map f l =
    let rec loop st = function
      | [] -> st
      | x::xs -> loop (List.rev_append (f x) st) xs
    in
    loop [] l

  TEST "TL.rev_concat_map" =
      rev_concat_map (fun x -> [x; x+1]) [1;2;3] = [4;3;3;2;2;1]

  let concat_map f l = List.rev (rev_concat_map f l)

  TEST "TL.concat_map" =
      concat_map (fun x -> [x; x+1]) [1;2;3] = [1;2;2;3;3;4]
end

let take n xs =
  let rec take_ n st xs =
    if n <= 0 then st
    else match xs with
    | [] -> st
    | x::xs -> take_ (n-1) (x::st) xs
  in
  List.rev (take_ n [] xs)

let split_at n xs =
  let rec split_at_ n st xs =
    if n <= 0 then st, xs
    else match xs with
    | [] -> st, []
    | x::xs -> split_at_ (n-1) (x::st) xs
  in
  let r, dropped = split_at_ n [] xs in
  List.rev r, dropped

TEST "split_at" =
      split_at 2 [] = ([], [])
      && split_at 2 [1] = ([1], [])
      && split_at 2 [1;2] = ([1;2], [])
      && split_at 2 [1;2;3] = ([1;2], [3])

let rec drop n xs =
  if n <= 0 then xs
  else match xs with
  | [] -> []
  | _::xs -> drop (n-1) xs

let span p xs =
  let rec span_ st = function
    | [] -> List.rev st, []
    | x::xs when p x -> span_ (x::st) xs
    | l -> List.rev st, l
  in
  span_ [] xs

let partition p xs =
  let rec part yes no = function
    | [] -> List.rev yes, List.rev no
    | x::xs when p x -> part (x::yes) no xs
    | x::xs -> part yes (x::no) xs
  in
  part [] [] xs

let partition_map f xs =
  let rec part yes no = function
    | [] -> List.rev yes, List.rev no
    | x::xs ->
        match f x with
        | Some y -> part (y::yes) no xs
        | None -> part yes (x::no) xs
  in
  part [] [] xs

let unique xs =
  let rec unique st = function
    | [] -> List.rev st
    | x::xs ->
        let st' = if List.mem x st then st else x::st in
        unique st' xs
  in
  unique [] xs

let is_unique key xs =
  let rec is_unique st = function
    | [] -> None
    | x::xs ->
        let k = key x in
        try
          let x' = List.assoc k st in
          Some (x,x')
        with Not_found -> is_unique ((k,x)::st) xs
  in
  is_unique [] xs

let rec intersperse x = function
  | [] -> []
  | [y] -> [y]
  | y::ys -> y::x::intersperse x ys

let rec last = function
  | [] -> failwith "Xlist.last"
  | [x] -> x
  | _::xs -> last xs

let assoc_all k =
  let rec aux rev = function
    | [] -> List.rev rev
    | (k',v)::kvs when k = k' -> aux (v::rev) kvs
    | _::kvs -> aux rev kvs
  in
  aux []

let assoc_opt k l = try Some (List.assoc k l) with Not_found -> None

let scani_left f acc xs =
  let rec scan acc pos xs = match pos, xs with
    | _pos, [] -> acc
    | pos, x::xs ->
        match f pos acc x with
        | `Continue acc -> scan acc (pos+1) xs
        | `Stop acc -> acc
  in
  scan acc 0 xs

let fold_left1 f = function
  | [] -> invalid_arg "fold_left1" (* check the stack trace to see the use site *)
  | x::xs -> List.fold_left f x xs

let sum xs = List.fold_left (+) 0 xs

let rev_group eq rev_xs =
  let rec grouping gs cur_group = function
    | [] ->
        begin match cur_group with
        | [] -> gs
        | _ -> cur_group :: gs
        end
    | x::xs ->
        match cur_group with
        | [] -> grouping gs [x] xs
        | y::_ ->
            if eq x y then grouping gs (x::cur_group) xs
            else grouping (cur_group::gs) [x] xs
  in
  grouping [] [] rev_xs

let group eq xs = rev_group eq (List.rev xs)
let sort_then_group compare xs = 
  rev_group (fun x y -> compare x y = 0) (List.sort (fun x y -> - (compare x y)) xs)

TEST "group" =
      group (=) [1;2;3] = [[1];[2];[3]]

TEST "group" =
      group (=) [1;2;3;3;4;4;5;5;6] = [[1];[2];[3;3];[4;4];[5;5];[6]]

TEST "group" =
      group (=) [1;2;3;3;4;4;5;5;6;6] = [[1];[2];[3;3];[4;4];[5;5];[6;6]]

TEST "group" =
      group (=) [1;1;2;3;3;4;4;5;5;6;6] = [[1;1];[2];[3;3];[4;4];[5;5];[6;6]]

TEST "group" =
      group (=) [1;1;2;3;3;4;4;5;5;1;1;6;6] = [[1;1];[2];[3;3];[4;4];[5;5];[1;1];[6;6]]

let splits_by max items =
  assert (max > 0);
  let rec loop rev items =
    match split_at max items with
    | [], [] -> rev
    | [], _ -> assert false
    | xs, [] -> xs::rev
    | xs, ys -> loop (xs::rev) ys
  in
  List.rev (loop [] items)

TEST "splits_by" =
      splits_by 2 [] = []
      && splits_by 2 [1] = [[1]]
      && splits_by 2 [1;2] = [[1;2]]
      && splits_by 2 [1;2;3] = [[1;2]; [3]]

let accum xsref x = xsref := x :: !xsref
let (+::=) = accum

module Infix = struct
  let (--) = (--)
  let (+::=) = (+::=)
end

module Pervasives = Infix