Source file type_system.ml

(**************************************************************************)
(*                                                                        *)
(*                 ACG development toolkit                                *)
(*                                                                        *)
(*                  Copyright 2008-2023 INRIA                             *)
(*                                                                        *)
(*  More information on "https://acg.loria.fr/"                     *)
(*  License: CeCILL, see the LICENSE file or "http://www.cecill.info"     *)
(*  Authors: see the AUTHORS file                                         *)
(*                                                                        *)
(*                                                                        *)
(*                                                                        *)
(*                                                                        *)
(*                                                                        *)
(**************************************************************************)

open UtilsLib
open Logic.Abstract_syntax
open Logic.Lambda

let get_location = function
  | Abstract_syntax.Var (_, l) -> l
  | Abstract_syntax.Const (_, l) -> l
  | Abstract_syntax.Abs (_, _, _, l) -> l
  | Abstract_syntax.LAbs (_, _, _, l) -> l
  | Abstract_syntax.App (_, _, l) -> l

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

module type SIG_ACCESS = sig
  exception Not_found

  type t

  val unfold_type_definition : int -> t -> Lambda.stype
  (** [unfold_type_definition id t] returns the actual type for the
      type defined by [id] as the identifier of a type definition in
      the signature [t]. Fails with "Bug" if [id] does not correspond
      to a type definition *)

  val expand_type : Lambda.stype -> t -> Lambda.stype
  val find_term : string -> t -> Lambda.term * Lambda.stype
  val pp_type : t -> Format.formatter -> Lambda.stype -> unit
  val pp_term : t -> Format.formatter -> Lambda.term -> unit
end

module Type_System = struct
  module Make (Signature : SIG_ACCESS) = struct
    exception Not_functional_type
    exception Functional_type of Abstract_syntax.abstraction
    exception Not_normal_term

    exception
      Vacuous_abstraction of
        (string * Abstract_syntax.location * Abstract_syntax.location)

    exception Non_empty_linear_context of (string * Abstract_syntax.location)

    exception
      Not_linear of (Abstract_syntax.location * Abstract_syntax.location)

    exception
      Type_mismatch of (Abstract_syntax.location * Lambda.stype * Lambda.stype)

    let decompose_functional_type ty sg =
      match Signature.expand_type ty sg with
      | Lambda.LFun (ty1, ty2) -> (ty1, ty2, Abstract_syntax.Linear)
      | Lambda.Fun (ty1, ty2) -> (ty1, ty2, Abstract_syntax.Non_linear)
      (*      | Lambda.DAtom i -> decompose_functional_type (Signature.unfold_type_definition i sg) sg *)
      | _ -> raise Not_functional_type

    (* [get_typing x loc typing_env] returns [l,t,e] where [l] is the
       level of [x] and [t] its type in the typing environment
       [typing_env] when [x] is a variable located at [loc]. [e] is the
       new environment where [x] as been marked as used at location
       [l]. If [x] has been used alread once (this is marked by the
       [Some l], 3rd projection in the association list), the function
       raises [Not_linear (l,loc)] where [l] is the location of the former
       usage of [x]. If [x] is not in the typing environment
       [typing_env], it raises Not_found *)

    let get_typing x loc lst =
      let rec get_typing_aux lst k =
        match lst with
        | [] -> raise Not_found
        | (s, (level, ty, None)) :: tl when s = x ->
            k (level, ty, (s, (level, ty, Some loc)) :: tl)
        | (s, (_, _, Some l)) :: _ when s = x -> raise (Not_linear (l, loc))
        | hd :: tl ->
            get_typing_aux tl (fun (level, ty, r) -> k (level, ty, hd :: r))
      in
      get_typing_aux lst (fun (level, ty, env) -> (level, ty, env))
      [@@warning "-32"]

    type typing_env_content =
      | Delay of (int -> int -> Abstract_syntax.location -> Lambda.term)
      | Eval of
          (int ->
          int ->
          Abstract_syntax.location ->
          Lambda.term * typing_env_content)

    type typing_environment = {
      linear_level : int;
      (* The depth of the term with respect to the
         number of linear abstraction. Starting at 0 *)
      level : int;
      (* The depth of the term with respect to the number
         of non linear abstraction. Starting at 0 *)
      env :
        (typing_env_content * Lambda.stype * Abstract_syntax.abstraction * int)
        Utils.StringMap.t;
      (* the last int parameter is the nomber of occurences of the variable with name the key of the map *)
      wrapper :
        (Abstract_syntax.location * Lambda.stype * Abstract_syntax.location)
        option;
      almost_linearity : bool;
    }

    let remove_lin_context env =
      Utils.StringMap.fold
        (fun k ((_, _, abs, _) as v) (l_acc, nl_acc) ->
          match abs with
          | Abstract_syntax.Linear -> (Utils.StringMap.add k v l_acc, nl_acc)
          | Abstract_syntax.Non_linear -> (l_acc, Utils.StringMap.add k v nl_acc))
        env
        (Utils.StringMap.empty, Utils.StringMap.empty)

    let insert_lin_var (ty : Lambda.stype) (env : typing_environment) =
      ( Eval
          (fun l_level _ loc ->
            (* let i = (l_level - 1 - env.linear_level) in
               let () = Printf.printf "Inserting variable %d - 1 - %d = %d\n%!" l_level env.linear_level i in *)
            ( Lambda.LVar (l_level - 1 - env.linear_level),
              Delay (fun _ _ l -> raise (Not_linear (l, loc))) )),
        ty,
        Abstract_syntax.Linear,
        0 )

    let insert_non_lin_var (ty : Lambda.stype) (env : typing_environment) =
      ( Delay (fun _ level _ -> Lambda.Var (level - 1 - env.level)),
        ty,
        Abstract_syntax.Non_linear,
        0 )

    let compute (k : typing_env_content) (env : typing_environment)
        (l : Abstract_syntax.location) =
      match k with
      | Delay f -> (f env.linear_level env.level l, Delay f)
      | Eval f -> f env.linear_level env.level l

    let get_binding x map =
      try Some (Utils.StringMap.find x map) with Not_found -> None

    let replace_binding x v map =
      match v with
      | None -> Utils.StringMap.remove x map
      | Some b -> Utils.StringMap.add x b map

    let typecheck (t : Abstract_syntax.term) (ty : Lambda.stype)
        (sg : Signature.t) : Lambda.term * bool =
      let local_expand ty = Signature.expand_type ty sg in
      let rec typecheck_aux t ty (tenv : typing_environment) =
        match t with
        | Abstract_syntax.Var (x, l) -> (
            try
              let f_var, var_type, lin, nb_occ =
                Utils.StringMap.find x tenv.env
              in
              let var, new_f = compute f_var tenv l in
              match ty with
              | None ->
                  ( var,
                    var_type,
                    {
                      tenv with
                      env =
                        Utils.StringMap.add x
                          (new_f, var_type, lin, nb_occ + 1)
                          tenv.env;
                    } )
              | Some l_ty when l_ty = local_expand var_type ->
                  ( var,
                    var_type,
                    {
                      tenv with
                      env =
                        Utils.StringMap.add x
                          (new_f, var_type, lin, nb_occ + 1)
                          tenv.env;
                    } )
              | Some l_ty -> raise (Type_mismatch (l, l_ty, var_type))
            with Not_found -> raise (Non_empty_linear_context (x, l)))
        | Abstract_syntax.Const (x, l) -> (
            try
              let l_term, l_type = Signature.find_term x sg in
              match ty with
              | None -> (l_term, l_type, tenv)
              | Some l_ty when local_expand l_type = l_ty ->
                  (l_term, l_type, tenv)
              | Some l_ty -> raise (Type_mismatch (l, l_ty, l_type))
            with Signature.Not_found -> failwith "Bug")
        | Abstract_syntax.LAbs (x, l_x, u, l_u) -> (
            match ty with
            | None -> raise Not_normal_term
            | Some l_ty -> (
                let b = get_binding x tenv.env in
                try
                  let ty1, ty2, lin = decompose_functional_type l_ty sg in
                  match lin with
                  | Abstract_syntax.Linear -> (
                      let u_term, _, new_typing_env =
                        typecheck_aux u
                          (Some (local_expand ty2))
                          {
                            tenv with
                            linear_level = tenv.linear_level + 1;
                            env =
                              Utils.StringMap.add x (insert_lin_var ty1 tenv)
                                tenv.env;
                          }
                      in
                      let f_var, _, _, nb_occ =
                        Utils.StringMap.find x new_typing_env.env
                      in
                      try
                        let _ = compute f_var new_typing_env l_x in
                        (* if the Not_linear exception is not raised, it means the variable
                           was not
                           used in
                           u_term *)
                        let () = assert (nb_occ = 0) in
                        raise (Vacuous_abstraction (x, l_x, get_location u))
                      with Not_linear _ ->
                        let () = assert (nb_occ = 1) in
                        ( Lambda.LAbs (x, u_term),
                          l_ty,
                          {
                            new_typing_env with
                            env = replace_binding x b new_typing_env.env;
                            linear_level = tenv.linear_level;
                          } ))
                  | Abstract_syntax.Non_linear as l -> raise (Functional_type l)
                with
                | Not_functional_type
                | Functional_type Abstract_syntax.Non_linear ->
                  Errors.(TypeErrors.emit (Type_l.IsUsed
                   (Format.asprintf "%a" (Signature.pp_type sg) l_ty, "\"'a → 'b\" (linear abstraction)")) ~loc:l_u)))
        | Abstract_syntax.Abs (x, _, u, l_u) -> (
            match ty with
            | None -> raise Not_normal_term
            | Some l_ty -> (
                let b = get_binding x tenv.env in
                try
                  let ty1, ty2, lin = decompose_functional_type l_ty sg in
                  match lin with
                  | Abstract_syntax.Non_linear ->
                      let u_term, _, new_typing_env =
                        typecheck_aux u
                          (Some (local_expand ty2))
                          {
                            tenv with
                            level = tenv.level + 1;
                            env =
                              Utils.StringMap.add x
                                (insert_non_lin_var ty1 tenv)
                                tenv.env;
                          }
                      in
                      let () =
                        Log.debug (fun m ->
                            m "Binding and occurrences:@,@[<v>  @[%a@]@]"
                              (fun fmt env ->
                                 Utils.StringMap.iter
                                   (fun var_name (_, ty, _, nb) ->
                                      Format.fprintf fmt "%s: \"%a\", %d" var_name
                                        (Signature.pp_type sg)
                                        ty
                                        nb)
                                   env)
                              new_typing_env.env)
                      in
                      let is_almost_linear =
                        match get_binding x new_typing_env.env with
                        | Some (_, var_type, _, nb_occ)
                          when nb_occ = 0
                               || nb_occ > 1
                                  && not (Lambda.is_atomic var_type (fun i -> Signature.unfold_type_definition i sg)) ->
                            let () =
                              Log.info (fun m ->
                                  m "var \"%s\", nb occ: %d, atomic type: %b" x
                                    nb_occ
                                    (Lambda.is_atomic var_type (fun i -> Signature.unfold_type_definition i sg)))
                            in
                            false
                        | None ->
                            failwith "Bug: the variable should be given a type"
                        | _ -> true
                      in
                      ( Lambda.Abs (x, u_term),
                        l_ty,
                        {
                          new_typing_env with
                          env = replace_binding x b new_typing_env.env;
                          level = tenv.level;
                          almost_linearity =
                            new_typing_env.almost_linearity && is_almost_linear;
                        } )
                  | Abstract_syntax.Linear as l -> raise (Functional_type l)
                with Not_functional_type | Functional_type _ ->
                  Errors.(
                    TypeErrors.emit (Type_l.IsUsed
                   (Format.asprintf "%a" (Signature.pp_type sg) l_ty, "\"'a ⇒ 'b\" (non-linear abstraction)")) ~loc:l_u)))
        | Abstract_syntax.App (u, v, l) -> (
            let u_term, u_type, new_typing_env =
              try typecheck_aux u None tenv
              with Not_normal_term ->
                Errors.(TypeErrors.emit Type_l.NotNormal ~loc:l)
            in
            let ty1, ty2, lin =
              try decompose_functional_type u_type sg
              with Not_functional_type ->
                let u_loc = get_location u in
                Errors.(
                  TypeErrors.emit
                    (Type_l.IsUsed
                      (Format.asprintf "%a" (Signature.pp_type sg) u_type, "\"'a -> 'b\" or \"'a => 'b\" in order to enable application"))
                    ~loc:u_loc)
            in
            let v_term, _, new_new_typing_env =
              match lin with
              | Abstract_syntax.Linear ->
                  typecheck_aux v (Some (local_expand ty1)) new_typing_env
              | Abstract_syntax.Non_linear -> (
                  let lin_env, non_lin_env =
                    remove_lin_context new_typing_env.env
                  in
                  let wrapper = (get_location u, u_type, get_location v) in
                  try
                    let v_t, v_ty, new_env (*{ wrapper = w; env = l_env; _ }*) =
                      typecheck_aux v
                        (Some (local_expand ty1))
                        {
                          new_typing_env with
                          env = non_lin_env;
                          wrapper = Some wrapper;
                        }
                    in
                    ( v_t,
                      v_ty,
                      {
                        new_env with
                        env =
                          Utils.StringMap.union
                            (fun _ _ _ -> failwith "Bug")
                            lin_env new_env.env;
                      } )
                  with Non_empty_linear_context (x, l) ->
                    let func_loc, func_st, v_loc =
                      match tenv.wrapper with None -> wrapper | Some a -> a
                    in
                    (*        let v_loc = get_location v in*)
                    Errors.(
                      TypeErrors.emit
                        (Type_l.NonEmptyContext (x, l, func_loc,
                                                 (Format.asprintf "%a" (Signature.pp_type sg) func_st)))
                      ~loc:v_loc))
            in
            match ty with
            | None -> (Lambda.App (u_term, v_term), ty2, new_new_typing_env)
            | Some l_ty when l_ty = local_expand ty2 ->
                (Lambda.App (u_term, v_term), l_ty, new_new_typing_env)
            | Some l_ty -> raise (Type_mismatch (l, l_ty, ty2)))
      in
      try
        let t_term, t_type, ({ almost_linearity = b; _ } : typing_environment) =
          typecheck_aux t
            (Some (local_expand ty))
            {
              linear_level = 0;
              level = 0;
              env = Utils.StringMap.empty;
              wrapper = None;
              almost_linearity = true;
            }
        in
        Log.debug (fun m ->
            m "@[Type-checked: @[%a:@[<2>@ %a@]@]@]" (Signature.pp_term sg)
              t_term (Signature.pp_type sg) t_type);
        (t_term, b)
      with
      | Type_mismatch (loc, t1, t2) ->
        Errors.(
          TypeErrors.emit
            (Type_l.IsUsed (Format.asprintf "%a" (Signature.pp_type sg) t1,
              Format.asprintf "\"%a\"" (Signature.pp_type sg) t2))
            ~loc)
      | Not_linear (loc1, loc2) ->
          Errors.(TypeErrors.emit (Type_l.TwoOccurrencesOfLinearVariable loc1) ~loc:loc2)
      | Vacuous_abstraction (x, l1, l2) ->
          Errors.(TypeErrors.emit (Type_l.VacuousAbstraction (x, l2)) ~loc:l1)
  end
end