(*
copyright (c) 2013, simon cruanes
all rights reserved.

redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.  redistributions in binary
form must reproduce the above copyright notice, this list of conditions and the
following disclaimer in the documentation and/or other materials provided with
the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*)

(** {1 Composable State Machines}

This module defines state machines that should help design applications
with a more explicit control of state (e.g. for networking applications. *)

type ('a, 's, 'b) t = 's -> 'a -> ('b * 's) option
(** transition function that fully describes an automaton *)

type ('a, 's, 'b) automaton = ('a, 's, 'b) t

(** {2 Basic Interface} *)

let empty _st _x = None

let id () x = Some (x,())

let repeat x () () = Some (x, ())

let get_state a state x = match a state x with
  | None -> None
  | Some (_, state') -> Some (state', state')

let next a s x = a s x

let scan a (st, prev) x =
  match a st x with
  | None -> None
  | Some (y,state') ->
      Some (y::prev, (state', y::prev))

let map_in f a state x = a state (f x)
let map_out f a state x = match a state x with
  | None -> None
  | Some (y, state') ->
      Some (f y, state')

exception ExitNest

let nest l =
  let rec eval (answers, res_states) l state x =
    match l, state with
    | [], [] ->
      Some (List.rev answers, List.rev res_states)
    | a::l', state::states' ->
        begin match a state x with
        | None -> raise ExitNest
        | Some (ans,state') ->
          eval (ans::answers, state'::res_states) l' states' x
        end
    | [], _
    | _, [] ->
      raise (Invalid_argument "CSM.next: list length mismatch")
  in
  fun state x ->
    try eval ([],[]) l state x
    with ExitNest -> None

let split a state x = match a state x with
  | None -> None
  | Some (y, state') -> Some ((y,y), state')

let unsplit merge a state x = match a state x with
  | None -> None
  | Some ((y,z), state') ->
      Some (merge y z, state')

let pair a1 a2 (s1,s2) (x1,x2) =
  match a1 s1 x1, a2 s2 x2 with
  | Some (y1,s1'), Some (y2, s2') ->
      Some ((y1,y2), (s1',s2'))
  | Some _, None
  | None, Some _
  | None, None -> None

let ( *** ) = pair

let first a state (x,keep) = match a state x with
  | None -> None
  | Some (y,state') ->
      Some ((y,keep), state')

let second a state (keep,x) = match a state x with
  | None -> None
  | Some (y,state') ->
      Some ((keep,y), state')

let (>>>) a1 a2 (s1, s2) x =
  match a1 s1 x with
  | None -> None
  | Some (y, s1') ->
      match a2 s2 y with
      | None -> None
      | Some (z, s2') ->
          Some (z, (s1', s2'))

let _flatmap_opt f o = match o with
  | None -> None
  | Some x -> f x

let append a1 a2 state x =
  match state with
  | `Left s1 ->
      _flatmap_opt (fun (y,s1) -> Some (y,`Left s1)) (a1 s1 x)
  | `Right s2 ->
      _flatmap_opt (fun (y,s2) -> Some (y,`Right s2)) (a2 s2 x)

let rec flatten (automata,state) x = match automata with
  | [] -> None
  | a::automata' ->
      match a state x with
      | None -> flatten (automata', state) x
      | Some (y, state') ->
          Some (y, (automata,state'))

let filter p a state x = match a state x with
  | None -> None
  | Some (y, state') ->
      if p y then Some (Some y, state') else Some (None, state')

type ('a, 'c, 's1, 's2) flat_map_state =
  ('s1 * (('a, 's2, 'c) t * 's2) option)

let rec flat_map f a state x =
  match state with
  | s1, None ->
      begin match a s1 x with
      | None -> None
      | Some (y, s1') ->
          let a2, s2 = f y in
          flat_map f a (s1', Some (a2,s2)) x
      end
  | s1, Some(a2,s2) ->
      begin match a2 s2 x with
      | None -> flat_map f a (s1, None) x
      | Some (z, s2') ->
          let state' = s1, Some (a2, s2') in
          Some (z, state')
      end

(** {2 Mutable Interface} *)

module Mut = struct
  type ('a, 's, 'b) t = {
    next : ('a, 's, 'b) automaton;
    mutable state : 's;
  } (** mutable automaton, with in-place modification *)

  let create a ~init =
    { next=a; state=init; }

  let next a x =
    match a.next a.state x with
    | None -> None
    | Some (y,state) ->
        a.state <- state;
        Some y

  let copy a = { a with state=a.state; }
end

(** {2 Instances} *)

module Int = struct
  let range j state () =
    if state > j then None
    else Some (state, state+1)
end