diff --git a/Cslib.lean b/Cslib.lean
index 982b94f5d..0b0e5a50e 100644
--- a/Cslib.lean
+++ b/Cslib.lean
@@ -58,8 +58,6 @@ public import Cslib.Computability.URM.StraightLine
 public import Cslib.Crypto.Protocols.PerfectSecrecy.Basic
 public import Cslib.Crypto.Protocols.PerfectSecrecy.Defs
 public import Cslib.Crypto.Protocols.PerfectSecrecy.Encryption
-public import Cslib.Crypto.Protocols.PerfectSecrecy.Internal.OneTimePad
-public import Cslib.Crypto.Protocols.PerfectSecrecy.Internal.PerfectSecrecy
 public import Cslib.Crypto.Protocols.PerfectSecrecy.OneTimePad
 public import Cslib.Crypto.Protocols.SecretSharing.Defs
 public import Cslib.Crypto.Protocols.SecretSharing.Scheme
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
index c5f5a29ef..26a154ee8 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Basic.lean
@@ -7,16 +7,20 @@ Authors: Samuel Schlesinger
 module
 
 public import Cslib.Crypto.Protocols.PerfectSecrecy.Defs
-public import Cslib.Crypto.Protocols.PerfectSecrecy.Internal.PerfectSecrecy
+import Mathlib.Probability.Distributions.Uniform
 
 /-!
 # Perfect Secrecy
 
 Characterisation theorems for perfect secrecy following
-[KatzLindell2020], Chapter 2.
+[KatzLindell2020], Chapter 2: the equivalence with message-ciphertext
+independence, the ciphertext indistinguishability characterization, and
+Shannon's key-space bound.
 
 ## Main results
 
+- `Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme.perfectlySecret_iff_indep`:
+  perfect secrecy is exactly message-ciphertext independence
 - `Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme.perfectlySecret_iff_ciphertextIndist`:
   ciphertext indistinguishability characterization ([KatzLindell2020], Lemma 2.5)
 - `Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme.perfectlySecret_keySpace_ge`:
@@ -31,21 +35,98 @@ Characterisation theorems for perfect secrecy following
 
 namespace Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme
 
+open Cslib.Probability.PMF
+
 universe u
 variable {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)
+    (m : M) (c : C) :
+    scheme.jointDist msgDist (m, c) = msgDist m * scheme.ciphertextDist m c :=
+  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) :
+    ∑' m, scheme.jointDist msgDist (m, c) = scheme.marginalCiphertextDist msgDist c :=
+  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) :
+    scheme.PerfectlySecret ↔
+    ∀ (msgDist : PMF M) (m : M) (c : C),
+      scheme.jointDist msgDist (m, c) =
+        msgDist m * scheme.marginalCiphertextDist msgDist c := by
+  refine ⟨fun h msgDist m c => ?_, fun h msgDist c hc => ?_⟩
+  · by_cases hc : (scheme.marginalCiphertextDist msgDist) c = 0
+    · have := ENNReal.tsum_eq_zero.mp
+        ((jointDist_tsum_fst scheme msgDist c).trans hc) m
+      rw [this, hc, mul_zero]
+    · have hne_top := PMF.apply_ne_top (scheme.marginalCiphertextDist msgDist) c
+      have := DFunLike.congr_fun (h msgDist c ((PMF.mem_support_iff _ _).mpr hc)) m
+      simp only [posteriorMsgDist_apply] at this
+      rw [← this, ENNReal.div_mul_cancel hc hne_top]
+  · ext m
+    simp only [posteriorMsgDist_apply]
+    rw [h msgDist m c, ENNReal.mul_div_cancel_right
+      ((PMF.mem_support_iff _ _).mp hc) (PMF.apply_ne_top _ c)]
+
 /-- 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) :
-    scheme.PerfectlySecret ↔ scheme.CiphertextIndist :=
-  ⟨PerfectSecrecy.ciphertextIndist_of_perfectlySecret scheme,
-   PerfectSecrecy.perfectlySecret_of_ciphertextIndist scheme⟩
+    scheme.PerfectlySecret ↔ scheme.CiphertextIndist := by
+  classical
+  refine ⟨fun h => ?_, fun h msgDist c hc =>
+    posteriorDist_eq_prior_of_outputIndist msgDist scheme.ciphertextDist h c hc⟩
+  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
+  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 := PMF.apply_ne_top μ m
+  exact (ENNReal.mul_right_inj hne hne_top).mp (by rw [← jointDist_eq]; exact h μ m c)
+
+/-- Ciphertext indistinguishability implies message-ciphertext independence. -/
+theorem indep_of_ciphertextIndist (scheme : EncScheme M K C)
+    (h : scheme.CiphertextIndist) (msgDist : PMF M) (m : M) (c : C) :
+    scheme.jointDist msgDist (m, c) =
+      msgDist m * scheme.marginalCiphertextDist msgDist c :=
+  (perfectlySecret_iff_indep scheme).mp
+    ((perfectlySecret_iff_ciphertextIndist scheme).mpr h) msgDist 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). -/
+private lemma encrypt_key_injective (scheme : EncScheme 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) :
+    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|`
+/-- Perfect secrecy requires `|K| ≥ |M|` — Shannon's theorem
 ([KatzLindell2020], Theorem 2.12). -/
 theorem perfectlySecret_keySpace_ge [Finite K]
     (scheme : EncScheme M K C) (h : scheme.PerfectlySecret) :
-    Nat.card K ≥ Nat.card M :=
-  PerfectSecrecy.shannonKeySpace scheme h
+    Nat.card K ≥ Nat.card M := by
+  classical
+  have hci := (perfectlySecret_iff_ciphertextIndist scheme).mp 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
+    intro m
+    exact (PMF.mem_support_bind_iff _ _ _).mp
+      (show c₀ ∈ (scheme.ciphertextDist m).support by rw [hci m m₀]; exact hc₀)
+  choose f hf_mem hf_enc using key_exists
+  exact Nat.card_le_card_of_injective f
+    (encrypt_key_injective scheme f c₀ hf_mem hf_enc)
 
 end Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Defs.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Defs.lean
index e7d824dd6..4c1d5dcfb 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Defs.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/Defs.lean
@@ -35,6 +35,8 @@ Core definitions for perfect secrecy following [KatzLindell2020], Chapter 2.
 
 namespace Cslib.Crypto.Protocols.PerfectSecrecy.EncScheme
 
+open Cslib.Probability.PMF
+
 universe u
 variable {M K C : Type u}
 
@@ -58,7 +60,7 @@ the marginal distribution. -/
 noncomputable def posteriorMsgDist (scheme : EncScheme M K C)
     (msgDist : PMF M) (c : C)
     (hc : c ∈ (scheme.marginalCiphertextDist msgDist).support) : PMF M :=
-  Cslib.Probability.PMF.posteriorDist msgDist scheme.ciphertextDist c hc
+  posteriorDist msgDist scheme.ciphertextDist c hc
 
 @[simp]
 theorem posteriorMsgDist_apply (scheme : EncScheme M K C)
@@ -66,7 +68,7 @@ theorem posteriorMsgDist_apply (scheme : EncScheme M K C)
     (hc : c ∈ (scheme.marginalCiphertextDist msgDist).support) (m : M) :
     scheme.posteriorMsgDist msgDist c hc m =
       scheme.jointDist msgDist (m, c) / scheme.marginalCiphertextDist msgDist c :=
-  rfl
+  posteriorDist_apply msgDist scheme.ciphertextDist c hc m
 
 /-- An encryption scheme is perfectly secret if the posterior message
 distribution equals the prior for every ciphertext with positive probability
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/OneTimePad.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/OneTimePad.lean
deleted file mode 100644
index 451e3253a..000000000
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/OneTimePad.lean
+++ /dev/null
@@ -1,38 +0,0 @@
-/-
-Copyright (c) 2026 Samuel Schlesinger. All rights reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Samuel Schlesinger
--/
-
-module
-
-public import Cslib.Init
-public import Mathlib.Probability.Distributions.Uniform
-
-/-!
-# One-Time Pad: Internal proofs
-
-The OTP ciphertext distribution is uniform regardless of message.
--/
-
-@[expose] public section
-
-namespace Cslib.Crypto.Protocols.PerfectSecrecy.OTP
-
--- TODO: upstream to Mathlib as a FinEnum instance for BitVec.
-instance bitVecFintype (n : ℕ) : Fintype (BitVec n) :=
-  Fintype.ofEquiv (Fin (2 ^ n))
-    ⟨BitVec.ofFin, BitVec.toFin, fun x => by simp, fun x => by simp⟩
-
--- TODO: upstream to Mathlib — general BitVec XOR cancellation lemma.
-/-- XOR by a fixed mask is self-inverse on `BitVec`: `c = k ^^^ m ↔ c ^^^ m = k`. -/
-lemma eq_xor_iff_xor_eq {l : ℕ} (c m k : BitVec l) :
-    (c = k ^^^ m) ↔ (c ^^^ m = k) := by grind
-
-/-- The ciphertext distribution of the OTP is uniform, regardless of the message. -/
-theorem otp_ciphertextDist_eq_uniform (l : ℕ) (m : BitVec l) :
-    (PMF.uniformOfFintype (BitVec l)).bind
-      (fun k => PMF.pure (k ^^^ m)) =
-    PMF.uniformOfFintype (BitVec l) := by simp [PMF.ext_iff, eq_xor_iff_xor_eq]
-
-end Cslib.Crypto.Protocols.PerfectSecrecy.OTP
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
deleted file mode 100644
index 66656e8ba..000000000
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/Internal/PerfectSecrecy.lean
+++ /dev/null
@@ -1,131 +0,0 @@
-/-
-Copyright (c) 2026 Samuel Schlesinger. All rights reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Samuel Schlesinger
--/
-
-module
-
-public import Cslib.Crypto.Protocols.PerfectSecrecy.Defs
-public import Mathlib.Probability.Distributions.Uniform
-
-/-!
-# Perfect Secrecy: Internal proofs
-
-Auxiliary lemmas for perfect secrecy:
-- Equivalence of the conditional-probability and independence formulations
-- Both directions of the ciphertext indistinguishability characterization
-  ([KatzLindell2020], Lemma 2.5)
-- Shannon's key-space bound ([KatzLindell2020], Theorem 2.12)
--/
-
-@[expose] public section
-
-namespace Cslib.Crypto.Protocols.PerfectSecrecy
-
-open PMF ENNReal
-
-universe u
-variable {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)
-    (m : M) (c : C) :
-    scheme.jointDist msgDist (m, c) = msgDist 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) :
-    ∑' 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) :
-    scheme.PerfectlySecret ↔
-    ∀ (msgDist : PMF M) (m : M) (c : C),
-      scheme.jointDist msgDist (m, c) =
-        msgDist m * scheme.marginalCiphertextDist msgDist c := by
-  constructor
-  · intro h msgDist m c
-    by_cases hc : (scheme.marginalCiphertextDist msgDist) c = 0
-    · have := ENNReal.tsum_eq_zero.mp
-        ((jointDist_tsum_fst scheme msgDist c).trans hc) m
-      rw [this, hc, mul_zero]
-    · have hne_top := ne_top_of_le_ne_top one_ne_top
-        (PMF.coe_le_one (scheme.marginalCiphertextDist msgDist) c)
-      have := DFunLike.congr_fun (h msgDist c ((PMF.mem_support_iff _ _).mpr hc)) m
-      simp only [EncScheme.posteriorMsgDist_apply] at this
-      rw [← this, ENNReal.div_mul_cancel hc hne_top]
-  · intro h msgDist c hc; ext m
-    simp only [EncScheme.posteriorMsgDist_apply]
-    rw [h msgDist m c, ENNReal.mul_div_cancel_right
-      ((PMF.mem_support_iff _ _).mp hc)
-      (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)
-    (h : scheme.CiphertextIndist)
-    (msgDist : PMF M) (m : M) (c : C) :
-    scheme.jointDist msgDist (m, c) =
-      msgDist 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
-  rw [PMF.bind_apply]
-  conv_rhs => arg 1; ext m'; rw [h m' m]
-  rw [ENNReal.tsum_mul_right, PMF.tsum_coe, one_mul]
-
-/-- Ciphertext indistinguishability implies perfect secrecy. -/
-theorem perfectlySecret_of_ciphertextIndist (scheme : EncScheme 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) :
-    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
-  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)
-
-/-- 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)
-    (f : M → K) (c₀ : C)
-    (hf_mem : ∀ m, f m ∈ scheme.gen.support)
-    (hf_enc : ∀ m, c₀ ∈ (scheme.enc (f m) m).support) :
-    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) :
-    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
-    intro m
-    exact (PMF.mem_support_bind_iff _ _ _).mp
-      (show c₀ ∈ (scheme.ciphertextDist m).support by rw [hci m m₀]; exact hc₀)
-  choose f hf_mem hf_enc using key_exists
-  exact Nat.card_le_card_of_injective f
-    (encrypt_key_injective scheme f c₀ hf_mem hf_enc)
-
-end Cslib.Crypto.Protocols.PerfectSecrecy
diff --git a/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean b/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
index ed53950f6..1df26b1af 100644
--- a/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
+++ b/Cslib/Crypto/Protocols/PerfectSecrecy/OneTimePad.lean
@@ -7,8 +7,9 @@ Authors: Samuel Schlesinger
 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.FinEnum
+import Cslib.Probability.PMF
+import Mathlib.Data.LawfulXor.Equiv
 
 /-!
 # One-Time Pad
@@ -34,18 +35,28 @@ The one-time pad (Vernam cipher) over `BitVec l`
 
 namespace Cslib.Crypto.Protocols.PerfectSecrecy
 
+open Cslib.Probability.PMF
+
 /-- 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]
+    simp [xor_cancel_left]
+
+/-- The ciphertext distribution of the OTP is uniform, regardless of the
+message: masking with a uniform key is the permutation `Equiv.xor` of the
+uniform distribution. -/
+theorem otp_ciphertextDist_eq_uniform (l : ℕ) (m : BitVec l) :
+    (otp l).ciphertextDist m = PMF.uniformOfFintype (BitVec l) := by
+  have h : (fun k : BitVec l => PMF.pure (k ^^^ m)) = (PMF.pure ∘ ⇑(Equiv.xor m)) :=
+    congrArg (PMF.pure ∘ ·) xor_right_eq
+  change (PMF.uniformOfFintype (BitVec l)).bind (fun k => PMF.pure (k ^^^ m)) = _
+  rw [h, PMF.bind_pure_comp, uniformOfFintype_map_equiv]
 
 /-- The one-time pad is perfectly secret ([KatzLindell2020], Theorem 2.10). -/
 theorem otp_perfectlySecret (l : ℕ) : (otp l).PerfectlySecret :=
-  (EncScheme.perfectlySecret_iff_ciphertextIndist _).mpr fun m₀ m₁ => by
-    simp only [EncScheme.ciphertextDist, otp]
-    exact (OTP.otp_ciphertextDist_eq_uniform l m₀).trans
-      (OTP.otp_ciphertextDist_eq_uniform l m₁).symm
+  (EncScheme.perfectlySecret_iff_ciphertextIndist _).mpr fun m₀ m₁ =>
+    (otp_ciphertextDist_eq_uniform l m₀).trans (otp_ciphertextDist_eq_uniform l m₁).symm
 
 end Cslib.Crypto.Protocols.PerfectSecrecy
diff --git a/Cslib/Crypto/Protocols/SecretSharing/Defs.lean b/Cslib/Crypto/Protocols/SecretSharing/Defs.lean
index e8788f08d..5da7d7096 100644
--- a/Cslib/Crypto/Protocols/SecretSharing/Defs.lean
+++ b/Cslib/Crypto/Protocols/SecretSharing/Defs.lean
@@ -41,6 +41,8 @@ namespace Cslib.Crypto.Protocols.SecretSharing
 
 namespace Scheme
 
+open Cslib.Probability.PMF
+
 variable {Secret Randomness Party Share : Type*}
 
 /-- The distribution of the full share assignment for one secret. -/
@@ -58,9 +60,8 @@ theorem viewDist_eq_of_not_authorized
     (scheme : Scheme Secret Randomness Party Share)
     {s : Finset Party} (hs : ¬ scheme.authorized s)
     (secret₀ secret₁ : Secret) :
-    scheme.viewDist s secret₀ = scheme.viewDist s secret₁ := by
-  unfold viewDist
-  exact scheme.view_indist s hs secret₀ secret₁
+    scheme.viewDist s secret₀ = scheme.viewDist s secret₁ :=
+  scheme.view_indist s hs secret₀ secret₁
 
 /-- The posterior distribution on secrets after observing the coalition view
 `v`. -/
@@ -68,8 +69,7 @@ noncomputable def posteriorSecretDist
     (scheme : Scheme Secret Randomness Party Share)
     (s : Finset Party) (secretDist : PMF Secret) (v : s → Share)
     (hv : v ∈ (secretDist.bind (scheme.viewDist s)).support) : PMF Secret :=
-  Cslib.Probability.PMF.posteriorDist
-    (p := secretDist) (f := scheme.viewDist s) v hv
+  posteriorDist (p := secretDist) (f := scheme.viewDist s) v hv
 
 @[simp]
 theorem posteriorSecretDist_apply
@@ -80,7 +80,7 @@ theorem posteriorSecretDist_apply
       (secretDist.bind fun secret' =>
         (scheme.viewDist s secret').bind fun v' => PMF.pure (secret', v')) (secret, v) /
         (secretDist.bind (scheme.viewDist s)) v :=
-  rfl
+  posteriorDist_apply secretDist (scheme.viewDist s) v hv secret
 
 /-- Perfect privacy for unauthorized coalitions: conditioning on a view does not
 change the prior on secrets. -/
@@ -93,11 +93,10 @@ def PerfectlyPrivate (scheme : Scheme Secret Randomness Party Share) : Prop :=
 /-- Every scheme has posterior privacy by definition of `Scheme`. -/
 theorem perfectlyPrivate
     (scheme : Scheme Secret Randomness Party Share) :
-    scheme.PerfectlyPrivate := by
-  intro s hs secretDist v hv
-  exact Cslib.Probability.PMF.posteriorDist_eq_prior_of_outputIndist
-    (p := secretDist) (f := scheme.viewDist s)
-    (fun secret₀ secret₁ => scheme.viewDist_eq_of_not_authorized hs secret₀ secret₁) v hv
+    scheme.PerfectlyPrivate :=
+  fun s hs secretDist v hv =>
+    posteriorDist_eq_prior_of_outputIndist secretDist (scheme.viewDist s)
+      (scheme.viewDist_eq_of_not_authorized hs) v hv
 
 end Scheme
 
diff --git a/Cslib/Crypto/Protocols/SecretSharing/Scheme.lean b/Cslib/Crypto/Protocols/SecretSharing/Scheme.lean
index 09e07ae2a..937ad0ba6 100644
--- a/Cslib/Crypto/Protocols/SecretSharing/Scheme.lean
+++ b/Cslib/Crypto/Protocols/SecretSharing/Scheme.lean
@@ -59,8 +59,7 @@ structure Scheme (Secret Randomness Party Share : Type*) where
   /-- Authorized coalitions. -/
   authorized : Finset Party → Prop
   /-- Authorization is monotone in the coalition. -/
-  authorized_mono :
-    ∀ {s t : Finset Party}, s ⊆ t → authorized s → authorized t
+  authorized_mono : Monotone authorized
   /-- Authorized coalitions reconstruct the secret from the restricted view. -/
   correct :
     ∀ (r : Randomness) (secret : Secret) (s : Finset Party),
@@ -98,9 +97,8 @@ theorem not_authorized_of_subset
     (scheme : Scheme Secret Randomness Party Share)
     {s t : Finset Party} (hst : s ⊆ t)
     (ht : ¬ scheme.authorized t) :
-    ¬ scheme.authorized s := by
-  intro hs
-  exact ht (scheme.authorized_mono hst hs)
+    ¬ scheme.authorized s :=
+  mt (scheme.authorized_mono hst) ht
 
 end Scheme
 
diff --git a/Cslib/Crypto/Protocols/SecretSharing/Shamir.lean b/Cslib/Crypto/Protocols/SecretSharing/Shamir.lean
index 0c227bdd8..c85bcb41d 100644
--- a/Cslib/Crypto/Protocols/SecretSharing/Shamir.lean
+++ b/Cslib/Crypto/Protocols/SecretSharing/Shamir.lean
@@ -63,6 +63,8 @@ noncomputable section
 
 namespace Cslib.Crypto.Protocols.SecretSharing.Shamir
 
+open Cslib.Probability.PMF
+
 variable {F Party : Type*} [Field F] [Fintype Party]
 
 /-- Public parameters for a finite Shamir secret-sharing instance. The threshold
@@ -116,21 +118,12 @@ structure TailSampler (params : Params F Party) where
   /-- Translating the coefficients does not change the distribution. -/
   map_add_eq_self : ∀ δ : Randomness params, gen.map (fun coeffs => coeffs + δ) = gen
 
-private def coeffTranslate {params : Params F Party} (δ : Randomness params) :
-    Randomness params ≃ Randomness params where
-  toFun coeffs := coeffs + δ
-  invFun coeffs := coeffs - δ
-  left_inv coeffs := by simp
-  right_inv coeffs := by simp
-
 /-- Uniform tail coefficients form the canonical privacy-compatible sampler. -/
 noncomputable def uniformTailSampler (params : Params F Party)
     [Fintype F] [Nonempty F] : TailSampler params where
   gen := PMF.uniformOfFintype (Randomness params)
   map_add_eq_self δ := by
-    simpa [coeffTranslate] using
-      (Cslib.Probability.PMF.uniformOfFintype_map_equiv
-        (coeffTranslate (params := params) δ))
+    simpa using uniformOfFintype_map_equiv (Equiv.addRight δ)
 
 private noncomputable def privacyCorrectionPolynomial
     (params : Params F Party) (s : Finset Party)
@@ -255,9 +248,7 @@ noncomputable def schemeWith (params : Params F Party) (sampler : TailSampler pa
     share := share params
     reconstruct := reconstruct params
     authorized := authorized params
-    authorized_mono := by
-      intro s u hsu hs
-      exact le_trans hs (Finset.card_le_card hsu)
+    authorized_mono := fun _ _ hsu hs => le_trans hs (Finset.card_le_card hsu)
     correct := by
       intro coeffs secretValue s hs
       have hdeg₀ :
diff --git a/Cslib/Probability/PMF.lean b/Cslib/Probability/PMF.lean
index 8393eb229..8c4c8a63d 100644
--- a/Cslib/Probability/PMF.lean
+++ b/Cslib/Probability/PMF.lean
@@ -28,7 +28,6 @@ the Mathlib module instead.
 - `Cslib.Probability.PMF.bind_pair_tsum_fst`: marginalizing over the first component
 - `Cslib.Probability.PMF.uniformOfFintype_map_equiv`:
   a uniform distribution is invariant under equivalence
-- `Cslib.Probability.PMF.posterior_hasSum`: posterior probabilities sum to 1
 - `Cslib.Probability.PMF.posteriorDist`: the posterior as a `PMF`
 - `Cslib.Probability.PMF.posteriorDist_eq_prior_of_outputIndist`:
   if the output distribution does not depend on the input, conditioning does
@@ -64,52 +63,28 @@ theorem bind_pair_tsum_fst (p : PMF α) (f : α → PMF β) (b : β) :
 theorem uniformOfFintype_map_equiv {γ : Type v} [Fintype α] [Fintype γ] [Nonempty α] [Nonempty γ]
     (e : α ≃ γ) :
     (PMF.uniformOfFintype α).map e = PMF.uniformOfFintype γ := by
-  classical
-  have hcard : Fintype.card α = Fintype.card γ := Fintype.card_congr e
   ext c
-  rw [PMF.map_apply, PMF.uniformOfFintype_apply, tsum_eq_single (e.symm c)]
-  · simp_rw [PMF.uniformOfFintype_apply]
-    simp [hcard]
-  · intro a ha
-    simp_rw [PMF.uniformOfFintype_apply]
-    split_ifs with h
-    · exfalso
-      apply ha
-      simpa using congrArg e.symm h.symm
-    · simp
-
-/-- 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 : β)
-    (hb : b ∈ (p.bind f).support) :
-    HasSum (fun a =>
-      (p.bind fun a' => (f a').bind fun b' => PMF.pure (a', b')) (a, b) /
-        (p.bind f) b) 1 := by
-  have hne := (PMF.mem_support_iff _ _).mp hb
-  have hne_top := ne_top_of_le_ne_top one_ne_top (PMF.coe_le_one (p.bind f) b)
-  have : ∑' a, (p.bind fun a' => (f a').bind fun b' => PMF.pure (a', b')) (a, b) /
-      (p.bind f) b = 1 := by
-    simp only [div_eq_mul_inv]
-    rw [ENNReal.tsum_mul_right, bind_pair_tsum_fst]
-    exact ENNReal.mul_inv_cancel hne hne_top
-  exact this ▸ ENNReal.summable.hasSum
+  rw [PMF.map_apply, tsum_eq_single (e.symm c)]
+  · simp [Fintype.card_congr e]
+  · exact fun a ha => ite_eq_right fun h => ha (by simp [h])
 
 /-- The posterior distribution `Pr[A = a | B = b]` as a `PMF`,
-given `a ← p`, `b ← f a`, and that `b` has positive marginal probability. -/
+given `a ← p`, `b ← f a`, and that `b` has positive marginal probability:
+the joint distribution's slice at `b`, normalized. -/
 noncomputable def posteriorDist (p : PMF α) (f : α → PMF β) (b : β)
     (hb : b ∈ (p.bind f).support) : PMF α :=
-  ⟨fun a =>
-    (p.bind fun a' => (f a').bind fun b' => PMF.pure (a', b')) (a, b) /
-      (p.bind f) b,
-   posterior_hasSum p f b hb⟩
+  PMF.normalize
+    (fun a => (p.bind fun a' => (f a').bind fun b' => PMF.pure (a', b')) (a, b))
+    (by rw [bind_pair_tsum_fst]; exact (PMF.mem_support_iff _ _).mp hb)
+    (by rw [bind_pair_tsum_fst]; exact PMF.apply_ne_top _ _)
 
 @[simp]
 theorem posteriorDist_apply (p : PMF α) (f : α → PMF β) (b : β)
     (hb : b ∈ (p.bind f).support) (a : α) :
     posteriorDist p f b hb a =
       (p.bind fun a' => (f a').bind fun b' => PMF.pure (a', b')) (a, b) /
-        (p.bind f) b :=
-  rfl
+        (p.bind f) b := by
+  rw [posteriorDist, PMF.normalize_apply, bind_pair_tsum_fst, div_eq_mul_inv]
 
 /-- If the output distribution of a channel does not depend on the input, then
 conditioning on any output with positive probability leaves the prior unchanged. -/
@@ -118,16 +93,10 @@ theorem posteriorDist_eq_prior_of_outputIndist (p : PMF α) (f : α → PMF β)
     (b : β) (hb : b ∈ (p.bind f).support) :
     posteriorDist p f b hb = p := by
   ext a
-  rw [posteriorDist_apply, bind_pair_apply, PMF.bind_apply]
-  have hf : ∀ a', f a' b = f a b := fun a' => by rw [h a' a]
-  simp_rw [hf]
-  rw [ENNReal.tsum_mul_right, PMF.tsum_coe, one_mul]
-  have hb' : (p.bind f) b ≠ 0 := (PMF.mem_support_iff _ _).mp hb
-  have hmarg : (p.bind f) b = f a b := by
-    rw [PMF.bind_apply]
-    simp_rw [hf]
-    rw [ENNReal.tsum_mul_right, PMF.tsum_coe, one_mul]
-  exact ENNReal.mul_div_cancel_right (hmarg ▸ hb')
-    (ne_top_of_le_ne_top ENNReal.one_ne_top (PMF.coe_le_one _ _))
+  have hbind : p.bind f = f a :=
+    (congrArg p.bind (funext fun a' => h a' a)).trans (PMF.bind_const p (f a))
+  rw [posteriorDist_apply, bind_pair_apply, hbind]
+  exact ENNReal.mul_div_cancel_right ((PMF.mem_support_iff _ _).mp (hbind ▸ hb))
+    (PMF.apply_ne_top _ _)
 
 end Cslib.Probability.PMF