Source file monadic2.ml

(**************************************************************************)
(*                                                                        *)
(*                                 OCaml                                  *)
(*                                                                        *)
(*                 Simon Cruanes                                          *)
(*                                                                        *)
(*   Copyright 2017 Institut National de Recherche en Informatique et     *)
(*     en Automatique.                                                    *)
(*                                                                        *)
(*                 Raphaël Proust                                         *)
(*                                                                        *)
(*   Copyright 2022 Nomadic Labs                                          *)
(*   Copyright 2025 Raphaël Proust                                        *)
(*                                                                        *)
(*   All rights reserved.  This file is distributed under the terms of    *)
(*   the GNU Lesser General Public License version 2.1, with the          *)
(*   special exception on linking described in the file LICENSE.          *)
(*                                                                        *)
(**************************************************************************)

(* This implementation is a refactoring and functorisation of the original code
   present in the OCaml {!Stdlib.Seq} module. *)

module type S2 = sig
  type ('a, 'e) mon
  type ('a, 'e) t = unit -> (('a, 'e) node, 'e) mon
  and ('a, 'e) node =
    | Nil
    | Cons of 'a * ('a, 'e) t

  include Sigs2.SEQMON2ALL
    with type ('a, 'e) mon := ('a, 'e) mon
    with type ('a, 'e) t := ('a, 'e) t

  module M
  : Sigs2.SEQMON2TRANSFORMERS
    with type ('a, 'e) mon := ('a, 'e) mon
    with type ('a, 'e) callermon := ('a, 'e) mon
    with type ('a, 'e) t := ('a, 'e) t

  module Make
    (Alt : Sigs2.MONAD2)
    (Glue : Sigs2.GLUE2
      with type ('a, 'e) x := ('a, 'e) Alt.t
      with type ('a, 'e) f := ('a, 'e) mon
      with type ('a, 'e) ret := ('a, 'e) mon)
  : Sigs2.SEQMON2TRANSFORMERS
      with type ('a, 'e) mon := ('a, 'e) mon
      with type ('a, 'e) callermon := ('a, 'e) Alt.t
      with type ('a, 'e) t := ('a, 'e) t

  module MakeTraversors
    (Alt : Sigs2.MONAD2)
    (Ret : Sigs2.MONAD2)
    (GlueAlt : Sigs2.GLUE2
      with type ('a, 'e) x := ('a, 'e) Alt.t
      with type ('a, 'e) f := ('a, 'e) Ret.t
      with type ('a, 'e) ret := ('a, 'e) Ret.t)
    (GlueMon : Sigs2.GLUE2
      with type ('a, 'e) x := ('a, 'e) mon
      with type ('a, 'e) f := ('a, 'e) Ret.t
      with type ('a, 'e) ret := ('a, 'e) Ret.t)
  : Sigs2.SEQMON2TRAVERSORS
      with type ('a, 'e) mon := ('a, 'e) Ret.t
      with type ('a, 'e) callermon := ('a, 'e) Alt.t
      with type ('a, 'e) t := ('a, 'e) t
end

module Make2
  (Mon: Sigs2.MONAD2)
: S2
   with type ('a, 'e) mon := ('a, 'e) Mon.t
= struct

  let ( let* ) = Mon.bind
  let return = Mon.return

  type ('a, 'e) t = unit -> (('a, 'e) node, 'e) Mon.t
  and ('a, 'e) node =
    | Nil
    | Cons of 'a * ('a, 'e) t

  let is_empty xs =
    let* n = xs () in
    match n with
    | Nil ->
        return true
    | Cons (_, _) ->
        return false

  let uncons xs =
    let* n = xs () in
    match n with
    | Cons (x, xs) ->
        return (Some (x, xs))
    | Nil ->
        return None

  let rec length_aux accu xs =
    let* n = xs () in
    match n with
    | Nil ->
        return accu
    | Cons (_, xs) ->
        length_aux (accu + 1) xs
  let[@inline] length xs = length_aux 0 xs

  let empty () = return Nil

  let cons x next () = return (Cons (x, next))

  let rec append seq1 seq2 () =
    let* n = seq1 () in
    match n with
    | Nil -> seq2()
    | Cons (x, next) -> return (Cons (x, append next seq2))

  let rec repeat x () =
    return (Cons (x, repeat x))

  let rec cycle_nonempty xs () =
    append xs (cycle_nonempty xs) ()

  let cycle xs () =
    let* n = xs () in
    match n with
    | Nil ->
        return Nil
    | Cons (x, xs') ->
        return (Cons (x, append xs' (cycle_nonempty xs)))

  let rec take_aux n xs =
    if n = 0 then
      empty
    else
      fun () ->
        let* node = xs () in
        match node with
        | Nil ->
            return Nil
        | Cons (x, xs) ->
            return (Cons (x, take_aux (n-1) xs))
  let take n xs =
    if n < 0 then invalid_arg "Seq.take";
    take_aux n xs

  let rec force_drop n xs =
    let* node = xs () in
    match node with
    | Nil ->
        return Nil
    | Cons (_, xs) ->
        let n = n - 1 in
        if n = 0 then
          xs()
        else
          force_drop n xs

  let drop n xs =
    if n < 0 then invalid_arg "Seq.drop"
    else if n = 0 then
      xs
    else
      fun () ->
        force_drop n xs

  let rec memoize xs =
    Suspension.memoize (fun () ->
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) ->
          return (Cons (x, memoize xs))
    )

  let rec once xs =
    Suspension.once (fun () ->
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) ->
          return (Cons (x, once xs))
    )

  let rec map_fst xys () =
    let* n = xys () in
    match n with
    | Nil ->
        return Nil
    | Cons ((x, _), xys) ->
        return (Cons (x, map_fst xys))

  let rec map_snd xys () =
    let* n = xys () in
    match n with
    | Nil ->
        return Nil
    | Cons ((_, y), xys) ->
        return (Cons (y, map_snd xys))

  let unzip xys =
    map_fst xys, map_snd xys

  let split =
    unzip

  let rec filter_map_beeg f seq () =
    let* n = seq () in
    match n with
    | Nil -> return Nil
    | Cons (x, next) ->
        let* o = f x in
        match o with
        | None -> filter_map_beeg f next ()
        | Some y ->  return (Cons (y, filter_map_beeg f next))

  let peel xss =
    unzip (filter_map_beeg uncons xss)

  let rec transpose xss () =
    let heads, tails = peel xss in
    let* b = is_empty heads in
    if b then begin
      return Nil
    end
    else
      return (Cons (heads, transpose tails))

  let rec concat seq () =
    let* n = seq () in
    match n with
    | Nil -> return Nil
    | Cons (x, next) ->
       append x (concat next) ()

  let rec zip xs ys () =
    let* n = xs () in
    match n with
    | Nil ->
        return Nil
    | Cons (x, xs) ->
        let* n = ys () in
        match n with
        | Nil ->
            return Nil
        | Cons (y, ys) ->
            return (Cons ((x, y), zip xs ys))

  let rec interleave xs ys () =
    let* n = xs () in
    match n with
    | Nil ->
        ys()
    | Cons (x, xs) ->
        return (Cons (x, interleave ys xs))

  let to_dispenser xs =
    let s = ref xs in
    fun () ->
      let* n = (!s) () in
      match n with
      | Nil ->
          return None
      | Cons (x, xs) ->
          s := xs;
          return (Some x)

  let rec ints i () =
    return (Cons (i, ints (i + 1)))

  let rec of_seq s () =
    match s () with
    | Seq.Nil ->  return Nil
    | Seq.Cons (x, s) -> return (Cons (x, of_seq s))

  module MakeTraversors
    (Alt : Sigs2.MONAD2)
    (Ret : Sigs2.MONAD2)
    (GlueAlt : Sigs2.GLUE2
      with type ('a, 'e) x := ('a, 'e) Alt.t
      with type ('a, 'e) f := ('a, 'e) Ret.t
      with type ('a, 'e) ret := ('a, 'e) Ret.t)
    (GlueMon : Sigs2.GLUE2
      with type ('a, 'e) x := ('a, 'e) Mon.t
      with type ('a, 'e) f := ('a, 'e) Ret.t
      with type ('a, 'e) ret := ('a, 'e) Ret.t)
  : Sigs2.SEQMON2TRAVERSORS
      with type ('a, 'e) mon := ('a, 'e) Ret.t
      with type ('a, 'e) callermon := ('a, 'e) Alt.t
      with type ('a, 'e) t := ('a, 'e) t
  = struct

    let return = Ret.return
    let ( let* ) = GlueMon.bind
    let ( let** ) = GlueAlt.bind

    let rec equal eq xs ys =
      (* TODO: support ( and* ) with Beeg.join *)
      let* xs = xs () in
      let* ys = ys () in
      match xs, ys with
      | Nil, Nil ->
          return true
      | Cons (x, xs), Cons (y, ys) ->
          let** e = eq x y in
          if e then
            equal eq xs ys
          else
            return false
      | Nil, Cons (_, _)
      | Cons (_, _), Nil ->
          return false

    let rec compare cmp xs ys =
      let* xs = xs () in
      let* ys = ys () in
      match xs, ys with
      | Nil, Nil ->
          return 0
      | Cons (x, xs), Cons (y, ys) ->
          let** c = cmp x y in
          if c <> 0 then return c else compare cmp xs ys
      | Nil, Cons (_, _) ->
          return (-1)
      | Cons (_, _), Nil ->
          return (+1)

    let rec fold_left f acc seq =
      let* n = seq () in
      match n with
      | Nil -> return acc
      | Cons (x, next) ->
          let** acc = f acc x in
          fold_left f acc next

    let rec iter f seq =
      let* n = seq () in
      match n with
      | Nil -> return ()
      | Cons (x, next) ->
          let** () = f x in
          iter f next

    let rec iteri_aux f i xs =
      let* n = xs () in
      match n with
      | Nil ->
          return ()
      | Cons (x, xs) ->
          let** () = f i x in
          iteri_aux f (i+1) xs
    let[@inline] iteri f xs = iteri_aux f 0 xs

    let rec fold_lefti_aux f accu i xs =
      let* n = xs () in
      match n with
      | Nil ->
          return accu
      | Cons (x, xs) ->
          let** accu = f accu i x in
          fold_lefti_aux f accu (i+1) xs
    let[@inline] fold_lefti f accu xs = fold_lefti_aux f accu 0 xs

    let rec for_all p xs =
      let* n = xs () in
      match n with
      | Nil ->
          return true
      | Cons (x, xs) ->
          let** b = p x in
          if b then
            for_all p xs
          else
            return false

    let rec exists p xs =
      let* n = xs () in
      match n with
      | Nil ->
          return false
      | Cons (x, xs) ->
          let** b = p x in
          if b then
            return true
          else
            exists p xs

    let rec find p xs =
      let* n = xs () in
      match n with
      | Nil ->
          return None
      | Cons (x, xs) ->
          let** b = p x in
          if b then return (Some x) else find p xs

    let rec find_map f xs =
      let* n = xs () in
      match n with
      | Nil ->
          return None
      | Cons (x, xs) ->
          let** o = f x in
          match o with
          | None ->
              find_map f xs
          | Some _ as result ->
              return result

    let rec iter2 f xs ys =
      let* n = xs () in
      match n with
      | Nil ->
          return ()
      | Cons (x, xs) ->
          let* n = ys () in
          match n with
          | Nil ->
              return ()
          | Cons (y, ys) ->
              let** () = f x y in
              iter2 f xs ys

    let rec fold_left2 f accu xs ys =
      let* n = xs () in
      match n with
      | Nil ->
          return accu
      | Cons (x, xs) ->
          let* n = ys () in
          match n with
          | Nil ->
              return accu
          | Cons (y, ys) ->
              let** accu = f accu x y in
              fold_left2 f accu xs ys

    let rec for_all2 f xs ys =
      let* n = xs () in
      match n with
      | Nil ->
          return true
      | Cons (x, xs) ->
          let* n = ys () in
          match n with
          | Nil ->
              return true
          | Cons (y, ys) ->
              let** b = f x y in
              if b then
                for_all2 f xs ys
              else
                return false

    let rec exists2 f xs ys =
      let* n = xs () in
      match n with
      | Nil ->
          return false
      | Cons (x, xs) ->
          let* n = ys () in
          match n with
          | Nil ->
              return false
          | Cons (y, ys) ->
              let** b = f x y in
              if b then
                return true
              else
                exists2 f xs ys

  end

  module Make
    (Alt : Sigs2.MONAD2)
    (Glue : Sigs2.GLUE2
      with type ('a, 'e) x := ('a, 'e) Alt.t
      with type ('a, 'e) f := ('a, 'e) Mon.t
      with type ('a, 'e) ret := ('a, 'e) Mon.t)
  : Sigs2.SEQMON2TRANSFORMERS
      with type ('a, 'e) mon := ('a, 'e) Mon.t
      with type ('a, 'e) callermon := ('a, 'e) Alt.t
      with type ('a, 'e) t := ('a, 'e) t
  = struct

    include MakeTraversors (Alt)(Mon)(Glue)(Mon)

    let return = Mon.return
    let ( let** ) = Glue.bind

    (* we define return at the end of the file to avoid shadowing the monad *)

    let rec map f seq () =
      let* n = seq () in
      match n with
      | Nil -> return Nil
      | Cons (x, next) ->
          let** y = f x in
          return (Cons (y, map f next))

    let rec flat_map f seq () =
      let* n = seq () in
      match n with
      | Nil -> return Nil
      | Cons (x, next) ->
          let** y = f x in
        append y (flat_map f next) ()

    let concat_map = flat_map

    let rec filter_map f seq () =
      let* n = seq () in
      match n with
      | Nil -> return Nil
      | Cons (x, next) ->
          let** o = f x in
          match o with
          | None -> filter_map f next ()
          | Some y ->  return (Cons (y, filter_map f next))

    let rec filter f seq () =
      let* n = seq () in
      match n with
      | Nil -> return Nil
      | Cons (x, next) ->
          let** b = f x in
          if b
          then return (Cons (x, filter f next))
          else filter f next ()

    let rec unfold f u () =
      let** o = f u in
      match o with
      | None -> return Nil
      | Some (x, u') -> return (Cons (x, unfold f u'))

    (* [init_aux f i j] is the sequence [f i, ..., f (j-1)]. *)

    let rec init_aux f i j () =
      if i < j then begin
        let** fi = f i in
        return (Cons (fi, init_aux f (i + 1) j))
      end
      else
        return Nil

    let init n f =
      if n < 0 then
        invalid_arg "Seq.init"
      else
        init_aux f 0 n

    let rec forever f () =
      let** x = f () in
      return (Cons (x, forever f))

    let rec iterate1 f x () =
      let** y = f x in
      return (Cons (y, iterate1 f y))

    let iterate f x =
      cons x (iterate1 f x)

    let rec mapi_aux f i xs () =
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) ->
          let** y = f i x in
          return (Cons (y, mapi_aux f (i+1) xs))

    let[@inline] mapi f xs = mapi_aux f 0 xs

    let rec tail_scan f s xs () =
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) ->
          let** s = f s x in
          return (Cons (s, tail_scan f s xs))

    let scan f s xs =
      cons s (tail_scan f s xs)

    let rec sorted_merge1l cmp x xs ys () =
      let* n = ys () in
      match n with
      | Nil ->
          return (Cons (x, xs))
      | Cons (y, ys) ->
          sorted_merge1 cmp x xs y ys

    and sorted_merge1r cmp xs y ys () =
      let* n = xs () in
      match n with
      | Nil ->
          return (Cons (y, ys))
      | Cons (x, xs) ->
          sorted_merge1 cmp x xs y ys

    and sorted_merge1 cmp x xs y ys =
      let** c = cmp x y in
      if c <= 0 then
        return (Cons (x, sorted_merge1r cmp xs y ys))
      else
        return (Cons (y, sorted_merge1l cmp x xs ys))

    let sorted_merge cmp xs ys () =
      let* xs = xs () in
      let* ys = ys () in
      match xs, ys with
        | Nil, Nil ->
            return Nil
        | Nil, c
        | c, Nil ->
            return c
        | Cons (x, xs), Cons (y, ys) ->
            sorted_merge1 cmp x xs y ys

    let rec take_while p xs () =
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) ->
          let** b = p x in
          if b then return (Cons (x, take_while p xs)) else return Nil

    let rec drop_while p xs () =
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) as node ->
          let** b = p x in
          if b then drop_while p xs () else return node

    let rec group eq xs () =
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) ->
          return (Cons (cons x (take_while (eq x) xs), group eq (drop_while (eq x) xs)))

    let rec map2 f xs ys () =
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) ->
          let* n = ys () in
          match n with
          | Nil ->
              return Nil
          | Cons (y, ys) ->
              let** z = f x y in
              return (Cons (z, map2 f xs ys))


    let rec filter_map_find_left_map f xs () =
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) ->
          let** fx = f x in
          match fx with
          | Either.Left y ->
              return (Cons (y, filter_map_find_left_map f xs))
          | Either.Right _ ->
              filter_map_find_left_map f xs ()

    let rec filter_map_find_right_map f xs () =
      let* n = xs () in
      match n with
      | Nil ->
          return Nil
      | Cons (x, xs) ->
          let** fx = f x in
          match fx with
          | Either.Left _ ->
              filter_map_find_right_map f xs ()
          | Either.Right z ->
              return (Cons (z, filter_map_find_right_map f xs))

    let partition_map f xs =
      filter_map_find_left_map f xs,
      filter_map_find_right_map f xs

    let partition p xs =
      filter p xs, filter (fun x -> Alt.bind (p x) (fun b -> Alt.return (not b))) xs

    let rec diagonals remainders xss () =
      let* n = xss () in
      match n with
      | Cons (xs, xss) ->
          begin
            let* n = xs () in
            match n with
          | Cons (x, xs) ->
              let heads, tails = peel remainders in
              return (Cons (cons x heads, diagonals (cons xs tails) xss))
          | Nil ->
              let heads, tails = peel remainders in
              return (Cons (heads, diagonals tails xss))
          end
      | Nil ->
          transpose remainders ()

    let diagonals xss =
      diagonals empty xss

    let map_product f xs ys =
      concat (diagonals (
        map
          (fun x ->
            Alt.return (fun () ->
            map
              (fun y -> f x y)
              ys
              ()
            )
          )
          xs
      ))

    let of_dispenser it =
      let rec c () =
        let** o = it () in
        match o with
        | None ->
            return Nil
        | Some x ->
            return (Cons (x, c))
      in
      c

  end

  include Make (Identity2) (Identity2)

  module M
  : Sigs2.SEQMON2TRANSFORMERS
      with type ('a, 'e) mon := ('a, 'e) Mon.t
      with type ('a, 'e) callermon := ('a, 'e) Mon.t
      with type ('a, 'e) t := ('a, 'e) t
  = Make(Mon)(Mon)

  (* defined here to avoid duplicating the code of [map_product] defined by
     [MakeTraversors(..)]'s inclusion *)
  let product xs ys =
    map_product (fun x y -> (x, y)) xs ys

  let return x () = return (Cons (x, empty))
end