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 u → Type*} [MonadLiftT n PMF] [Probability.HasUniformSelectFinset n]
+  {M K C : Type u}
 
 /-- 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) (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,31 @@ 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*) [MonadLiftT m PMF]
+    (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
+  correct : ∀ key, key ∈ PMF.support gen → ∀ message ciphertext,
+    ciphertext ∈ PMF.support (enc key message) → 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)
+noncomputable def EncScheme.ofPure.{u} {m : Type u → Type*}
+    [Monad m] [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 Message Key Ciphertext where
+    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
+  correct key _ message ciphertext hc := by
+    have : ciphertext = enc key message := by simpa 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 u → Type*} [MonadLiftT n PMF] {M K C : Type u}
 
 /-- 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,54 @@ 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 [Probability.HasUniformSelectFinset 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 := Probability.HasUniformSelectFinset.uniformSelectFinset _ hs
+  have hμ : (μ : PMF M) = PMF.uniformOfFinset _ hs := by simp [μ]
   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)
+  rw [Probability.HasUniformSelectFinset.liftM_uniformSelectFinset] at hne_top
+  refine (ENNReal.mul_right_inj hne 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)
     (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₂))
 
 /-- Perfect secrecy requires `|K| ≥ |M|` (Shannon's theorem). -/
-theorem shannonKeySpace [Finite K]
-    (scheme : EncScheme M K C) (h : scheme.PerfectlySecret) :
+theorem shannonKeySpace [Finite K] [Probability.HasUniformSelectFinset n]
+    (scheme : EncScheme n M K C) (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

@@ -37,7 +37,7 @@ namespace Cslib.Crypto.Protocols.PerfectSecrecy
 /-- 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) :=
+    EncScheme PMF (BitVec l) (BitVec l) (BitVec l) :=
   .ofPure (PMF.uniformOfFintype _) (· ^^^ ·) (· ^^^ ·) fun k m => by
     simp [← BitVec.xor_assoc]
 

Cslib/Probability/HasUniform.lean

@@ -0,0 +1,55 @@
+/-
+Copyright (c) 2026 Devon Tuma. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Devon Tuma
+-/
+
+module
+
+public import Cslib.Init
+public import Mathlib.Probability.Distributions.Uniform
+
+/-!
+# Monads with Uniform Selection
+
+This file defines general typeclasses for asserting that a monad `m` with an embedding into `PMF`
+can model computations that lift to uniform distributions in `PMF`.
+
+## Main Definitions
+
+- `Cslib.Probability.HasUniformSelectFinset`: monad supports uniform finset selection
+- `Cslib.Probability.HasUniformSelectFintype`: monad supports uniform finite types
+
+-/
+
+@[expose] public section
+
+namespace Cslib.Probability
+
+universe u v
+
+/-- The monad `m` has a way to model uniform selection over non-empty finsets. -/
+class HasUniformSelectFinset (m : Type u → Type*) [MonadLiftT m PMF] where
+  uniformSelectFinset {α : Type u} (s : Finset α) (hs : s.Nonempty) : m α
+  liftM_uniformSelectFinset {α : Type u} (s : Finset α) (hs : s.Nonempty) :
+    (uniformSelectFinset s hs : PMF α) = PMF.uniformOfFinset s hs
+
+attribute [simp] HasUniformSelectFinset.liftM_uniformSelectFinset
+
+noncomputable instance : HasUniformSelectFinset PMF where
+  uniformSelectFinset := PMF.uniformOfFinset
+  liftM_uniformSelectFinset _ _ := rfl
+
+/-- The monad `m` has a way to model uniform selection over non-empty finsets. -/
+class HasUniformSelectFintype (m : Type u → Type*) [MonadLiftT m PMF] where
+  uniformSelectFintype (α : Type u) [Fintype α] [Nonempty α] : m α
+  liftM_uniformSelectFinset (α : Type u) [Fintype α] [Nonempty α] :
+    (uniformSelectFintype α : PMF α) = PMF.uniformOfFintype α
+
+attribute [simp] HasUniformSelectFinset.liftM_uniformSelectFinset
+
+noncomputable instance : HasUniformSelectFintype PMF where
+  uniformSelectFintype := PMF.uniformOfFintype
+  liftM_uniformSelectFinset _ _ := rfl
+
+end Cslib.Probability

Cslib/Probability/PMF.lean

@@ -130,4 +130,7 @@ theorem posteriorDist_eq_prior_of_outputIndist (p : PMF α) (f : α → PMF β)
   exact ENNReal.mul_div_cancel_right (hmarg ▸ hb')
     (ne_top_of_le_ne_top ENNReal.one_ne_top (PMF.coe_le_one _ _))
 
+@[simp]
+lemma monad_pure_eq_pure (x : α) : (pure x : PMF α) = PMF.pure x := rfl
+
 end Cslib.Probability.PMF