Provide fresh randomized reply transformations

src/Commons/Definitions/ReplyTransform.lean

@@ -1,14 +1,17 @@
 module
 
 public import Commons.Definitions.Ideal
-public import RandomSystems.Typed.QueryRouting
+public import RandomSystems.Typed.ReplyTransform
 
 /-!
 # Transforming evaluation replies
 
 `mapReplies f q` forwards an `eval(())` query once and returns `f x` on the inside reply `x`.
 The per-port allowance `q` is unchanged. `mapReplies_smul_memorylessSource` identifies
-the attached source with the source of the mapped probability law.
+the attached source with the source of the mapped probability law. `sampleReplies P f q`
+additionally draws `k ∼ P` independently on each call and returns `f k x`;
+`sampleReplies_smul_memorylessSource` gives the fresh reply law
+`(prodProbDist P Q).map (fun kx => f kx.1 kx.2)` on a source with law `Q`.
 
 ```text
 caller ── eval(()) ──► [mapReplies f] ── eval(()) ──► source P
@@ -42,4 +45,27 @@ lemma mapReplies_smul_memorylessSource
   rw [MemorylessSource.eq_ofMemoryless, MemorylessSource.eq_ofMemoryless]
   exact TypedConverter.mapReplies_smul_ofMemoryless _ _ _
 
+/-- Draw `k ∼ P` afresh, request one inside reply `x`, and return `f k x`. -/
+noncomputable def sampleReplies
+    -- `X`, `Y` carry inside and outside replies; `K` carries the auxiliary sample.
+    {X Y K : Type} [Fintype X] [Fintype Y] [Fintype K]
+    -- `P` is sampled independently on every call; `f` computes the outside reply.
+    (P : ProbDist K) (f : K → X → Y)
+    -- `q` gives the same evaluation allowance on both public interfaces.
+    (q : Evaluation.Port → ℕ) :
+    TypedConverter (Evaluation.perPort Unit Y q) (Evaluation.perPort Unit X q) :=
+  TypedConverter.sampleReplies Evaluation.Port (fun _ => Unit) (fun _ => X) q P (fun _ => f)
+
+/-- Independent fresh `k ∼ P` and `x ∼ Q` give fresh replies with the mapped product law. -/
+lemma sampleReplies_smul_memorylessSource
+    -- `X`, `Y`, `K` carry resource replies, outside replies and auxiliary samples.
+    {X Y K : Type} [Fintype X] [Fintype Y] [Fintype K]
+    -- `P`, `Q` are the two independent laws; `f` combines their draws and `q` bounds calls.
+    (P : ProbDist K) (f : K → X → Y) (Q : ProbDist X) (q : Evaluation.Port → ℕ) :
+    sampleReplies P f q • MemorylessSource.perPort X Q (budget := q) =
+      MemorylessSource.perPort Y ((prodProbDist P Q).map (fun kx => f kx.1 kx.2))
+        (budget := q) := by
+  rw [MemorylessSource.eq_ofMemoryless, MemorylessSource.eq_ofMemoryless]
+  exact TypedConverter.sampleReplies_smul_ofMemoryless _ _ _ _
+
 end Commons.Evaluation

src/RandomSystems.lean

@@ -112,6 +112,7 @@ public import RandomSystems.Typed.ProbProgramAttachment
 public import RandomSystems.Typed.Relabel
 public import RandomSystems.Typed.RecordingGuard
 public import RandomSystems.Typed.QueryRouting
+public import RandomSystems.Typed.ReplyTransform
 public import RandomSystems.Typed.SamplingAlgebra
 public import RandomSystems.Typed.Source
 public import RandomSystems.Typed.SystemCoherence

src/RandomSystems/Typed/ReplyTransform.lean

@@ -0,0 +1,133 @@
+module
+
+public import RandomSystems.Typed.QueryRouting
+
+/-!
+# Fresh randomized reply transformations
+
+`TypedConverter.sampleReplies` draws `k ∼ P`, forwards the current input once, and returns
+`f i k z` on the inside reply `z`. It assembles `sampleReplyProgram` with `TypedSystem.source`;
+the source and inside resource are independent. Each invocation consumes another source value.
+
+```text
+caller ── i(x) ──► [sampleReplyProgram]
+                    1. source P: () ↦ k
+                    2. resource: i(x) ↦ z
+caller ◄─ f i k z ──┘
+```
+
+`sampleReplyProgram_bounded` limits each invocation to two calls. Port preservation assigns
+one call to the inside resource's matching port. `sampleReplyProgram_memorylessReplyLaw`
+computes one reply law; `sampleReplies_smul_ofMemoryless` lifts it to the entire system.
+
+-/
+
+@[expose] public section
+
+namespace SystemAlgebra
+
+open DSL Probability Probability.Distribution
+
+-- `I` indexes queries `X i` and inside replies `Z i`; `K` carries fresh randomness and `Y` replies.
+variable {I : Type} {X Z : I → Type} {Y K : Type}
+
+/-- On `i(x)`, query the auxiliary source, then inside port `i`, and return `f i k z`.
+The state is `Unit`; `k` is used only within this invocation. -/
+noncomputable def sampleReplyProgram
+    -- `f i` combines a fresh value `k : K` and inside reply `z : Z i` into the outside reply.
+    (f : ∀ i, K → Z i → Y) :
+    Program Unit I (I ⊕ Unit) X (fun _ => Y) (Sum.rec X (fun _ => Unit))
+      (Sum.rec Z (fun _ => K)) := fun _ i x =>
+  callProg (.inr ()) () fun k =>
+    callProg (.inl i) x fun z => returnProg ((), f i k z)
+
+/-- One source call followed by one resource call uses at most two inside queries. -/
+lemma sampleReplyProgram_bounded (f : ∀ i, K → Z i → Y) :
+    (sampleReplyProgram (X := X) f).Bounded 2 := by
+  intro s i x
+  exact callBound_call _ _ _ (fun _ => callBound_call _ _ _ (fun _ => callBound_return 0 _))
+
+/-- An invocation at `i` calls resource port `inl i` once and no other resource port. -/
+lemma sampleReplyProgram_portPreserving (f : ∀ i, K → Z i → Y) :
+    (sampleReplyProgram (X := X) f).PortPreserving (fun i => some (.inl i)) := by
+  intro s i x
+  -- The source port is outside the protected resource ports.
+  apply preservesAt_call_other
+  · intro k
+    -- The sole resource call matches `i`; returning makes no further query.
+    apply preservesAt_call_self
+    · intro z; exact avoids_return _ _
+    · intro j hj; exact Sum.inl.inj (Option.some.inj hj)
+  · intro j hj; cases hj
+
+/-- An invocation returns the pushforward of the independent pair `(k,z) ∼ P × ν i x`. -/
+lemma sampleReplyProgram_memorylessReplyLaw
+    -- `P` supplies the fresh auxiliary value; `f i` combines it with the resource reply.
+    (P : ProbDist K) (f : ∀ i, K → Z i → Y)
+    -- `ν i x` is the resource's fresh reply law; `s` is the unique state at query `i(x)`.
+    (ν : ∀ i, X i → ProbDist (Z i)) (s : Unit) (i : I) (x : X i) :
+    (sampleReplyProgram f).memorylessReplyLaw
+        (RandomSystem.parallelLaw ν (fun _ _ => P)) s i x 2 [] =
+      ((prodProbDist P (ν i x)).map (fun kz => f i kz.1 kz.2)).1 := by
+  -- Read the source value and resource reply, then push their joint law through `f i`.
+  simp only [Program.memorylessReplyLaw, sampleReplyProgram, callProg, returnProg,
+    List.nil_append, List.cons_append, ↓reduceDIte, RandomSystem.parallelLaw,
+    ProbDist.map_val, prodProbDist_val]
+  rw [prod_eq_bindK, fTransform_bindK]
+  apply bindK_congr
+  intro k _
+  rw [fTransform_fTransform, bindK_single_comp]
+  rfl
+
+namespace TypedConverter
+
+-- These typed components use finite ports and alphabets, with allowance `q i` at each port.
+variable [Fintype I] [∀ i, Fintype (X i)] [∀ i, Fintype (Z i)] [Fintype Y] [Fintype K]
+
+/-- The two-call program on per-port domains, before supplying its independent source.
+The resource keeps allowance `q i`; the source allows the total call bound `2 * ∑ i, q i`. -/
+noncomputable def sampleRepliesProgram
+    -- `I`, `X`, `Z` specify the resource ports, queries and replies.
+    (I : Type) [Fintype I] (X Z : I → Type) [∀ i, Fintype (X i)] [∀ i, Fintype (Z i)]
+    -- `q i` bounds outside and resource calls at `i`; `f i` combines replies with fresh `K`.
+    (q : I → ℕ) (f : ∀ i, K → Z i → Y) :
+    TypedConverter (Interface.portBudget I X (fun _ => Y) q)
+      (Interface.tensor (Interface.portBudget I X Z q)
+        (Interface.source K (2 * (Interface.portBudget I X Z q).bound))) :=
+  ofPreservingProgram () (sampleReplyProgram f) (sampleReplyProgram_bounded f)
+    (fun i => some (.inl i)) (sampleReplyProgram_portPreserving f) q (fun _ hd => hd.2)
+    (fun ys hne hl hc => Interface.perPort_tensor_admits
+      (A := Interface.portBudget I X Z q)
+      (B := Interface.source K (2 * (Interface.portBudget I X Z q).bound)) q
+      (fun _ hn hq => ⟨hn, hq⟩) (fun _ hn hq => ⟨hn, hq⟩)
+      ys hne hl (fun i => hc i _ rfl))
+
+/-- Draw `k ∼ P` afresh, forward `i(x)` once, then return `f i k z` on resource reply `z`.
+Both public interfaces permit `q i` calls at port `i`. -/
+noncomputable def sampleReplies
+    -- `I`, `X`, `Z` specify the resource ports, queries and replies.
+    (I : Type) [Fintype I] (X Z : I → Type) [∀ i, Fintype (X i)] [∀ i, Fintype (Z i)]
+    -- `q` supplies per-port allowances; `P` supplies independent auxiliary values in `K`.
+    (q : I → ℕ) (P : ProbDist K) (f : ∀ i, K → Z i → Y) :
+    TypedConverter (Interface.portBudget I X (fun _ => Y) q) (Interface.portBudget I X Z q) :=
+  comp (sampleRepliesProgram I X Z q f)
+    (rightContext _ (TypedSystem.source P (2 * (Interface.portBudget I X Z q).bound)))
+
+/-- Fresh resource law `ν i x` gives fresh outside law `(k,z) ↦ f i k z` under `P × ν i x`. -/
+lemma sampleReplies_smul_ofMemoryless
+    -- `q i` bounds calls; `P` is the auxiliary law and `f i` computes the outside reply.
+    (q : I → ℕ) (P : ProbDist K) (f : ∀ i, K → Z i → Y)
+    -- `ν i x` is the inside reply law, independent of previous replies and source draws.
+    (ν : ∀ i, X i → ProbDist (Z i)) :
+    sampleReplies I X Z q P f • TypedSystem.ofMemoryless (Interface.portBudget I X Z q) ν =
+      TypedSystem.ofMemoryless (Interface.portBudget I X (fun _ => Y) q)
+        (fun i x => (prodProbDist P (ν i x)).map (fun kz => f i kz.1 kz.2)) := by
+  -- Attachment places the independent source beside the resource, then runs the program.
+  rw [sampleReplies, comp_smul, attach_rightContext,
+    TypedSystem.source_eq_ofMemoryless, TypedSystem.parallel_ofMemoryless]
+  -- The invocation law is state-independent, so it identifies every cumulative probability.
+  exact ofPreservingProgram_smul_ofMemoryless _ _ _ _ _ _ _ _ rfl
+    (sampleReplyProgram_memorylessReplyLaw P f ν)
+
+end TypedConverter
+end SystemAlgebra

src/Tests.lean

@@ -47,6 +47,7 @@ public import Tests.DSL.Inlining
 public import Tests.DSL.ParallelOutside
 public import Tests.DSL.ProbabilityNotation
 public import Tests.DSL.Semantics
+public import Tests.FreshReplyTransform
 public import Tests.NormalizedProbability
 public import Tests.ProbabilityNotation
 public import Tests.ProbabilityTablePatterns

src/Tests/FreshReplyTransform.lean

@@ -0,0 +1,87 @@
+module
+
+public import Commons.Definitions.ReplyTransform
+
+/-!
+# Fresh replies across repeated calls
+
+`freshBits` draws an auxiliary uniform bit on every invocation. Its two replies `false, true`
+have probability `1/4`. `retainedBits` chooses one bit for both replies and gives that transcript
+probability zero. The source contract therefore checks sampling scope across invocations.
+-/
+
+@[expose] public section
+
+namespace Tests.FreshReplyTransform
+
+open Commons Commons.Evaluation SystemAlgebra Probability Probability.Distribution
+
+/-- Two `eval(())` queries, each returning a bit. -/
+abbrev bitInterface : Interface := Evaluation.perPort Unit Bool (fun _ => 2)
+
+/-- The transcript of two evaluations returning `false` and then `true`. -/
+def mixedReplies : List
+    ((Σ i, bitInterface.X i) × (Σ i, bitInterface.Y i)) :=
+  [(⟨.eval, ()⟩, ⟨.eval, false⟩), (⟨.eval, ()⟩, ⟨.eval, true⟩)]
+
+/-- Each evaluation draws a fresh fair bit and ignores the inside source's independent reply. -/
+noncomputable def freshBits : TypedSystem bitInterface :=
+  sampleReplies 𝒰[Bool] (fun b (_ : Bool) => b) (fun _ => 2) •
+    MemorylessSource.perPort Bool 𝒰[Bool] (budget := fun _ => 2)
+
+/-- Sample a fair bit once and return that same bit on both evaluations. -/
+noncomputable def retainedBits : TypedSystem bitInterface :=
+  TypedSystem.sample 𝒰[Bool] (fun b => TypedSystem.ofFunction fun _ _ => b)
+
+/-- A fresh reply law `P` assigns the product `P(false) * P(true)` to two replies. -/
+lemma two_reply_probability (P : ProbDist Bool) :
+    (TypedSystem.ofMemoryless bitInterface (fun _ _ => P)).1 mixedReplies =
+      P.1 false * P.1 true := by
+  have d1 : bitInterface.inputDomain [] ⟨.eval, ()⟩ := by
+    constructor
+    · exact List.cons_ne_nil _ _
+    · intro i; cases i
+      norm_num [restrict_cons_self, restrict_nil]
+  have d2 : bitInterface.inputDomain [(⟨.eval, ()⟩, ⟨.eval, false⟩)] ⟨.eval, ()⟩ := by
+    constructor
+    · exact List.cons_ne_nil _ _
+    · intro i; cases i
+      norm_num [restrict_cons_self, restrict_nil]
+  -- Each admitted reply multiplies the preceding transcript mass by its fresh probability.
+  change RandomSystem.memoryless bitInterface.inputDomain (RandomSystem.portLaw (fun _ _ => P))
+    ([(⟨.eval, ()⟩, ⟨.eval, false⟩)] ++ [(⟨.eval, ()⟩, ⟨.eval, true⟩)]) = _
+  rw [RandomSystem.memoryless_snoc, ite_eq_left d2, RandomSystem.portLaw_apply_mk]
+  change RandomSystem.memoryless bitInterface.inputDomain (RandomSystem.portLaw (fun _ _ => P))
+    ([] ++ [(⟨.eval, ()⟩, ⟨.eval, false⟩)]) * _ = _
+  rw [RandomSystem.memoryless_snoc, ite_eq_left d1, RandomSystem.portLaw_apply_mk,
+    RandomSystem.mass_nil, one_mul]
+
+/-- The public component contract makes the auxiliary bits independent across calls. -/
+lemma freshBits_probability : freshBits.1 mixedReplies = 1 / 4 := by
+  -- Discarding the independent inside reply leaves the auxiliary uniform law.
+  have hp : (prodProbDist 𝒰[Bool] 𝒰[Bool]).map Prod.fst = 𝒰[Bool] := by
+    apply Subtype.ext
+    change fTransform Prod.fst (prod (uniform Bool) (uniform Bool)) = uniform Bool
+    rw [fTransform_fst_prod, weight_uniform, one_smul]
+  rw [freshBits, sampleReplies_smul_memorylessSource, hp, MemorylessSource.eq_ofMemoryless]
+  rw [two_reply_probability]
+  change uniform Bool false * uniform Bool true = 1 / 4
+  norm_num [uniform_apply]
+
+/-- A retained bit cannot produce two different replies. -/
+lemma retainedBits_probability : retainedBits.1 mixedReplies = 0 := by
+  simp only [retainedBits, TypedSystem.sample_mass, TypedSystem.ofFunction_mass]
+  simp [mixedReplies]
+
+/-- Fresh sampling cannot be replaced by one draw shared by all invocations. -/
+example : freshBits ≠ retainedBits := by
+  intro he
+  have hm := congrArg (fun R : TypedSystem bitInterface => R.1 mixedReplies) he
+  rw [freshBits_probability, retainedBits_probability] at hm
+  norm_num at hm
+
+#print axioms TypedConverter.sampleReplies_smul_ofMemoryless
+#print axioms sampleReplies_smul_memorylessSource
+#print axioms freshBits_probability
+
+end Tests.FreshReplyTransform