From 95f758d240a08fe56ab4260e76ceec32c8ab6fc9 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 26 Jul 2018 16:07:30 -0700
Subject: [PATCH] feat(library/system/io_interface): improve `iterate`
 interface

The new version is better for lifting.
---
 library/system/io.lean           | 12 ++++++------
 library/system/io_interface.lean |  2 +-
 src/library/vm/vm_io.cpp         | 10 +++++-----
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/library/system/io.lean b/library/system/io.lean
index 45da8a31b9..c45095b5bd 100644
--- a/library/system/io.lean
+++ b/library/system/io.lean
@@ -31,15 +31,15 @@ monad_io_is_alternative io_core
 io_core io.error α
 
 namespace io
-/- Remark: the following definitions can be generalized and defined for any (m : Type -> Type -> Type)
+/- Remark: the following definitions can be generalized and defined for any (m : Type -> Type -> Type 1)
    that implements the required type classes. However, the generalized versions are very inconvenient to use,
    (example: `#eval io.put_str "hello world"` does not work because we don't have enough information to infer `m`.).
 -/
-def iterate {e α} (a : α) (f : α → io_core e (option α)) : io_core e α :=
-monad_io.iterate e α a f
+def iterate {e} {α β : Type} (a : α) (f : α → io_core e (sum α β)) : io_core e β :=
+monad_io.iterate e α β a f
 
 def forever {e} (a : io_core e unit) : io_core e unit :=
-iterate () $ λ _, a >> return (some ())
+iterate () $ λ _, a >> return (sum.inl ())
 
 -- TODO(Leo): delete after we merge #1881
 def catch {e₁ e₂ α} (a : io_core e₁ α) (b : e₁ → io_core e₂ α) : io_core e₂ α :=
@@ -140,11 +140,11 @@ def read_to_end (h : handle) : io string :=
 iterate "" $ λ r,
   do done ← is_eof h,
     if done
-    then return none
+    then return (sum.inr r) -- stop
     else do
       -- HACK: use less efficient `get_line` while `read` is broken
       c ← get_line h,
-      return $ some (r ++ c)
+      return $ sum.inl (r ++ c) -- continue
 
 def read_file (fname : string) (bin := ff) : io string :=
 do h ← mk_file_handle fname io.mode.read bin,
diff --git a/library/system/io_interface.lean b/library/system/io_interface.lean
index bf472a4f52..e59b37890c 100644
--- a/library/system/io_interface.lean
+++ b/library/system/io_interface.lean
@@ -36,7 +36,7 @@ class monad_io (m : Type → Type → Type 1) :=
 -- TODO(Leo): use monad_except after it is merged
 (catch    : Π e₁ e₂ α, m e₁ α → (e₁ → m e₂ α) → m e₂ α)
 (fail     : Π e α, e → m e α)
-(iterate  : Π e α, α → (α → m e (option α)) → m e α)
+(iterate  : Π e (α β : Type), α → (α → m e (sum α β)) → m e β)
 -- Primitive Types
 (handle   : Type)
 
diff --git a/src/library/vm/vm_io.cpp b/src/library/vm/vm_io.cpp
index 9f9d9f665f..a16ab73e4d 100644
--- a/src/library/vm/vm_io.cpp
+++ b/src/library/vm/vm_io.cpp
@@ -394,8 +394,8 @@ static vm_obj io_fail(vm_obj const &, vm_obj const &, vm_obj const & e, vm_obj c
     return mk_io_failure(e);
 }
 
-/* (iterate  : Π e α, α → (α → m e (option α)) → m e α) */
-static vm_obj io_iterate(vm_obj const &, vm_obj const &, vm_obj const & a, vm_obj const & fn, vm_obj const &) {
+/* (iterate  : Π e (α β : Type), α → (α → m e (sum α β)) → m e β) */
+static vm_obj io_iterate(vm_obj const &, vm_obj const &, vm_obj const &, vm_obj const & a, vm_obj const & fn, vm_obj const &) {
     vm_obj r = a;
     while (true) {
         vm_obj p = invoke(fn, r, mk_vm_unit());
@@ -403,10 +403,10 @@ static vm_obj io_iterate(vm_obj const &, vm_obj const &, vm_obj const & a, vm_ob
             return p;
         } else {
             vm_obj v = cfield(p, 0);
-            if (is_none(v)) {
-                return mk_io_result(r);
+            if (cidx(v) == 1) {
+                return mk_io_result(cfield(v, 0));
             } else {
-                r = get_some_value(v);
+                r = cfield(v, 0);
             }
         }
     }