Pre-computing Results of Oracle Queries

This file defines a function QueryImpl.withPregen that modifies a query implementation to take in a list of pre-chosen outputs to use when answering queries. Uses StateT so that consumed seed values are threaded to subsequent queries.

Note that ordering is subtle, for example so.withCaching.withPregen will first check for seeds and not cache the result if one is found, while so.withPregen.withCaching checks the cache first, and include seed values into the cache after returning them.

def QueryImpl.withPregen {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] {m : Type u → Type v} [Monad m] (so : QueryImpl spec m) : QueryImpl spec (StateT spec.QuerySeed m)

Modify a QueryImpl to check for pregenerated responses for oracle queries first. If a seed value is available for the query, it is used and the seed is consumed (the tail replaces the head). When no seed remains, falls back to the underlying implementation.

@[simp]

theorem QueryImpl.withPregen_apply {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] {m : Type u → Type v} [Monad m] (so : QueryImpl spec m) (t : spec.Domain) : so.withPregen t = StateT.mk fun (seed : spec.QuerySeed) => match seed t with | u :: us => pure (u, Function.update seed t us) | [] => (fun (x : spec.Range t) => (x, seed)) <$> so t

def seededOracle {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] : QueryImpl spec (StateT spec.QuerySeed (OracleComp spec))

Use pregenerated oracle responses for queries, falling back to the real oracle when the seed is exhausted. Seed consumption is tracked via StateT.

theorem seededOracle.probEvent_liftComp_uniformSample_eq_of_eq {ι : Type} {spec : OracleSpec ι} [(i : ι) → SampleableType (spec.Range i)] [unifSpec ⊂ₒ spec] [unifSpec.LawfulSubSpec spec] [spec.Fintype] [spec.Inhabited] {i : ι} (u₀ : spec.Range i) : (probEvent (OracleComp.liftComp ($ᵗ spec.Range i) spec) fun (u : spec.Range i) => u₀ = u) = (↑(Fintype.card (spec.Range i)))⁻¹

The probability that a lifted uniform sample equals a fixed value is (card α)⁻¹.

@[simp]

theorem seededOracle.apply_eq {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] (t : spec.Domain) : seededOracle t = StateT.mk fun (seed : spec.QuerySeed) => match seed t with | u :: us => pure (u, Function.update seed t us) | [] => (fun (x : spec.Range t) => (x, seed)) <$> liftM (OracleQuery.query t)

theorem seededOracle.run_bind_query_eq_pop {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] {α : Type u} (t : spec.Domain) (mx : spec.Range t → OracleComp spec α) (seed : spec.QuerySeed) : (do let u ← seededOracle t simulateQ seededOracle (mx u)).run seed = match seed.pop t with | none => do let u ← liftM (OracleQuery.query t) (simulateQ seededOracle (mx u)).run seed | some (u, seed') => (simulateQ seededOracle (mx u)).run seed'

@[simp]

theorem seededOracle.probOutput_generateSeed_bind_simulateQ_bind {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [(i : ι₀) → SampleableType (spec₀.Range i)] [unifSpec ⊂ₒ spec₀] [unifSpec.LawfulSubSpec spec₀] [spec₀.Fintype] [spec₀.Inhabited] (qc : ι₀ → ℕ) (js : List ι₀) {α β : Type} (oa : OracleComp spec₀ α) (ob : α → OracleComp spec₀ β) (y : β) : Pr[= y | do let seed ← OracleComp.liftComp (OracleComp.generateSeed spec₀ qc js) spec₀ let x ← Prod.fst <$> (simulateQ seededOracle oa).run seed ob x] = Pr[= y | oa >>= ob]

@[simp]

theorem seededOracle.probOutput_generateSeed_bind_map_simulateQ {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [(i : ι₀) → SampleableType (spec₀.Range i)] [unifSpec ⊂ₒ spec₀] [unifSpec.LawfulSubSpec spec₀] [spec₀.Fintype] [spec₀.Inhabited] (qc : ι₀ → ℕ) (js : List ι₀) {α β : Type} (oa : OracleComp spec₀ α) (f : α → β) (y : β) : Pr[= y | do let seed ← OracleComp.liftComp (OracleComp.generateSeed spec₀ qc js) spec₀ f <$> Prod.fst <$> (simulateQ seededOracle oa).run seed] = Pr[= y | f <$> oa]

theorem seededOracle.evalDist_liftComp_uniformSample_bind_simulateQ_run'_addValue {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [(j : ι₀) → SampleableType (spec₀.Range j)] [unifSpec ⊂ₒ spec₀] [unifSpec.LawfulSubSpec spec₀] [spec₀.Fintype] [spec₀.Inhabited] (σ : spec₀.QuerySeed) (i : ι₀) {α : Type} (oa : OracleComp spec₀ α) : (evalDist do let u ← OracleComp.liftComp ($ᵗ spec₀.Range i) spec₀ (simulateQ seededOracle oa).run' (σ.addValue i u)) = evalDist ((simulateQ seededOracle oa).run' σ)

Adding a uniform value at index i to a seed does not change the distribution of running a computation with the seeded oracle. This is because the extra value replaces what would otherwise be a fresh uniform oracle response.

theorem seededOracle.evalDist_liftComp_replicate_uniformSample_bind_simulateQ_run'_addValues {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [(j : ι₀) → SampleableType (spec₀.Range j)] [unifSpec ⊂ₒ spec₀] [unifSpec.LawfulSubSpec spec₀] [spec₀.Fintype] [spec₀.Inhabited] (i : ι₀) {α : Type} (oa : OracleComp spec₀ α) (n : ℕ) (σ : spec₀.QuerySeed) : (evalDist do let us ← (OracleComp.replicate n ($ᵗ spec₀.Range i)).liftComp spec₀ (simulateQ seededOracle oa).run' (σ.addValues us)) = evalDist ((simulateQ seededOracle oa).run' σ)

theorem seededOracle.evalDist_liftComp_generateSeed_bind_simulateQ_run'_takeAtIndex {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [(i : ι₀) → SampleableType (spec₀.Range i)] [unifSpec ⊂ₒ spec₀] [unifSpec.LawfulSubSpec spec₀] [spec₀.Fintype] [spec₀.Inhabited] (qc : ι₀ → ℕ) (js : List ι₀) (i₀ : ι₀) (k : ℕ) {α : Type} (oa : OracleComp spec₀ α) : (evalDist do let seed ← OracleComp.liftComp (OracleComp.generateSeed spec₀ qc js) spec₀ (simulateQ seededOracle oa).run' (seed.takeAtIndex i₀ k)) = evalDist oa

Truncating the seed at oracle i₀ to only the first k entries does not change the distribution when averaging over seeds from generateSeed.

@[simp]

theorem seededOracle.probOutput_generateSeed_bind_map_simulateQ_takeAtIndex {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [(i : ι₀) → SampleableType (spec₀.Range i)] [unifSpec ⊂ₒ spec₀] [unifSpec.LawfulSubSpec spec₀] [spec₀.Fintype] [spec₀.Inhabited] (qc : ι₀ → ℕ) (js : List ι₀) (i₀ : ι₀) (k : ℕ) {α β : Type} (oa : OracleComp spec₀ α) (f : α → β) (y : β) : Pr[= y | do let seed ← OracleComp.liftComp (OracleComp.generateSeed spec₀ qc js) spec₀ f <$> (simulateQ seededOracle oa).run' (seed.takeAtIndex i₀ k)] = Pr[= y | f <$> oa]

theorem seededOracle.tsum_probOutput_generateSeed_weight_takeAtIndex {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [(i : ι₀) → SampleableType (spec₀.Range i)] [unifSpec ⊂ₒ spec₀] [unifSpec.LawfulSubSpec spec₀] [spec₀.Fintype] [spec₀.Inhabited] (qc : ι₀ → ℕ) (js : List ι₀) (i₀ : ι₀) (k : ℕ) {α : Type} (oa : OracleComp spec₀ α) (x : α) (h : spec₀.QuerySeed → ENNReal) : ∑' (σ : spec₀.QuerySeed), Pr[= σ | OracleComp.generateSeed spec₀ qc js] * (h (σ.takeAtIndex i₀ k) * Pr[= x | (simulateQ seededOracle oa).run' σ]) = ∑' (σ : spec₀.QuerySeed), Pr[= σ | OracleComp.generateSeed spec₀ qc js] * (h (σ.takeAtIndex i₀ k) * Pr[= x | (simulateQ seededOracle oa).run' (σ.takeAtIndex i₀ k)])

Weighted takeAtIndex faithfulness: a prefix-dependent weight h preserves the faithfulness equality between full-seed and truncated-seed simulation. This generalizes the basic takeAtIndex faithfulness by allowing multiplication by an arbitrary function of the seed prefix σ.takeAtIndex i₀ k.

theorem seededOracle.isPerIndexQueryBound_run'_of_seedCoverage {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] {α : Type u} {oa : OracleComp spec α} {qb residual : ι → ℕ} {seed : spec.QuerySeed} (hqb : oa.IsPerIndexQueryBound qb) (hcover : ∀ (t : ι), qb t - residual t ≤ (seed t).length) : ((simulateQ seededOracle oa).run' seed).IsPerIndexQueryBound residual

If a pre-generated seed already supplies all but residual t of the answers allowed by the structural per-index query bound qb, then running oa against seededOracle can make at most those residual live queries.

This theorem is the core replay-cost statement for seeded simulations: a large enough seed turns most oracle interactions into deterministic table lookups, leaving only the uncovered suffix of the computation as genuine oracle queries.

theorem seededOracle.isPerIndexQueryBound_run'_zero {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] {α : Type u} {oa : OracleComp spec α} {qb : ι → ℕ} {seed : spec.QuerySeed} (hqb : oa.IsPerIndexQueryBound qb) (hcover : ∀ (t : ι), qb t ≤ (seed t).length) : ((simulateQ seededOracle oa).run' seed).IsPerIndexQueryBound 0

A seed that covers the full structural query bound eliminates all live oracle queries.

theorem seededOracle.isPerIndexQueryBound_run'_takeAtIndex {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] {α : Type u} {oa : OracleComp spec α} {qb : ι → ℕ} {seed : spec.QuerySeed} {i : ι} {k : ℕ} (hqb : oa.IsPerIndexQueryBound qb) (hcover : ∀ (t : ι), qb t ≤ (seed t).length) (hk : k ≤ qb i) : ((simulateQ seededOracle oa).run' (seed.takeAtIndex i k)).IsPerIndexQueryBound (Function.update 0 i (qb i - k))

If the seed stores only the first k answers for oracle i, then the replay can make live queries only to i, and at most qb i - k of them remain.

This is the structural query-bound form of the usual forking-lemma intuition: after rewinding to the k-th query to oracle i, every earlier answer is fixed by the prefix seed, so only the suffix after the fork point can still hit the live oracle.

theorem seededOracle.isPerIndexQueryBound_run'_takeAtIndex_addValue {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] {α : Type u} {oa : OracleComp spec α} {qb : ι → ℕ} {seed : spec.QuerySeed} {i : ι} (hqb : oa.IsPerIndexQueryBound qb) (hcover : ∀ (t : ι), qb t ≤ (seed t).length) (s : Fin (qb i + 1)) (u : spec.Range i) : ((simulateQ seededOracle oa).run' ((seed.takeAtIndex i ↑s).addValue i u)).IsPerIndexQueryBound (Function.update 0 i (qb i - (↑s + 1)))

After rewinding to query index s and appending one fresh answer at oracle i, the replayed run can still make live queries only to i, with at most qb i - (s + 1) such queries left.