SPRTwoHop

instance instNeZeroMlkemModulus : NeZero MLKEM.modulus := ⟨by decide⟩

instance instSampleableTypeRq : SampleableType MLKEM.Rq :=
  inferInstanceAs (SampleableType (Vector MLKEM.Coeff MLKEM.ringDegree))

instance instSampleableTypeTq : SampleableType MLKEM.Tq :=
  SampleableType.ofEquiv (α := MLKEM.Rq)
    ⟨fun f => ⟨f⟩, fun t => t.coeffs, fun _ => rfl, fun _ => rfl⟩

theorem LearningWithErrors.advantage_eq_boolDistAdvantage {Sample Secret Output : Type}
    [Add Output] (problem : LearningWithErrors.Problem Sample Secret Output)
    (adv : LearningWithErrors.Adversary problem) :
    LearningWithErrors.advantage problem adv =
      (LearningWithErrors.game0 problem adv).boolDistAdvantage
        (LearningWithErrors.game1 problem adv) := by

def pkOf (rho : Seed32) (tHat : TqVec params.k) : EncapsulationKey params encoding :=
  { tHatEncoded := encoding.byteEncode12Vec tHat, rho := rho }

def ctOf (u : RqVec params.k) (v : Rq) : Ciphertext params encoding :=
  { uEncoded := encoding.byteEncodeDUVec (encoding.compressDU u)
    vEncoded := encoding.byteEncodeDV (encoding.compressDV v) }

def sampleEta1Hat (offset : ℕ) : ProbComp (TqVec params.k) := do
  let sigma ← $ᵗ Seed32
  pure (ring.nttVec (prims.sampleVecEta1 sigma offset))

def sampleNoise : ProbComp (RqVec params.k × Rq) := do
  let c₁ ← $ᵗ Coins
  let c₂ ← $ᵗ Coins
  pure (prims.sampleVecEta2 c₁ params.k, prims.prfEta2 c₂ (2 * params.k))

def idealKeygen : ProbComp (EncapsulationKey params encoding) := do
  let rho ← $ᵗ Seed32
  let s ← sampleEta1Hat ring encoding prims 0
  let e ← sampleEta1Hat ring encoding prims params.k
  pure (pkOf encoding rho (ring.matVecMul (prims.publicMatrix rho) s + e))

def idealEncrypt (ek : EncapsulationKey params encoding) :
    ProbComp (Ciphertext params encoding) := do
  let y ← sampleEta1Hat ring encoding prims 0
  let ne ← sampleNoise encoding prims
  let m ← $ᵗ Message
  pure (ctOf encoding
    (ring.invNTTVec (ring.matTransposeVecMul (prims.publicMatrix ek.rho) y) + ne.1)
    (ring.invNTT (ring.dot (encoding.byteDecode12Vec ek.tHatEncoded) y) + ne.2
      + encoding.decompress1 (encoding.byteDecode1 m)))

def uniformKeygen : ProbComp (EncapsulationKey params encoding) := do
  let rho ← $ᵗ Seed32
  let t ← $ᵗ (TqVec params.k)
  pure (pkOf encoding rho t)

def randCt : ProbComp (Ciphertext params encoding) := do
  let u ← $ᵗ (RqVec params.k)
  let v ← $ᵗ Rq
  let m ← $ᵗ Message
  pure (ctOf encoding u (v + encoding.decompress1 (encoding.byteDecode1 m)))

def idealBranch (adv : StealthScheme.UnlinkAdv (EncapsulationKey params encoding)
      (Ciphertext params encoding))
    (kb : ProbComp (EncapsulationKey params encoding))
    (ct : EncapsulationKey params encoding → ProbComp (Ciphertext params encoding))
    (b : Bool) : ProbComp Bool := do
  let pkb ← kb
  let c ← ct pkb
  let pkOther ← idealKeygen ring encoding prims
  adv (if b then pkOther else pkb) (if b then pkb else pkOther) c

def keyHopProblem :
    LearningWithErrors.Problem Seed32 (TqVec params.k) (TqVec params.k) where
  sampleChallenge := $ᵗ Seed32
  sampleSecret := sampleEta1Hat ring encoding prims 0
  sampleError := sampleEta1Hat ring encoding prims params.k
  noiseless := fun s rho => ring.matVecMul (prims.publicMatrix rho) s
  sampleUniform := $ᵗ (TqVec params.k)

def ctHopProblem :
    LearningWithErrors.Problem (Seed32 × TqVec params.k) (TqVec params.k)
      (RqVec params.k × Rq) where
  sampleChallenge := do
    let rho ← $ᵗ Seed32
    let t ← $ᵗ (TqVec params.k)
    pure (rho, t)
  sampleSecret := sampleEta1Hat ring encoding prims 0
  sampleError := sampleNoise encoding prims
  noiseless := fun y ch =>
    (ring.invNTTVec (ring.matTransposeVecMul (prims.publicMatrix ch.1) y),
      ring.invNTT (ring.dot ch.2 y))
  sampleUniform := do
    let u ← $ᵗ (RqVec params.k)
    let v ← $ᵗ Rq
    pure (u, v)

def keyHopAdv (b : Bool) : LearningWithErrors.Adversary (keyHopProblem ring encoding prims) :=
  fun chal => do
    let c ← idealEncrypt ring encoding prims (pkOf encoding chal.1 chal.2)
    let pkOther ← idealKeygen ring encoding prims
    adv (if b then pkOther else pkOf encoding chal.1 chal.2)
      (if b then pkOf encoding chal.1 chal.2 else pkOther) c

def ctHopAdv (b : Bool) : LearningWithErrors.Adversary (ctHopProblem ring encoding prims) :=
  fun chal => do
    let m ← $ᵗ Message
    let pkOther ← idealKeygen ring encoding prims
    adv (if b then pkOther else pkOf encoding chal.1.1 chal.1.2)
      (if b then pkOf encoding chal.1.1 chal.1.2 else pkOther)
      (ctOf encoding chal.2.1 (chal.2.2 + encoding.decompress1 (encoding.byteDecode1 m)))

theorem idealEncrypt_pkOf (rho : Seed32) (t : TqVec params.k) :
    idealEncrypt ring encoding prims (pkOf encoding rho t) = (do
      let y ← sampleEta1Hat ring encoding prims 0
      let ne ← sampleNoise encoding prims
      let m ← $ᵗ Message
      pure (ctOf encoding
        (ring.invNTTVec (ring.matTransposeVecMul (prims.publicMatrix rho) y) + ne.1)
        (ring.invNTT (ring.dot t y) + ne.2
          + encoding.decompress1 (encoding.byteDecode1 m)))) := by

theorem game0_keyHopProblem (b : Bool) :
    LearningWithErrors.game0 (keyHopProblem ring encoding prims)
      (keyHopAdv ring encoding prims adv b) =
      idealBranch ring encoding prims adv (idealKeygen ring encoding prims)
        (idealEncrypt ring encoding prims) b := by

theorem game1_keyHopProblem (b : Bool) :
    LearningWithErrors.game1 (keyHopProblem ring encoding prims)
      (keyHopAdv ring encoding prims adv b) =
      idealBranch ring encoding prims adv (uniformKeygen encoding)
        (idealEncrypt ring encoding prims) b := by

theorem game0_ctHopProblem (b : Bool) :
    LearningWithErrors.game0 (ctHopProblem ring encoding prims)
      (ctHopAdv ring encoding prims adv b) =
      idealBranch ring encoding prims adv (uniformKeygen encoding)
        (idealEncrypt ring encoding prims) b := by

theorem game1_ctHopProblem (b : Bool) :
    LearningWithErrors.game1 (ctHopProblem ring encoding prims)
      (ctHopAdv ring encoding prims adv b) =
      idealBranch ring encoding prims adv (uniformKeygen encoding)
        (fun _ => randCt encoding) b := by

def KEM.anonBranchReordered (b : Bool) : ProbComp Bool := do
  let pkb ← Prod.fst <$> kem.keygen
  let c ← Prod.fst <$> kem.encaps pkb
  let pkOther ← Prod.fst <$> kem.keygen
  adv (if b then pkOther else pkb) (if b then pkb else pkOther) c

theorem KEM.evalDist_anonSetup_bind_anonBranch (b : Bool) :
    𝒟[kem.anonSetup >>= kem.anonBranch adv b] = 𝒟[kem.anonBranchReordered adv b] := by

theorem KEM.boolDistAdvantage_anonBranch_anonBranchReordered (b : Bool) :
    (kem.anonSetup >>= kem.anonBranch adv b).boolDistAdvantage
      (kem.anonBranchReordered adv b) = 0 := by

noncomputable def primitiveIdealization (b : Bool) : ℝ :=
  (KEM.anonBranchReordered (MLKEM.asKEMScheme ring encoding prims) adv b).boolDistAdvantage
    (idealBranch ring encoding prims adv (idealKeygen ring encoding prims)
      (idealEncrypt ring encoding prims) b)

noncomputable def simulatorGap (b : Bool) : ℝ :=
  (idealBranch ring encoding prims adv (uniformKeygen encoding)
      (fun _ => randCt encoding) b).boolDistAdvantage
    (KEM.anonSetup (MLKEM.asKEMScheme ring encoding prims) >>= KEM.simBranch sim adv)

noncomputable def encodingRegularity (b : Bool) : ℝ :=
  (idealBranch ring encoding prims adv (uniformKeygen encoding)
      (fun _ => randCt encoding) b).boolDistAdvantage
    (idealBranch ring encoding prims adv (uniformKeygen encoding) (fun _ => sim) b)

noncomputable def keyRestoration (b : Bool) : ℝ :=
  (idealBranch ring encoding prims adv (uniformKeygen encoding) (fun _ => sim) b).boolDistAdvantage
    (KEM.anonSetup (MLKEM.asKEMScheme ring encoding prims) >>= KEM.simBranch sim adv)

theorem simulatorGap_le (b : Bool) :
    simulatorGap ring encoding prims adv sim b ≤
      encodingRegularity ring encoding prims adv sim b
        + keyRestoration ring encoding prims adv sim b :=
  ProbComp.boolDistAdvantage_triangle _ _ _

theorem sprAdv_le_two_hop_decomposition (b : Bool) :
    KEM.sprAdv (MLKEM.asKEMScheme ring encoding prims) sim adv b ≤
      primitiveIdealization ring encoding prims adv b
        + LearningWithErrors.advantage (keyHopProblem ring encoding prims)
            (keyHopAdv ring encoding prims adv b)
        + LearningWithErrors.advantage (ctHopProblem ring encoding prims)
            (ctHopAdv ring encoding prims adv b)
        + simulatorGap ring encoding prims adv sim b := by

def mlkem768KeyHopProblem :
    LearningWithErrors.Problem Seed32 (TqVec mlkem768.k) (TqVec mlkem768.k) :=
  keyHopProblem concreteNTTRingOps mlkem768Encoding mlkem768Primitives

def mlkem768CtHopProblem :
    LearningWithErrors.Problem (Seed32 × TqVec mlkem768.k) (TqVec mlkem768.k)
      (RqVec mlkem768.k × Rq) :=
  ctHopProblem concreteNTTRingOps mlkem768Encoding mlkem768Primitives

theorem mlkem768_sprAdv_le_two_hop_decomposition (b : Bool) :
    mlkem768KEM.sprAdv mlkem768UniformCiphertext adv768 b ≤
      primitiveIdealization concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768 b
        + LearningWithErrors.advantage mlkem768KeyHopProblem
            (keyHopAdv concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768 b)
        + LearningWithErrors.advantage mlkem768CtHopProblem
            (ctHopAdv concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768 b)
        + simulatorGap concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768
            mlkem768UniformCiphertext b :=
  sprAdv_le_two_hop_decomposition concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768
    mlkem768UniformCiphertext b

def primitiveIdealizationBound (epsilonPrim : ℝ) : Prop :=
  ∀ b, primitiveIdealization ring encoding prims adv b ≤ epsilonPrim

def encodingRegularityBound (epsilonEnc : ℝ) : Prop :=
  ∀ b, encodingRegularity ring encoding prims adv sim b ≤ epsilonEnc

def keyRestorationAdv (b : Bool) :
    LearningWithErrors.Adversary (keyHopProblem ring encoding prims) :=
  fun chal => do
    let simCt ← sim
    let pkOther ← idealKeygen ring encoding prims
    adv (if b then pkOther else pkOf encoding chal.1 chal.2)
      (if b then pkOf encoding chal.1 chal.2 else pkOther) simCt
 
omit [DecidableEq encoding.EncodedTHat] [DecidableEq encoding.EncodedU]
  [DecidableEq encoding.EncodedV] in

theorem game0_keyRestorationAdv (b : Bool) :
    LearningWithErrors.game0 (keyHopProblem ring encoding prims)
      (keyRestorationAdv ring encoding prims adv sim b) =
      idealBranch ring encoding prims adv (idealKeygen ring encoding prims)
        (fun _ => sim) b := by

theorem game1_keyRestorationAdv (b : Bool) :
    LearningWithErrors.game1 (keyHopProblem ring encoding prims)
      (keyRestorationAdv ring encoding prims adv sim b) =
      idealBranch ring encoding prims adv (uniformKeygen encoding)
        (fun _ => sim) b := by

noncomputable def keyIdealization (b : Bool) : ℝ :=
  (idealBranch ring encoding prims adv (idealKeygen ring encoding prims)
      (fun _ => sim) b).boolDistAdvantage
    (KEM.anonSetup (MLKEM.asKEMScheme ring encoding prims) >>= KEM.simBranch sim adv)

theorem keyRestoration_le_mlwe_add_keyIdealization (b : Bool) :
    keyRestoration ring encoding prims adv sim b ≤
      LearningWithErrors.advantage (keyHopProblem ring encoding prims)
        (keyRestorationAdv ring encoding prims adv sim b)
        + keyIdealization ring encoding prims adv sim b := by

def keyRestorationMLWE (epsilonRestore : ℝ) : Prop :=
  ∀ b, keyRestoration ring encoding prims adv sim b ≤ epsilonRestore

theorem sprAdv_le_mlwe (b : Bool)
    (epsilonPrim epsilonEnc epsilonRestore : ℝ)
    (hPrim : ∀ b, primitiveIdealization ring encoding prims adv b ≤ epsilonPrim)
    (hEnc : ∀ b, encodingRegularity ring encoding prims adv sim b ≤ epsilonEnc)
    (hRestore : ∀ b, keyRestoration ring encoding prims adv sim b ≤ epsilonRestore) :
    KEM.sprAdv (MLKEM.asKEMScheme ring encoding prims) sim adv b ≤
      epsilonPrim
        + LearningWithErrors.advantage (keyHopProblem ring encoding prims)
            (keyHopAdv ring encoding prims adv b)
        + LearningWithErrors.advantage (ctHopProblem ring encoding prims)
            (ctHopAdv ring encoding prims adv b)
        + epsilonEnc
        + epsilonRestore := by

theorem mlkem768_sprAdv_le_mlwe (b : Bool)
    (epsilonPrim epsilonEnc epsilonRestore : ℝ)
    (hPrim : ∀ b, primitiveIdealization concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768 b ≤ epsilonPrim)
    (hEnc : ∀ b, encodingRegularity concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768
      mlkem768UniformCiphertext b ≤ epsilonEnc)
    (hRestore : ∀ b, keyRestoration concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768
      mlkem768UniformCiphertext b ≤ epsilonRestore) :
    mlkem768KEM.sprAdv mlkem768UniformCiphertext adv768 b ≤
      epsilonPrim
        + LearningWithErrors.advantage mlkem768KeyHopProblem
            (keyHopAdv concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768 b)
        + LearningWithErrors.advantage mlkem768CtHopProblem
            (ctHopAdv concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768 b)
        + epsilonEnc
        + epsilonRestore :=
  sprAdv_le_mlwe concreteNTTRingOps mlkem768Encoding mlkem768Primitives adv768
    mlkem768UniformCiphertext b epsilonPrim epsilonEnc epsilonRestore hPrim hEnc hRestore