(* Open Source License *)
(* Copyright (c) 2021 Nomadic Labs. <contact@nomadic-labs.com> *)

(* Permission is hereby granted, free of charge, to any person obtaining a *)
(* copy of this software and associated documentation files (the "Software"), *)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense, *)
(* and/or sell copies of the Software, and to permit persons to whom the *)
(* Software is furnished to do so, subject to the following conditions: *)

(* The above copyright notice and this permission notice shall be included *)
(* in all copies or substantial portions of the Software. *)

(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER *)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER *)
(* DEALINGS IN THE SOFTWARE. *)

Require Import NArith.
Require Import ZArith.
Require Import Bool.
Require Import String.
Require Import List.
Require Import Sorted.
Require Import syntax.
Require Import semantics.
Require Import comparable.
Require Import error.
Require Import Lia.
Require Import util.
Require Import utils.
Require Import fa12_interface.
Require Import fa12_spec.

Module FA12SumOfBalances (C : ContractContext)
                         (FT : FA12Types)
                         (FD : FA12Data C FT)
                         (FA : FA12Axioms C FT FD).

  Module fa12Spec := FA12Spec C FT FD FA.
  Import fa12Spec.

  Import FT FD FA.
  Import self fa12Notations.
  Import semantics.

  Notation positive sto
    := (fun owner => negb (getBalance sto owner =? 0))%N.

  Lemma addBalances_setAllowance sto owner spender amount l :
    fold_right (addBalance (setAllowance sto owner spender amount)) 0%N l =
    fold_right (addBalance sto) 0%N l.
  Proof.
    cbn; induction l as [ | owner' l IHl]; cbn; auto.
    rewrite IHl, getBalance_setAllowance; auto.
  Qed.

  Lemma addBalances_setBalance sto owner amount l :
    ~ In owner l ->
    fold_right (addBalance (setBalance sto owner amount)) 0%N l =
    fold_right (addBalance sto) 0%N l.
  Proof.
    cbn; induction l as [ | owner' l IHl]; cbn; auto.
    intro H_not_in.
    rewrite IHl, getBalance_setBalance_neq; auto.
  Qed.

  Lemma balance_leq_addBalances sto owner l :
    In owner l ->
    (getBalance sto owner <= fold_right (addBalance sto) 0 l)%N.
  Proof.
    cbn; induction l as [ | owner' l Hl]; cbn.
    1: intuition.
    intros [H | H].
    - subst; lia.
    - apply Hl in H; lia.
  Qed.

  Lemma sumOfBalances_positive sto l :
    sumOfBalances sto l = sumOfBalances sto (cset_subset l (positive sto)).
  Proof.
    destruct l as [l Hl]; cbn; clear Hl.
    induction l as [ | k l IHl]; cbn; auto.
    destruct (negb _) eqn:eq; cbn.
    all: rewrite IHl; auto.
    rewrite negb_false_iff, N.eqb_eq in eq.
    rewrite eq.
    lia.
  Qed.

  Lemma sumOfBalances_equal {sto} {l1} l2 :
    cset_subset l1 (positive sto) = cset_subset l2 (positive sto) ->
    sumOfBalances sto l1 = sumOfBalances sto l2.
  Proof.
    intro H.
    rewrite (sumOfBalances_positive _ l1), (sumOfBalances_positive _ l2).
    rewrite H; auto.
  Qed.

  Lemma balance_leq_sumOfAllBalances sto owner :
    (getBalance sto owner <= sumOfAllBalances sto)%N.
  Proof.
    pose proof (getOwners_mem sto owner) as H_in.
    destruct (getOwners _) as [l Hl]; cbn.
    pose proof (balance_leq_addBalances sto owner l) as H.
    destruct (getBalanceOpt sto owner) as [balance | ] eqn:H_balance.
    2: lia.
    apply H.
    apply Is_true_eq_left in H_in.
    apply (cset_list_in_sorted_correct address) in H_in; auto.
  Qed.

  Lemma sumOfAllBalances_setAllowance sto owner spender amount :
    sumOfAllBalances (setAllowance sto owner spender amount) =
    sumOfAllBalances sto.
  Proof.
    rewrite (sumOfBalances_equal (getOwners sto)).
    - destruct getOwners as [l Hl]; cbn.
      apply addBalances_setAllowance.
    - apply (cset_extensionality address); intro owner'.
      rewrite !(cset_mem_subset address), !getOwners_mem.
      pose proof (getBalance_setAllowance sto owner' owner spender amount) as H.
      destruct getBalanceOpt as [balance | ]; cbn.
      all: destruct getBalanceOpt as [balance' | ]; cbn.
      all: subst; split; intuition.
  Qed.

  Opaque compare.

  Lemma sumOfAllBalances_setBalance sto owner amount :
    sumOfAllBalances (setBalance sto owner amount) =
    Z.abs_N (Z.of_N (sumOfAllBalances sto) - Z.of_N (getBalance sto owner) + Z.of_N amount).
  Proof.
    rewrite (sumOfBalances_equal (cset_insert address owner (getOwners sto))).
    - pose proof (getOwners_mem sto owner) as H_in; revert H_in.
      destruct getOwners as [l Hl]; cbn.
      destruct (getBalanceOpt sto owner) as [balance | ] eqn:H_balance.
      all: induction Hl as [ | owner' l H_sorted IHl H_forall]; cbn.
      1: discriminate.
      2: rewrite getBalance_setBalance_eq; lia.
      all: intro H_in.
      all: destruct compare eqn:eq; cbn.
      + rewrite compare_eq_iff in eq.
        subst.
        rewrite getBalance_setBalance_eq.
        rewrite H_balance, addBalances_setBalance.
        1: cbn; lia.
        apply cset_not_in_sorted.
        apply SSorted_cons; auto.
      + discriminate.
      + rewrite getBalance_setBalance_neq.
        2: apply compare_diff; auto.
        rewrite IHl; auto.
        apply Is_true_eq_left in H_in.
        apply (cset_list_in_sorted_correct address) in H_in; auto.
        pose proof (balance_leq_addBalances sto owner l H_in) as H.
        rewrite H_balance in H.
        destruct (getBalanceOpt sto owner') as [balance' | ]; cbn.
        all: cbn in H; lia.
      + discriminate.
      + rewrite getBalance_setBalance_eq.
        pose proof (addBalances_setBalance sto owner amount (owner'::l)) as H.
        rewrite<- (cset_list_in_sorted_correct address) in H.
        2: apply SSorted_cons; auto.
        cbn in H; rewrite eq in H.
        rewrite H.
        all: cbn; auto; lia.
      + rewrite IHl; auto.
        rewrite getBalance_setBalance_neq.
        2: apply compare_diff; auto.
        destruct (getBalanceOpt sto owner') as [balance' | ].
        all: lia.
    - apply (cset_extensionality address); intro owner'.
      destruct (compare_equal owner owner') eqn:eq.
      all: rewrite ?eqb_eq, ?eqb_neq in eq; subst.
      all: rewrite !(cset_mem_subset address).
      1: rewrite (cset_mem_insert_eq address).
      2: rewrite (cset_mem_insert_neq address); auto.
      all: rewrite !getOwners_mem.
      2: pose proof (getBalance_setBalance_neq sto owner owner' amount eq) as H.
      all: destruct getBalanceOpt as [balance | ]; cbn.
      3: destruct getBalanceOpt as [balance' | ]; cbn.
      all: subst; intuition.
  Qed.

  Theorem sumOfAllBalances_constant
          (env : @proto_env self_type)
          (p : data fa12_parameter_ty)
          (sto : data storage_ty)
          (ret_ops : data (list operation))
          (ret_sto : data storage_ty) :
    main env p sto ret_ops ret_sto ->
    sumOfAllBalances ret_sto = sumOfAllBalances sto.
  Proof.
    destruct p as [[p | p] | [p | [p | p]]]; cbn.
    - (* %transfer *)
      destruct p as [from [to amount]]; cbn.
      intros [from_balance' [H_from_balance' H]].
      destruct H as [[_ H] | [current_allowance' [_ H]]].
      all: subst.
      2: rewrite sumOfAllBalances_setAllowance.
      all: destruct (compare_equal from to) eqn:eq.
      all: rewrite ?eqb_eq, ?eqb_neq in eq; try symmetry in eq; subst.
      all: rewrite !sumOfAllBalances_setBalance.
      all: try rewrite !getBalance_setBalance_eq.
      all: try rewrite !getBalance_setBalance_neq; auto.
      all: pose proof (balance_leq_sumOfAllBalances sto from) as H_from_balance.
      all: destruct getBalanceOpt as [from_balance | ]; cbn.
      all: lia.
    - (* %appprove *)
      destruct p as [spender new_allowance]; cbn.
      intros [[_ H] | [_ H]].
      all: subst; auto.
      rewrite sumOfAllBalances_setAllowance; auto.
    - (* %getAllowance *)
      destruct p as [[sender spender] contr]; cbn.
      intros [H _].
      subst; auto.
    - (* %getBalance *)
      destruct p as [owner contr]; cbn.
      intros [H _].
      subst; auto.
    - (* %getTotalSupply *)
      destruct p as [[] contr]; cbn.
      intros [H _].
      subst; auto.
  Qed.

End FA12SumOfBalances.

Module FA12SpecIffImpl (C : ContractContext).

  Module semantics := Semantics C.
  Import semantics.

  Section FA12SpecIffImpl.

    Context (parameter_ty : type).
    Context (storage_ty : type).
    Context (fuel_bound : Datatypes.nat).

    Notation input_ty := (pair parameter_ty storage_ty :: nil).
    Notation output_ty := (pair (list operation) storage_ty :: nil).

    Context (impl : Datatypes.nat -> stack input_ty -> M (stack output_ty)).
    Context (spec : data parameter_ty -> data storage_ty -> data (list operation) -> data storage_ty -> Prop).

    Context (spec_if_impl : forall fuel p sto ret_ops ret_sto,
      fuel_bound <= fuel ->
      impl fuel ((p, sto), tt) = Return ((ret_ops, ret_sto), tt) ->
      spec p sto ret_ops ret_sto).

    Context (spec_unique : forall p sto ret_ops1 ret_ops2 ret_sto1 ret_sto2,
      spec p sto ret_ops1 ret_sto1 ->
      spec p sto ret_ops2 ret_sto2 ->
      ret_sto1 = ret_sto2 /\ ret_ops1 = ret_ops2).

    Context (spec_success : forall fuel p sto ret_ops ret_sto,
      fuel_bound <= fuel ->
      spec p sto ret_ops ret_sto ->
      precond (impl fuel ((p, sto), tt)) (fun _ => Logic.True)).

    Lemma spec_iff_impl fuel p sto ret_ops ret_sto :
      fuel_bound <= fuel ->
      impl fuel ((p, sto), tt) = Return ((ret_ops, ret_sto), tt) <->
      spec p sto ret_ops ret_sto.
    Proof.
      intros H_fuel.
      split.
      1: eapply spec_if_impl; auto.
      intro H_spec.
      pose proof H_spec as H_spec'.
      eapply spec_success in H_spec'.
      2: exact H_fuel.
      destruct impl as [e | [[ret_ops' ret_sto'] []]] eqn:H_eval.
      all: destruct H_spec'.
      apply spec_if_impl in H_eval; auto.
      eapply spec_unique in H_eval.
      2: exact H_spec.
      destruct H_eval as [Hret_ops Hret_sto].
      subst; auto.
    Qed.

  End FA12SpecIffImpl.

End FA12SpecIffImpl.