Require Import Crypto.Spec.Encoding.
Require Bedrock.Word Crypto.Util.WordUtil.
Require Coq.ZArith.Znumtheory Coq.ZArith.BinInt.
Require Coq.Numbers.Natural.Peano.NPeano.
Require Crypto.CompleteEdwardsCurve.CompleteEdwardsCurveTheorems.

Local Infix "^" := NPeano.pow.
Local Infix "mod" := NPeano.modulo (at level 40, no associativity).
Local Infix "++" := Word.combine.

Generalizable All Variables.
Section EdDSA.
  Class EdDSA (* <https://eprint.iacr.org/2015/677.pdf> *)
        {E Eeq Eadd Ezero Eopp} {EscalarMult} (* the underllying elliptic curve operations *)

        {b : nat} (* public keys are k bits, signatures are 2*k bits *)
        {H : forall {n}, Word.word n -> Word.word (b + b)} (* main hash function *)
        {c : nat} (* cofactor E = 2^c *)
        {n : nat} (* secret keys are (n+1) bits *)
        {l : nat} (* order of the subgroup of E generated by B *)

        {B : E} (* base point *)

        {PointEncoding : canonical encoding of E as Word.word b} (* wire format *)
        {FlEncoding : canonical encoding of { n | n < l } as Word.word b}
    :=
      {
        EdDSA_group:@Algebra.group E Eeq Eadd Ezero Eopp;

        EdDSA_c_valid : c = 2 \/ c = 3;

        EdDSA_n_ge_c : n >= c;
        EdDSA_n_le_b : n <= b;

        EdDSA_B_not_identity : B <> Ezero;

        EdDSA_l_prime : Znumtheory.prime (BinInt.Z.of_nat l);
        EdDSA_l_odd : l > 2;
        EdDSA_l_order_B : EscalarMult l B = Ezero
      }.
  Global Existing Instance EdDSA_group.

  Context `{prm:EdDSA}.

  Local Infix "=" := Eeq.
  Local Coercion Word.wordToNat : Word.word >-> nat.
  Local Notation secretkey := (Word.word b) (only parsing).
  Local Notation publickey := (Word.word b) (only parsing).
  Local Notation signature := (Word.word (b + b)) (only parsing).

  Local Arguments H {n} _.
  Local Notation wfirstn n w := (@WordUtil.wfirstn n _ w _) (only parsing).

  Require Import Omega.
  Obligation Tactic := simpl; intros; try apply NPeano.Nat.mod_upper_bound; destruct prm; omega.

  Program Definition curveKey (sk:secretkey) : nat :=
    let x := wfirstn n (H sk) in (* hash the key, use first "half" for secret scalar *)
    let x := x - (x mod (2^c)) in (* it is implicitly 0 mod (2^c) *)
             x + 2^n. (* and the high bit is always set *)

  Local Infix "+" := Eadd.
  Local Infix "*" := EscalarMult.

  Definition prngKey (sk:secretkey) : Word.word b := Word.split2 b b (H sk).
  Definition public (sk:secretkey) : publickey := enc (curveKey sk*B).

  Program Definition sign (A_:publickey) sk {n} (M : Word.word n) :=
    let r : nat := H (prngKey sk ++ M) in (* secret nonce *)
    let R : E := r * B in (* commitment to nonce *)
    let s : nat := curveKey sk in (* secret scalar *)
    let S : {n|n<l} := exist _ ((r + H (enc R ++ public sk ++ M) * s) mod l) _ in
        enc R ++ enc S.

  (* For a [n]-bit [message] from public key [A_], validity of a signature [R_ ++ S_] *)
  Inductive valid {n:nat} : Word.word n -> publickey -> signature -> Prop :=
    ValidityRule : forall (message:Word.word n) (A:E) (R:E) (S:nat) S_lt_l,
      S * B = R + (H (enc R ++ enc A ++ message) mod l) * A
      -> valid message (enc A) (enc R ++ enc (exist _ S S_lt_l)).
End EdDSA.