Query Implementations with State Monads

This file gives lemmas about QueryImpl spec m when m is something like StateT σ n.

TODO: should generalize things to MonadState once laws for it exist.

noncomputable def QueryImpl.mapStateTBase {ι₀ : Type u_1} {ι₁ : Type u_2} {spec₀ : OracleSpec ι₀} {spec₁ : OracleSpec ι₁} {m : Type u → Type v} [Monad m] {σ : Type u} (outer : QueryImpl spec₁ m) (inner : QueryImpl spec₀ (StateT σ (OracleComp spec₁))) : QueryImpl spec₀ (StateT σ m)

Push an outer oracle interpretation through the base monad of a StateT-valued query implementation.

theorem QueryImpl.simulateQ_mapStateTBase_run {ι₀ : Type u_1} {ι₁ : Type u_2} {spec₀ : OracleSpec ι₀} {spec₁ : OracleSpec ι₁} {m : Type u → Type v} [Monad m] [LawfulMonad m] {σ : Type u} (outer : QueryImpl spec₁ m) (inner : QueryImpl spec₀ (StateT σ (OracleComp spec₁))) {α : Type u} (oa : OracleComp spec₀ α) (s : σ) : simulateQ outer ((simulateQ inner oa).run s) = (simulateQ (outer.mapStateTBase inner) oa).run s

Running a StateT handler and then interpreting its base oracle computations is the same as first mapping the handler's base through the outer interpreter.

theorem QueryImpl.simulateQ_mapStateTBase_run' {ι₀ : Type u_1} {ι₁ : Type u_2} {spec₀ : OracleSpec ι₀} {spec₁ : OracleSpec ι₁} {m : Type u → Type v} [Monad m] [LawfulMonad m] {σ : Type u} (outer : QueryImpl spec₁ m) (inner : QueryImpl spec₀ (StateT σ (OracleComp spec₁))) {α : Type u} (oa : OracleComp spec₀ α) (s : σ) : simulateQ outer ((simulateQ inner oa).run' s) = (simulateQ (outer.mapStateTBase inner) oa).run' s

Output-only corollary of simulateQ_mapStateTBase_run.

def QueryImpl.parallelStateT {ι₁ : Type u_1} {ι₂ : Type u_2} {spec₁ : OracleSpec ι₁} {spec₂ : OracleSpec ι₂} {m : Type (max u_3 u_4) → Type u_5} [Functor m] {σ₁ σ₂ : Type (max u_3 u_4)} (impl₁ : QueryImpl spec₁ (StateT σ₁ m)) (impl₂ : QueryImpl spec₂ (StateT σ₂ m)) : QueryImpl (spec₁ + spec₂) (StateT (σ₁ × σ₂) m)

Given implementations for oracles in spec₁ and spec₂ in terms of state monads for two different contexts σ₁ and σ₂, implement the combined set spec₁ + spec₂ in terms of a combined σ₁ × σ₂ state.

def QueryImpl.flattenStateT {ι : Type u_1} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] {σ τ : Type u} (impl : QueryImpl spec (StateT σ (StateT τ m))) : QueryImpl spec (StateT (σ × τ) m)

Reassociate a nested state transformer into one product state.

The outer state is the first component of the product; the inner/base state is the second component. This is the state-transformer analogue of reassociating handler stacks into an explicit joint state before applying projection lemmas.

@[simp]

theorem QueryImpl.flattenStateT_liftTarget_apply_run {ι : Type u_1} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] [LawfulMonad m] {σ τ : Type u} (impl : QueryImpl spec (StateT τ m)) (t : spec.Domain) (s : σ) (q : τ) : ((liftTarget (StateT σ (StateT τ m)) impl).flattenStateT t).run (s, q) = (fun (y : spec.Range t × τ) => (y.1, s, y.2)) <$> (impl t).run q

def QueryImpl.piStateT {τ : Type} [DecidableEq τ] {ι : τ → Type u_1} {spec : (t : τ) → OracleSpec (ι t)} {m : Type (max u_2 u_3) → Type u_4} [Monad m] {σ : τ → Type (max u_2 u_3)} (impl : (t : τ) → QueryImpl (spec t) (StateT (σ t) m)) : QueryImpl (OracleSpec.sigma spec) (StateT ((t : τ) → σ t) m)

Indexed version of QueryImpl.parallelStateT. Note that m cannot vary with t. dtumad: The Function.update thing is nice but forces DecidableEq.

def QueryImpl.withBadFlag {ι : Type u_1} {spec : OracleSpec ι} {m : Type (max u_2 u_3) → Type u_4} [Functor m] {σ : Type (max u_2 u_3)} (impl : QueryImpl spec (StateT σ m)) : QueryImpl spec (StateT (σ × Bool) m)

Lift a stateful query implementation to a (state × Bool)-stateful version that threads the boolean (bad) flag unchanged. The output value and updated state come from the underlying impl; the second Bool component is preserved verbatim across each query.

def QueryImpl.withBadUpdate {ι : Type u_1} {spec : OracleSpec ι} {m : Type u_2 → Type u_3} [Functor m] {σ : Type u_2} (impl : QueryImpl spec (StateT σ m)) (f : (t : spec.Domain) → σ → spec.Range t → Bool) : QueryImpl spec (StateT (σ × Bool) m)

Lift a stateful query implementation to a (state × Bool)-stateful version that OR-updates the boolean (bad) flag with a predicate f evaluated on the pre-state and produced output. The flag is monotone: if it was already true, it stays true.

@[simp]

theorem QueryImpl.withBadFlag_apply_run {ι : Type u_1} {spec : OracleSpec ι} {m : Type (max u_2 u_3) → Type u_4} [Functor m] {σ : Type (max u_2 u_3)} (impl : QueryImpl spec (StateT σ m)) (t : spec.Domain) (s : σ) (b : Bool) : (impl.withBadFlag t).run (s, b) = (fun (vs : spec.Range t × σ) => (vs.1, vs.2, b)) <$> (impl t).run s

Run-shape of withBadFlag: the lifted implementation maps the underlying run by tagging each (value, state) pair with the unchanged bad flag b.

@[simp]

theorem QueryImpl.withBadUpdate_apply_run {ι : Type u_1} {spec : OracleSpec ι} {m : Type (max u_2 u_3) → Type u_4} [Functor m] {σ : Type (max u_2 u_3)} (impl : QueryImpl spec (StateT σ m)) (f : (t : spec.Domain) → σ → spec.Range t → Bool) (t : spec.Domain) (s : σ) (b : Bool) : (impl.withBadUpdate f t).run (s, b) = (fun (vs : spec.Range t × σ) => (vs.1, vs.2, b || f t s vs.1)) <$> (impl t).run s

Run-shape of withBadUpdate: the lifted implementation maps the underlying run by appending the OR-updated bad flag b || f t s vs.1.

theorem OracleComp.StateT_run_simulateQ_eq_map_run'_simulateQ {ι : Type u_1} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] [LawfulMonad m] {σ : Type u} (so : QueryImpl spec (StateT σ m)) {α : Type u} [Subsingleton σ] (oa : OracleComp spec α) (s s' : σ) : (simulateQ so oa).run s = (fun (x : α) => (x, s')) <$> (simulateQ so oa).run' s

If the state type is Subsingleton, then we can represent simulation in terms of simulate', adding back any state at the end of the computation.

theorem OracleComp.StateT_run'_simulateQ_eq_self {ι : Type u_1} {spec : OracleSpec ι} {σ α : Type u} (so : QueryImpl spec (StateT σ (OracleComp spec))) (h : ∀ (t : spec.Domain) (s : σ), (so t).run' s = query t) (oa : OracleComp spec α) (s : σ) : (simulateQ so oa).run' s = oa

theorem OracleComp.liftM_run_StateT {m : Type u → Type v} [Monad m] {σ α : Type u} (x : m α) (s : σ) : (liftM x).run s = do let a ← x pure (a, s)

theorem OracleComp.simulateQ_flattenStateT_run {ι : Type u_1} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] [LawfulMonad m] {σ τ : Type u} (impl : QueryImpl spec (StateT σ (StateT τ m))) {α : Type u} (oa : OracleComp spec α) (s : σ) (q : τ) : (simulateQ impl.flattenStateT oa).run (s, q) = do let __discr ← ((simulateQ impl oa).run s).run q match __discr with | ((a, s'), q') => pure (a, s', q')

Running a computation under a flattened nested-state implementation is the same as running the original nested computation and reassociating the final states into a product.

theorem OracleComp.simulateQ_flattenStateT_run' {ι : Type u_1} {spec : OracleSpec ι} {m : Type u → Type v} [Monad m] [LawfulMonad m] {σ τ : Type u} (impl : QueryImpl spec (StateT σ (StateT τ m))) {α : Type u} (oa : OracleComp spec α) (s : σ) (q : τ) : (simulateQ impl.flattenStateT oa).run' (s, q) = (Prod.fst <$> (simulateQ impl oa).run s).run' q

Output-only corollary of simulateQ_flattenStateT_run.

theorem OracleComp.simulateQ_mapStateTBase_run_eq_map_flattenStateT {ι₀ : Type u_2} {ι₁ : Type u_3} {spec₀ : OracleSpec ι₀} {spec₁ : OracleSpec ι₁} {m : Type u → Type v} [Monad m] [LawfulMonad m] {σ τ : Type u} (outer : QueryImpl spec₁ (StateT τ m)) (inner : QueryImpl spec₀ (StateT σ (OracleComp spec₁))) {α : Type u} (oa : OracleComp spec₀ α) (s : σ) (q : τ) : (simulateQ outer ((simulateQ inner oa).run s)).run q = (fun (z : α × σ × τ) => ((z.1, z.2.1), z.2.2)) <$> (simulateQ (outer.mapStateTBase inner).flattenStateT oa).run (s, q)

Running an adversary-side StateT handler under an outer stateful interpreter produces the same distribution as the flattened product-state handler, up to reassociating ((output, localState), outerState) and (output, (localState, outerState)).