val bn_lshift_add_early_stop:
    #t:limb_t
  -> #aLen:size_nat
  -> #bLen:size_nat
  -> a:lbignum t aLen
  -> b:lbignum t bLen
  -> i:nat{i + bLen <= aLen} ->
  tuple2 (carry t) (lbignum t aLen)
let bn_lshift_add_early_stop #t #aLen #bLen a b i =
  let r = sub a i bLen in
  let c, r' = bn_add r b in
  let a' = update_sub a i bLen r' in
  c, a'
val bn_lshift_add_early_stop:
    #t:limb_t
  -> #aLen:size_nat
  -> #bLen:size_nat
  -> a:lbignum t aLen
  -> b:lbignum t bLen
  -> i:nat{i + bLen <= aLen} ->
  tuple2 (carry t) (lbignum t aLen)
let bn_lshift_add_early_stop #t #aLen #bLen a b i =
let r = sub a i bLen in
let c, r' = bn_add r b in
let a' = update_sub a i bLen r' in
c, a'
Hacl.Spec.Bignum.Karatsuba.fst.checked
prims.fst.checked Lib.Sequence.fsti.checked Lib.LoopCombinators.fsti.checked Lib.IntTypes.fsti.checked Hacl.Spec.Lib.fst.checked Hacl.Spec.Karatsuba.Lemmas.fst.checked Hacl.Spec.Bignum.Squaring.fst.checked Hacl.Spec.Bignum.Multiplication.fst.checked Hacl.Spec.Bignum.Lib.fst.checked Hacl.Spec.Bignum.Definitions.fst.checked Hacl.Spec.Bignum.Base.fst.checked Hacl.Spec.Bignum.Addition.fst.checked FStar.UInt64.fsti.checked FStar.UInt32.fsti.checked FStar.Pervasives.Native.fst.checked FStar.Pervasives.fsti.checked FStar.Mul.fst.checked FStar.Math.Lemmas.fst.checked FStar.Calc.fsti.checked
Hacl.Spec.Bignum.Karatsuba.fst
Hacl.Spec.Bignum.Definitions.limb_t Lib.IntTypes.size_nat Hacl.Spec.Bignum.Definitions.lbignum Prims.nat Prims.b2t Prims.op_LessThanOrEqual Prims.op_Addition Hacl.Spec.Bignum.Base.carry FStar.Pervasives.Native.Mktuple2 Lib.Sequence.lseq Hacl.Spec.Bignum.Definitions.limb Prims.l_and Prims.eq2 Lib.Sequence.sub Prims.l_Forall Prims.l_or Prims.op_LessThan FStar.Seq.Base.index Lib.Sequence.to_seq Lib.Sequence.index Lib.Sequence.update_sub FStar.Pervasives.Native.tuple2 Hacl.Spec.Bignum.Addition.bn_add FStar.Seq.Base.seq FStar.Seq.Base.slice
module Hacl.Spec.Bignum.Karatsuba

open FStar.Mul

open Lib.IntTypes
open Lib.Sequence
open Lib.LoopCombinators

open Hacl.Spec.Bignum.Definitions
open Hacl.Spec.Bignum.Base
open Hacl.Spec.Bignum.Lib
open Hacl.Spec.Lib

open Hacl.Spec.Bignum.Addition
open Hacl.Spec.Bignum.Multiplication
open Hacl.Spec.Bignum.Squaring

module K = Hacl.Spec.Karatsuba.Lemmas

#reset-options "--z3rlimit 50 --fuel 0 --ifuel 0"

inline_for_extraction noextract
let bn_mul_threshold = 32

(* this carry means nothing but the sign of the result *)
val bn_sign_abs:
    #t:limb_t
  -> #aLen:size_nat
  -> a:lbignum t aLen
  -> b:lbignum t aLen ->
  tuple2 (carry t) (lbignum t aLen)

let bn_sign_abs #t #aLen a b =
  let c0, t0 = bn_sub a b in
  let c1, t1 = bn_sub b a in
  let res = map2 (mask_select (uint #t 0 -. c0)) t1 t0 in
  c0, res


val bn_sign_abs_lemma:
    #t:limb_t
  -> #aLen:size_nat
  -> a:lbignum t aLen
  -> b:lbignum t aLen ->
  Lemma (let c, res = bn_sign_abs a b in
   bn_v res == K.abs (bn_v a) (bn_v b) /\
   v c == (if bn_v a < bn_v b then 1 else 0))

let bn_sign_abs_lemma #t #aLen a b =
  let s, r = K.sign_abs (bn_v a) (bn_v b) in

  let c0, t0 = bn_sub a b in
  bn_sub_lemma a b;
  assert (bn_v t0 - v c0 * pow2 (bits t * aLen) == bn_v a - bn_v b);

  let c1, t1 = bn_sub b a in
  bn_sub_lemma b a;
  assert (bn_v t1 - v c1 * pow2 (bits t * aLen) == bn_v b - bn_v a);

  let mask = uint #t 0 -. c0 in
  assert (v mask == (if v c0 = 0 then 0 else v (ones t SEC)));
  let res = map2 (mask_select mask) t1 t0 in
  lseq_mask_select_lemma t1 t0 mask;
  assert (bn_v res == (if v mask = 0 then bn_v t0 else bn_v t1));

  bn_eval_bound a aLen;
  bn_eval_bound b aLen;
  bn_eval_bound t0 aLen;
  bn_eval_bound t1 aLen

  // if bn_v a < bn_v b then begin
  //   assert (v mask = v (ones U64 SEC));
  //   assert (bn_v res == bn_v b - bn_v a);
  //   assert (bn_v res == r /\ v c0 = 1) end
  // else begin
  //   assert (v mask = 0);
  //   assert (bn_v res == bn_v a - bn_v b);
  //   assert (bn_v res == r /\ v c0 = 0) end;
  // assert (bn_v res == r /\ v c0 == (if bn_v a < bn_v b then 1 else 0))


val bn_middle_karatsuba:
    #t:limb_t
  -> #aLen:size_nat
  -> c0:carry t
  -> c1:carry t
  -> c2:carry t
  -> t01:lbignum t aLen
  -> t23:lbignum t aLen ->
  limb t & lbignum t aLen

let bn_middle_karatsuba #t #aLen c0 c1 c2 t01 t23 =
  let c_sign = c0 ^. c1 in
  let c3, t45 = bn_sub t01 t23 in let c3 = c2 -. c3 in
  let c4, t67 = bn_add t01 t23 in let c4 = c2 +. c4 in
  let mask = uint #t 0 -. c_sign in
  let t45 = map2 (mask_select mask) t67 t45 in
  let c5 = mask_select mask c4 c3 in
  c5, t45


val sign_lemma: #t:limb_t -> c0:carry t -> c1:carry t ->
  Lemma (v (c0 ^. c1) == (if v c0 = v c1 then 0 else 1))
let sign_lemma #t c0 c1 =
  logxor_spec c0 c1;
  match t with
  | U32 ->
    assert_norm (UInt32.logxor 0ul 0ul == 0ul);
    assert_norm (UInt32.logxor 0ul 1ul == 1ul);
    assert_norm (UInt32.logxor 1ul 0ul == 1ul);
    assert_norm (UInt32.logxor 1ul 1ul == 0ul)
  | U64 ->
    assert_norm (UInt64.logxor 0uL 0uL == 0uL);
    assert_norm (UInt64.logxor 0uL 1uL == 1uL);
    assert_norm (UInt64.logxor 1uL 0uL == 1uL);
    assert_norm (UInt64.logxor 1uL 1uL == 0uL)


val bn_middle_karatsuba_lemma:
    #t:limb_t
  -> #aLen:size_nat
  -> c0:carry t
  -> c1:carry t
  -> c2:carry t
  -> t01:lbignum t aLen
  -> t23:lbignum t aLen ->
  Lemma (let (c, res) = bn_middle_karatsuba c0 c1 c2 t01 t23 in
    let c3, t45 = bn_sub t01 t23 in let c3' = c2 -. c3 in
    let c4, t67 = bn_add t01 t23 in let c4' = c2 +. c4 in

    if v c0 = v c1 then
      v c == v c3' /\ bn_v res == bn_v t45
    else
      v c == v c4' /\ bn_v res == bn_v t67)

let bn_middle_karatsuba_lemma #t #aLen c0 c1 c2 t01 t23 =
  let lp = bn_v t01 + v c2 * pow2 (bits t * aLen) - bn_v t23 in
  let rp = bn_v t01 + v c2 * pow2 (bits t * aLen) + bn_v t23 in

  let c_sign = c0 ^. c1 in
  sign_lemma c0 c1;
  assert (v c_sign == (if v c0 = v c1 then 0 else 1));

  let c3, t45 = bn_sub t01 t23 in let c3' = c2 -. c3 in
  let c4, t67 = bn_add t01 t23 in let c4' = c2 +. c4 in

  let mask = uint #t 0 -. c_sign in
  let t45' = map2 (mask_select mask) t67 t45 in
  lseq_mask_select_lemma t67 t45 mask;
  //assert (bn_v t45' == (if v mask = 0 then bn_v t45 else bn_v t67));
  let c5 = mask_select mask c4' c3' in
  mask_select_lemma mask c4' c3'
  //assert (v c5 == (if v mask = 0 then v c3' else v c4'));


val bn_middle_karatsuba_eval_aux:
    #t:limb_t
  -> #aLen:size_nat
  -> a0:lbignum t (aLen / 2)
  -> a1:lbignum t (aLen / 2)
  -> b0:lbignum t (aLen / 2)
  -> b1:lbignum t (aLen / 2)
  -> res:lbignum t aLen
  -> c2:carry t
  -> c3:carry t -> Lemma
  (requires
    bn_v res + (v c2 - v c3) * pow2 (bits t * aLen) ==
    bn_v a0 * bn_v b1 + bn_v a1 * bn_v b0)
  (ensures  0 <= v c2 - v c3 /\ v c2 - v c3 <= 1)

let bn_middle_karatsuba_eval_aux #t #aLen a0 a1 b0 b1 res c2 c3 =
  bn_eval_bound res aLen


val bn_middle_karatsuba_eval:
    #t:limb_t
  -> #aLen:size_nat
  -> a0:lbignum t (aLen / 2)
  -> a1:lbignum t (aLen / 2)
  -> b0:lbignum t (aLen / 2)
  -> b1:lbignum t (aLen / 2)
  -> c2:carry t
  -> t01:lbignum t aLen
  -> t23:lbignum t aLen ->
  Lemma
  (requires
   (let t0 = K.abs (bn_v a0) (bn_v a1) in
    let t1 = K.abs (bn_v b0) (bn_v b1) in
    bn_v t01 + v c2 * pow2 (bits t * aLen) == bn_v a0 * bn_v b0 + bn_v a1 * bn_v b1 /\
    bn_v t23 == t0 * t1))
  (ensures
   (let c0, t0 = bn_sign_abs a0 a1 in
    let c1, t1 = bn_sign_abs b0 b1 in

    let c, res = bn_middle_karatsuba c0 c1 c2 t01 t23 in
    bn_v res + v c * pow2 (bits t * aLen) == bn_v a0 * bn_v b1 + bn_v a1 * bn_v b0))

let bn_middle_karatsuba_eval #t #aLen a0 a1 b0 b1 c2 t01 t23 =
  let pbits = bits t in
  let c0, t0 = bn_sign_abs a0 a1 in
  bn_sign_abs_lemma a0 a1;
  assert (bn_v t0 == K.abs (bn_v a0) (bn_v a1));
  assert (v c0 == (if bn_v a0 < bn_v a1 then 1 else 0));

  let c1, t1 = bn_sign_abs b0 b1 in
  bn_sign_abs_lemma b0 b1;
  assert (bn_v t1 == K.abs (bn_v b0) (bn_v b1));
  assert (v c1 == (if bn_v b0 < bn_v b1 then 1 else 0));

  let c, res = bn_middle_karatsuba c0 c1 c2 t01 t23 in
  bn_middle_karatsuba_lemma c0 c1 c2 t01 t23;

  let c3, t45 = bn_sub t01 t23 in let c3' = c2 -. c3 in
  let c4, t67 = bn_add t01 t23 in let c4' = c2 +. c4 in

  if v c0 = v c1 then begin
    assert (bn_v a0 * bn_v b0 + bn_v a1 * bn_v b1 - bn_v t0 * bn_v t1 == bn_v a0 * bn_v b1 + bn_v a1 * bn_v b0);
    assert (v c2 * pow2 (pbits * aLen) + bn_v t01 - bn_v t23 == bn_v a0 * bn_v b1 + bn_v a1 * bn_v b0);

    assert (v c == v c3' /\ bn_v res == bn_v t45);
    //assert (v c = (v c2 - v c3) % pow2 pb);
    bn_sub_lemma t01 t23;
    assert (bn_v res - v c3 * pow2 (pbits * aLen) == bn_v t01 - bn_v t23);
    Math.Lemmas.distributivity_sub_left (v c2) (v c3) (pow2 (pbits * aLen));
    assert (bn_v res + (v c2 - v c3) * pow2 (pbits * aLen) == v c2 * pow2 (pbits * aLen) + bn_v t01 - bn_v t23);
    bn_middle_karatsuba_eval_aux a0 a1 b0 b1 res c2 c3;
    Math.Lemmas.small_mod (v c2 - v c3) (pow2 pbits);
    assert (bn_v res + v c * pow2 (pbits * aLen) == v c2 * pow2 (pbits * aLen) + bn_v t01 - bn_v t23);
    () end
  else begin
    assert (bn_v a0 * bn_v b0 + bn_v a1 * bn_v b1 + bn_v t0 * bn_v t1 == bn_v a0 * bn_v b1 + bn_v a1 * bn_v b0);
    assert (v c2 * pow2 (pbits * aLen) + bn_v t01 + bn_v t23 == bn_v a0 * bn_v b1 + bn_v a1 * bn_v b0);

    assert (v c == v c4' /\ bn_v res == bn_v t67);
    //assert (v c = v c2 + v c4);
    bn_add_lemma t01 t23;
    assert (bn_v res + v c4 * pow2 (pbits * aLen) == bn_v t01 + bn_v t23);
    Math.Lemmas.distributivity_add_left (v c2) (v c4) (pow2 (pbits * aLen));
    Math.Lemmas.small_mod (v c2 + v c4) (pow2 pbits);
    assert (bn_v res + v c * pow2 (pbits * aLen) == v c2 * pow2 (pbits * aLen) + bn_v t01 + bn_v t23);
    () end


val bn_lshift_add:
    #t:limb_t
  -> #aLen:size_nat
  -> a:lbignum t aLen
  -> b1:limb t
  -> i:nat{i + 1 <= aLen} ->
  tuple2 (carry t) (lbignum t aLen)

let bn_lshift_add #t #aLen a b1 i =
  let r = sub a i (aLen - i) in
  let c, r' = bn_add1 r b1 in
  let a' = update_sub a i (aLen - i) r' in
  c, a'


val bn_lshift_add_lemma:
    #t:limb_t
  -> #aLen:size_nat
  -> a:lbignum t aLen
  -> b1:limb t
  -> i:nat{i + 1 <= aLen} ->
  Lemma (let c, res = bn_lshift_add a b1 i in
    bn_v res + v c * pow2 (bits t * aLen) == bn_v a + v b1 * pow2 (bits t * i))

let bn_lshift_add_lemma #t #aLen a b1 i =
  let pbits = bits t in
  let r = sub a i (aLen - i) in
  let c, r' = bn_add1 r b1 in
  let a' = update_sub a i (aLen - i) r' in
  let p = pow2 (pbits * aLen) in

  calc (==) {
    bn_v a' + v c * p;
    (==) { bn_update_sub_eval a r' i }
    bn_v a - bn_v r * pow2 (pbits * i) + bn_v r' * pow2 (pbits * i) + v c * p;
    (==) { bn_add1_lemma r b1 }
    bn_v a - bn_v r * pow2 (pbits * i) + (bn_v r + v b1 - v c * pow2 (pbits * (aLen - i))) * pow2 (pbits * i) + v c * p;
    (==) { Math.Lemmas.distributivity_add_left (bn_v r) (v b1 - v c * pow2 (pbits * (aLen - i))) (pow2 (pbits * i)) }
    bn_v a + (v b1 - v c * pow2 (pbits * (aLen - i))) * pow2 (pbits * i) + v c * p;
    (==) { Math.Lemmas.distributivity_sub_left (v b1) (v c * pow2 (pbits * (aLen - i))) (pow2 (pbits * i)) }
    bn_v a + v b1 * pow2 (pbits * i) - (v c * pow2 (pbits * (aLen - i))) * pow2 (pbits * i) + v c * p;
    (==) { Math.Lemmas.paren_mul_right (v c) (pow2 (pbits * (aLen - i))) (pow2 (pbits * i));
           Math.Lemmas.pow2_plus (pbits * (aLen - i)) (pbits * i) }
    bn_v a + v b1 * pow2 (pbits * i);
    }


val bn_lshift_add_early_stop:
    #t:limb_t
  -> #aLen:size_nat
  -> #bLen:size_nat
  -> a:lbignum t aLen
  -> b:lbignum t bLen
  -> i:nat{i + bLen <= aLen} ->
  tuple2 (carry t) (lbignum t aLen)
val bn_lshift_add_early_stop:
    #t:limb_t
  -> #aLen:size_nat
  -> #bLen:size_nat
  -> a:lbignum t aLen
  -> b:lbignum t bLen
  -> i:nat{i + bLen <= aLen} ->
  tuple2 (carry t) (lbignum t aLen)

    a: Hacl.Spec.Bignum.Definitions.lbignum t aLen ->
    b: Hacl.Spec.Bignum.Definitions.lbignum t bLen ->
    i: Prims.nat{i + bLen <= aLen}
  -> Hacl.Spec.Bignum.Base.carry t * Hacl.Spec.Bignum.Definitions.lbignum t aLen