Library MetaRocq.Erasure.EUnboxing

From Stdlib Require Import Utf8 Program.
From MetaRocq.Utils Require Import utils.
From MetaRocq.Common Require Import config Kernames Primitive BasicAst EnvMap.
From MetaRocq.Erasure Require Import EPrimitive EAst EAstUtils EInduction EArities
    ELiftSubst ESpineView EGlobalEnv EWellformed EEnvMap
    EWcbvEval EEtaExpanded ECSubst EWcbvEvalEtaInd EProgram.

Local Open Scope string_scope.
Set Asymmetric Patterns.
Import MRMonadNotation.

From Equations Require Import Equations.
Set Equations Transparent.
Local Set Keyed Unification.
From Stdlib Require Import ssreflect ssrbool.

We assume Prop </= Type and universes are checked correctly in the following.
Local Existing Instance extraction_checker_flags.

Ltac introdep := let H := fresh in intros H; depelim H.

#[global]
Hint Constructors eval : core.

Import MRList (map_InP, map_InP_spec).

Equations safe_hd {A} (l : list A) (nnil : l nil) : A :=
| [] | nnil := False_rect _ (nnil eq_refl)
| hd :: _ | _nnil := hd.

Definition inspect {A : Type} (x : A) : {y : A | x = y} :=
  @exist _ (fun yx = y) x eq_refl.

Section unbox.
  Context (Σ : GlobalContextMap.t).

  Section Def.

  Definition unboxable (idecl : one_inductive_body) (cdecl : constructor_body) :=
    (#|idecl.(ind_ctors)| == 1) && (cdecl.(cstr_nargs) == 1).

  Equations is_unboxable (kn : inductive) (c : nat) : bool :=
    | kn | 0 with GlobalContextMap.lookup_constructor Σ kn 0 := {
      | Some (mdecl, idecl, cdecl) := unboxable idecl cdecl
      | None := false }
    | kn | S _ := false.

  Notation " t 'eqn:' h " := (exist t h) (only parsing, at level 12).

  Opaque is_unboxable.

  Equations get_unboxable_case_branch (ind : inductive) (brs : list (list name × term)) : option (name × term) :=
  get_unboxable_case_branch ind [([brna], bbody)] := Some (brna, bbody);
  get_unboxable_case_branch ind _ := None.

  Lemma get_unboxable_case_branch_spec {ind : inductive} {brs : list (list name × term)} {brna bbody} :
    get_unboxable_case_branch ind brs = Some (brna, bbody)
    brs = [([brna], bbody)].
  Proof.
    funelim (get_unboxable_case_branch ind brs) ⇒ //.
    now intros [= <- <-].
  Qed.

  Equations unbox (t : term) : term :=
    | tRel iEAst.tRel i
    | tEvar ev argsEAst.tEvar ev (map unbox args)
    | tLambda na MEAst.tLambda na (unbox M)
    | tApp u vEAst.tApp (unbox u) (unbox v)
    | tLetIn na b b'EAst.tLetIn na (unbox b) (unbox b')
    | tCase ind c brs with inspect (is_unboxable ind.1 0) ⇒
      { | true eqn:unb with inspect (get_unboxable_case_branch ind.1 (map (on_snd unbox) brs)) := {
          | Some (brna, bbody) eqn:heqbrEAst.tLetIn brna (unbox c) bbody
          | None eqn:heqbr := EAst.tCase ind (unbox c) (map (on_snd unbox) brs) }
        | false eqn:unb := EAst.tCase ind (unbox c) (map (on_snd unbox) brs) }
    | tProj p c with inspect (is_unboxable p.(proj_ind) 0) := {
        | true eqn:unbunbox c
        | false eqn:unbEAst.tProj p (unbox c) }
    | tConstruct ind n args with inspect (is_unboxable ind n) ⇒ {
       unbox (tConstruct ind n [a]) (true eqn:unb) ⇒ unbox a ;
       unbox (tConstruct ind n args) isunbtConstruct ind n (map unbox args) }
    | tFix mfix idx
      let mfix' := map (map_def unbox) mfix in
      EAst.tFix mfix' idx
    | tCoFix mfix idx
      let mfix' := map (map_def unbox) mfix in
      EAst.tCoFix mfix' idx
    | tBoxEAst.tBox
    | tVar nEAst.tVar n
    | tConst nEAst.tConst n
    | tPrim pEAst.tPrim (map_prim unbox p)
    | tLazy tEAst.tLazy (unbox t)
    | tForce tEAst.tForce (unbox t).

  End Def.

  Lemma map_repeat {A B} (f : A B) x n : map f (repeat x n) = repeat (f x) n.
  Proof using Type.
    now induction n; simpl; auto; rewrite IHn.
  Qed.

  Lemma map_unbox_repeat_box n : map unbox (repeat tBox n) = repeat tBox n.
  Proof using Type. now rewrite map_repeat. Qed.

  Arguments eqb : simpl never.


End unbox.

#[universes(polymorphic)] Global Hint Rewrite @map_primIn_spec @map_InP_spec : unbox.
Tactic Notation "simp_eta" "in" hyp(H) := simp isEtaExp in H; rewrite -?isEtaExp_equation_1 in H.
Ltac simp_eta := simp isEtaExp; rewrite -?isEtaExp_equation_1.
Tactic Notation "simp_unbox" "in" hyp(H) := simp unbox in H; rewrite -?unbox_equation_1 in H.
Ltac simp_unbox := simp unbox; rewrite -?unbox_equation_1.

Definition unbox_constant_decl Σ cb :=
  {| cst_body := option_map (unbox Σ) cb.(cst_body) |}.

Definition unbox_inductive_decl idecl :=
  {| ind_finite := idecl.(ind_finite); ind_npars := 0; ind_bodies := idecl.(ind_bodies) |}.

Definition unbox_decl Σ d :=
  match d with
  | ConstantDecl cbConstantDecl (unbox_constant_decl Σ cb)
  | InductiveDecl ideclInductiveDecl (unbox_inductive_decl idecl)
  end.

Definition unbox_env Σ :=
  map (on_snd (unbox_decl Σ)) Σ.(GlobalContextMap.global_decls).

Definition unbox_program (p : eprogram_env) : eprogram :=
  (unbox_env p.1, unbox p.1 p.2).

Import EGlobalEnv.

Lemma lookup_env_unbox Σ kn :
  lookup_env (unbox_env Σ) kn =
  option_map (unbox_decl Σ) (lookup_env Σ.(GlobalContextMap.global_decls) kn).
Proof.
  unfold unbox_env.
  destruct Σ as [Σ map repr wf]; cbn.
  induction Σ at 2 4; simpl; auto.
  case: eqb_spec ⇒ //.
Qed.

Lemma lookup_constructor_unbox {Σ kn c} :
  lookup_constructor (unbox_env Σ) kn c =
  match lookup_constructor Σ.(GlobalContextMap.global_decls) kn c with
  | Some (mdecl, idecl, cdecl)Some (unbox_inductive_decl mdecl, idecl, cdecl)
  | NoneNone
  end.
Proof.
  rewrite /lookup_constructor /lookup_inductive /lookup_minductive.
  rewrite lookup_env_unbox.
  destruct lookup_env as [[]|] ⇒ /= //.
  do 2 destruct nth_error ⇒ //.
Qed.

Lemma lookup_constructor_pars_args_unbox (Σ : GlobalContextMap.t) i n npars nargs :
  EGlobalEnv.lookup_constructor_pars_args Σ i n = Some (npars, nargs)
  EGlobalEnv.lookup_constructor_pars_args (unbox_env Σ) i n = Some (0, nargs).
Proof.
  rewrite /EGlobalEnv.lookup_constructor_pars_args. rewrite lookup_constructor_unbox //=.
  destruct EGlobalEnv.lookup_constructor ⇒ //. destruct p as [[] ?] ⇒ //=. now intros [= <- <-].
Qed.

Lemma is_propositional_unbox (Σ : GlobalContextMap.t) ind :
  match inductive_isprop_and_pars Σ.(GlobalContextMap.global_decls) ind with
  | Some (prop, npars)
    inductive_isprop_and_pars (unbox_env Σ) ind = Some (prop, 0)
  | None
    inductive_isprop_and_pars (unbox_env Σ) ind = None
  end.
Proof.
  rewrite /inductive_isprop_and_pars /lookup_inductive /lookup_minductive.
  rewrite lookup_env_unbox.
  destruct lookup_env; simpl; auto.
  destruct g; simpl; auto. destruct nth_error ⇒ //.
Qed.

Lemma is_propositional_cstr_unbox {Σ : GlobalContextMap.t} {ind c} :
  match constructor_isprop_pars_decl Σ.(GlobalContextMap.global_decls) ind c with
  | Some (prop, npars, cdecl)
    constructor_isprop_pars_decl (unbox_env Σ) ind c = Some (prop, 0, cdecl)
  | None
  constructor_isprop_pars_decl (unbox_env Σ) ind c = None
  end.
Proof.
  rewrite /constructor_isprop_pars_decl /lookup_constructor /lookup_inductive /lookup_minductive.
  rewrite lookup_env_unbox.
  destruct lookup_env; simpl; auto.
  destruct g; simpl; auto. do 2 destruct nth_error ⇒ //.
Qed.

Lemma lookup_inductive_pars_unbox {efl : EEnvFlags} {Σ : GlobalContextMap.t} ind :
  wf_glob Σ
   pars, lookup_inductive_pars Σ ind = Some pars
  EGlobalEnv.lookup_inductive_pars (unbox_env Σ) ind = Some 0.
Proof.
  rewrite /lookup_inductive_parswf pars.
  rewrite /lookup_inductive /lookup_minductive.
  rewrite (lookup_env_unbox _ ind).
  destruct lookup_env as [[decl|]|] ⇒ //=.
Qed.

Arguments eval {wfl}.

Arguments isEtaExp : simpl never.

Lemma isEtaExp_mkApps {Σ} {f u} : isEtaExp Σ (tApp f u)
  let (hd, args) := decompose_app (tApp f u) in
  match construct_viewc hd with
  | view_construct kn c block_args
    args [] f = mkApps hd (remove_last args) u = last args u
    isEtaExp_app Σ kn c #|args| && forallb (isEtaExp Σ) args && is_nil block_args
  | view_other _ discr
    [&& isEtaExp Σ hd, forallb (isEtaExp Σ) args, isEtaExp Σ f & isEtaExp Σ u]
  end.
Proof.
  move/(isEtaExp_mkApps Σ f [u]).
  cbn -[decompose_app]. destruct decompose_app eqn:da.
  destruct construct_viewc eqn:cv ⇒ //.
  intros →.
  pose proof (decompose_app_inv da).
  pose proof (decompose_app_notApp _ _ _ da).
  destruct l using rev_case. cbn. intuition auto. solve_discr. noconf H.
  rewrite mkApps_app in H. noconf H.
  rewrite remove_last_app last_last. intuition auto.
  destruct l; cbn in *; congruence.
  pose proof (decompose_app_inv da).
  pose proof (decompose_app_notApp _ _ _ da).
  destruct l using rev_case. cbn. intuition auto. destruct t ⇒ //.
  rewrite mkApps_app in H. noconf H.
  move⇒ /andP[] etat. rewrite forallb_app ⇒ /andP[] etal /=.
  rewrite andb_true_retaa. rewrite etaa andb_true_r.
  rewrite etat etal. cbn. rewrite andb_true_r.
  eapply isEtaExp_mkApps_intro; auto; solve_all.
Qed.

Lemma decompose_app_tApp_split f a hd args :
  decompose_app (tApp f a) = (hd, args) f = mkApps hd (remove_last args) a = last args a.
Proof.
  unfold decompose_app. cbn.
  move/decompose_app_rec_inv' ⇒ [n [napp [hskip heq]]].
  rewrite -(firstn_skipn n args).
  rewrite -hskip. rewrite last_last; split ⇒ //.
  rewrite heq. f_equal.
  now rewrite remove_last_app.
Qed.