module type Monoid = sig
  type t

  val unit : t
  val op : t -> t -> t
  val eq : t -> t -> bool
  val uid : int (* For testing equality of monoids *)
end

module type MonoidElem = sig
  module Monoid : Monoid

  val elem : Monoid.t
end

type monoid_elem = (module MonoidElem)

module type MonoidWithPacker = sig
  include Monoid

  val pack : t -> monoid_elem
end

module AddPacker (M : Monoid) : MonoidWithPacker with type t = M.t = struct
  include M

  let pack t =
    let module Elem = struct
      module Monoid = M

      let elem = t
    end
    in
    (module Elem : MonoidElem)
  ;;
end

module IntMonoid = AddPacker (struct
  type t = int

  let unit = 0
  let op = ( + )
  let eq t1 t2 = t1 = t2
  let uid = Hashtbl.hash "IntMonoid"
end)

module StrMonoid = AddPacker (struct
  type t = string

  let unit = "0"
  let op = ( ^ )
  let eq t1 t2 = t1 = t2
  let uid = Hashtbl.hash "StrMonoid"
end)

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 = monoid_elem DList.t
  (* Invariants : there are no neutral elements, and consecutive elements live in distinct monoids *)

  let unit = []

  let rec op t1 t2 =
    match DList.uncons_last t1, DList.uncons_first t2 with
    | NoLast, _ -> t2
    | _, NoFirst -> t1
    | RestAndLast (es1, e1), FirstAndRest (e2, es2) ->
      let module E1 = (val e1 : MonoidElem) in
      let module E2 = (val e2 : MonoidElem) in
      if E1.Monoid.uid != E2.Monoid.uid
      then DList.append t1 t2
      else (
        let elem = E1.Monoid.op E1.elem (Obj.magic E2.elem) in
        (* Remove Obj.magic? *)
        if E1.Monoid.eq elem E1.Monoid.unit
        then op es1 es2
        else
          let module E1MonoidWithPacker = AddPacker (E1.Monoid) in
          let e = E1MonoidWithPacker.pack elem in
          op (op es1 [ e ]) es2)
  ;;

  let rec eq t1 t2 =
    List.equal (fun e1 e2 ->
      let module E1 = (val e1 : MonoidElem) in
      let module E2 = (val e2 : MonoidElem) in
      E1.Monoid.uid = E2.Monoid.uid && E1.Monoid.eq E1.elem (Obj.magic E2.elem)) t1 t2
  ;;
end

let test () =
  let int i = [ IntMonoid.pack i ] in
  let str s = [ StrMonoid.pack 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
;;