From 3c9092ba3ca1a369c84c7e64203e61c766c78fa5 Mon Sep 17 00:00:00 2001
From: Devon Tuma <dtumad@gmail.com>
Date: Wed, 20 May 2026 12:24:21 -0500
Subject: [PATCH 1/6] feat: Generalize encryption schemes over arbitrary monads

---
 Cslib.lean                                    |  1 +
 .../Protocols/PerfectSecrecy/Basic.lean       |  4 +-
 .../Crypto/Protocols/PerfectSecrecy/Defs.lean | 24 ++++----
 .../Protocols/PerfectSecrecy/Encryption.lean  | 29 +++++-----
 .../Internal/PerfectSecrecy.lean              | 52 ++++++++++--------
 .../Protocols/PerfectSecrecy/OneTimePad.lean  |  2 +-
 Cslib/Probability/HasUniform.lean             | 55 +++++++++++++++++++
 Cslib/Probability/PMF.lean                    |  3 +
 8 files changed, 119 insertions(+), 51 deletions(-)
 create mode 100644 Cslib/Probability/HasUniform.lean

diff --git a/Cslib.lean b/Cslib.lean
index b504973c3..c3ccc596e 100644
--- a/Cslib.lean
+++ b/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
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
index c5f5a29ef..bfd637422 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
@@ -36,7 +36,7 @@ variable {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 PMF M K C) :
     scheme.PerfectlySecret ↔ scheme.CiphertextIndist :=
   ⟨PerfectSecrecy.ciphertextIndist_of_perfectlySecret scheme,
    PerfectSecrecy.perfectlySecret_of_ciphertextIndist scheme⟩
@@ -44,7 +44,7 @@ theorem perfectlySecret_iff_ciphertextIndist (scheme : EncScheme M K C) :
 /-- 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 PMF M K C) (h : scheme.PerfectlySecret) :
     Nat.card K ≥ Nat.card M :=
   PerfectSecrecy.shannonKeySpace scheme h
 
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Defs.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Defs.lean
index e7d824dd6..7531b4b31 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Defs.lean
+++ b/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
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean
index a6896bcdf..0d49676f2 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean
+++ b/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
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
index 66656e8ba..bd8c27b11 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
+++ b/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,55 @@ 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 to fix the monad of the encryption scheme to `PMF` to define `μ`,
+this could be avoided with a class to assert a monad has uniform selction. -/
+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₀)
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
index ed53950f6..9275a7257 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
+++ b/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]
 
diff --git a/Cslib/Probability/HasUniform.lean b/Cslib/Probability/HasUniform.lean
new file mode 100644
index 000000000..0dbbcc56c
--- /dev/null
+++ b/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
diff --git a/Cslib/Probability/PMF.lean b/Cslib/Probability/PMF.lean
index d20be22be..f7dd57c72 100644
--- a/Cslib/Probability/PMF.lean
+++ b/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

From 8e4475e00c9b526082aaf86a30c7c239226a335a Mon Sep 17 00:00:00 2001
From: Devon Tuma <dtumad@gmail.com>
Date: Wed, 20 May 2026 12:26:18 -0500
Subject: [PATCH 2/6] stale comment

---
 .../Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean      | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
index bd8c27b11..15a6d1b96 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
@@ -85,8 +85,7 @@ theorem perfectlySecret_of_ciphertextIndist (scheme : EncScheme n M K C)
     indep_of_ciphertextIndist scheme h msgDist m c)
 
 /-- Perfect secrecy implies ciphertext indistinguishability.
-Note we need to fix the monad of the encryption scheme to `PMF` to define `μ`,
-this could be avoided with a class to assert a monad has uniform selction. -/
+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

From 771ecc72b6c8ea054d4ba2dfd1e352c1b803010e Mon Sep 17 00:00:00 2001
From: Devon Tuma <dtumad@gmail.com>
Date: Wed, 20 May 2026 12:34:29 -0500
Subject: [PATCH 3/6] missed file generalizations

---
 Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
index bfd637422..cc2c02b79 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
@@ -32,11 +32,12 @@ 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 PMF 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⟩
@@ -44,7 +45,7 @@ theorem perfectlySecret_iff_ciphertextIndist (scheme : EncScheme PMF M K C) :
 /-- Perfect secrecy requires `|K| ≥ |M|`
 ([KatzLindell2020], Theorem 2.12). -/
 theorem perfectlySecret_keySpace_ge [Finite K]
-    (scheme : EncScheme PMF 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
 

From 35b6f4caf7e1a7ecd3dd828691f6491e61f6b3f1 Mon Sep 17 00:00:00 2001
From: Devon Tuma <dtumad@gmail.com>
Date: Sat, 23 May 2026 17:31:52 -0500
Subject: [PATCH 4/6] Start making an example better

---
 .../Protocols/PerfectSecrecy/Basic.lean       |  6 +-
 .../Protocols/PerfectSecrecy/Encryption.lean  | 31 ++++---
 .../Internal/PerfectSecrecy.lean              | 25 +++---
 .../Protocols/PerfectSecrecy/OneTimePad.lean  | 59 +++++++++++--
 Cslib/Probability/HasUniform.lean             | 88 ++++++++++++++-----
 Cslib/Probability/PMF.lean                    | 35 ++++++++
 6 files changed, 192 insertions(+), 52 deletions(-)

diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
index cc2c02b79..7640982db 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
@@ -32,8 +32,8 @@ Characterisation theorems for perfect secrecy following
 namespace Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme
 
 universe u
-variable {n : Type u → Type*} [MonadLiftT n PMF] [Probability.HasUniformSelectFinset n]
-  {M K C : Type u}
+variable {n : Type → Type*} [Monad n] [MonadLiftT n PMF] [LawfulMonadLiftT n PMF]
+  [Probability.HasUniformBitVec n] {M K C : Type}
 
 /-- A scheme is perfectly secret iff the ciphertext distribution is
 independent of the plaintext ([KatzLindell2020], Lemma 2.5). -/
@@ -45,7 +45,7 @@ theorem perfectlySecret_iff_ciphertextIndist (scheme : EncScheme n M K C) :
 /-- Perfect secrecy requires `|K| ≥ |M|`
 ([KatzLindell2020], Theorem 2.12). -/
 theorem perfectlySecret_keySpace_ge [Finite K]
-    (scheme : EncScheme n 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
 
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean
index 0d49676f2..06d021980 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean
@@ -34,7 +34,7 @@ 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 (m : Type u → Type*) [MonadLiftT m PMF]
+structure EncScheme (m : Type u → Type*)
     (Message Key Ciphertext : Type u) where
   /-- Probabilistic key generation. -/
   gen : m Key
@@ -42,23 +42,32 @@ structure EncScheme (m : Type u → Type*) [MonadLiftT m PMF]
   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 ∈ PMF.support gen → ∀ message ciphertext,
-    ciphertext ∈ PMF.support (enc key message) → dec key ciphertext = message
+
+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) 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
 where decryption is a left inverse of encryption for every 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)) :
+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 := pure (enc key message)
   dec := dec
-  correct key _ message ciphertext hc := by
-    have : ciphertext = enc key message := by simpa using hc
+
+@[reducible]
+def 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
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
index 15a6d1b96..e30dd033b 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
@@ -27,7 +27,7 @@ namespace Cslib.Crypto.Protocols.PerfectSecrecy
 open PMF ENNReal
 
 universe u
-variable {n : Type u → Type*} [MonadLiftT n PMF] {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 n M K C) (msgDist : n M)
@@ -86,39 +86,42 @@ theorem perfectlySecret_of_ciphertextIndist (scheme : EncScheme n M K C)
 
 /-- 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]
+theorem ciphertextIndist_of_perfectlySecret [Monad n] [LawfulMonadLiftT n PMF]
+    [Probability.HasUniformBitVec 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 ..⟩
-  let μ : n M := Probability.HasUniformSelectFinset.uniformSelectFinset _ hs
-  have hμ : (μ : PMF M) = PMF.uniformOfFinset _ hs := by simp [μ]
+  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)
-  rw [Probability.HasUniformSelectFinset.liftM_uniformSelectFinset] at hne_top
-  refine (ENNReal.mul_right_inj hne hne_top).mp ?_
+  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 n 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 ∈ 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] [Probability.HasUniformSelectFinset n]
-    (scheme : EncScheme n M K C) (h : scheme.PerfectlySecret) :
+theorem shannonKeySpace [Finite K] [Monad n] [LawfulMonadLiftT n PMF]
+    [Probability.HasUniformBitVec 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
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
index 9275a7257..e294e5e3e 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
+++ b/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,63 @@ The one-time pad (Vernam cipher) over `BitVec l`
 
 namespace Cslib.Crypto.Protocols.PerfectSecrecy
 
+variable (n : Type → Type*) [Monad n] [MonadLiftT n PMF] [LawfulMonadLiftT n PMF]
+  [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 PMF (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 : ℕ) : (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 : ℕ) : (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.HasUniformBitVec.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 executable_otp
+
+/-- Monad to model computations with access to uniform bitvector selection -/
+private abbrev UniformBitVecM : Type → Type 1 :=
+  FreeM (PFunctor.Obj {A := ℕ, B := BitVec})
+
+/-- Currently this can't be made computable because we use `PMF` for probabilities.
+Executing `run_otp -/
+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 _
+
+private instance : Probability.HasUniformBitVec UniformBitVecM where
+  uniformBitVec n := FreeM.lift ⟨n, id⟩
+  liftM_uniformBitVec n := by
+    simp [liftM, monadLift, MonadLift.monadLift, bind, PMF.map_id]; rfl
+
+private instance : MonadLift UniformBitVecM IO where
+  monadLift := FreeM.liftM fun | ⟨n, f⟩ => do
+    let m ← IO.rand 0 (2 ^ n)
+    return f (BitVec.ofNat n m)
+
+private def foo (n : ℕ) : UniformBitVecM (BitVec n) := do
+  let x ← Probability.HasUniformBitVec.uniformBitVec n
+  return x
+
+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')
+
+end executable_otp
+
 end Cslib.Crypto.Protocols.PerfectSecrecy
diff --git a/Cslib/Probability/HasUniform.lean b/Cslib/Probability/HasUniform.lean
index 0dbbcc56c..f13095288 100644
--- a/Cslib/Probability/HasUniform.lean
+++ b/Cslib/Probability/HasUniform.lean
@@ -7,7 +7,9 @@ Authors: Devon Tuma
 module
 
 public import Cslib.Init
+public import Cslib.Probability.PMF
 public import Mathlib.Probability.Distributions.Uniform
+public import Mathlib.Data.FinEnum
 
 /-!
 # Monads with Uniform Selection
@@ -17,8 +19,9 @@ 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
+- `Cslib.Probability.HasUniformBitVec`: monad supports uniform `BitVec n`
+- `Cslib.Probability.HasUniformFinset`: monad supports uniform selection from a non-empty finset
+- `Cslib.Probability.HasUniformFintype`: monad supports uniform selection from a finite type
 
 -/
 
@@ -28,28 +31,71 @@ 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
+/-- The monad `m` has a way to model uniform selection over `BitVec n`. -/
+class HasUniformBitVec (m : Type → Type*) [MonadLiftT m PMF] where
+  uniformBitVec (n : ℕ) : m (BitVec n)
+  liftM_uniformBitVec (n : ℕ) : (uniformBitVec n : PMF _) = PMF.uniformOfFintype (BitVec n)
 
-attribute [simp] HasUniformSelectFinset.liftM_uniformSelectFinset
+attribute [simp] HasUniformBitVec.liftM_uniformBitVec
 
-noncomputable instance : HasUniformSelectFinset PMF where
-  uniformSelectFinset := PMF.uniformOfFinset
-  liftM_uniformSelectFinset _ _ := rfl
+/-- A uniform `Bool` derived from a single uniform bit. -/
+def uniformBool {m : Type → Type*} [Monad m] [MonadLiftT m PMF]
+    [HasUniformBitVec m] : m Bool :=
+  do return (← HasUniformBitVec.uniformBitVec 1) == 0
 
 /-- 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
+class HasUniformFinset (m : Type u → Type*) [MonadLiftT m PMF] where
+  uniformFinset {α : Type u} (s : Finset α) (hs : s.Nonempty) : m α
+  liftM_uniformFinset {α : Type u} (s : Finset α) (hs : s.Nonempty) :
+    (uniformFinset s hs : PMF α) = PMF.uniformOfFinset s hs
+
+attribute [simp] HasUniformFinset.liftM_uniformFinset
+
+instance {m} [MonadLiftT m PMF] [HasUniformFinset m] : HasUniformBitVec m where
+  uniformBitVec n :=
+    HasUniformFinset.uniformFinset (Finset.univ : Finset (BitVec n)) Finset.univ_nonempty
+  liftM_uniformBitVec _ := by
+    rw [HasUniformFinset.liftM_uniformFinset]; rfl
+
+/-- The monad `m` has a way to model uniform selection over inhabited finite types. -/
+class HasUniformFintype (m : Type u → Type*) [MonadLiftT m PMF] where
+  uniformFintype (α : Type u) [Fintype α] [Nonempty α] : m α
+  liftM_uniformFintype (α : Type u) [Fintype α] [Nonempty α] :
+    (uniformFintype α : PMF α) = PMF.uniformOfFintype α
+
+attribute [simp] HasUniformFintype.liftM_uniformFintype
+
+instance {m} [Monad m] [MonadLiftT m PMF] [LawfulMonadLiftT m PMF]
+    [HasUniformFintype m] : HasUniformFinset m where
+  uniformFinset {α} s hs :=
+    letI : Nonempty {x // x ∈ s} := hs.coe_sort
+    HasUniformFintype.uniformFintype {x // x ∈ s} >>= fun x => pure x.val
+  liftM_uniformFinset {α} s hs := by
+    letI : Nonempty {x // x ∈ s} := hs.coe_sort
+    simp only [liftM_bind, liftM_pure, HasUniformFintype.liftM_uniformFintype]
+    change (PMF.uniformOfFintype {x // x ∈ s}).bind (PMF.pure ∘ Subtype.val) = _
+    rw [PMF.bind_pure_comp]
+    exact Cslib.Probability.PMF.map_subtypeVal_uniformOfFintype s hs
+
+noncomputable instance : HasUniformFintype PMF where
+  uniformFintype := PMF.uniformOfFintype
+  liftM_uniformFintype _ _ := rfl
+
+/-- The equivalence between `BitVec 1` and `Bool` via the predicate `· == 0`. -/
+private def bitVecOneEquivBool : BitVec 1 ≃ Bool where
+  toFun bv := bv == 0
+  invFun b := bif b then 0 else 1
+  left_inv := by decide
+  right_inv := by decide
+
+/-- The PMF coercion of `uniformBool` is the uniform distribution on `Bool`. -/
+@[simp]
+theorem liftM_uniformBool {m : Type → Type*} [Monad m] [MonadLiftT m PMF]
+    [LawfulMonadLiftT m PMF] [HasUniformBitVec m] :
+    (liftM (uniformBool : m Bool) : PMF Bool) = PMF.uniformOfFintype Bool := by
+  simp only [uniformBool, liftM_bind, liftM_pure, HasUniformBitVec.liftM_uniformBitVec]
+  change (PMF.uniformOfFintype (BitVec 1)).bind (PMF.pure ∘ bitVecOneEquivBool) = _
+  rw [PMF.bind_pure_comp]
+  exact Cslib.Probability.PMF.uniformOfFintype_map_equiv bitVecOneEquivBool
 
 end Cslib.Probability
diff --git a/Cslib/Probability/PMF.lean b/Cslib/Probability/PMF.lean
index f7dd57c72..9faaaf952 100644
--- a/Cslib/Probability/PMF.lean
+++ b/Cslib/Probability/PMF.lean
@@ -78,6 +78,41 @@ theorem uniformOfFintype_map_equiv {γ : Type v} [Fintype α] [Fintype γ] [None
       simpa using congrArg e.symm h.symm
     · simp
 
+/-- Pushing forward a uniform distribution on the subtype `↥s` along `Subtype.val`
+gives the uniform distribution on the underlying finset `s`. -/
+theorem map_subtypeVal_uniformOfFintype (s : Finset α) (hs : s.Nonempty) :
+    haveI : Nonempty {x // x ∈ s} := hs.coe_sort
+    (PMF.uniformOfFintype {x // x ∈ s}).map Subtype.val = PMF.uniformOfFinset s hs := by
+  classical
+  haveI : Nonempty {x // x ∈ s} := hs.coe_sort
+  ext a
+  rw [PMF.map_apply, PMF.uniformOfFinset_apply]
+  by_cases ha : a ∈ s
+  · rw [tsum_eq_single ⟨a, ha⟩, if_pos rfl, PMF.uniformOfFintype_apply,
+      Fintype.card_coe, if_pos ha]
+    intro x hx
+    rw [if_neg (fun h => hx (Subtype.ext h.symm))]
+  · rw [if_neg ha, ENNReal.tsum_eq_zero]
+    intro x
+    rw [if_neg]
+    rintro rfl; exact ha x.2
+
+/-- Binding a uniform distribution on `Bool` against a two-branch `pure` gives the
+uniform distribution on the (possibly singleton) two-element finset. -/
+theorem uniformOfFintype_bool_bind_ite [DecidableEq α] (a b : α) :
+    ((PMF.uniformOfFintype Bool).bind fun x => PMF.pure (bif x then a else b)) =
+      PMF.uniformOfFinset {a, b} (Finset.insert_nonempty a {b}) := by
+  ext m
+  rw [PMF.bind_apply, PMF.uniformOfFinset_apply,
+    tsum_bool, PMF.uniformOfFintype_apply, Fintype.card_bool]
+  by_cases hab : a = b
+  · subst hab
+    by_cases hm : m = a <;>
+      simp [hm, PMF.pure_apply, ENNReal.inv_two_add_inv_two]
+  · by_cases h₀ : m = a <;> by_cases h₁ : m = b
+    all_goals simp_all [Finset.card_pair (Ne.symm fun h => hab h.symm),
+      PMF.pure_apply, ENNReal.inv_two_add_inv_two]
+
 /-- Posterior probabilities `joint(a, b) / marginal(b)` sum to 1
 when `b` is in the support of the marginal. -/
 theorem posterior_hasSum (p : PMF α) (f : α → PMF β) (b : β)

From 6122a3655ea5dcbe09d5fa2189af4e9be1702570 Mon Sep 17 00:00:00 2001
From: Devon Tuma <dtumad@gmail.com>
Date: Sat, 23 May 2026 18:49:02 -0500
Subject: [PATCH 5/6] Working executable implementation

---
 .../Protocols/PerfectSecrecy/Basic.lean       |   2 +-
 .../Internal/PerfectSecrecy.lean              |   4 +-
 .../Protocols/PerfectSecrecy/OneTimePad.lean  |  38 +++---
 Cslib/Probability/HasUniform.lean             | 114 +++++++++++-------
 Cslib/Probability/PMF.lean                    |   2 +
 5 files changed, 97 insertions(+), 63 deletions(-)

diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
index 7640982db..4eafc8f5e 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
@@ -33,7 +33,7 @@ namespace Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme
 
 universe u
 variable {n : Type → Type*} [Monad n] [MonadLiftT n PMF] [LawfulMonadLiftT n PMF]
-  [Probability.HasUniformBitVec n] {M K C : Type}
+  [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). -/
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
index e30dd033b..8594d5fc4 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
@@ -87,7 +87,7 @@ theorem perfectlySecret_of_ciphertextIndist (scheme : EncScheme n M K C)
 /-- 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.HasUniformBitVec n] [Probability.LawfulUniformBitVec n]
     (scheme : EncScheme n M K C) (h : scheme.PerfectlySecret) :
     scheme.CiphertextIndist := by
   classical
@@ -120,7 +120,7 @@ lemma encrypt_key_injective (scheme : EncScheme n M K C) [scheme.Correct]
 
 /-- Perfect secrecy requires `|K| ≥ |M|` (Shannon's theorem). -/
 theorem shannonKeySpace [Finite K] [Monad n] [LawfulMonadLiftT n PMF]
-    [Probability.HasUniformBitVec n]
+    [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
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
index e294e5e3e..25e6bce11 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
@@ -36,8 +36,7 @@ The one-time pad (Vernam cipher) over `BitVec l`
 
 namespace Cslib.Crypto.Protocols.PerfectSecrecy
 
-variable (n : Type → Type*) [Monad n] [MonadLiftT n PMF] [LawfulMonadLiftT n PMF]
-  [Probability.HasUniformBitVec n]
+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). -/
@@ -45,26 +44,26 @@ def otp (l : ℕ) : EncScheme n (BitVec l) (BitVec l) (BitVec l) :=
   let gen : n (BitVec l) := Probability.HasUniformBitVec.uniformBitVec l
   .ofPure gen (· ^^^ ·) (· ^^^ ·)
 
-instance (l : ℕ) : (otp n l).Correct :=
+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 n 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, EncScheme.ofPure,
-      Probability.HasUniformBitVec.liftM_uniformBitVec, liftM_pure,
+      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 executable_otp
+section computable_otp
 
-/-- Monad to model computations with access to uniform bitvector selection -/
+/-- Monad to model computations with access to uniform bitvector selection. -/
 private abbrev UniformBitVecM : Type → Type 1 :=
   FreeM (PFunctor.Obj {A := ℕ, B := BitVec})
 
-/-- Currently this can't be made computable because we use `PMF` for probabilities.
-Executing `run_otp -/
+/-- 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
 
@@ -72,20 +71,20 @@ 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
-    simp [liftM, monadLift, MonadLift.monadLift, bind, PMF.map_id]; rfl
+    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
-    let m ← IO.rand 0 (2 ^ n)
-    return f (BitVec.ofNat n m)
-
-private def foo (n : ℕ) : UniformBitVecM (BitVec n) := do
-  let x ← Probability.HasUniformBitVec.uniformBitVec n
-  return x
+  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
@@ -93,6 +92,9 @@ private def run_otp (message : BitVec l) :
   let message' := (otp UniformBitVecM l).dec key ciphertext
   return (message, key, ciphertext, message')
 
-end executable_otp
+-- #eval run_otp 0#6
+-- #eval run_otp 32#6
+
+end computable_otp
 
 end Cslib.Crypto.Protocols.PerfectSecrecy
diff --git a/Cslib/Probability/HasUniform.lean b/Cslib/Probability/HasUniform.lean
index f13095288..8eee3c718 100644
--- a/Cslib/Probability/HasUniform.lean
+++ b/Cslib/Probability/HasUniform.lean
@@ -14,14 +14,20 @@ public import Mathlib.Data.FinEnum
 /-!
 # 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`.
+This file defines general typeclasses for asserting that a monad `m` supports uniform sampling
+operations, along with companion `Lawful*` classes that, given an embedding into `PMF`, witness
+that the operations lift to the corresponding uniform distributions.
+
+Splitting the operation from its correctness proof keeps the sampling classes computable; only
+the `Lawful*` classes (which mention `PMF`) need to be noncomputable.
 
 ## Main Definitions
 
-- `Cslib.Probability.HasUniformBitVec`: monad supports uniform `BitVec n`
-- `Cslib.Probability.HasUniformFinset`: monad supports uniform selection from a non-empty finset
-- `Cslib.Probability.HasUniformFintype`: monad supports uniform selection from a finite type
+- `Cslib.Probability.HasUniformBitVec` / `LawfulUniformBitVec`: monad supports uniform `BitVec n`
+- `Cslib.Probability.HasUniformFinset` / `LawfulUniformFinset`: monad supports uniform selection
+  from a non-empty finset
+- `Cslib.Probability.HasUniformFintype` / `LawfulUniformFintype`: monad supports uniform selection
+  from an inhabited finite type
 
 -/
 
@@ -31,71 +37,95 @@ namespace Cslib.Probability
 
 universe u v
 
+section uniformBitvec
+
 /-- The monad `m` has a way to model uniform selection over `BitVec n`. -/
-class HasUniformBitVec (m : Type → Type*) [MonadLiftT m PMF] where
+class HasUniformBitVec (m : Type → Type*) where
   uniformBitVec (n : ℕ) : m (BitVec n)
-  liftM_uniformBitVec (n : ℕ) : (uniformBitVec n : PMF _) = PMF.uniformOfFintype (BitVec n)
 
-attribute [simp] HasUniformBitVec.liftM_uniformBitVec
+/-- `LawfulUniformBitVec` witnesses that `HasUniformBitVec.uniformBitVec` lifts to the uniform
+distribution on `BitVec n`. -/
+class LawfulUniformBitVec (m : Type → Type*) [MonadLiftT m PMF] [HasUniformBitVec m] : Prop where
+  liftM_uniformBitVec (n : ℕ) :
+    (liftM (HasUniformBitVec.uniformBitVec n : m (BitVec n)) : PMF (BitVec n)) =
+      PMF.uniformOfFintype (BitVec n)
+
+attribute [simp] LawfulUniformBitVec.liftM_uniformBitVec
 
 /-- A uniform `Bool` derived from a single uniform bit. -/
-def uniformBool {m : Type → Type*} [Monad m] [MonadLiftT m PMF]
-    [HasUniformBitVec m] : m Bool :=
+def uniformBool {m : Type → Type*} [Monad m] [HasUniformBitVec m] : m Bool :=
   do return (← HasUniformBitVec.uniformBitVec 1) == 0
 
+@[simp]
+theorem liftM_uniformBool {m : Type → Type*} [Monad m] [MonadLiftT m PMF]
+    [LawfulMonadLiftT m PMF] [HasUniformBitVec m] [LawfulUniformBitVec m] :
+    (liftM (uniformBool : m Bool) : PMF Bool) = PMF.uniformOfFintype Bool := by
+  simp only [uniformBool, liftM_bind, LawfulUniformBitVec.liftM_uniformBitVec, liftM_pure]
+  let : BitVec 1 ≃ Bool := ⟨(· == 0), (bif · then 0 else 1), by decide, by decide⟩
+  exact Cslib.Probability.PMF.uniformOfFintype_map_equiv this
+
+end uniformBitvec
+
+section uniformFinset
+
 /-- The monad `m` has a way to model uniform selection over non-empty finsets. -/
-class HasUniformFinset (m : Type u → Type*) [MonadLiftT m PMF] where
+class HasUniformFinset (m : Type u → Type*) where
   uniformFinset {α : Type u} (s : Finset α) (hs : s.Nonempty) : m α
+
+/-- `LawfulUniformFinset` witnesses that `HasUniformFinset.uniformFinset` lifts to the uniform
+distribution on the underlying finset. -/
+class LawfulUniformFinset (m : Type u → Type*) [MonadLiftT m PMF] [HasUniformFinset m] : Prop where
   liftM_uniformFinset {α : Type u} (s : Finset α) (hs : s.Nonempty) :
-    (uniformFinset s hs : PMF α) = PMF.uniformOfFinset s hs
+    (liftM (HasUniformFinset.uniformFinset s hs : m α) : PMF α) = PMF.uniformOfFinset s hs
 
-attribute [simp] HasUniformFinset.liftM_uniformFinset
+attribute [simp] LawfulUniformFinset.liftM_uniformFinset
 
-instance {m} [MonadLiftT m PMF] [HasUniformFinset m] : HasUniformBitVec m where
+instance {m} [HasUniformFinset m] : HasUniformBitVec m where
   uniformBitVec n :=
     HasUniformFinset.uniformFinset (Finset.univ : Finset (BitVec n)) Finset.univ_nonempty
-  liftM_uniformBitVec _ := by
-    rw [HasUniformFinset.liftM_uniformFinset]; rfl
+
+instance {m} [MonadLiftT m PMF] [HasUniformFinset m] [LawfulUniformFinset m] :
+    LawfulUniformBitVec m where
+  liftM_uniformBitVec n := by
+    change (liftM (HasUniformFinset.uniformFinset
+      (Finset.univ : Finset (BitVec n)) Finset.univ_nonempty : m (BitVec n)) : PMF _) = _
+    rw [LawfulUniformFinset.liftM_uniformFinset]; rfl
+
+end uniformFinset
+
+section uniformFintype
 
 /-- The monad `m` has a way to model uniform selection over inhabited finite types. -/
-class HasUniformFintype (m : Type u → Type*) [MonadLiftT m PMF] where
+class HasUniformFintype (m : Type u → Type*) where
   uniformFintype (α : Type u) [Fintype α] [Nonempty α] : m α
+
+/-- `LawfulUniformFintype` witnesses that `HasUniformFintype.uniformFintype` lifts to the uniform
+distribution on the underlying finite type. -/
+class LawfulUniformFintype (m : Type u → Type*) [MonadLiftT m PMF]
+    [HasUniformFintype m] : Prop where
   liftM_uniformFintype (α : Type u) [Fintype α] [Nonempty α] :
-    (uniformFintype α : PMF α) = PMF.uniformOfFintype α
+    (liftM (HasUniformFintype.uniformFintype α : m α) : PMF α) = PMF.uniformOfFintype α
 
-attribute [simp] HasUniformFintype.liftM_uniformFintype
+attribute [simp] LawfulUniformFintype.liftM_uniformFintype
 
-instance {m} [Monad m] [MonadLiftT m PMF] [LawfulMonadLiftT m PMF]
-    [HasUniformFintype m] : HasUniformFinset m where
-  uniformFinset {α} s hs :=
+instance {m} [Monad m] [HasUniformFintype m] : HasUniformFinset m where
+  uniformFinset {_} s hs :=
     letI : Nonempty {x // x ∈ s} := hs.coe_sort
     HasUniformFintype.uniformFintype {x // x ∈ s} >>= fun x => pure x.val
+
+instance {m} [Monad m] [MonadLiftT m PMF] [LawfulMonadLiftT m PMF]
+    [HasUniformFintype m] [LawfulUniformFintype m] : LawfulUniformFinset m where
   liftM_uniformFinset {α} s hs := by
-    letI : Nonempty {x // x ∈ s} := hs.coe_sort
-    simp only [liftM_bind, liftM_pure, HasUniformFintype.liftM_uniformFintype]
-    change (PMF.uniformOfFintype {x // x ∈ s}).bind (PMF.pure ∘ Subtype.val) = _
-    rw [PMF.bind_pure_comp]
+    simp only [HasUniformFinset.uniformFinset, liftM_bind,
+      LawfulUniformFintype.liftM_uniformFintype, liftM_pure, PMF.monad_pure_eq_pure]
     exact Cslib.Probability.PMF.map_subtypeVal_uniformOfFintype s hs
 
 noncomputable instance : HasUniformFintype PMF where
   uniformFintype := PMF.uniformOfFintype
-  liftM_uniformFintype _ _ := rfl
 
-/-- The equivalence between `BitVec 1` and `Bool` via the predicate `· == 0`. -/
-private def bitVecOneEquivBool : BitVec 1 ≃ Bool where
-  toFun bv := bv == 0
-  invFun b := bif b then 0 else 1
-  left_inv := by decide
-  right_inv := by decide
+noncomputable instance : LawfulUniformFintype PMF where
+  liftM_uniformFintype _ _ := rfl
 
-/-- The PMF coercion of `uniformBool` is the uniform distribution on `Bool`. -/
-@[simp]
-theorem liftM_uniformBool {m : Type → Type*} [Monad m] [MonadLiftT m PMF]
-    [LawfulMonadLiftT m PMF] [HasUniformBitVec m] :
-    (liftM (uniformBool : m Bool) : PMF Bool) = PMF.uniformOfFintype Bool := by
-  simp only [uniformBool, liftM_bind, liftM_pure, HasUniformBitVec.liftM_uniformBitVec]
-  change (PMF.uniformOfFintype (BitVec 1)).bind (PMF.pure ∘ bitVecOneEquivBool) = _
-  rw [PMF.bind_pure_comp]
-  exact Cslib.Probability.PMF.uniformOfFintype_map_equiv bitVecOneEquivBool
+end uniformFintype
 
 end Cslib.Probability
diff --git a/Cslib/Probability/PMF.lean b/Cslib/Probability/PMF.lean
index 9faaaf952..4be3f5613 100644
--- a/Cslib/Probability/PMF.lean
+++ b/Cslib/Probability/PMF.lean
@@ -168,4 +168,6 @@ theorem posteriorDist_eq_prior_of_outputIndist (p : PMF α) (f : α → PMF β)
 @[simp]
 lemma monad_pure_eq_pure (x : α) : (pure x : PMF α) = PMF.pure x := rfl
 
+lemma monad_bind_eq_bind {α β} (p : PMF α) (q : α → PMF β) : p >>= q = p.bind q := rfl
+
 end Cslib.Probability.PMF

From b59374a306dbe5eebc8f95777b00b381a4e8aafe Mon Sep 17 00:00:00 2001
From: Devon Tuma <dtumad@gmail.com>
Date: Sat, 23 May 2026 19:24:36 -0500
Subject: [PATCH 6/6] linting errors

---
 Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean | 9 ++++-----
 Cslib/Probability/HasUniform.lean                     | 3 +++
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean
index 06d021980..58b021a97 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Encryption.lean
@@ -47,12 +47,11 @@ 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) where
+    (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
-where decryption is a left inverse of encryption for every key. -/
+/-- 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
@@ -60,8 +59,8 @@ def EncScheme.ofPure {Message Key Ciphertext : Type u} (gen : m Key)
   enc key message := pure (enc key message)
   dec := dec
 
-@[reducible]
-def EncScheme.ofPure.Correct [MonadLiftT m PMF] [LawfulMonadLiftT m PMF]
+/-- `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)) :
diff --git a/Cslib/Probability/HasUniform.lean b/Cslib/Probability/HasUniform.lean
index 8eee3c718..3ad2fa41c 100644
--- a/Cslib/Probability/HasUniform.lean
+++ b/Cslib/Probability/HasUniform.lean
@@ -41,6 +41,7 @@ section uniformBitvec
 
 /-- The monad `m` has a way to model uniform selection over `BitVec n`. -/
 class HasUniformBitVec (m : Type → Type*) where
+  /-- Select a random bitvector of length `n`. -/
   uniformBitVec (n : ℕ) : m (BitVec n)
 
 /-- `LawfulUniformBitVec` witnesses that `HasUniformBitVec.uniformBitVec` lifts to the uniform
@@ -70,6 +71,7 @@ section uniformFinset
 
 /-- The monad `m` has a way to model uniform selection over non-empty finsets. -/
 class HasUniformFinset (m : Type u → Type*) where
+  /-- Select a uniform element from the elements of `s : Finset α`. -/
   uniformFinset {α : Type u} (s : Finset α) (hs : s.Nonempty) : m α
 
 /-- `LawfulUniformFinset` witnesses that `HasUniformFinset.uniformFinset` lifts to the uniform
@@ -97,6 +99,7 @@ section uniformFintype
 
 /-- The monad `m` has a way to model uniform selection over inhabited finite types. -/
 class HasUniformFintype (m : Type u → Type*) where
+  /-- Select a random element of the finite and nonempty type `α`. -/
   uniformFintype (α : Type u) [Fintype α] [Nonempty α] : m α
 
 /-- `LawfulUniformFintype` witnesses that `HasUniformFintype.uniformFintype` lifts to the uniform