(* ------------------------------------------------------------------------ *)
(* Adapted from Celsius project : https://github.com/clementblaudeau/celsius *)

Require Export Syntax.
From Stdlib Require Import List.
Import ListNotations.
Require Export Notations LibTactics Tactics.
Require Export ssreflect ssrbool Stdlib.Sets.Finite_sets_facts.

(* ------------------------------------------------------------------------ *)

Definition gget {X: Type} (l : list X)  : Loc -> option X := nth_error l.
Definition runtime_getObj (l : heap)    : Loc -> option Obj := nth_error l.
Definition getVal (l : list value)  : Loc -> option value := nth_error l.
Definition runtime_getVal (rΓ : r_env): Loc -> option value := nth_error (rΓ.(vars)).
Definition static_getType (sΓ: s_env): Loc -> option qualified_type := nth_error sΓ.

Fixpoint mapM {A B : Type} (f : A -> option B) (l : list A) : option (list B) :=
  match l with
  | [] => Some []
  | x :: xs =>
      match f x, mapM f xs with
      | Some y, Some ys => Some (y :: ys)
      | _, _ => None
      end
  end.

Definition static_getType_list (sΓ: s_env): list Loc -> option (list qualified_type) :=
  fun l => mapM (fun x => static_getType sΓ x) l.

Lemma static_getType_list_preserves_length : forall sΓ args argtypes,
  static_getType_list sΓ args = Some argtypes ->
  dom argtypes = dom args.
Proof.
  intros sΓ args argtypes Hstatic.
  unfold static_getType_list in Hstatic.
  generalize dependent argtypes.
  induction args as [|a args' IH]; intros argtypes Hstatic.
  - (* Base case: args = [] *)
    simpl in Hstatic.
    injection Hstatic as Hstatic.
    subst argtypes.
    reflexivity.
  - (* Inductive case: args = a :: args' *)
    simpl in Hstatic.
    destruct (static_getType sΓ a) as [T|] eqn:HstaticT; [|discriminate].
    destruct (mapM (static_getType sΓ) args') as [Ts|] eqn:HstaticTs; [|discriminate].
    injection Hstatic as Hstatic.
    subst argtypes.
    simpl.
    f_equal.
    apply IH.
    reflexivity.
Qed.

Definition runtime_lookup_list (rΓ: r_env): list Loc -> option (list value) :=
  fun l => mapM (fun x => runtime_getVal rΓ x) l.

(* ------------------------------------------------------------------------ *)

Local Hint Unfold runtime_getObj runtime_getVal static_getType: updates.
(* ------------------------------------------------------------------------ *)

Fixpoint forallb2 {A B} (f : A -> B -> bool) (l1 : list A) (l2 : list B) : bool :=
  match l1, l2 with
  | [], [] => true
  | x1 :: xs1, x2 :: xs2 => f x1 x2 && forallb2 f xs1 xs2
  | _, _ => false
  end.

Fixpoint update {X : Type} (position : nat) (value : X) (l : list X) : list X :=
  match l, position with
  | [], _ => []
  | _::t, 0 => value :: t
  | h::l', S n => h :: update n value l'
  end.
Notation "[ x ↦  v ] l" := (update x v l) (at level 0).

Definition update_r_env_value (rΓ : r_env) (l : Loc) (v : value) : r_env :=
  match rΓ with
  {| vars := vars; |} =>
      {| vars := update l v vars;|}
  end.

Lemma update_same :
  forall X p v (l: list X),
    p < (length l) ->
    (nth_error [p ↦ v]l p) = Some v.
Proof.
  intros X p v; generalize dependent p.
  induction p; steps; eauto with lia.
Qed.
Global Hint Resolve update_same: updates.

Lemma runtime_getObj_update_same:
  forall h l O, l < dom h -> runtime_getObj ([l ↦ O]h) l = Some O.
Proof.
  eauto using update_same.
Qed.

Lemma runtime_getVal_update_same :
  forall rΓ l v
         (Hl : l < dom rΓ.(vars)),
    runtime_getVal (update_r_env_value rΓ l v) l = Some v.
Proof.
  intros rΓ l v Hl.
  unfold runtime_getVal, update_r_env_value.
  destruct rΓ as [vars]; simpl in *.
  (* Goal now is nth_error ([l ↦ v] vars) l = Some v *)
  apply update_same.
  exact Hl.
Qed.

Lemma update_diff :
  forall X p p' v (l: list X),
    p <> p' ->
    (nth_error [p ↦ v]l p') = (nth_error l p').
Proof.
  induction p; intros; destruct l ; destruct p' => //.
  simpl; eauto.
Qed.
Global Hint Resolve update_diff: updates.

Lemma runtime_getObj_update_diff:
  forall h l l' O,
    l <> l' ->
    runtime_getObj ([l ↦ O]h) l' = runtime_getObj h l'.
Proof.
  eauto using update_diff.
Qed.

Lemma runtime_getVal_update_diff:
  forall rΓ l l' v,
    l <> l' ->
    runtime_getVal (update_r_env_value rΓ l v) l' = runtime_getVal rΓ l'.
Proof.
  intros rΓ l l' v Hneq.
  unfold runtime_getVal, update_r_env_value.
  destruct rΓ as [vars]; simpl.
  (* nth_error ([l ↦ v] vars) l' = nth_error vars l' *)
  eapply update_diff; exact Hneq.
Qed.

Lemma update_length :
  forall X p v (l: list X),
    length ([p ↦ v]l) = length l.
Proof.
  intros X.
  induction p; steps.
Qed.
Global Hint Resolve update_length: updates.

Lemma gget_dom :
  forall {X : Type} (l : list X) (C : Loc) x,
    gget l C = Some x -> C < length l.
Proof.
  unfold gget.
  intros. eapply nth_error_Some; eauto.
Qed.
Global Hint Resolve gget_dom: updates.

Lemma gget_not_dom :
  forall {X : Type} (l : list X) (C : Loc),
    gget l C = None -> C >= length l.
Proof.
  unfold gget.
  intros. eapply nth_error_None; eauto.
Qed.
Global Hint Resolve gget_not_dom: updates.

Lemma runtime_getObj_dom:
  forall l O h, runtime_getObj h l = Some O -> l < dom h.
Proof.
  unfold runtime_getObj.
  intros; eapply nth_error_Some; eauto.
Qed.
Global Hint Resolve runtime_getObj_dom: updates.

Lemma runtime_getObj_not_dom:
  forall l h, runtime_getObj h l = None -> l >= dom h.
Proof.
  unfold runtime_getObj.
  intros; eapply nth_error_None; eauto.
Qed.
Global Hint Resolve runtime_getObj_dom: updates.

Lemma getVal_dom:
  forall l v ρ, getVal ρ l = Some v -> l < dom ρ.
Proof.
  unfold getVal.
  intros; eapply nth_error_Some; eauto.
Qed.
Global Hint Resolve getVal_dom: updates.

Lemma getVal_not_dom:
  forall l ρ, getVal ρ l = None -> l >= dom ρ.
Proof.
  unfold getVal.
  intros; eapply nth_error_None; eauto.
Qed.
Global Hint Resolve getVal_not_dom: updates.

Lemma runtime_getVal_dom:
  forall l v ρ, runtime_getVal ρ l = Some v -> l < dom ρ.(vars).
Proof.
  unfold runtime_getVal.
  intros; eapply nth_error_Some; eauto.
Qed.
Global Hint Resolve runtime_getVal_dom: updates.

Lemma runtime_getVal_not_dom:
  forall l ρ, runtime_getVal ρ l = None -> l >= dom ρ.(vars).
Proof.
  unfold runtime_getVal.
  intros; eapply nth_error_None; eauto.
Qed.
Global Hint Resolve runtime_getVal_not_dom: updates.

Lemma static_getType_dom:
  forall l O sΓ, static_getType sΓ l = Some O -> l < dom sΓ.
Proof.
  unfold static_getType.
  intros; eapply nth_error_Some; eauto.
Qed.
Global Hint Resolve static_getType_dom: updates.

Lemma static_getType_not_dom:
  forall l sΓ, static_getType sΓ l = None -> l >= dom sΓ.
Proof.
  unfold static_getType.
  intros; eapply nth_error_None; eauto.
Qed.
Global Hint Resolve static_getType_not_dom: updates.

(* ------------------------------------------------------------------------ *)

Lemma runtime_getObj_Some : forall σ l,
    l < dom σ ->
    exists C ω, runtime_getObj σ l = Some (mkObj C ω).
Proof.
  intros.
  destruct (runtime_getObj σ l) as [[C ω]|] eqn:?.
  exists C, ω; auto.
  exfalso. eapply nth_error_None in Heqo. lia.
Qed.


(* ------------------------------------------------------------------------ *)

(* ------------------------------------------------------------------------ *)

Lemma runtime_getObj_last :
  forall σ C ρ,
    runtime_getObj (σ++[mkObj C ρ]) (dom σ) = Some (mkObj C ρ).
Proof.
  induction σ; steps.
Qed.
Global Hint Resolve runtime_getObj_last: core.

Lemma runtime_getObj_last2 :
  forall σ C ρ l,
    l < (dom σ) ->
    runtime_getObj (σ++[ mkObj C ρ ]) l = runtime_getObj σ l.
Proof.
  induction σ; simpl; intros; try lia.
  destruct l;
    steps;
    eauto with lia.
Qed.

Lemma static_getType_last :
  forall Σ T,
    static_getType (Σ++[T]) (dom Σ) = Some T.
Proof.
  induction Σ; steps.
Qed.
Global Hint Resolve static_getType_last: core.


Lemma runtime_getVal_last :
  forall (rΓ : r_env) (v : value),
    runtime_getVal (set_vars rΓ (vars rΓ ++ [v])) (dom rΓ.(vars)) = Some v.
Proof.
  intros rΓ v.
  unfold runtime_getVal. simpl.
  rewrite nth_error_app2.
  - reflexivity.
  - rewrite Nat.sub_diag. reflexivity.
Qed.
Global Hint Resolve runtime_getVal_last: core.

Lemma runtime_getVal_last2 :
  forall (rΓ : r_env) (x : Loc) (v : value)
    (Hdom : x < dom rΓ.(vars)),
    runtime_getVal (set_vars rΓ (vars rΓ ++ [v])) x = runtime_getVal rΓ x.
Proof.
  intros rΓ x v Hdom.
  unfold runtime_getVal. simpl.
  apply nth_error_app1.
  exact Hdom.
Qed.

Lemma Forall_nth_error {A} (P : A -> Prop) (l : list A) (n : nat) (x : A) :
  Forall P l ->
  nth_error l n = Some x ->
  P x.
Proof.
  revert l. induction n as [|n IH]; intros l HForall Hnth.
  - destruct l as [|y l']; simpl in *; [discriminate|].
    inversion Hnth; subst. inversion HForall; subst. assumption.
  - destruct l as [|y l']; simpl in *; [discriminate|].
    inversion HForall; subst.
    apply IH with (l := l'); assumption.
Qed.

(* ------------------------------------------------------------------------ *)

Lemma gget_Some : forall {A : Type} (l : list A) (n : nat),
  n < length l ->
  exists x, gget l n = Some x.
Proof.
  intros A l n H.
  unfold gget.
  destruct (nth_error l n) as [x|] eqn:Hnth.
  - exists x. reflexivity.
  - exfalso.
    apply nth_error_None in Hnth.
    lia.
Qed.

(* Find a class declaration in the class table *)
Definition find_class (CT : class_table) (C : class_name) : option class_def :=
    gget CT C.

Lemma find_class_dom : forall CT C x,
  find_class CT C = Some x -> C < dom CT.
Proof.
  intros. unfold find_class in H. apply gget_dom in H. exact H.
Qed.

Lemma find_class_Some : forall CT C,
  C < dom CT -> exists def, find_class CT C = Some def.
Proof.
  intros CT C H.
  unfold find_class.
  apply gget_Some.
  exact H.
Qed.

(* Class bound look up in the class table  *)
Definition bound (CT : class_table) (C : class_name) : option q_c :=
  match find_class CT C with
  | Some decl => Some (class_qualifier (signature decl))
  | None => None
  end.

Lemma Forall_update : forall {A : Type} (P : A -> Prop) (l : list A) (n : nat) (v : A),
  Forall P l ->
  P v ->
  n < length l ->
  Forall P (update n v l).
Proof.
  intros A P l n v H_forall H_v H_bound.
  generalize dependent n.
  induction l as [|h t IH]; intros n H_bound.
  - simpl in H_bound. lia.
  - destruct n as [|n'].
    + simpl. constructor; [exact H_v | inversion H_forall; exact H2].
    + simpl. constructor; [inversion H_forall; exact H1 | apply IH; [inversion H_forall; exact H2 | simpl in H_bound; lia]].
Qed.

Lemma nth_error_Some_exists : forall {A : Type} (l : list A) (i : nat),
  i < length l ->
  exists x, nth_error l i = Some x.
Proof.
  intros A l i H.
  destruct (nth_error l i) as [x|] eqn:Hnth.
  - exists x. reflexivity.
  - exfalso.
    apply nth_error_None in Hnth.
    lia.
Qed.

Lemma Forall2_nth_error : forall {A B : Type} (P : A -> B -> Prop) (l1 : list A) (l2 : list B) (i : nat) (a : A) (b : B),
  Forall2 P l1 l2 ->
  nth_error l1 i = Some a ->
  nth_error l2 i = Some b ->
  P a b.
Proof.
  intros A B P l1 l2 i a b H Ha Hb.
  generalize dependent i.
  induction H; intros i Ha Hb.
  - (* Empty lists *)
    destruct i; discriminate Ha.
  - (* Non-empty lists *)
    destruct i as [|i'].
    + (* i = 0 *)
      simpl in Ha, Hb.
      injection Ha as Ha_eq. injection Hb as Hb_eq.
      subst. exact H.
    + (* i = S i' *)
      simpl in Ha, Hb.
      eapply IHForall2; eauto.
Qed.

Lemma Forall2_update : forall A B (P : A -> B -> Prop) (la : list A) (lb : list B) f a,
  Forall2 P la lb ->
  f < length lb ->
  (forall b, nth_error lb f = Some b -> P a b) ->
  Forall2 P (update f a la) lb.
Proof.
  intros A B P la lb f a H_forall2 H_bound H_prop.
  generalize dependent f.
  induction H_forall2; intros f H_bound H_prop.
  - (* Empty lists *)
    simpl in H_bound. lia.
  - (* Non-empty lists *)
    destruct f as [|f'].
    + (* f = 0 *)
      simpl. constructor.
      * apply H_prop. simpl. reflexivity.
      * exact H_forall2.
    + (* f = S f' *)
      simpl. constructor.
      * exact H.
      * apply IHH_forall2.
        -- simpl in H_bound. lia.
        -- intros b Hnth. apply H_prop. simpl. exact Hnth.
Qed.