(*  Title:      Pure/envir.ML

    Author:     Lawrence C Paulson, Cambridge University Computer Laboratory

Free-form environments.  The type of a term variable / sort of a type variable is

part of its name.  The lookup function must apply type substitutions,

since they may change the identity of a variable.

*)

signature ENVIR =

sig

  type tenv = (typ * term) Vartab.table

  datatype env = Envir of {maxidx: int, tenv: tenv, tyenv: Type.tyenv}

  val maxidx_of: env -> int

  val term_env: env -> tenv

  val type_env: env -> Type.tyenv

  val is_empty: env -> bool

  val empty: int -> env

  val merge: env * env -> env

  val insert_sorts: env -> sort list -> sort list

  val genvars: string -> env * typ list -> env * term list

  val genvar: string -> env * typ -> env * term

  val lookup: env * (indexname * typ) -> term option

  val lookup': tenv * (indexname * typ) -> term option

  val update: ((indexname * typ) * term) * env -> env

  val above: env -> int -> bool

  val vupdate: ((indexname * typ) * term) * env -> env

  val norm_type_same: Type.tyenv -> typ Same.operation

  val norm_types_same: Type.tyenv -> typ list Same.operation

  val norm_type: Type.tyenv -> typ -> typ

  val norm_term_same: env -> term Same.operation

  val norm_term: env -> term -> term

  val beta_norm: term -> term

  val head_norm: env -> term -> term

  val eta_contract: term -> term

  val beta_eta_contract: term -> term

  val fastype: env -> typ list -> term -> typ

  val subst_type_same: Type.tyenv -> typ Same.operation

  val subst_term_types_same: Type.tyenv -> term Same.operation

  val subst_term_same: Type.tyenv * tenv -> term Same.operation

  val subst_type: Type.tyenv -> typ -> typ

  val subst_term_types: Type.tyenv -> term -> term

  val subst_term: Type.tyenv * tenv -> term -> term

  val expand_atom: typ -> typ * term -> term

  val expand_term: (term -> (typ * term) option) -> term -> term

  val expand_term_frees: ((string * typ) * term) list -> term -> term

end;

structure Envir: ENVIR =

struct

(** datatype env **)

(*Updating can destroy environment in 2 ways!

   (1) variables out of range

   (2) circular assignments

*)

type tenv = (typ * term) Vartab.table;

datatype env = Envir of

 {maxidx: int,          (*upper bound of maximum index of vars*)

  tenv: tenv,           (*assignments to Vars*)

  tyenv: Type.tyenv};   (*assignments to TVars*)

fun make_env (maxidx, tenv, tyenv) =

  Envir {maxidx = maxidx, tenv = tenv, tyenv = tyenv};

fun maxidx_of (Envir {maxidx, ...}) = maxidx;

fun term_env (Envir {tenv, ...}) = tenv;

fun type_env (Envir {tyenv, ...}) = tyenv;

fun is_empty env =

  Vartab.is_empty (term_env env) andalso

  Vartab.is_empty (type_env env);

(* build env *)

fun empty maxidx = make_env (maxidx, Vartab.empty, Vartab.empty);

fun merge

   (Envir {maxidx = maxidx1, tenv = tenv1, tyenv = tyenv1},

    Envir {maxidx = maxidx2, tenv = tenv2, tyenv = tyenv2}) =

  make_env (Int.max (maxidx1, maxidx2),

    Vartab.merge (op =) (tenv1, tenv2),

    Vartab.merge (op =) (tyenv1, tyenv2));

(*NB: type unification may invent new sorts*)  (* FIXME tenv!? *)

val insert_sorts = Vartab.fold (fn (_, (_, T)) => Sorts.insert_typ T) o type_env;

(*Generate a list of distinct variables.

  Increments index to make them distinct from ALL present variables. *)

fun genvars name (Envir {maxidx, tenv, tyenv}, Ts) : env * term list =

  let

    fun genvs (_, [] : typ list) : term list = []

      | genvs (n, [T]) = [Var ((name, maxidx + 1), T)]

      | genvs (n, T :: Ts) =

          Var ((name ^ radixstring (26, "a" , n), maxidx + 1), T)

            :: genvs (n + 1, Ts);

  in (Envir {maxidx = maxidx + 1, tenv = tenv, tyenv = tyenv}, genvs (0, Ts)) end;

(*Generate a variable.*)

fun genvar name (env, T) : env * term =

  let val (env', [v]) = genvars name (env, [T])

  in (env', v) end;

fun var_clash xi T T' =

  raise TYPE ("Variable " ^ quote (Term.string_of_vname xi) ^ " has two distinct types",

    [T', T], []);

fun lookup_check check tenv (xi, T) =

  (case Vartab.lookup tenv xi of

    NONE => NONE

  | SOME (U, t) => if check (T, U) then SOME t else var_clash xi T U);

(* When dealing with environments produced by matching instead *)

(* of unification, there is no need to chase assigned TVars.   *)

(* In this case, we can simply ignore the type substitution    *)

(* and use = instead of eq_type.                               *)

fun lookup' (tenv, p) = lookup_check (op =) tenv p;

fun lookup2 (tyenv, tenv) =

  lookup_check (Type.eq_type tyenv) tenv;

fun lookup (Envir {tenv, tyenv, ...}, p) = lookup2 (tyenv, tenv) p;

fun update (((xi, T), t), Envir {maxidx, tenv, tyenv}) =

  Envir {maxidx = maxidx, tenv = Vartab.update_new (xi, (T, t)) tenv, tyenv = tyenv};

(*Determine if the least index updated exceeds lim*)

fun above (Envir {tenv, tyenv, ...}) lim =

  (case Vartab.min_key tenv of SOME (_, i) => i > lim | NONE => true) andalso

  (case Vartab.min_key tyenv of SOME (_, i) => i > lim | NONE => true);

(*Update, checking Var-Var assignments: try to suppress higher indexes*)

fun vupdate ((aU as (a, U), t), env as Envir {tyenv, ...}) =

  (case t of

    Var (nT as (name', T)) =>

      if a = name' then env     (*cycle!*)

      else if Term_Ord.indexname_ord (a, name') = LESS then

        (case lookup (env, nT) of  (*if already assigned, chase*)

          NONE => update ((nT, Var (a, T)), env)

        | SOME u => vupdate ((aU, u), env))

      else update ((aU, t), env)

  | _ => update ((aU, t), env));

(** beta normalization wrt. environment **)

(*Chases variables in env.  Does not exploit sharing of variable bindings

  Does not check types, so could loop.*)

local

fun norm_type0 tyenv : typ Same.operation =

  let

    fun norm (Type (a, Ts)) = Type (a, Same.map norm Ts)

      | norm (TFree _) = raise Same.SAME

      | norm (TVar v) =

          (case Type.lookup tyenv v of

            SOME U => Same.commit norm U

          | NONE => raise Same.SAME);

  in norm end;

fun norm_term1 tenv : term Same.operation =

  let

    fun norm (Var v) =

          (case lookup' (tenv, v) of

            SOME u => Same.commit norm u

          | NONE => raise Same.SAME)

      | norm (Abs (a, T, body)) = Abs (a, T, norm body)

      | norm (Abs (_, _, body) $ t) = Same.commit norm (subst_bound (t, body))

      | norm (f $ t) =

          ((case norm f of

             Abs (_, _, body) => Same.commit norm (subst_bound (t, body))

           | nf => nf $ Same.commit norm t)

          handle Same.SAME => f $ norm t)

      | norm _ = raise Same.SAME;

  in norm end;

fun norm_term2 tenv tyenv : term Same.operation =

  let

    val normT = norm_type0 tyenv;

    fun norm (Const (a, T)) = Const (a, normT T)

      | norm (Free (a, T)) = Free (a, normT T)

      | norm (Var (xi, T)) =

          (case lookup2 (tyenv, tenv) (xi, T) of

            SOME u => Same.commit norm u

          | NONE => Var (xi, normT T))

      | norm (Abs (a, T, body)) =

          (Abs (a, normT T, Same.commit norm body)

            handle Same.SAME => Abs (a, T, norm body))

      | norm (Abs (_, _, body) $ t) = Same.commit norm (subst_bound (t, body))

      | norm (f $ t) =

          ((case norm f of

             Abs (_, _, body) => Same.commit norm (subst_bound (t, body))

           | nf => nf $ Same.commit norm t)

          handle Same.SAME => f $ norm t)

      | norm _ = raise Same.SAME;

  in norm end;

in

fun norm_type_same tyenv T =

  if Vartab.is_empty tyenv then raise Same.SAME

  else norm_type0 tyenv T;

fun norm_types_same tyenv Ts =

  if Vartab.is_empty tyenv then raise Same.SAME

  else Same.map (norm_type0 tyenv) Ts;

fun norm_type tyenv T = norm_type_same tyenv T handle Same.SAME => T;

fun norm_term_same (Envir {tenv, tyenv, ...}) =

  if Vartab.is_empty tyenv then norm_term1 tenv

  else norm_term2 tenv tyenv;

fun norm_term envir t = norm_term_same envir t handle Same.SAME => t;

fun beta_norm t = if Term.has_abs t then norm_term (empty 0) t else t;

end;

(*Put a term into head normal form for unification.*)

fun head_norm env =

  let

    fun norm (Var v) =

        (case lookup (env, v) of

          SOME u => head_norm env u

        | NONE => raise Same.SAME)

      | norm (Abs (a, T, body)) = Abs (a, T, norm body)

      | norm (Abs (_, _, body) $ t) = Same.commit norm (subst_bound (t, body))

      | norm (f $ t) =

          (case norm f of

            Abs (_, _, body) => Same.commit norm (subst_bound (t, body))

          | nf => nf $ t)

      | norm _ = raise Same.SAME;

  in Same.commit norm end;

(*Eta-contract a term (fully)*)

local

fun decr lev (Bound i) = if i >= lev then Bound (i - 1) else raise Same.SAME

  | decr lev (Abs (a, T, body)) = Abs (a, T, decr (lev + 1) body)

  | decr lev (t $ u) = (decr lev t $ decrh lev u handle Same.SAME => t $ decr lev u)

  | decr _ _ = raise Same.SAME

and decrh lev t = (decr lev t handle Same.SAME => t);

fun eta (Abs (a, T, body)) =

    ((case eta body of

        body' as (f $ Bound 0) =>

          if Term.is_dependent f then Abs (a, T, body')

          else decrh 0 f

     | body' => Abs (a, T, body')) handle Same.SAME =>

        (case body of

          f $ Bound 0 =>

            if Term.is_dependent f then raise Same.SAME

            else decrh 0 f

        | _ => raise Same.SAME))

  | eta (t $ u) = (eta t $ Same.commit eta u handle Same.SAME => t $ eta u)

  | eta _ = raise Same.SAME;

in

fun eta_contract t =

  if Term.has_abs t then Same.commit eta t else t;

val beta_eta_contract = eta_contract o beta_norm;

end;

(*finds type of term without checking that combinations are consistent

  Ts holds types of bound variables*)

fun fastype (Envir {tyenv, ...}) =

  let

    val funerr = "fastype: expected function type";

    fun fast Ts (f $ u) =

          (case Type.devar tyenv (fast Ts f) of

            Type ("fun", [_, T]) => T

          | TVar v => raise TERM (funerr, [f $ u])

          | _ => raise TERM (funerr, [f $ u]))

      | fast Ts (Const (_, T)) = T

      | fast Ts (Free (_, T)) = T

      | fast Ts (Bound i) =

          (nth Ts i handle General.Subscript => raise TERM ("fastype: Bound", [Bound i]))

      | fast Ts (Var (_, T)) = T

      | fast Ts (Abs (_, T, u)) = T --> fast (T :: Ts) u;

  in fast end;

(** plain substitution -- without variable chasing **)

local

fun subst_type0 tyenv = Term_Subst.map_atypsT_same

  (fn TVar v =>

        (case Type.lookup tyenv v of

          SOME U => U

        | NONE => raise Same.SAME)

    | _ => raise Same.SAME);

fun subst_term1 tenv = Term_Subst.map_aterms_same

  (fn Var v =>

        (case lookup' (tenv, v) of

          SOME u => u

        | NONE => raise Same.SAME)

    | _ => raise Same.SAME);

fun subst_term2 tenv tyenv : term Same.operation =

  let

    val substT = subst_type0 tyenv;

    fun subst (Const (a, T)) = Const (a, substT T)

      | subst (Free (a, T)) = Free (a, substT T)

      | subst (Var (xi, T)) =

          (case lookup' (tenv, (xi, T)) of

            SOME u => u

          | NONE => Var (xi, substT T))

      | subst (Bound _) = raise Same.SAME

      | subst (Abs (a, T, t)) =

          (Abs (a, substT T, Same.commit subst t)

            handle Same.SAME => Abs (a, T, subst t))

      | subst (t $ u) = (subst t $ Same.commit subst u handle Same.SAME => t $ subst u);

  in subst end;

in

fun subst_type_same tyenv T =

  if Vartab.is_empty tyenv then raise Same.SAME

  else subst_type0 tyenv T;

fun subst_term_types_same tyenv t =

  if Vartab.is_empty tyenv then raise Same.SAME

  else Term_Subst.map_types_same (subst_type0 tyenv) t;

fun subst_term_same (tyenv, tenv) =

  if Vartab.is_empty tenv then subst_term_types_same tyenv

  else if Vartab.is_empty tyenv then subst_term1 tenv

  else subst_term2 tenv tyenv;

fun subst_type tyenv T = subst_type_same tyenv T handle Same.SAME => T;

fun subst_term_types tyenv t = subst_term_types_same tyenv t handle Same.SAME => t;

fun subst_term envs t = subst_term_same envs t handle Same.SAME => t;

end;

(** expand defined atoms -- with local beta reduction **)

fun expand_atom T (U, u) =

  subst_term_types (Type.raw_match (U, T) Vartab.empty) u

    handle Type.TYPE_MATCH => raise TYPE ("expand_atom: ill-typed replacement", [T, U], [u]);

fun expand_term get =

  let

    fun expand tm =

      let

        val (head, args) = Term.strip_comb tm;

        val args' = map expand args;

        fun comb head' = Term.list_comb (head', args');

      in

        (case head of

          Abs (x, T, t) => comb (Abs (x, T, expand t))

        | _ =>

          (case get head of

            SOME def => Term.betapplys (expand_atom (Term.fastype_of head) def, args')

          | NONE => comb head))

      end;

  in expand end;

fun expand_term_frees defs =

  let

    val eqs = map (fn ((x, U), u) => (x, (U, u))) defs;

    val get = fn Free (x, _) => AList.lookup (op =) eqs x | _ => NONE;

  in expand_term get end;

end;