(*****************************************************************************)
(*                                                                           *)
(* Open Source License                                                       *)
(* Copyright (c) 2022 Nomadic Labs. <contact@nomadic-labs.com>               *)
(*                                                                           *)
(* Permission is hereby granted, free of charge, to any person obtaining a   *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
(* and/or sell copies of the Software, and to permit persons to whom the     *)
(* Software is furnished to do so, subject to the following conditions:      *)
(*                                                                           *)
(* The above copyright notice and this permission notice shall be included   *)
(* in all copies or substantial portions of the Software.                    *)
(*                                                                           *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
(* DEALINGS IN THE SOFTWARE.                                                 *)
(*                                                                           *)
(*****************************************************************************)

type (_, _, _) arity =
  | Zero_arity : ('elt, 'elt, unit) arity
  | Succ_arity : ('elt, 'b, 'a) arity -> ('elt, 'elt -> 'b, int * 'a) arity

let arity_0 = Zero_arity

let arity_1 = Succ_arity arity_0

let arity_2 = Succ_arity arity_1

let arity_3 = Succ_arity arity_2

module type Model_impl = sig
  type arg_type

  val name : Namespace.t

  val takes_saturation_reprs : bool

  module Def (X : Costlang.S) : sig
    type model_type

    val arity : (X.size, model_type, arg_type) arity

    val model : model_type X.repr
  end
end

module type Instantiated = sig
  type 'a repr

  type size

  type arg_type

  type model_type

  val arity : (size, model_type, arg_type) arity

  val model : arg_type -> size repr
end

type 'arg model = (module Model_impl with type arg_type = 'arg)

module type App = sig
  type t

  val applied : t
end

module type Applied = functor (X : Costlang.S) ->
  App with type t = X.size X.repr

type applied = (module Applied)

type packed_model = Model : _ model -> packed_model

type _ t =
  | Abstract : {conv : 'workload -> 'arg; model : 'arg model} -> 'workload t
  | Aggregate : {
      model : 'workload -> applied;
      sub_models : packed_model list;
    }
      -> 'workload t

let pp_packed_model ppf (Model model) =
  let module Model = (val model) in
  let module Pp = Model.Def (Costlang.Pp) in
  Format.fprintf ppf "@[<v2>%a:@ %s@]" Namespace.pp Model.name Pp.model

let pp ppf = function
  | Abstract {model; _} ->
      Format.fprintf ppf "@[<v2>Abstract@ %a@]" pp_packed_model (Model model)
  | Aggregate {sub_models; _} ->
      Format.fprintf
        ppf
        "@[<v2>Aggregate with submodels:@ @[%a@]@]"
        (Format.pp_print_list pp_packed_model)
        sub_models

let apply_model : 'arg -> 'arg model -> applied =
  fun (type e) (elim : e) ((module Impl) : e model) ->
   let module Applied (X : Costlang.S) = struct
     include Impl.Def (X)

     type t = X.size X.repr

     let rec apply :
         type a b c.
         (int -> c X.repr) -> (c, a, b) arity -> a X.repr -> b -> c X.repr =
      fun conv arity f arg ->
       match arity with
       | Zero_arity -> f
       | Succ_arity ar ->
           let arg, rest = arg in
           apply conv ar (X.app f (conv arg)) rest

     let applied = apply X.int arity model elim
   end in
   ((module Applied) : applied)

module Instantiate (X : Costlang.S) (M : Model_impl) :
  Instantiated
    with type 'a repr = 'a X.repr
     and type size = X.size
     and type arg_type = M.arg_type = struct
  type 'a repr = 'a X.repr

  type size = X.size

  include M
  include Def (X)

  let rec apply :
      type a b c.
      (int -> c X.repr) -> (c, a, b) arity -> a X.repr -> b -> c X.repr =
   fun conv arity f arg ->
    match arity with
    | Zero_arity -> f
    | Succ_arity ar ->
        let arg, rest = arg in
        apply conv ar (X.app f (conv arg)) rest

  let model elim = apply X.int arity model elim
end

let set_takes_saturation_reprs (type a) b ((module Model) : a model) : a model =
  let module Model' = struct
    include Model

    let takes_saturation_reprs = b
  end in
  (module Model')

let make ?(takes_saturation_reprs = false) ~conv model =
  let model = set_takes_saturation_reprs takes_saturation_reprs model in
  Abstract {conv; model}

let make_aggregated ~model ~sub_models = Aggregate {model; sub_models}

let apply model workload =
  match model with
  | Abstract {conv; model} -> apply_model (conv workload) model
  | Aggregate {model; _} -> model workload

let force_aggregated ~model =
  match model with
  | Aggregate _ -> model
  | Abstract {conv = _; model = model2} ->
      Aggregate {model = apply model; sub_models = [Model model2]}

let add_aggregated_models :
    ('w1 -> applied) -> ('w2 -> applied) -> 'w1 * 'w2 -> applied =
 fun m1 m2 (w1, w2) ->
  let (module M1) = m1 w1 in
  let (module M2) = m2 w2 in
  let module M (X : Costlang.S) = struct
    type t = X.size X.repr

    let applied =
      let (module M1 : App with type t = X.size X.repr) = (module M1 (X)) in
      let (module M2 : App with type t = X.size X.repr) = (module M2 (X)) in
      X.(M1.applied + M2.applied)
  end in
  (module M : Applied)

let add_model m1 m2 =
  let m1 = force_aggregated ~model:m1 in
  let m2 = force_aggregated ~model:m2 in
  match (m1, m2) with
  | ( Aggregate {model = m1; sub_models = l1},
      Aggregate {model = m2; sub_models = l2} ) ->
      Aggregate {model = add_aggregated_models m1 m2; sub_models = l1 @ l2}
  | _ -> assert false (* impossible *)

let precompose : type a b. (a -> b) -> b t -> a t =
 fun f model ->
  match model with
  | Abstract {conv; model} ->
      let conv x = conv (f x) in
      Abstract {conv; model}
  | Aggregate {model; sub_models} ->
      Aggregate {model = (fun x -> model (f x)); sub_models}

let get_free_variable_set (type a) (model : a model) =
  let module M = (val model) in
  let module T0 = Costlang.Fold_constants (Costlang.Free_variables) in
  let module T1 = Costlang.Beta_normalize (T0) in
  let module R = M.Def (T1) in
  T0.prj @@ T1.prj R.model

(* No workload application.  For [Aggregate _], only extract
   the free variables of the [sub_models].
*)
let get_free_variable_set_of_t =
  let get_free_variables_of_packed_model (Model (module Model) : packed_model) =
    let module M = Model.Def (Costlang.Free_variables) in
    M.model
  in
  function
  | Abstract {model; _} -> get_free_variables_of_packed_model (Model model)
  | Aggregate {sub_models; _} ->
      List.fold_left
        (fun acc packed_model ->
          Free_variable.Set.union acc
          @@ get_free_variables_of_packed_model packed_model)
        Free_variable.Set.empty
        sub_models

let get_free_variable_set_applied (type workload) (model : workload t)
    (workload : workload) =
  (* If a parameter is fixed to 0 in the workload data, the application
      of the workload can eliminate free variables multiplied
      by the parameter.

      The typical example is the intercept case where some parameters
      tend to be fixed to 0.  This may not work when the intercept point
      is not at "zero"s.

      It is unfortunate that we need to apply workload data to a model to
      know which variables can be optimized out.  We may be able to do it
      without workload, but it seems not an easy task.
  *)
  let applied = apply model workload in
  let module M = (val applied) in
  let module T0 = Costlang.Fold_constants (Costlang.Free_variables) in
  let module T1 = Costlang.Beta_normalize (T0) in
  let module R = M (T1) in
  T0.prj @@ T1.prj R.applied

(* -------------------------------------------------------------------------- *)
(* Commonly used models *)

let zero =
  let module M = struct
    type arg_type = unit

    let name = Namespace.root "zero"

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size

      let arity = arity_0

      let model = int 0
    end
  end in
  (module M : Model_impl with type arg_type = unit)

let unknown_const1 ~name ~const =
  let module M = struct
    type arg_type = unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size

      let arity = arity_0

      let model = free ~name:const
    end
  end in
  (module M : Model_impl with type arg_type = unit)

let unknown_const1_skip1 ~name ~const =
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model = lam ~name:"size" @@ fun (_ : size repr) -> free ~name:const
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let unknown_const1_skip2 ~name ~const =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun (_ : size repr) ->
        lam ~name:"size2" @@ fun (_ : size repr) -> free ~name:const
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let linear ~name ~coeff =
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model = lam ~name:"size" @@ fun size -> free ~name:coeff * size
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let affine ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size ->
        free ~name:intercept + (free ~name:coeff * size)
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let affine_offset ~name ~intercept ~coeff ~offset =
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size ->
        free ~name:intercept + (free ~name:coeff * sat_sub size (int offset))
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let quadratic ~name ~coeff =
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size -> free ~name:coeff * (size * size)
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let nlogn ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size ->
        free ~name:intercept + (free ~name:coeff * (size * log2 (int 1 + size)))
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let nsqrtn_const ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size ->
        free ~name:intercept + (free ~name:coeff * (size * sqrt size))
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let logn ~name ~coeff =
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size -> free ~name:coeff * log2 (int 1 + size)
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let linear_sum ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept + (free ~name:coeff * (size1 + size2))
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let linear_sat_sub ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept + (free ~name:coeff * sat_sub size1 size2)
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let linear_max ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept + (free ~name:coeff * max size1 size2)
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let linear_min ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept + (free ~name:coeff * min size1 size2)
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let linear_min_offset ~name ~intercept ~coeff ~offset =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept
        + (free ~name:coeff * sat_sub (min size1 size2) (int offset))
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let linear_mul ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept + (free ~name:coeff * (size1 * size2))
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let bilinear ~name ~coeff1 ~coeff2 =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        (free ~name:coeff1 * size1) + (free ~name:coeff2 * size2)
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let bilinear_affine ~name ~intercept ~coeff1 ~coeff2 =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept
        + (free ~name:coeff1 * size1)
        + (free ~name:coeff2 * size2)
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let affine_skip1 ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun (_size1 : size repr) ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept + (free ~name:coeff * size2)
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let nlogm ~name ~intercept ~coeff =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept
        + (free ~name:coeff * (size1 * log2 (int 1 + size2)))
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let n_plus_logm ~name ~intercept ~linear_coeff ~log_coeff =
  let module M = struct
    type arg_type = int * (int * unit)

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size

      let arity = arity_2

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        free ~name:intercept
        + (free ~name:linear_coeff * size1)
        + (free ~name:log_coeff * log2 (int 1 + size2))
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * unit))

let trilinear ~name ~coeff1 ~coeff2 ~coeff3 =
  let module M = struct
    type arg_type = int * (int * (int * unit))

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size -> size -> size

      let arity = arity_3

      let model =
        lam ~name:"size1" @@ fun size1 ->
        lam ~name:"size2" @@ fun size2 ->
        lam ~name:"size3" @@ fun size3 ->
        (free ~name:coeff1 * size1)
        + (free ~name:coeff2 * size2)
        + (free ~name:coeff3 * size3)
    end
  end in
  (module M : Model_impl with type arg_type = int * (int * (int * unit)))

let breakdown ~name ~coeff1 ~coeff2 ~break =
  assert (0 <= break) ;

  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size ->
        (free ~name:coeff1 * min (int break) size)
        + (free ~name:coeff2 * sat_sub size (int break))
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let breakdown2 ~name ~coeff1 ~coeff2 ~coeff3 ~break1 ~break2 =
  assert (0 <= break1 && break1 <= break2) ;
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size ->
        (free ~name:coeff1 * min (int break1) size)
        + (free ~name:coeff2 * sat_sub (min (int break2) size) (int break1))
        + (free ~name:coeff3 * sat_sub size (int break2))
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let breakdown2_const ~name ~coeff1 ~coeff2 ~coeff3 ~const ~break1 ~break2 =
  assert (0 <= break1 && break1 <= break2) ;
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size ->
        (free ~name:coeff1 * min (int break1) size)
        + (free ~name:coeff2 * sat_sub (min (int break2) size) (int break1))
        + (free ~name:coeff3 * sat_sub size (int break2))
        + free ~name:const
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

let breakdown2_const_offset ~name ~coeff1 ~coeff2 ~coeff3 ~const ~break1 ~break2
    ~offset =
  assert (0 <= break1 && break1 <= break2) ;
  let module M = struct
    type arg_type = int * unit

    let name = name

    let takes_saturation_reprs = false

    module Def (X : Costlang.S) = struct
      open X

      type model_type = size -> size

      let arity = arity_1

      let model =
        lam ~name:"size" @@ fun size ->
        let_ ~name:"size" (sat_sub size (int offset)) @@ fun size ->
        (free ~name:coeff1 * min (int break1) size)
        + (free ~name:coeff2 * sat_sub (min (int break2) size) (int break1))
        + (free ~name:coeff3 * sat_sub size (int break2))
        + free ~name:const
    end
  end in
  (module M : Model_impl with type arg_type = int * unit)

module type Binary_operation = sig
  module Def (X : Costlang.S) : sig
    val op : X.size X.repr -> X.size X.repr -> X.size X.repr
  end
end

module Synthesize
    (B : Binary_operation)
    (X : Model_impl)
    (Y : Model_impl with type arg_type = X.arg_type) (Names : sig
      val name : Namespace.t

      val x_label : string

      val y_label : string
    end) : Model_impl with type arg_type = X.arg_type = struct
  type arg_type = X.arg_type

  let name = Names.name

  (* Use X's configuration *)
  let takes_saturation_reprs = X.takes_saturation_reprs

  module Def (C : Costlang.S) = struct
    module Args = X.Def (Costlang.Arg_names)
    module BinOp = B.Def (C)
    module X = X.Def (C)
    module Y = Y.Def (C)

    type model_type = X.model_type

    let arity = X.arity

    let rec synthesize :
        type a args_model x_model y_model.
        (Costlang.Arg_names.size, args_model, a) arity ->
        args_model Costlang.Arg_names.repr ->
        (C.size, x_model, a) arity ->
        x_model C.repr ->
        (C.size, y_model, a) arity ->
        y_model C.repr ->
        x_model C.repr =
     fun arg_arity args arity1 term1 arity2 term2 ->
      match arg_arity with
      | Zero_arity ->
          (* These bindings of Zero_arity are necessary for type checking *)
          let Zero_arity = arity1 in
          let Zero_arity = arity2 in
          let open C in
          let_ ~name:Names.x_label term1 @@ fun term1 ->
          let_ ~name:Names.y_label term2 @@ fun term2 -> BinOp.op term1 term2
      | Succ_arity arg_arity ->
          let (Succ_arity arity1) = arity1 in
          let (Succ_arity arity2) = arity2 in
          let open C in
          lam ~name:(Costlang.Arg_names.arg_name args) @@ fun arg ->
          synthesize
            arg_arity
            (Costlang.Arg_names.unwrap_size args)
            arity1
            (app term1 arg)
            arity2
            (app term2 arg)

    let model = synthesize Args.arity Args.model X.arity X.model Y.arity Y.model
  end
end

let synthesize (type a) ~name ~binop ~x_label ~x_model ~y_label ~y_model =
  let (module B : Binary_operation) = binop in
  let ((module X) : a model) = x_model in
  let ((module Y) : a model) = y_model in
  let module M =
    Synthesize (B) (X) (Y)
      (struct
        let x_label = x_label

        let y_label = y_label

        let name = name
      end)
  in
  ((module M) : a model)