Add State.runViaIO that interprets state in IOEff

Author: paulcadman

Effects/IOEff.lean

@@ -1,6 +1,5 @@
 import Effects.Container
 import Effects.Prog
-import Effects.State
 
 open scoped Container
 
@@ -51,25 +50,22 @@ end IOEff
 
 section Examples
 
-def logIO {effs : List Effect} [IOEff ∈ effs] (msg : String) : Prog effs Unit :=
+def logIO
+  {effs : List Effect}
+  [IOEff ∈ effs]
+  (msg : String)
+  : Prog effs Unit :=
   IOEff.embed (IO.println msg)
 
-def exStateReaderIO : Prog [State Nat, IOEff] Unit := do
-  let (n : Nat) ← State.get
-  logIO s!"n={n}"
-  State.put (n + 1)
-  logIO "updated state"
-
-def runStateReaderIO : IO Unit :=
-  exStateReaderIO
-    |> State.eval 1
-    |> IOEff.run
+def exLog : Prog [IOEff] Unit := do
+  logIO "a"
+  logIO "b"
 
 /--
-info: n=1
-updated state
+info: a
+b
 -/
 #guard_msgs in
-#eval runStateReaderIO
+#eval exLog |> IOEff.run
 
 end Examples

Effects/State.lean

@@ -1,5 +1,6 @@
 import Effects.Container
 import Effects.Prog
+import Effects.IOEff
 
 open scoped Container
 
@@ -67,16 +68,110 @@ def run
 def eval (s : S) (p : Prog (State S :: effs) α) : Prog effs α :=
   Prod.snd <$> run s p
 
+def runViaIO'
+  [IOEff ∈ effs]
+  (p : Prog (State S :: effs) α) :
+  IO.Ref S → Prog effs (S × α) :=
+  Prog.foldP
+    (P := fun _ => IO.Ref S → Prog effs (S × α))
+    (var0 := fun x => fun iost =>
+      IOEff.embed iost.get |>.bind (fun st => pure (st, x)))
+    (varS := id)
+    (op := fun ⟨c, k⟩ =>
+      match c with
+      | .inl StateOps.getOp => fun iost =>
+        IOEff.embed iost.get |>.bind (fun st => k st iost)
+      | .inl (StateOps.putOp s') => fun iost =>
+        IOEff.embed (iost.set s') |>.bind (fun _ => k () iost)
+      | .inr s => fun iost =>
+        Prog.op ⟨s, fun p => k p iost⟩)
+    (scp := fun ⟨c, k⟩ =>
+      match c with
+      | .inl x => nomatch x
+      | .inr s => fun iost => Prog.scp ⟨s, fun p => ProgN.varS (k p iost)⟩)
+    p
+
+def runViaIO
+  [IOEff ∈ effs]
+  (s : IO.Ref S)
+  (p : Prog (State S :: effs) α)
+  : Prog effs (S × α) :=
+  runViaIO' p s
+
+def evalViaIO
+  [IOEff ∈ effs]
+  (s : IO.Ref S)
+  (p : Prog (State S :: effs) α)
+  : Prog effs α :=
+  Prod.snd <$> runViaIO' p s
+
 end State
 
 section Examples
 
 open State
 
-def tick {effs} [State Nat ∈ effs] : Prog effs Unit := do
-  let i ← get
-  put (1 + i)
-
-#guard Prog.run (State.run 0 (do tick; tick)) = (2, ())
+def ticks
+  {effs}
+  [State Nat ∈ effs]
+  : Prog effs Nat := do
+  let n ← get
+  put (n + 1)
+  let m ← get
+  pure (n + 2 * m)
+
+
+def runTicksPure : Nat × Nat :=
+  ticks
+  |> State.run 0
+  |> Prog.run
+
+#guard runTicksPure = (1, 2)
+
+def runTicksIO : IO (Nat × Nat) := do
+  let ref ← IO.mkRef 0
+  ticks
+  |> State.runViaIO ref
+  |> IOEff.run
+
+/--
+info: (1, 2)
+-/
+#guard_msgs in
+#eval runTicksIO
+
+def twoState
+  {effs}
+  [State Bool ∈ effs]
+  [State Nat ∈ effs]
+  : Prog effs Unit := do
+  let b ← get
+  let n ← get
+  put (n + 1)
+  let n ← get
+  put (b && n > 5)
+
+def runTwoStatePure : Nat × Bool × Unit :=
+  twoState
+  |> State.run true
+  |> State.run 4
+  |> Prog.run
+
+#guard runTwoStatePure = (5, false, ())
+
+def runTwoStateIO : IO (Nat × Nat × Bool × Unit) := do
+  let ref : IO.Ref Nat ← IO.mkRef 4
+  let res ← twoState
+    |> State.run true
+    |> State.runViaIO ref
+    |> IOEff.run
+  -- Also check that the IO.Ref was updated
+  (← ref.get, res) |> pure
+
+/--
+info: (5, 5, false, ())
+-/
+#guard_msgs in
+#eval runTwoStateIO
 
 end Examples