(* This file is part of the Catala compiler, a specification language for tax
   and social benefits computation rules. Copyright (C) 2023 Inria, contributor:
   Denis Merigoux <denis.merigoux@inria.fr>

   Licensed under the Apache License, Version 2.0 (the "License"); you may not
   use this file except in compliance with the License. You may obtain a copy of
   the License at

   http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
   WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
   License for the specific language governing permissions and limitations under
   the License. *)

open Shared_ast
open Ast
open Catala_utils

type option_instance = {
  name : EnumName.t;
  some_cons : EnumConstructor.t;
  some_typ : naked_typ;
  none_cons : EnumConstructor.t;
}

type tuple_instance = {
  name : StructName.t;
  fields : (StructField.t * naked_typ) list;
}

type array_instance = {
  name : StructName.t;
  len_field : StructField.t;
  content_field : StructField.t;
  content_typ : naked_typ;
}

type monomorphized_instances = {
  (* The keys are the types inside the [TOption] (before monomorphization). *)
  options : option_instance Type.Map.t;
  (* The keys are the [TTuple] types themselves (before monomorphization). *)
  tuples : tuple_instance Type.Map.t;
  (* The keys are the types inside the [TArray] (before monomorphization). *)
  arrays : array_instance Type.Map.t;
}

let collect_monomorphized_instances (prg : typed program) :
    monomorphized_instances =
  let option_instances_counter = ref 0 in
  let tuple_instances_counter = ref 0 in
  let array_instances_counter = ref 0 in
  let rec collect_typ acc typ =
    match Mark.remove typ with
    | TTuple args when List.for_all (fun t -> Mark.remove t <> TAny) args ->
      let new_acc =
        {
          acc with
          tuples =
            Type.Map.update typ
              (fun monomorphized_name ->
                match monomorphized_name with
                | Some e -> Some e
                | None ->
                  incr tuple_instances_counter;
                  Some
                    {
                      fields =
                        List.mapi
                          (fun i arg ->
                            ( StructField.fresh
                                ("elt_" ^ string_of_int i, Pos.no_pos),
                              Mark.remove arg ))
                          args;
                      name =
                        StructName.fresh []
                          ( "tuple_" ^ string_of_int !option_instances_counter,
                            Pos.no_pos );
                    })
              acc.tuples;
        }
      in
      List.fold_left collect_typ new_acc args
    | TArray t ->
      let new_acc =
        {
          acc with
          arrays =
            Type.Map.update t
              (fun monomorphized_name ->
                match monomorphized_name with
                | Some e -> Some e
                | None ->
                  incr array_instances_counter;
                  Some
                    {
                      len_field = StructField.fresh ("length", Pos.no_pos);
                      content_field = StructField.fresh ("content", Pos.no_pos);
                      content_typ = Mark.remove t;
                      name =
                        StructName.fresh []
                          ( "array_" ^ string_of_int !array_instances_counter,
                            Pos.no_pos );
                    })
              acc.arrays;
        }
      in
      collect_typ new_acc t
    | TDefault t -> collect_typ acc t
    | TArrow (args, ret) ->
      List.fold_left collect_typ (collect_typ acc ret) args
    | TOption t when Mark.remove t <> TAny ->
      let new_acc =
        {
          acc with
          options =
            Type.Map.update t
              (fun monomorphized_name ->
                match monomorphized_name with
                | Some e -> Some e
                | None ->
                  incr option_instances_counter;
                  Some
                    {
                      some_cons =
                        EnumConstructor.fresh
                          ( "Some_" ^ string_of_int !option_instances_counter,
                            Pos.no_pos );
                      none_cons =
                        EnumConstructor.fresh
                          ( "None_" ^ string_of_int !option_instances_counter,
                            Pos.no_pos );
                      some_typ = Mark.remove t;
                      name =
                        EnumName.fresh []
                          ( "option_" ^ string_of_int !option_instances_counter,
                            Pos.no_pos );
                    })
              acc.options;
        }
      in
      collect_typ new_acc t
    | TStruct _ | TEnum _ | TAny | TClosureEnv | TLit _ -> acc
    | TOption _ | TTuple _ ->
      raise
        (Message.CompilerError
           (Message.Content.add_position
              (Message.Content.to_internal_error
                 (Message.Content.of_message (fun fmt ->
                      Format.fprintf fmt
                        "Some types in tuples or option have not been resolved \
                         by the typechecking before monomorphization.")))
              (Mark.get typ)))
  in
  let rec collect_expr e acc =
    Expr.shallow_fold collect_expr e (collect_typ acc (Expr.ty e))
  in
  let acc =
    Scope.fold_left
      ~init:
        {
          options = Type.Map.empty;
          tuples = Type.Map.empty;
          arrays = Type.Map.empty;
        }
      ~f:(fun acc item _ ->
        match item with
        | Topdef (_, typ, e) -> collect_typ (collect_expr e acc) typ
        | ScopeDef (_, body) ->
          let _, body = Bindlib.unbind body.scope_body_expr in
          Scope.fold_left_lets ~init:acc
            ~f:(fun acc { scope_let_typ; scope_let_expr; _ } _ ->
              collect_typ (collect_expr scope_let_expr acc) scope_let_typ)
            body)
      prg.code_items
  in
  EnumName.Map.fold
    (fun _ constructors acc ->
      EnumConstructor.Map.fold
        (fun _ t acc -> collect_typ acc t)
        constructors acc)
    prg.decl_ctx.ctx_enums
    (StructName.Map.fold
       (fun _ fields acc ->
         StructField.Map.fold (fun _ t acc -> collect_typ acc t) fields acc)
       prg.decl_ctx.ctx_structs acc)

let rec monomorphize_typ
    (monomorphized_instances : monomorphized_instances)
    (typ : typ) : typ =
  match Mark.remove typ with
  | TStruct _ | TEnum _ | TAny | TClosureEnv | TLit _ -> typ
  | TArray t1 ->
    TStruct (Type.Map.find t1 monomorphized_instances.arrays).name, Mark.get typ
  | TDefault t1 ->
    TDefault (monomorphize_typ monomorphized_instances t1), Mark.get typ
  | TArrow (t1s, t2) ->
    ( TArrow
        ( List.map (monomorphize_typ monomorphized_instances) t1s,
          monomorphize_typ monomorphized_instances t2 ),
      Mark.get typ )
  | TTuple _ ->
    ( TStruct (Type.Map.find typ monomorphized_instances.tuples).name,
      Mark.get typ )
  | TOption t1 ->
    TEnum (Type.Map.find t1 monomorphized_instances.options).name, Mark.get typ

let is_some c =
  EnumConstructor.equal Expr.some_constr c
  ||
  (assert (EnumConstructor.equal Expr.none_constr c);
   false)

let rec monomorphize_expr
    (monomorphized_instances : monomorphized_instances)
    (e0 : typed expr) : typed expr boxed =
  let ty0 = Expr.ty e0 in
  (* Keys in [monomorphized_instances] are before monomorphization, so collect
     this top-down *)
  let f_expr = monomorphize_expr monomorphized_instances in
  let f_ty = monomorphize_typ monomorphized_instances in
  (* Proceed bottom-up: apply first to the sub-terms *)
  let e = Expr.map ~f:f_expr ~typ:f_ty ~op:Fun.id e0 in
  let m = Mark.get e in
  let map_box f = Expr.Box.app1 e (fun e -> f (Mark.remove e)) m in
  map_box
  @@ function
  | ETuple args ->
    let tuple_instance = Type.Map.find ty0 monomorphized_instances.tuples in
    EStruct
      {
        name = tuple_instance.name;
        fields =
          StructField.Map.of_list
          @@ List.map2
               (fun (tuple_field, _) arg -> tuple_field, arg)
               tuple_instance.fields args;
      }
  | ETupleAccess { e; index; _ } ->
    (* The type of the tuple needs to be recovered from the untransformed
       expr *)
    let tup_ty =
      match e0 with ETupleAccess { e; _ }, _ -> Expr.ty e | _ -> assert false
    in
    let tuple_instance = Type.Map.find tup_ty monomorphized_instances.tuples in
    EStructAccess
      {
        name = tuple_instance.name;
        e;
        field = fst (List.nth tuple_instance.fields index);
      }
  | EMatch { name; e; cases } when EnumName.equal name Expr.option_enum ->
    let option_instance = Type.Map.find ty0 monomorphized_instances.options in
    EMatch
      {
        name = option_instance.name;
        e;
        cases =
          EnumConstructor.Map.fold
            (fun c ->
              EnumConstructor.Map.add
                (if is_some c then option_instance.some_cons
                 else option_instance.none_cons))
            cases EnumConstructor.Map.empty;
      }
  | EInj { name; e; cons } when EnumName.equal name Expr.option_enum ->
    let option_instance =
      Type.Map.find
        (match Mark.remove ty0 with TOption t -> t | _ -> assert false)
        monomorphized_instances.options
    in
    EInj
      {
        name = option_instance.name;
        e;
        cons =
          (if is_some cons then option_instance.some_cons
           else option_instance.none_cons);
      }
  | EArray elts as e ->
    let elt_ty =
      match Mark.remove ty0 with TArray t -> t | _ -> assert false
    in
    let array_instance = Type.Map.find elt_ty monomorphized_instances.arrays in
    EStruct
      {
        name = array_instance.name;
        fields =
          StructField.Map.of_list
            [
              ( array_instance.len_field,
                ( ELit (LInt (Runtime.integer_of_int (List.length elts))),
                  Expr.with_ty m (TLit TInt, Expr.mark_pos m) ) );
              ( array_instance.content_field,
                (e, Expr.with_ty m (TArray (f_ty elt_ty), Expr.mark_pos m)) );
            ];
      }
  | e -> e

let program (prg : typed program) :
    typed program * Scopelang.Dependency.TVertex.t list =
  let monomorphized_instances = collect_monomorphized_instances prg in
  (* First we remove the polymorphic option type *)
  let prg =
    {
      prg with
      decl_ctx =
        {
          prg.decl_ctx with
          ctx_enums =
            EnumName.Map.remove Expr.option_enum prg.decl_ctx.ctx_enums;
        };
    }
  in
  (* Then we replace all hardcoded types and expressions with the monomorphized
     instances *)
  let prg =
    {
      prg with
      decl_ctx =
        {
          prg.decl_ctx with
          ctx_enums =
            EnumName.Map.map
              (EnumConstructor.Map.map
                 (monomorphize_typ monomorphized_instances))
              prg.decl_ctx.ctx_enums;
          ctx_structs =
            StructName.Map.map
              (StructField.Map.map (monomorphize_typ monomorphized_instances))
              prg.decl_ctx.ctx_structs;
        };
    }
  in
  (* Then we augment the [decl_ctx] with the monomorphized instances *)
  let prg =
    {
      prg with
      decl_ctx =
        {
          prg.decl_ctx with
          ctx_enums =
            Type.Map.fold
              (fun _ (option_instance : option_instance) (ctx_enums : enum_ctx) ->
                EnumName.Map.add option_instance.name
                  (EnumConstructor.Map.add option_instance.none_cons
                     (TLit TUnit, Pos.no_pos)
                     (EnumConstructor.Map.singleton option_instance.some_cons
                        (monomorphize_typ monomorphized_instances
                           (option_instance.some_typ, Pos.no_pos))))
                  ctx_enums)
              monomorphized_instances.options prg.decl_ctx.ctx_enums;
          ctx_structs =
            Type.Map.fold
              (fun _ (tuple_instance : tuple_instance)
                   (ctx_structs : struct_ctx) ->
                StructName.Map.add tuple_instance.name
                  (List.fold_left
                     (fun acc (field, typ) ->
                       StructField.Map.add field
                         (monomorphize_typ monomorphized_instances
                            (typ, Pos.no_pos))
                         acc)
                     StructField.Map.empty tuple_instance.fields)
                  ctx_structs)
              monomorphized_instances.tuples
              (Type.Map.fold
                 (fun _ (array_instance : array_instance)
                      (ctx_structs : struct_ctx) ->
                   StructName.Map.add array_instance.name
                     (StructField.Map.add array_instance.content_field
                        ( TArray
                            (monomorphize_typ monomorphized_instances
                               (array_instance.content_typ, Pos.no_pos)),
                          Pos.no_pos )
                        (StructField.Map.singleton array_instance.len_field
                           (TLit TInt, Pos.no_pos)))
                     ctx_structs)
                 monomorphized_instances.arrays prg.decl_ctx.ctx_structs);
        };
    }
  in
  let code_items =
    Bindlib.unbox
    @@ Scope.map
         ~f:(fun code_item ->
           match code_item with
           | Topdef (name, typ, e) -> Bindlib.box (Topdef (name, typ, e))
           | ScopeDef (name, body) ->
             let s_var, scope_body = Bindlib.unbind body.scope_body_expr in
             Bindlib.box_apply
               (fun scope_body_expr ->
                 ScopeDef (name, { body with scope_body_expr }))
               (Bindlib.bind_var s_var
                  (Scope.map_exprs_in_lets ~varf:Fun.id
                     ~transform_types:(monomorphize_typ monomorphized_instances)
                     ~f:(monomorphize_expr monomorphized_instances)
                     scope_body)))
         ~varf:Fun.id prg.code_items
  in
  ( { prg with code_items },
    Scopelang.Dependency.check_type_cycles prg.decl_ctx.ctx_structs
      prg.decl_ctx.ctx_enums )