open! Core

module type Monoid = sig
  type t

  val t : t Type_equal.Id.t
  val unit : t
  val op : t -> t -> t
  val eq : t -> t -> bool
end
module type S = Monoid

module Monoid = struct
  type 'a t = (module Monoid with type t = 'a)

  let same (type a1 a2) (monoid1 : a1 t) (monoid2 : a2 t) =
    let module Monoid1 = (val monoid1) in
    let module Monoid2 = (val monoid2) in
    Type_equal.Id.same_witness Monoid1.t Monoid2.t
  ;;
end

module MonoidElement = struct
  type t = Constructor : 'a Monoid.t * 'a -> t

  (* (* Alternative definition *)
  module type MonoidElement = sig
    module Monoid : Monoid
    val element : Monoid.t
  end
  type t = (module MonoidElement)
  *)

  let is_unit (Constructor (monoid, element)) =
    let module Monoid1 = (val monoid) in
    Monoid1.eq element Monoid1.unit
  ;;

  let eq (Constructor (monoid1, element1)) (Constructor (monoid2, element2)) =
    let module Monoid1 = (val monoid1) in
    let module Monoid2 = (val monoid2) in
    match Type_equal.Id.same_witness Monoid1.t Monoid2.t with
    | None -> false
    | Some eq_wit ->
      let element1_in_2 = Type_equal.conv eq_wit element1 in
      Monoid2.eq element1_in_2 element2
  ;;
end

module IntMonoid = struct
  type t = int

  let t = Type_equal.Id.create ~name:"IntMonoid.t" [%sexp_of: _]
  let unit = 0
  let op = ( + )
  let eq = Int.equal
end

let int_monoid : int Monoid.t = (module IntMonoid)

module StrMonoid = struct
  type t = string

  let t = Type_equal.Id.create ~name:"StrMonoid.t" [%sexp_of: _]
  let unit = "0"
  let op = ( ^ )
  let eq = String.equal
end

let str_monoid : string Monoid.t = (module StrMonoid)

module DList = struct
  type 'a t = 'a list

  let append l1 l2 = List.append l1 l2

  type 'a first_and_rest =
    | NoFirst
    | FirstAndRest of 'a * 'a t

  type 'a rest_and_last =
    | NoLast
    | RestAndLast of 'a t * 'a

  let uncons_first l : 'a first_and_rest =
    match l with
    | [] -> NoFirst
    | h :: t -> FirstAndRest (h, t)
  ;;

  let uncons_last l =
    match List.rev l with
    | [] -> NoLast
    | h :: t -> RestAndLast (t, h)
  ;;
end

module FreeProductOfAllMonoids = struct
  type t = MonoidElement.t DList.t
  (* Invariants : there are no neutral elements, and consecutive elements live in distinct monoids *)

  let unit = []

  type reduction_result =
    | Irreducible
    | Reduced of MonoidElement.t

  let maybe_reduce_aux (type a1 a2) (monoid1, element1) (monoid2, element2) =
    let module Monoid1 = (val monoid1 : Monoid with type t = a1) in
    let module Monoid2 = (val monoid2 : Monoid with type t = a2) in
    let element1_is_unit = Monoid1.eq element1 Monoid1.unit in
    let element2_is_unit = Monoid2.eq element2 Monoid2.unit in
    let same_monoid = Type_equal.Id.same_witness Monoid1.t Monoid2.t in
    match element1_is_unit, element2_is_unit, same_monoid with
    | true, _, _ -> Reduced (MonoidElement.Constructor (monoid2, element2))
    | _, true, _ -> Reduced (MonoidElement.Constructor (monoid1, element1))
    | _, _, Some eq ->
      let element1_in_2 = Type_equal.conv eq element1 in
      let element = Monoid2.op element1_in_2 element2 in
      Reduced (MonoidElement.Constructor (monoid2, element))
    | false, false, None -> Irreducible
  ;;

  let maybe_reduce x y =
    match x, y with
    | ( MonoidElement.Constructor (monoid1, element1)
      , MonoidElement.Constructor (monoid2, element2) ) ->
      maybe_reduce_aux (monoid1, element1) (monoid2, element2)
  ;;

  let rec op (t1 : t) (t2 : t) =
    match DList.uncons_last t1, DList.uncons_first t2 with
    | NoLast, _ -> t2
    | _, NoFirst -> t1
    | RestAndLast (es1, e1), FirstAndRest (e2, es2) ->
      (match maybe_reduce e1 e2 with
       | Irreducible -> DList.append t1 t2
       | Reduced e ->
         let ( $ ) = op in
         es1 $ [ e ] $ es2)
  ;;

  let rec eq t1 t2 = List.equal MonoidElement.eq t1 t2
end

let test () =
  let int i = [ MonoidElement.Constructor (int_monoid, i) ] in
  let str s = [ MonoidElement.Constructor (str_monoid, s) ] in
  let ( $ ) = FreeProductOfAllMonoids.op in
  let ok_str =
    if FreeProductOfAllMonoids.eq
         (str "a" $ int 1 $ int 1 $ int (-2) $ str "b")
         (str "ab")
    then "OK"
    else "Not OK"
  in
  print_endline ok_str
;;