Counting Queries Made by a Computation

This file defines a simulation oracle countingOracle for counting the number of queries made while running the computation. The count is represented by a function from oracle indices to counts, allowing each oracle to be tracked individually.

Tracking individually is not necessary, but gives tighter security bounds in some cases. It also allows for generating things like seed values for a computation more tightly.

def QueryImpl.withCost {ι : Type u} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] {ω : Type u} [Monoid ω] (so : QueryImpl spec m) (costFn : spec.Domain → ω) : QueryImpl spec (WriterT ω m)

Wrap an oracle implementation to accumulate cost in a WriterT ω layer. The cost function costFn assigns a cost value to each oracle query. Cost is accumulated before the implementation runs, so failed queries are still costed.

Side-channel trace instrumentation. withCost doubles as automatic side-channel instrumentation: given any OracleComp spec α and a cost function costFn : spec.Domain → ω, (simulateQ (base.withCost costFn) oa).run produces m (α × ω) without modifying the computation's source code. The cost function encodes the observation model:

• fun _ => () (constant-time): every query looks identical to the observer.
• fun t => queryLabel t (typed): the observer sees which oracle was queried but not the arguments or results.

This is the preferred approach when the observation model is "every oracle query is visible." For observations at non-query points (e.g. pure-computation timing), use the explicit observe / runObs API from ObservationOracle.

@[simp]

theorem QueryImpl.withCost_apply {ι : Type u} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] {ω : Type u} [Monoid ω] (so : QueryImpl spec m) (costFn : spec.Domain → ω) (t : spec.Domain) : so.withCost costFn t = do tell (costFn t) liftM (so t)

theorem QueryImpl.fst_map_run_withCost {ι : Type u} {spec : OracleSpec ι} {α : Type u} {m : Type u → Type v} [Monad m] {ω : Type u} [Monoid ω] [LawfulMonad m] (so : QueryImpl spec m) (costFn : spec.Domain → ω) (mx : OracleComp spec α) : Prod.fst <$> (simulateQ (so.withCost costFn) mx).run = simulateQ so mx

theorem QueryImpl.probFailure_run_simulateQ_withCost {ι : Type u} {spec : OracleSpec ι} {α : Type u} {m : Type u → Type v} [Monad m] {ω : Type u} [Monoid ω] [LawfulMonad m] [HasEvalSPMF m] (so : QueryImpl spec m) (costFn : spec.Domain → ω) (mx : OracleComp spec α) : Pr[⊥ | (simulateQ (so.withCost costFn) mx).run] = Pr[⊥ | simulateQ so mx]

Cost-tracking preserves failure probability: for any base monad m with HasEvalSPMF, wrapping an oracle implementation with withCost does not change the probability of failure.

theorem QueryImpl.NeverFail_run_simulateQ_withCost_iff {ι : Type u} {spec : OracleSpec ι} {α : Type u} {m : Type u → Type v} [Monad m] {ω : Type u} [Monoid ω] [LawfulMonad m] [HasEvalSPMF m] (so : QueryImpl spec m) (costFn : spec.Domain → ω) (mx : OracleComp spec α) : NeverFail (simulateQ (so.withCost costFn) mx).run ↔ NeverFail (simulateQ so mx)

@[simp]

theorem QueryImpl.run_simulateQ_withCost_const_one {ι : Type u} {spec : OracleSpec ι} {α : Type u} {m : Type u → Type v} [Monad m] {ω : Type u} [Monoid ω] [LawfulMonad m] (so : QueryImpl spec m) (mx : OracleComp spec α) : (simulateQ (so.withCost fun (x : spec.Domain) => 1) mx).run = (fun (x : α) => (x, 1)) <$> simulateQ so mx

When every query costs the monoid identity 1, the trace is always 1, so withCost is a no-op up to pairing with 1.

EvalDist Bridge for withCost

These lemmas connect the result-marginal distribution of a withCost-instrumented computation to the distribution of the uninstrumented computation, enabling direct probability-level reasoning about traced computations.

theorem QueryImpl.evalDist_fst_run_withCost {ι : Type u} {spec : OracleSpec ι} {α : Type u} {m : Type u → Type v} [Monad m] {ω : Type u} [Monoid ω] [LawfulMonad m] [HasEvalSPMF m] (so : QueryImpl spec m) (costFn : spec.Domain → ω) (mx : OracleComp spec α) : evalDist (Prod.fst <$> (simulateQ (so.withCost costFn) mx).run) = evalDist (simulateQ so mx)

theorem QueryImpl.probOutput_fst_run_withCost {ι : Type u} {spec : OracleSpec ι} {α : Type u} {m : Type u → Type v} [Monad m] {ω : Type u} [Monoid ω] [LawfulMonad m] [HasEvalSPMF m] (so : QueryImpl spec m) (costFn : spec.Domain → ω) (mx : OracleComp spec α) (x : α) : Pr[= x | Prod.fst <$> (simulateQ (so.withCost costFn) mx).run] = Pr[= x | simulateQ so mx]

theorem QueryImpl.support_fst_run_withCost {ι : Type u} {spec : OracleSpec ι} {α : Type u} {m : Type u → Type v} [Monad m] {ω : Type u} [Monoid ω] [LawfulMonad m] [HasEvalSPMF m] (so : QueryImpl spec m) (costFn : spec.Domain → ω) (mx : OracleComp spec α) : support (Prod.fst <$> (simulateQ (so.withCost costFn) mx).run) = support (simulateQ so mx)

def QueryImpl.withCounting {ι : Type u} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] [DecidableEq ι] (so : QueryImpl spec m) : QueryImpl spec (WriterT (OracleSpec.QueryCount ι) m)

Wrap an oracle implementation to count queries in a WriterT (QueryCount ι) layer. Counting happens before the implementation runs, so failed queries are still counted. This is a special case of withCost where the cost function is QueryCount.single.

@[simp]

theorem QueryImpl.withCounting_apply {ι : Type u} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] [DecidableEq ι] (so : QueryImpl spec m) (t : spec.Domain) : so.withCounting t = do tell (OracleSpec.QueryCount.single t) liftM (so t)

theorem QueryImpl.withCounting_eq_withCost {ι : Type u} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] [DecidableEq ι] (so : QueryImpl spec m) : so.withCounting = so.withCost fun (x : spec.Domain) => OracleSpec.QueryCount.single x

theorem QueryImpl.fst_map_run_withCounting {ι : Type u} {spec : OracleSpec ι} {α : Type u} {m : Type u → Type v} [Monad m] [DecidableEq ι] [LawfulMonad m] (so : QueryImpl spec m) (mx : OracleComp spec α) : Prod.fst <$> (simulateQ so.withCounting mx).run = simulateQ so mx

def costOracle {ι : Type u} {spec : OracleSpec ι} {ω : Type u} [Monoid ω] (costFn : spec.Domain → ω) : QueryImpl spec (WriterT ω (OracleComp spec))

Oracle with arbitrary cost tracking. The cost is accumulated in a WriterT ω layer while preserving the original oracle behavior.

def countingOracle {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] : QueryImpl spec (WriterT (OracleSpec.QueryCount ι) (OracleComp spec))

Oracle for counting the number of queries made by a computation. The count is stored as a function from oracle indices to counts, to give finer grained information about the count.

theorem countingOracle_eq_costOracle {ι : Type u} {spec : OracleSpec ι} [DecidableEq ι] : countingOracle = costOracle fun (x : spec.Domain) => OracleSpec.QueryCount.single x

@[simp]

theorem costOracle.fst_map_run_simulateQ {ι : Type u} {spec : OracleSpec ι} {α ω : Type u} [Monoid ω] (costFn : spec.Domain → ω) (oa : OracleComp spec α) : Prod.fst <$> (simulateQ (costOracle costFn) oa).run = oa

@[simp]

theorem costOracle.evalDist_fst_run_simulateQ {ι : Type u} {spec : OracleSpec ι} {α ω : Type u} [Monoid ω] [spec.Fintype] [spec.Inhabited] (costFn : spec.Domain → ω) (oa : OracleComp spec α) : evalDist (Prod.fst <$> (simulateQ (costOracle costFn) oa).run) = evalDist oa

@[simp]

theorem costOracle.probOutput_fst_run_simulateQ {ι : Type u} {spec : OracleSpec ι} {α ω : Type u} [Monoid ω] [spec.Fintype] [spec.Inhabited] (costFn : spec.Domain → ω) (oa : OracleComp spec α) (x : α) : Pr[= x | Prod.fst <$> (simulateQ (costOracle costFn) oa).run] = Pr[= x | oa]

@[simp]

theorem costOracle.support_run_simulateQ {ι : Type u} {spec : OracleSpec ι} {α ω : Type u} [Monoid ω] [spec.Fintype] [spec.Inhabited] (costFn : spec.Domain → ω) (oa : OracleComp spec α) : support (Prod.fst <$> (simulateQ (costOracle costFn) oa).run) = support oa

@[simp]

theorem countingOracle.fst_map_run_simulateQ {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) : Prod.fst <$> (simulateQ countingOracle oa).run = oa

@[simp]

theorem countingOracle.run_simulateQ_bind_fst {ι : Type u} {spec : OracleSpec ι} {α β : Type u} [DecidableEq ι] (oa : OracleComp spec α) (ob : α → OracleComp spec β) : (do let x ← (simulateQ countingOracle oa).run ob x.1) = oa >>= ob

@[simp]

theorem countingOracle.probFailure_run_simulateQ {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [spec₀.Fintype] [spec₀.Inhabited] {α : Type} (oa : OracleComp spec₀ α) : Pr[⊥ | (simulateQ countingOracle oa).run] = Pr[⊥ | oa]

Specialization of QueryImpl.probFailure_run_simulateQ_withCost to countingOracle.

@[simp]

theorem countingOracle.NeverFail_run_simulateQ_iff {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [spec₀.Fintype] [spec₀.Inhabited] {α : Type} (oa : OracleComp spec₀ α) : NeverFail (simulateQ countingOracle oa).run ↔ NeverFail oa

Specialization of QueryImpl.NeverFail_run_simulateQ_withCost_iff to countingOracle.

@[simp]

theorem countingOracle.probEvent_fst_run_simulateQ {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [spec₀.Fintype] [spec₀.Inhabited] {α : Type} (oa : OracleComp spec₀ α) (p : α → Prop) : (probEvent (simulateQ countingOracle oa).run fun (z : α × OracleSpec.QueryCount ι₀) => p z.1) = probEvent oa p

@[simp]

theorem countingOracle.probOutput_fst_map_run_simulateQ {ι₀ : Type} {spec₀ : OracleSpec ι₀} [DecidableEq ι₀] [spec₀.Fintype] [spec₀.Inhabited] {α : Type} (oa : OracleComp spec₀ α) (x : α) : Pr[= x | Prod.fst <$> (simulateQ countingOracle oa).run] = Pr[= x | oa]

def countingOracle.simulate {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) (qc : OracleSpec.QueryCount ι) : OracleComp spec (α × OracleSpec.QueryCount ι)

Compatibility helper matching old state-style counting semantics: simulate with zero initial count, then offset by qc.

theorem countingOracle.run_simulateT_eq_run_simulateT_zero {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) (qc : OracleSpec.QueryCount ι) : simulate oa qc = (Prod.map id fun (x : OracleSpec.QueryCount ι) => qc + x) <$> simulate oa 0

theorem countingOracle.support_simulate {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) (qc : OracleSpec.QueryCount ι) : support (simulate oa qc) = (Prod.map id fun (x : OracleSpec.QueryCount ι) => qc + x) '' support (simulate oa 0)

We can always reduce simulation with counting to start with 0, and add the initial count back at the end.

theorem countingOracle.mem_support_simulate_iff {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) (qc : OracleSpec.QueryCount ι) (z : α × OracleSpec.QueryCount ι) : z ∈ support (simulate oa qc) ↔ ∃ (qc' : OracleSpec.QueryCount ι), (z.1, qc') ∈ support (simulate oa 0) ∧ qc + qc' = z.2

Reduce membership in support of simulation with counting to simulation starting from 0.

theorem countingOracle.mem_support_simulate_iff_of_le {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) (qc : OracleSpec.QueryCount ι) (z : α × OracleSpec.QueryCount ι) (hz : qc ≤ z.2) : z ∈ support (simulate oa qc) ↔ (z.1, z.2 - qc) ∈ support (simulate oa 0)

theorem countingOracle.le_of_mem_support_simulate {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] {oa : OracleComp spec α} {qc : OracleSpec.QueryCount ι} {z : α × OracleSpec.QueryCount ι} (h : z ∈ support (simulate oa qc)) : qc ≤ z.2

theorem countingOracle.mem_support_snd_map_simulate_iff {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) (qc qc' : OracleSpec.QueryCount ι) : qc' ∈ support ((fun (z : α × OracleSpec.QueryCount ι) => z.2) <$> simulate oa qc) ↔ ∃ (qc'' : OracleSpec.QueryCount ι) (x : α), (x, qc'') ∈ support (simulate oa 0) ∧ qc + qc'' = qc'

theorem countingOracle.mem_support_snd_map_simulate_iff_of_le {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) {qc qc' : OracleSpec.QueryCount ι} (hqc : qc ≤ qc') : qc' ∈ support ((fun (z : α × OracleSpec.QueryCount ι) => z.2) <$> simulate oa qc) ↔ qc' - qc ∈ support ((fun (z : α × OracleSpec.QueryCount ι) => z.2) <$> simulate oa 0)

theorem countingOracle.le_of_mem_support_snd_map_simulate {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] {oa : OracleComp spec α} {qc qc' : OracleSpec.QueryCount ι} (h : qc' ∈ support ((fun (z : α × OracleSpec.QueryCount ι) => z.2) <$> simulate oa qc)) : qc ≤ qc'

theorem countingOracle.sub_mem_support_snd_map_simulate {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] {oa : OracleComp spec α} {qc qc' : OracleSpec.QueryCount ι} (h : qc' ∈ support ((fun (z : α × OracleSpec.QueryCount ι) => z.2) <$> simulate oa qc)) : qc' - qc ∈ support ((fun (z : α × OracleSpec.QueryCount ι) => z.2) <$> simulate oa 0)

theorem countingOracle.add_mem_support_simulate {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] {oa : OracleComp spec α} {qc : OracleSpec.QueryCount ι} {z : α × OracleSpec.QueryCount ι} (hz : z ∈ support (simulate oa qc)) (qc' : OracleSpec.QueryCount ι) : (z.1, z.2 + qc') ∈ support (simulate oa (qc + qc'))

@[simp]

theorem countingOracle.add_right_mem_support_simulate_iff {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) (qc qc' : OracleSpec.QueryCount ι) (x : α) : (x, qc + qc') ∈ support (simulate oa qc) ↔ (x, qc') ∈ support (simulate oa 0)

@[simp]

theorem countingOracle.add_left_mem_support_simulate_iff {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (oa : OracleComp spec α) (qc qc' : OracleSpec.QueryCount ι) (x : α) : (x, qc' + qc) ∈ support (simulate oa qc) ↔ (x, qc') ∈ support (simulate oa 0)

theorem countingOracle.mem_support_simulate_pure_iff {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (x : α) (qc : OracleSpec.QueryCount ι) (z : α × OracleSpec.QueryCount ι) : z ∈ support (simulate (pure x) qc) ↔ z = (x, qc)

theorem countingOracle.apply_ne_zero_of_mem_support_simulate_queryBind {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] {t : spec.Domain} {oa : spec.Range t → OracleComp spec α} {qc : OracleSpec.QueryCount ι} {z : α × OracleSpec.QueryCount ι} (hz : z ∈ support (simulate (liftM (OracleQuery.query t) >>= oa) qc)) : z.2 t ≠ 0

theorem countingOracle.exists_mem_support_of_mem_support_simulate_queryBind {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] {t : spec.Domain} {oa : spec.Range t → OracleComp spec α} {qc : OracleSpec.QueryCount ι} {z : α × OracleSpec.QueryCount ι} (hz : z ∈ support (simulate (liftM (OracleQuery.query t) >>= oa) qc)) : ∃ (u : spec.Range t), (z.1, Function.update z.2 t (z.2 t - 1)) ∈ support (simulate (oa u) qc)

theorem countingOracle.mem_support_simulate_queryBind_iff {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] (t : spec.Domain) (oa : spec.Range t → OracleComp spec α) (qc : OracleSpec.QueryCount ι) (z : α × OracleSpec.QueryCount ι) : z ∈ support (simulate (liftM (OracleQuery.query t) >>= oa) qc) ↔ z.2 t ≠ 0 ∧ ∃ (u : spec.Range t), (z.1, Function.update z.2 t (z.2 t - 1)) ∈ support (simulate (oa u) qc)

theorem countingOracle.exists_mem_support_of_mem_support {ι : Type u} {spec : OracleSpec ι} {α : Type u} [DecidableEq ι] {oa : OracleComp spec α} {x : α} (hx : x ∈ support oa) (qc : OracleSpec.QueryCount ι) : ∃ (qc' : OracleSpec.QueryCount ι), (x, qc') ∈ support (simulate oa qc)