feat(Crypto): Generalize encryption schemes over arbitrary monads

Cslib.lean

@@ -139,4 +139,5 @@ public import Cslib.Logics.Modal.Denotation
 public import Cslib.Logics.Propositional.Defs
 public import Cslib.Logics.Propositional.NaturalDeduction.Basic
 public import Cslib.MachineLearning.PACLearning.Defs
+public import Cslib.Probability.HasUniform
 public import Cslib.Probability.PMF

Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean

@@ -32,19 +32,20 @@ Characterisation theorems for perfect secrecy following
 namespace Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme
 
 universe u
-variable {M K C : Type u}
+variable {n : Type → Type*} [Monad n] [MonadLiftT n PMF] [LawfulMonadLiftT n PMF]
+  [Probability.HasUniformBitVec n] [Probability.LawfulUniformBitVec n] {M K C : Type}
 
 /-- A scheme is perfectly secret iff the ciphertext distribution is
 independent of the plaintext ([KatzLindell2020], Lemma 2.5). -/
-theorem perfectlySecret_iff_ciphertextIndist (scheme : EncScheme M K C) :
+theorem perfectlySecret_iff_ciphertextIndist (scheme : EncScheme n M K C) :
     scheme.PerfectlySecret ↔ scheme.CiphertextIndist :=
   ⟨PerfectSecrecy.ciphertextIndist_of_perfectlySecret scheme,
    PerfectSecrecy.perfectlySecret_of_ciphertextIndist scheme⟩
 
 /-- Perfect secrecy requires `|K| ≥ |M|`
 ([KatzLindell2020], Theorem 2.12). -/
 theorem perfectlySecret_keySpace_ge [Finite K]
-    (scheme : EncScheme M K C) (h : scheme.PerfectlySecret) :
+    (scheme : EncScheme n M K C) [scheme.Correct] (h : scheme.PerfectlySecret) :
     Nat.card K ≥ Nat.card M :=
   PerfectSecrecy.shannonKeySpace scheme h
 

Cslib/Crypto/Protocols/PerfectSecrecy/Defs.lean

@@ -36,33 +36,33 @@ Core definitions for perfect secrecy following [KatzLindell2020], Chapter 2.
 namespace Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme
 
 universe u
-variable {M K C : Type u}
+variable {n : Type u → Type*} [MonadLiftT n PMF] {M K C : Type u}
 
 /-- The distribution of `Enc_K(m)` when `K ← Gen`. -/
-noncomputable def ciphertextDist (scheme : EncScheme M K C) (m : M) : PMF C := do
+noncomputable def ciphertextDist (scheme : EncScheme n M K C) (m : M) : PMF C := do
   scheme.enc (← scheme.gen) m
 
 /-- Joint distribution of `(M, C)` given a message prior. -/
-noncomputable def jointDist (scheme : EncScheme M K C) (msgDist : PMF M) : PMF (M × C) := do
+noncomputable def jointDist (scheme : EncScheme n M K C) (msgDist : n M) : PMF (M × C) := do
   let m ← msgDist
   return (m, ← scheme.ciphertextDist m)
 
 /-- Marginal ciphertext distribution given a message prior. -/
-noncomputable def marginalCiphertextDist (scheme : EncScheme M K C)
-    (msgDist : PMF M) : PMF C := do
+noncomputable def marginalCiphertextDist (scheme : EncScheme n M K C)
+    (msgDist : n M) : PMF C := do
   scheme.ciphertextDist (← msgDist)
 
 /-- The posterior message distribution `Pr[M | C = c]` as a probability
 distribution, given a message prior and a ciphertext in the support of
 the marginal distribution. -/
-noncomputable def posteriorMsgDist (scheme : EncScheme M K C)
-    (msgDist : PMF M) (c : C)
+noncomputable def posteriorMsgDist (scheme : EncScheme n M K C)
+    (msgDist : n M) (c : C)
     (hc : c ∈ (scheme.marginalCiphertextDist msgDist).support) : PMF M :=
   Cslib.Probability.PMF.posteriorDist msgDist scheme.ciphertextDist c hc
 
 @[simp]
-theorem posteriorMsgDist_apply (scheme : EncScheme M K C)
-    (msgDist : PMF M) (c : C)
+theorem posteriorMsgDist_apply (scheme : EncScheme n M K C)
+    (msgDist : n M) (c : C)
     (hc : c ∈ (scheme.marginalCiphertextDist msgDist).support) (m : M) :
     scheme.posteriorMsgDist msgDist c hc m =
       scheme.jointDist msgDist (m, c) / scheme.marginalCiphertextDist msgDist c :=
@@ -71,14 +71,14 @@ theorem posteriorMsgDist_apply (scheme : EncScheme M K C)
 /-- An encryption scheme is perfectly secret if the posterior message
 distribution equals the prior for every ciphertext with positive probability
 ([KatzLindell2020], Definition 2.3). -/
-def PerfectlySecret (scheme : EncScheme M K C) : Prop :=
-  ∀ (msgDist : PMF M) (c : C)
+def PerfectlySecret (scheme : EncScheme n M K C) : Prop :=
+  ∀ (msgDist : n M) (c : C)
     (hc : c ∈ (scheme.marginalCiphertextDist msgDist).support),
     scheme.posteriorMsgDist msgDist c hc = msgDist
 
 /-- Ciphertext indistinguishability: the ciphertext distribution is the same
 for all messages ([KatzLindell2020], Lemma 2.5). -/
-def CiphertextIndist (scheme : EncScheme M K C) : Prop :=
+def CiphertextIndist (scheme : EncScheme n M K C) : Prop :=
   ∀ m₀ m₁ : M, scheme.ciphertextDist m₀ = scheme.ciphertextDist m₁
 
 end Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme

Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean

@@ -1,13 +1,12 @@
 /-
 Copyright (c) 2026 Samuel Schlesinger. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Samuel Schlesinger
+Authors: Samuel Schlesinger, Devon Tuma
 -/
 
 module
 
-public import Cslib.Init
-public import Mathlib.Probability.ProbabilityMassFunction.Monad
+public import Cslib.Probability.PMF
 
 /-!
 # Private-Key Encryption Schemes (Information-Theoretic)
@@ -35,27 +34,39 @@ namespace Cslib.Crypto.Protocols.PerfectSecrecy
 A private-key encryption scheme over message space `M`, key space `K`,
 and ciphertext space `C` ([KatzLindell2020], Definition 2.1).
 -/
-structure EncScheme (Message Key Ciphertext : Type*) where
+structure EncScheme (m : Type u → Type*)
+    (Message Key Ciphertext : Type u) where
   /-- Probabilistic key generation. -/
-  gen : PMF Key
+  gen : m Key
   /-- (Possibly randomized) encryption. -/
-  enc (key : Key) (message : Message) : PMF Ciphertext
+  enc (key : Key) (message : Message) : m Ciphertext
   /-- Deterministic decryption. -/
   dec (key : Key) (ciphertext : Ciphertext) : Message
-  /-- Decryption inverts encryption for all keys in the support of `gen`. -/
-  correct : ∀ key, key ∈ gen.support → ∀ message ciphertext,
-    ciphertext ∈ (enc key message).support → dec key ciphertext = message
 
-/-- Build an encryption scheme from deterministic pure encryption/decryption
-where decryption is a left inverse of encryption for every key. -/
-noncomputable def EncScheme.ofPure.{u} {Message Key Ciphertext : Type u} (gen : PMF Key)
-    (enc : Key → Message → Ciphertext) (dec : Key → Ciphertext → Message)
-    (h : ∀ key, Function.LeftInverse (dec key) (enc key)) :
-    EncScheme Message Key Ciphertext where
+variable {m : Type u → Type*} [Monad m]
+
+/-- Decryption inverts encryption for all keys in the support of `gen`. -/
+class EncScheme.Correct [MonadLiftT m PMF] {Message Key Ciphertext : Type u}
+    (scheme : EncScheme m Message Key Ciphertext) : Prop where
+  dec_enc : ∀ key, key ∈ PMF.support scheme.gen → ∀ message ciphertext,
+    ciphertext ∈ PMF.support (scheme.enc key message) → scheme.dec key ciphertext = message
+
+/-- Build an encryption scheme from deterministic pure encryption/decryption. -/
+def EncScheme.ofPure {Message Key Ciphertext : Type u} (gen : m Key)
+    (enc : Key → Message → Ciphertext) (dec : Key → Ciphertext → Message) :
+    EncScheme m Message Key Ciphertext where
   gen := gen
-  enc key message := PMF.pure (enc key message)
+  enc key message := pure (enc key message)
   dec := dec
-  correct key _ message _ hc := by
-    rw [PMF.mem_support_pure_iff] at hc; subst hc; exact h key message
+
+/-- `EncScheme.ofPure` is correct if decryption is a left inverse of encryption for every key. -/
+lemma EncScheme.ofPure.Correct [MonadLiftT m PMF] [LawfulMonadLiftT m PMF]
+    {Message Key Ciphertext : Type u} (gen : m Key)
+    (enc : Key → Message → Ciphertext) (dec : Key → Ciphertext → Message)
+    (h : ∀ key, Function.LeftInverse (dec key) (enc key)) :
+    (EncScheme.ofPure gen enc dec).Correct where
+  dec_enc key _ message ciphertext hc := by
+    have : ciphertext = enc key message := by simpa [EncScheme.ofPure] using hc
+    subst this; exact h key message
 
 end Cslib.Crypto.Protocols.PerfectSecrecy

Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean

@@ -1,12 +1,13 @@
 /-
 Copyright (c) 2026 Samuel Schlesinger. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Samuel Schlesinger
+Authors: Samuel Schlesinger, Devon Tuma
 -/
 
 module
 
 public import Cslib.Crypto.Protocols.PerfectSecrecy.Defs
+public import Cslib.Probability.HasUniform
 public import Mathlib.Probability.Distributions.Uniform
 
 /-!
@@ -26,26 +27,26 @@ namespace Cslib.Crypto.Protocols.PerfectSecrecy
 open PMF ENNReal
 
 universe u
-variable {M K C : Type u}
+variable {n : Type → Type*} [MonadLiftT n PMF] {M K C : Type}
 
 /-- The joint distribution at `(m, c)` equals `msgDist m * ciphertextDist m c`. -/
-theorem jointDist_eq (scheme : EncScheme M K C) (msgDist : PMF M)
+theorem jointDist_eq (scheme : EncScheme n M K C) (msgDist : n M)
     (m : M) (c : C) :
-    scheme.jointDist msgDist (m, c) = msgDist m * scheme.ciphertextDist m c :=
+    scheme.jointDist msgDist (m, c) = (msgDist : PMF M) m * scheme.ciphertextDist m c :=
   Cslib.Probability.PMF.bind_pair_apply msgDist scheme.ciphertextDist m c
 
 /-- Summing the joint distribution over messages gives the marginal ciphertext distribution. -/
-theorem jointDist_tsum_fst (scheme : EncScheme M K C) (msgDist : PMF M) (c : C) :
+theorem jointDist_tsum_fst (scheme : EncScheme n M K C) (msgDist : n M) (c : C) :
     ∑' m, scheme.jointDist msgDist (m, c) = scheme.marginalCiphertextDist msgDist c :=
   Cslib.Probability.PMF.bind_pair_tsum_fst msgDist scheme.ciphertextDist c
 
 /-- Perfect secrecy is equivalent to message-ciphertext independence.
 The two formulations are related by multiplying/dividing by `marginal(c)`. -/
-theorem perfectlySecret_iff_indep (scheme : EncScheme M K C) :
+theorem perfectlySecret_iff_indep (scheme : EncScheme n M K C) :
     scheme.PerfectlySecret ↔
-    ∀ (msgDist : PMF M) (m : M) (c : C),
+    ∀ (msgDist : n M) (m : M) (c : C),
       scheme.jointDist msgDist (m, c) =
-        msgDist m * scheme.marginalCiphertextDist msgDist c := by
+        (msgDist : PMF M) m * scheme.marginalCiphertextDist msgDist c := by
   constructor
   · intro h msgDist m c
     by_cases hc : (scheme.marginalCiphertextDist msgDist) c = 0
@@ -64,11 +65,11 @@ theorem perfectlySecret_iff_indep (scheme : EncScheme M K C) :
       (ne_top_of_le_ne_top one_ne_top (PMF.coe_le_one _ c))]
 
 /-- Ciphertext indistinguishability implies message-ciphertext independence. -/
-theorem indep_of_ciphertextIndist (scheme : EncScheme M K C)
+theorem indep_of_ciphertextIndist (scheme : EncScheme n M K C)
     (h : scheme.CiphertextIndist)
-    (msgDist : PMF M) (m : M) (c : C) :
+    (msgDist : n M) (m : M) (c : C) :
     scheme.jointDist msgDist (m, c) =
-      msgDist m * scheme.marginalCiphertextDist msgDist c := by
+      (msgDist : PMF M) m * scheme.marginalCiphertextDist msgDist c := by
   rw [jointDist_eq]; congr 1
   change scheme.ciphertextDist m c =
     PMF.bind msgDist (fun m' => scheme.ciphertextDist m') c
@@ -77,50 +78,57 @@ theorem indep_of_ciphertextIndist (scheme : EncScheme M K C)
   rw [ENNReal.tsum_mul_right, PMF.tsum_coe, one_mul]
 
 /-- Ciphertext indistinguishability implies perfect secrecy. -/
-theorem perfectlySecret_of_ciphertextIndist (scheme : EncScheme M K C)
+theorem perfectlySecret_of_ciphertextIndist (scheme : EncScheme n M K C)
     (h : scheme.CiphertextIndist) :
     scheme.PerfectlySecret :=
   (perfectlySecret_iff_indep scheme).mpr (fun msgDist m c =>
     indep_of_ciphertextIndist scheme h msgDist m c)
 
-/-- Perfect secrecy implies ciphertext indistinguishability. -/
-theorem ciphertextIndist_of_perfectlySecret (scheme : EncScheme M K C)
-    (h : scheme.PerfectlySecret) :
+/-- Perfect secrecy implies ciphertext indistinguishability.
+Note we need `n` to support uniform selection for the proof to work -/
+theorem ciphertextIndist_of_perfectlySecret [Monad n] [LawfulMonadLiftT n PMF]
+    [Probability.HasUniformBitVec n] [Probability.LawfulUniformBitVec n]
+    (scheme : EncScheme n M K C) (h : scheme.PerfectlySecret) :
     scheme.CiphertextIndist := by
   classical
   rw [perfectlySecret_iff_indep] at h
   intro m₀ m₁; ext c
   have hs : ({m₀, m₁} : Finset M).Nonempty := ⟨m₀, Finset.mem_insert_self ..⟩
-  set μ := PMF.uniformOfFinset _ hs
+  let μ : n M := do return bif (← Probability.uniformBool) then m₀ else m₁
+  have hμ : (μ : PMF M) = PMF.uniformOfFinset _ hs := by
+    simp only [μ, liftM_bind, liftM_pure, Cslib.Probability.liftM_uniformBool]
+    exact Cslib.Probability.PMF.uniformOfFintype_bool_bind_ite m₀ m₁
   suffices key : ∀ m ∈ ({m₀, m₁} : Finset M),
       scheme.ciphertextDist m c = scheme.marginalCiphertextDist μ c by
     exact (key m₀ (by simp)).trans (key m₁ (by simp)).symm
   intro m hm
   have hne := (PMF.mem_support_uniformOfFinset_iff hs m).mpr hm
   have hne_top := ne_top_of_le_ne_top one_ne_top (PMF.coe_le_one μ m)
-  exact (ENNReal.mul_right_inj hne hne_top).mp (by rw [← jointDist_eq]; exact h μ m c)
+  refine (ENNReal.mul_right_inj hne (hμ ▸ hne_top)).mp ?_
+  exact (hμ ▸ jointDist_eq scheme _ m c).symm.trans (hμ ▸ h μ m c)
 
 /-- If each message maps to a key that encrypts it to a common ciphertext,
 then the key assignment is injective (by correctness of decryption). -/
-lemma encrypt_key_injective (scheme : EncScheme M K C)
+lemma encrypt_key_injective (scheme : EncScheme n M K C) [scheme.Correct]
     (f : M → K) (c₀ : C)
-    (hf_mem : ∀ m, f m ∈ scheme.gen.support)
-    (hf_enc : ∀ m, c₀ ∈ (scheme.enc (f m) m).support) :
+    (hf_mem : ∀ m, f m ∈ PMF.support scheme.gen)
+    (hf_enc : ∀ m, c₀ ∈ PMF.support (scheme.enc (f m) m)) :
     Function.Injective f :=
   fun m₁ m₂ heq =>
-    (scheme.correct _ (hf_mem m₁) m₁ c₀ (hf_enc m₁)).symm.trans
-      (heq ▸ scheme.correct _ (hf_mem m₂) m₂ c₀ (hf_enc m₂))
+    (EncScheme.Correct.dec_enc _ (hf_mem m₁) m₁ c₀ (hf_enc m₁)).symm.trans
+      (heq ▸ EncScheme.Correct.dec_enc _ (hf_mem m₂) m₂ c₀ (hf_enc m₂))
 
 /-- Perfect secrecy requires `|K| ≥ |M|` (Shannon's theorem). -/
-theorem shannonKeySpace [Finite K]
-    (scheme : EncScheme M K C) (h : scheme.PerfectlySecret) :
+theorem shannonKeySpace [Finite K] [Monad n] [LawfulMonadLiftT n PMF]
+    [Probability.HasUniformBitVec n] [Probability.LawfulUniformBitVec n]
+    (scheme : EncScheme n M K C) [scheme.Correct] (h : scheme.PerfectlySecret) :
     Nat.card K ≥ Nat.card M := by
   classical
   have hci := ciphertextIndist_of_perfectlySecret scheme h
   by_cases hM : IsEmpty M; · simp
   obtain ⟨m₀⟩ := not_isEmpty_iff.mp hM
   obtain ⟨c₀, hc₀⟩ := (scheme.ciphertextDist m₀).support_nonempty
-  have key_exists : ∀ m, ∃ k ∈ scheme.gen.support, c₀ ∈ (scheme.enc k m).support := by
+  have key_exists : ∀ m, ∃ k ∈ PMF.support scheme.gen, c₀ ∈ PMF.support (scheme.enc k m) := by
     intro m
     exact (PMF.mem_support_bind_iff _ _ _).mp
       (show c₀ ∈ (scheme.ciphertextDist m).support by rw [hci m m₀]; exact hc₀)

Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean

@@ -9,6 +9,8 @@ module
 public import Cslib.Crypto.Protocols.PerfectSecrecy.Basic
 public import Cslib.Crypto.Protocols.PerfectSecrecy.Internal.OneTimePad
 public import Mathlib.Probability.Distributions.Uniform
+public import Mathlib.Data.PFunctor.Univariate.Basic
+public import Cslib.Foundations.Control.Monad.Free
 
 /-!
 # One-Time Pad
@@ -34,18 +36,65 @@ The one-time pad (Vernam cipher) over `BitVec l`
 
 namespace Cslib.Crypto.Protocols.PerfectSecrecy
 
+variable (n : Type → Type*) [Monad n] [Probability.HasUniformBitVec n]
+
 /-- The one-time pad over `l`-bit strings. Encryption and decryption
 are XOR ([KatzLindell2020], Construction 2.9). -/
-noncomputable def otp (l : ℕ) :
-    EncScheme (BitVec l) (BitVec l) (BitVec l) :=
-  .ofPure (PMF.uniformOfFintype _) (· ^^^ ·) (· ^^^ ·) fun k m => by
-    simp [← BitVec.xor_assoc]
+def otp (l : ℕ) : EncScheme n (BitVec l) (BitVec l) (BitVec l) :=
+  let gen : n (BitVec l) := Probability.HasUniformBitVec.uniformBitVec l
+  .ofPure gen (· ^^^ ·) (· ^^^ ·)
+
+instance (l : ℕ) [MonadLiftT n PMF] [LawfulMonadLiftT n PMF] : (otp n l).Correct :=
+  EncScheme.ofPure.Correct _ _ _ fun k m => by simp [← BitVec.xor_assoc]
 
 /-- The one-time pad is perfectly secret ([KatzLindell2020], Theorem 2.10). -/
-theorem otp_perfectlySecret (l : ℕ) : (otp l).PerfectlySecret :=
+theorem otp_perfectlySecret (l : ℕ) [MonadLiftT n PMF] [LawfulMonadLiftT n PMF]
+    [Probability.LawfulUniformBitVec n] : (otp n l).PerfectlySecret :=
   (EncScheme.perfectlySecret_iff_ciphertextIndist _).mpr fun m₀ m₁ => by
-    simp only [EncScheme.ciphertextDist, otp]
+    simp only [EncScheme.ciphertextDist, otp, EncScheme.ofPure,
+      Probability.LawfulUniformBitVec.liftM_uniformBitVec, liftM_pure,
+      Probability.PMF.monad_pure_eq_pure]
     exact (OTP.otp_ciphertextDist_eq_uniform l m₀).trans
       (OTP.otp_ciphertextDist_eq_uniform l m₁).symm
 
+section computable_otp
+
+/-- Monad to model computations with access to uniform bitvector selection. -/
+private abbrev UniformBitVecM : Type → Type 1 :=
+  FreeM (PFunctor.Obj {A := ℕ, B := BitVec})
+
+/-- Semantics assigning distributions on `BitVec n` to selection queries. -/
+private noncomputable instance : MonadLift UniformBitVecM PMF where
+  monadLift := FreeM.liftM fun | ⟨n, f⟩ => (PMF.uniformOfFintype (BitVec n)).map f
+
+private instance : LawfulMonadLiftT UniformBitVecM PMF where
+  monadLift_pure := FreeM.liftM_pure _
+  monadLift_bind := FreeM.liftM_bind _
+
+/-- The node `n` with identity as the continuation models uniform bitvector selection. -/
+private instance : Probability.HasUniformBitVec UniformBitVecM where
+  uniformBitVec n := FreeM.lift ⟨n, id⟩
+
+private noncomputable instance : Probability.LawfulUniformBitVec UniformBitVecM where
+  liftM_uniformBitVec n := by
+    change (FreeM.lift ⟨n, id⟩ : UniformBitVecM (BitVec n)).liftM _ = _
+    simp [PMF.map_id, Probability.PMF.monad_bind_eq_bind]; rfl
+
+/-- Interpret the free monad in `IO` to allow execution. -/
+private instance : MonadLift UniformBitVecM IO where
+  monadLift := FreeM.liftM fun | ⟨n, f⟩ => do return f (BitVec.ofNat n (← IO.rand 0 (2 ^ n)))
+
+/-- Takes a message and runs the full `otp` algorithm, returning all the values in a tuple. -/
+private def run_otp (message : BitVec l) :
+    UniformBitVecM (BitVec l × BitVec l × BitVec l × BitVec l) := do
+  let key ← (otp UniformBitVecM l).gen
+  let ciphertext ← (otp UniformBitVecM l).enc key message
+  let message' := (otp UniformBitVecM l).dec key ciphertext
+  return (message, key, ciphertext, message')
+
+-- #eval run_otp 0#6
+-- #eval run_otp 32#6
+
+end computable_otp
+
 end Cslib.Crypto.Protocols.PerfectSecrecy