From Coq Require Import Permutation Sorting.Sorted Sorting.Mergesort.
Require Import Coqlib Maps List Integers HashedSet.
Import ListNotations.
Require Import AST.
Require Import Op Registers.
Require Import RTL RTLtyping.
Require Import CounterMeasures.
Require Import Errors.
Open Scope error_monad_scope.

Definition transf_sig sig :=
  mksignature (Xptr::Xint::sig_args sig) (sig_res sig) (sig_cc sig).

Parameter ipcfc_pre: ident.

Fixpoint hash_list m r (l: list ident) h :=
  match l with
  | [] => h
  | hd::tl =>
      let h := Int.xor h (Int.repr (Zpos hd)) in
      let h := Int.mul h m in
      let h := Int.xor h (Int.shl h r) in
      hash_list m r tl h
  end.

Definition hash_entry_sig (qi: qualident) :=
  hash_list (Int.repr 0x5bd1e995) (Int.repr 15) (proj1_sig qi) (Int.repr 0).
Definition hash_exit_sig (qi: qualident) :=
  hash_list (Int.repr 0x1b44c290) (Int.repr 15) (proj1_sig qi) (Int.repr 0).

Opaque hash_entry_sig hash_exit_sig.

Definition harden_name qi :=
  match prefix_qualident ipcfc_pre qi with
  | inl qi => OK (proj1_sig qi)
  | inr _ => Error (MSG "Cannot harden" :: errcode_of_qualident qi)
  end.

Lemma harden_name_not_twice : forall qi1 qi2 qi3,
    harden_name qi1 = OK qi2 ->
    ~harden_name qi2 = OK qi3.

Lemma harden_name_qualident_prefix : forall id id',
    harden_name id = Errors.OK id' ->
    qualident_prefixed id' ipcfc_pre.

Lemma harden_name_inj : forall id1 id2 id',
    harden_name id1 = Errors.OK id' ->
    harden_name id2 = Errors.OK id' ->
    id1 = id2.

Section transfer_function.

Variable protected : QITree.t unit.
Variable fid : qualident.

Section transfer_instruction.

Variable fid' : qualident.
Variable gsr gsrv rts tmp1 : reg.
Variable abort : node.

Definition transf_call' {e} sig id' args r : seqinstr e -> seqinstr e :=
  fun seq =>
    Ssinstr (Iassert (Ccomp Ceq) [AA_load Mint32 (Aindexed' Int.zero) [gsr]; AA_reg gsrv]);;
    Ssinstr (Iop (Oxorimm (hash_entry_sig id')) [gsrv] tmp1);;
    Ssinstr (Icall (transf_sig sig) (inr id') (gsr::tmp1::args) r);;
    Ssinstr (Iload TRAP Mint32 (Aindexed' Int.zero) [gsr] tmp1);;
    Ssinstr (Iop (Oxorimm (hash_exit_sig id')) [tmp1] tmp1);;
    Ssmerge1 (Ccomp Cne) [tmp1; gsrv] Secatch (Some false);;
    Ssinstr (Istore Mint32 (Aindexed' Int.zero) [gsr] gsrv);; seq.

Definition transf_call sig ros args r (_ : node) :=
  match ros with
  | inr id =>
      do id' <- harden_name id;
      match protected!.id with
      | Some _ => OK (transf_call' sig id' args r;; Seinstr Inop)
      | None => OK (Seinstr (Icall sig ros args r))
      end
  | _ => OK (Seinstr (Icall sig ros args r))
  end.

Definition transf_return o :=
  OK (Ssinstr (Iop Oxor [rts; gsrv] rts);;
      Ssinstr (Iop (Oxorimm (hash_exit_sig fid')) [rts] rts);;
      Ssinstr (Istore Mint32 (Aindexed' Int.zero) [gsr] rts);;
      Sereturn (Ireturn o)).

End transfer_instruction.

Section transfer_function_init_prot.

Definition check_params params sig :=
  if Nat.eqb (length params) (length (sig_args sig))
  then check_params_norepet params
  else Error (msg "Incorrect parameter length").

Definition transf_init_prot f :=
  let maxr := max_reg_function f in
  let gsr := Pos.succ maxr in
  let gsrv := Pos.succ gsr in
  let rts := Pos.succ gsrv in
  let tmp1 := Pos.succ rts in
  let tmp2 := Pos.succ tmp1 in
  let tmp3 := Pos.succ tmp2 in
  let oret := match sig_res (fn_sig f) with Xvoid => None | _ => Some tmp3 end in
  do fid' <- harden_name fid;
  do d <- check_params (fn_params f) (fn_sig f);
  OK (Ssinstr (Iop (Ointconst (hash_entry_sig fid)) [] gsrv);;
      Ssinstr (Istore Mint32 (Ainstack Ptrofs.zero) [] gsrv);;
      Ssinstr (Iop (Oaddrstack Ptrofs.zero) [] gsr);;
      transf_call' gsr gsrv tmp1 (fn_sig f) fid' (fn_params f) tmp3;;
      Sereturn (Ireturn oret)).

Definition transf_function_init_prot (f : function) : Errors.res function :=
  Errors.bind
    (transf_init_prot f)
    (fun seq =>
       OK (mkfunction
             (fn_sig f) (fn_params f) (Memdata.size_chunk Mint32)
             (transf_replace_function seq f)
             (fn_entrypoint f) (fn_attr f))).

End transfer_function_init_prot.

Section transfer_function_prot.

Definition prot_prologue gsr gsrv rts fid' :=
  Ssinstr (Iload TRAP Mint32 (Aindexed' Int.zero) [gsr] gsrv);;
  Ssinstr (Iop Oxor [gsrv; rts] gsrv);;
  Ssmerge1 (Ccompimm Cne (hash_entry_sig fid')) [gsrv] Secatch (Some false);;
  Seinstr (Istore Mint32 (Aindexed' Int.zero) [gsr] gsrv).

Definition transf_function_prot (f : function) : Errors.res function :=
  let maxr := max_reg_function f in
  let gsr := Pos.succ maxr in
  let gsrv := Pos.succ gsr in
  let rts := Pos.succ gsrv in
  let tmp1 := Pos.succ rts in
  match
      (do fid' <- harden_name fid;
       transf_partial_function
         (prot_prologue gsr gsrv rts fid')
         transf_partial_nop_default transf_partial_op_default
         transf_partial_load_default transf_partial_store_default
         (transf_call gsr gsrv tmp1)
         (fun _ _ _ => Errors.Error (msg "IPCFC: tail calls are not allowed"))
         transf_partial_builtin_default
         transf_partial_cond_default transf_partial_jumptable_default
         (transf_return fid' gsr gsrv rts)
         f)
  with
  | OK (code, entry) => OK (mkfunction (transf_sig (fn_sig f)) (gsr::rts::fn_params f) (fn_stacksize f) code entry (fn_attr f))
  | Error e => Error e
  end.

End transfer_function_prot.

End transfer_function.


Open Scope error_monad_scope.

Definition check_norepet (ids : list qualident) :=
  (fix aux l t :=
     match l with
     | [] => OK tt
     | hd::tl => match QITree.get hd t with
               | Some _ => Error (MSG "IPCFC: Duplicate qualident " :: errcode_of_qualident hd)
               | None => aux tl (QITree.set hd tt t)
               end
     end) ids QITree.empty.

Definition check_exit_sigs {F V} (ids : list (_ * globdef (AST.fundef F) V)) :=
  mmap (fun '(id, def) =>
          match def with
          | Gfun (Internal _) =>
              if Int.eq (hash_exit_sig id) Int.zero
              then Error (msg "IPCFC: exit sig is zero (change hash function)")
              else OK tt
          | _ => OK tt
          end) ids.

Module SigOrder <: Orders.TotalLeBool.
  Definition t : Type := (qualident * int).
  Definition leb (x y: t) := negb (Int.lt (snd y) (snd x)).

  Lemma leb_total : forall x y,
      leb x y = true \/ leb y x = true.

  Lemma leb_trans : forall x y z,
      leb x y = true ->
      leb y z = true ->
      leb x z = true.

  Lemma leb_antisym : forall x y,
      leb x y = true ->
      leb y x = true ->
      Int.eq (snd x) (snd y) = true.
End SigOrder.

Module IntSort := Sort(SigOrder).

Fixpoint check_norepet_int (acc: qualident * int) l :=
  match l with
  | [] => OK tt
  | (id, sig)::tl =>
      if Int.eq sig (snd acc)
      then Error (MSG "IPCFC: Duplicate entry sig for " :: errcode_of_qualident (fst acc) ++ MSG " and " :: errcode_of_qualident id)
      else check_norepet_int (id, sig) tl
  end.

Definition check_entry_sigs {F V} (ids : list (_ * globdef (AST.fundef F) V)) :=
  let ids := map_filter (fun '(id, def) =>
                           match def with
                           | Gfun (Internal _) => Some id
                           | _ => None
                           end) ids in
  let sigs := IntSort.sort (List.map (fun id => (id, hash_entry_sig id)) ids) in
  match sigs with [] => OK tt | hd::tl => check_norepet_int hd tl end.

Definition protected (p: RTL.program) :=
  let public := QITree_Properties.of_list (List.map (fun x => (x, tt)) (prog_public p)) in
  QITree_Properties.of_list
    (map_filter (fun '(qi, gd) =>
                   match gd with
                   | Gfun (Internal f) =>
                       let cc := f.(fn_sig).(sig_cc) in
                       match public!.qi, cc.(cc_vararg) with
                       | Some _, None => Some (qi, tt)
                       | _, _ => None
                       end
                   | _ => None
                   end) (prog_defs p)).

Definition transf_program (p: program) :=
  let defs := prog_defs p in
  let defids := List.map fst defs in
  let prot := protected p in
  do tt <- Errors.mmap (fun id =>
                         if qualident_eq_prefix id ipcfc_pre
                         then Error (MSG "Bad qualident " :: errcode_of_qualident id ++ msg " before ipcfc.")
                         else OK tt) defids;
  let defm := prog_defmap p in
  do tt <- Errors.mmap (fun id =>
                         match defm!.id with
                         | Some _ => OK tt
                         | None => Error (errcode_of_qualident id ++ msg " is marked public but not declared.")
                         end) (prog_public p);
  do tt <- check_norepet defids;
  do tt <- type_program p;
  do defs <- Errors.mmap (fun '(id, d) =>
                           match d, prot!.id with
                           | Gfun (Internal f), Some _ =>
                               do id' <- harden_name id;
                               do f1 <- transf_function_init_prot id f;
                               do f2 <- transf_function_prot prot id f;
                               OK (id, Gfun (Internal f1), Some (id', Gfun (Internal f2)))
                           | _, _ => OK ((id, d), None)
                           end) defs;
  let defs := List.map fst defs ++ map_filter snd defs in
  do tt <- check_exit_sigs defs;
  do tt <- check_entry_sigs defs;
  OK (mkprogram defs (prog_public p) (prog_main p)).

Strong match property

Inductive match_def (p tp: RTL.program) (id: qualident) : Prop :=
| match_def_none1 :
    (prog_defmap p)!.id = None ->
    (prog_defmap tp)!.id = None ->
    match_def p tp id
| match_def_none2 : forall id' f f',
    (prog_defmap p)!.id = None ->
    harden_name id' = Errors.OK id ->
    (prog_defmap p)!.id' = Some (Gfun (Internal f)) ->
    (protected p)!.id' = Some tt ->
    transf_function_prot (protected p) id' f = OK f' ->
    (prog_defmap tp)!.id = Some (Gfun (Internal f')) ->
    match_def p tp id
| match_def_var : forall v,
    (prog_defmap p)!.id = Some (Gvar v) ->
    (prog_defmap tp)!.id = Some (Gvar v) ->
    match_def p tp id
| match_def_unprot: forall f,
    (protected p)!.id = None ->
    (prog_defmap p)!.id = Some (Gfun f) ->
    (prog_defmap tp)!.id = Some (Gfun f) ->
    match_def p tp id
| match_def_prot: forall id' f f1 f2,
    (protected p)!.id = Some tt ->
    harden_name id = Errors.OK id' ->
    (prog_defmap p)!.id = Some (Gfun (Internal f)) ->
    transf_function_init_prot id f = Errors.OK f1 ->
    (prog_defmap tp)!.id = Some (Gfun (Internal f1)) ->
    transf_function_prot (protected p) id f = OK f2 ->
    (prog_defmap tp)!.id' = Some (Gfun (Internal f2)) ->
    match_def p tp id.

Record match_prog (p tp: program) : Prop := {
    match_prog_main: tp.(prog_main) = p.(prog_main);
    match_prog_public: tp.(prog_public) = p.(prog_public);
    match_prog_public_in: forall id, In id (prog_public p) -> In id (map fst (prog_defs p));
    match_prog_def: forall id, match_def p tp id;
    match_prog_inv: forall id id' tf,
      harden_name id' = Errors.OK id ->
      (prog_defmap tp)!.id = Some (Gfun (Internal tf)) ->
      exists f, (prog_defmap p)!.id' = Some (Gfun (Internal f));
    match_prog_wt: wt_program p;
    match_prog_norepet1: list_norepet (prog_defs_names p);
    match_prog_norepet2: list_norepet (prog_defs_names tp);
    match_prog_npref: forall id, In id (map fst (prog_defs p)) -> ~qualident_prefixed id ipcfc_pre;
    match_prog_sigs_nzero: forall id f,
      In (id, (Gfun (Internal f))) (prog_defs tp) -> hash_exit_sig id <> Int.zero;
    match_prog_sigs_inj: forall id1 f1 id2 f2,
      In (id1, (Gfun (Internal f1))) (prog_defs tp) ->
      In (id2, (Gfun (Internal f2))) (prog_defs tp) ->
      hash_entry_sig id1 = hash_entry_sig id2 ->
      id1 = id2
  }.

Lemma check_norepet_correct: forall ids,
    check_norepet ids = OK tt ->
    list_norepet ids.

Lemma transf_program_defs_norepet : forall prog tprog,
    transf_program prog = Errors.OK tprog ->
    list_norepet (map fst (prog_defs tprog)).

Lemma protected_internal : forall prog id,
    (protected prog)!.id = Some tt ->
    exists f, In (id, Gfun (Internal f)) (prog_defs prog).

Corollary external_unprot : forall prog id e,
    list_norepet (prog_defs_names prog) ->
    In (id, Gfun (External e)) (prog_defs prog) ->
    (protected prog)!.id = None.

Lemma check_exit_sigs_correct {F V} : forall ids tt,
    @check_exit_sigs F V ids = OK tt ->
    forall id f, In (id, Gfun (Internal f)) ids -> hash_exit_sig id <> Int.zero.

Lemma check_norepet_int_correct : forall tl hd tt,
    check_norepet_int hd tl = OK tt ->
    StronglySorted (fun x y => SigOrder.leb x y = true) (hd::tl) ->
    list_norepet (map snd (hd::tl)).

Lemma check_entry_sigs_correct {F V} : forall ids,
    @check_entry_sigs F V ids = OK tt ->
    forall id1 f1 id2 f2,
      In (id1, (Gfun (Internal f1))) ids ->
      In (id2, (Gfun (Internal f2))) ids ->
      hash_entry_sig id1 = hash_entry_sig id2 ->
      id1 = id2.

Lemma transf_program_match:
  forall prog tprog, transf_program prog = Errors.OK tprog -> match_prog prog tprog.