Source file typeInference.ml

open UtilsLib.Utils
open Lambda

module Log = UtilsLib.Xlog.Make (struct
  let name = "TypeInference"
end)

module Value = struct
  type t = int
  type value = Lambda.stype

  let unfold stype _ =
    match stype with
    | Lambda.Atom _ -> None
    | Lambda.LFun (a, b) -> Some (1, [ a; b ])
    | Lambda.Fun (a, b) ->
        Some (1, [ a; b ])
        (* CHECK TODO: Shouldn't
           there be a difference
           in the 2
           constructors? *)
    | _ -> failwith "Bug: No type inference on these types"

  let pp fmt ty = Format.fprintf fmt "%s" (Lambda.raw_type_to_string ty)
end

module UF = VarUnionFind.UF (Value)

module Type = struct
  exception Not_typable

  type typing_env = {
    l_level : int;
    nl_level : int;
    lvar_typing : int IntMap.t;
    nlvar_typing : int IntMap.t;
    const_typing : (int * int) IntMap.t;
    (* maps the occurrence position, which is unique,
       to a pair consisting of the type variable and the
       constant identifier *)
    cst_nbr : int;
    type_equations : UF.t;
  }

  let empty_env =
    {
      l_level = 0;
      nl_level = 0;
      lvar_typing = IntMap.empty;
      nlvar_typing = IntMap.empty;
      const_typing = IntMap.empty;
      cst_nbr = 0;
      type_equations = UF.empty;
    }

  let type_equation_log _ eq =
    Log.debug (fun m -> m "type equation:@,@[<v>  @[%a@]@]" UF.pp eq)
  (*    UF.log_content Logs.Debug eq *)

  let rec inference_aux level t ty_var env =
    let prefix = String.make (level * 3) ' ' in
    Log.debug (fun m ->
        m "%sType inference of %s (currently %d)." prefix
          (Lambda.raw_to_string t) ty_var);
    Log.debug (fun m ->
        m "%sEquations are:@,@[<v>  @[%a@]@]" prefix UF.pp env.type_equations);
    (*    type_equation_log prefix env.type_equations; *)
    let ty, new_env =
      match t with
      | Lambda.Var i -> (
          try
            let ty_in_env =
              IntMap.find (env.nl_level - i - 1) env.nlvar_typing
            in
            Log.debug (fun m ->
                m
                  "%sAdding an equation (variable found in the environment) \
                   %d<-->%d"
                  prefix ty_var ty_in_env);
            let new_eq = UF.union ty_var ty_in_env env.type_equations in
            (ty_var, { env with type_equations = new_eq })
          with Not_found ->
            let new_var, new_eq = UF.generate_new_var env.type_equations in
            Log.debug (fun m ->
                m "%sAdding a new variable %d and an equation" prefix new_var);
            ( new_var,
              {
                env with
                nlvar_typing = IntMap.add i new_var env.nlvar_typing;
                type_equations = new_eq;
              } ))
      | Lambda.LVar i -> (
          try
            let ty_in_env = IntMap.find (env.l_level - i - 1) env.lvar_typing in
            Log.debug (fun m ->
                m
                  "%sAdding an equation (Lvariable found in the environment) \
                   %d<-->%d"
                  prefix ty_var ty_in_env);
            let new_eq = UF.union ty_var ty_in_env env.type_equations in
            (ty_var, { env with type_equations = new_eq })
          with Not_found ->
            let new_var, new_eq = UF.generate_new_var env.type_equations in
            Log.debug (fun m ->
                m "%sAdding a new Lvariable %d and an equation" prefix new_var);
            ( new_var,
              {
                env with
                lvar_typing = IntMap.add i new_var env.lvar_typing;
                type_equations = new_eq;
              } ))
      | Lambda.Const i ->
          (* Each occurence of a constants is considered as a new free
             variables *)
          let new_var, new_eq = UF.generate_new_var env.type_equations in
          let new_eq = UF.union ty_var new_var new_eq in
          ( new_var,
            {
              env with
              type_equations = new_eq;
              const_typing =
                IntMap.add (env.cst_nbr + 1) (new_var, i) env.const_typing;
              cst_nbr = env.cst_nbr + 1;
            } )
      | Lambda.DConst _ ->
          failwith
            "Bug: there should not remain any defined constant when  computing \
             the principal type"
      | Lambda.Abs (_x, t) ->
          Log.debug (fun m -> m "%sType inference of an abstraction:" prefix);
          let alpha, new_eq = UF.generate_new_var env.type_equations in
          Log.debug (fun m ->
              m "%sAdded a variable at %d. Equations are:" prefix alpha);
          let () = type_equation_log prefix new_eq in
          let beta, new_eq = UF.generate_new_var new_eq in
          Log.debug (fun m ->
              m "%sAdded a variable at %d. Equations are:" prefix beta);
          let () = type_equation_log prefix new_eq in
          let new_const, new_eq =
            UF.generate_new_constr new_eq (1, [ alpha; beta ])
          in
          Log.debug (fun m ->
              m "%sAdded new const at %d. Equations are:" prefix new_const);
          let () = type_equation_log prefix new_eq in
          Log.debug (fun m ->
              m "%sPreparing a Union %d %d." prefix ty_var new_const);
          let new_eq = UF.union ty_var new_const new_eq in
          Log.debug (fun m ->
              m "%sAdded a varibale at %d. Equations are:" prefix beta);
          type_equation_log prefix new_eq;
          let _, new_env =
            inference_aux (level + 1) t beta
              {
                env with
                nl_level = env.nl_level + 1;
                nlvar_typing = IntMap.add env.nl_level alpha env.nlvar_typing;
                type_equations = new_eq;
              }
          in
          let _is_cyclic, new_eq = UF.cyclic ty_var new_env.type_equations in
          ( ty_var,
            {
              env with
              type_equations = new_eq;
              const_typing = new_env.const_typing;
              cst_nbr = new_env.cst_nbr;
            } )
      | Lambda.LAbs (_x, t) ->
          Log.debug (fun m ->
              m "%sType inference of a linear abstraction:" prefix);
          let alpha, new_eq = UF.generate_new_var env.type_equations in
          let beta, new_eq = UF.generate_new_var new_eq in
          let new_const, new_eq =
            UF.generate_new_constr new_eq (1, [ alpha; beta ])
          in
          let new_eq = UF.union ty_var new_const new_eq in
          let _, new_env =
            inference_aux (level + 1) t beta
              {
                env with
                l_level = env.l_level + 1;
                lvar_typing = IntMap.add env.l_level alpha env.lvar_typing;
                type_equations = new_eq;
              }
          in
          let _is_cyclic, new_eq = UF.cyclic ty_var new_env.type_equations in
          ( ty_var,
            {
              env with
              type_equations = new_eq;
              const_typing = new_env.const_typing;
              cst_nbr = new_env.cst_nbr;
            } )
      (*	ty_var,{new_env with type_equations=new_eq;lvar_typing=env.lvar_typing} *)
      (*	ty_var,{new_env with type_equations=new_eq} *)
      | Lambda.App (t, u) ->
          let u_type, new_eq = UF.generate_new_var env.type_equations in
          let t_type, new_eq =
            UF.generate_new_constr new_eq (1, [ u_type; ty_var ])
          in
          Log.debug (fun m ->
              m "%sType inference of the parameter in an application:" prefix);
          let _u_type, new_env =
            inference_aux (level + 1) u u_type
              { env with type_equations = new_eq }
          in
          Log.debug (fun m ->
              m "%sType inference of the functor in an application:" prefix);
          let _t_type, new_env = inference_aux (level + 1) t t_type new_env in
          (ty_var, new_env)
      | _ -> failwith "Bug: No principal typing algorithm for these types"
    in
    let is_cyclic, new_eq = UF.cyclic ty new_env.type_equations in
    if is_cyclic then raise Not_typable
    else (ty, { new_env with type_equations = new_eq })

  let rec build_type i type_eq =
    let (i, v), type_eq = UF.find i type_eq in
    match v with
    | UF.Link_to j when j = i -> Lambda.Atom (-i)
    | UF.Link_to _ ->
        failwith
          "Bug: when UF.find returns a Link_to, it should be a Link_to itself"
    | UF.Value _ ->
        failwith
          "Bug: when performing type inference for principal typing, no type \
           constant should appear"
    | UF.Constr (_c, [ alpha; beta ]) ->
        let alpha' = build_type alpha type_eq in
        let beta' = build_type beta type_eq in
        Lambda.Fun (alpha', beta')
    | UF.Constr _ ->
        failwith
          "Bug: when performing type inference for principal typing, the only \
           allowd type construction is the arrow"

  let inference t =
    try
      let vars = UF.empty in
      let ty, vars = UF.generate_new_var vars in
      let ty, env =
        inference_aux 0 t ty { empty_env with type_equations = vars }
      in
      ( build_type ty env.type_equations,
        IntMap.map
          (fun (ty, i) -> (Lambda.Const i, build_type ty env.type_equations))
          env.const_typing )
    with UF.Union_Failure -> raise Not_typable
end