Alternative tailRecM implementation (#57)

Avoids async bouncing by observing synchronous effects and looping.

Author: natefaubion

src/Control/Monad/Aff.js

@@ -293,3 +293,64 @@ exports._liftEff = function (nonCanceler, e) {
     return nonCanceler;
   };
 }
+
+exports._tailRecM = function (isLeft, f, a) {
+  return function(success, error) {
+    return function go(acc) {
+      var result, status, canceler;
+
+      // Observes synchronous effects using a flag.
+      //   status = 0 (unresolved status)
+      //   status = 1 (synchronous effect)
+      //   status = 2 (asynchronous effect)
+      while (true) {
+        status = 0;
+        canceler = f(acc)(function(v) {
+          // If the status is still unresolved, we have observed a
+          // synchronous effect. Otherwise, the status will be `2`.
+          if (status === 0) {
+            // Store the result for further synchronous processing.
+            result = v;
+            status = 1;
+          } else {
+            // When we have observed an asynchronous effect, we use normal
+            // recursion. This is safe because we will be on a new stack.
+            if (isLeft(v)) {
+              go(v.value0);
+            } else {
+              try {
+                success(v.value0);
+              } catch (err) {
+                error(err);
+              }
+            }
+          }
+        }, error);
+
+        // If the status has already resolved to `1` by our Aff handler, then
+        // we have observed a synchronous effect. Otherwise it will still be
+        // `0`.
+        if (status === 1) {
+          // When we have observed a synchronous effect, we merely swap out the
+          // accumulator and continue the loop, preserving stack.
+          if (isLeft(result)) {
+            acc = result.value0;
+            continue;
+          } else {
+            try {
+              success(result.value0);
+            } catch (err) {
+              error(err);
+            }
+          }
+        } else {
+          // If the status has not resolved yet, then we have observed an
+          // asynchronous effect.
+          status = 2;
+        }
+        return canceler;
+      }
+
+    }(a);
+  };
+};

src/Control/Monad/Aff.purs

@@ -33,7 +33,7 @@ import Control.Monad.Rec.Class (MonadRec, tailRecM)
 import Control.MonadPlus (MonadPlus)
 import Control.Plus (Plus)
 
-import Data.Either (Either(..), either)
+import Data.Either (Either(..), either, isLeft)
 import Data.Foldable (Foldable, foldl)
 import Data.Function (Fn2(), Fn3(), runFn2, runFn3)
 import Data.Monoid (Monoid, mempty)
@@ -191,14 +191,7 @@ instance alternativeAff :: Alternative (Aff e)
 instance monadPlusAff :: MonadPlus (Aff e)
 
 instance monadRecAff :: MonadRec (Aff e) where
-  tailRecM f a = go 0 f a
-    where
-    go size f a = do
-      e <- f a
-      case e of
-        Left a' | size < 100 -> go (size + 1) f a'
-                | otherwise -> later (tailRecM f a')
-        Right b -> pure b
+  tailRecM f a = runFn3 _tailRecM isLeft f a
 
 instance monadContAff :: MonadCont (Aff e) where
   callCC f = makeAff (\eb cb -> runAff eb cb (f \a -> makeAff (\_ _ -> cb a)))
@@ -234,3 +227,5 @@ foreign import _attempt :: forall e a. Fn3 (forall x y. x -> Either x y) (forall
 foreign import _runAff :: forall e a. Fn3 (Error -> Eff e Unit) (a -> Eff e Unit) (Aff e a) (Eff e Unit)
 
 foreign import _liftEff :: forall e a. Fn2 (Canceler e) (Eff e a) (Aff e a)
+
+foreign import _tailRecM :: forall e a b. Fn3 (Either a b -> Boolean) (a -> Aff e (Either a b)) a (Aff e b)

test/Test/Main.purs

@@ -123,13 +123,29 @@ test_cancelPar = do
   log (if v then "Success: Canceling composite of two Par succeeded"
                    else "Failure: Canceling composite of two Par failed")
 
-loop :: forall eff. Int -> Aff (console :: CONSOLE | eff) Unit
-loop n = tailRecM go n
+test_syncTailRecM :: TestAVar Unit
+test_syncTailRecM = do
+  v <- makeVar' false
+  _ <- forkAff $ tailRecM go { n: 1000000, v }
+  b <- takeVar v
+  log (if b then "Success: Synchronous tailRecM resolved synchronously"
+            else "Failure: Synchronous tailRecM resolved asynchronously")
   where
-  go 0 = do
-    log "Done!"
-    return (Right unit)
-  go n = return (Left (n - 1))
+  go { n = 0, v } = do
+    modifyVar (const true) v
+    pure (Right 0)
+  go { n, v } = pure (Left { n: n - 1, v })
+
+loopAndBounce :: forall eff. Int -> Aff (console :: CONSOLE | eff) Unit
+loopAndBounce n = do
+  res <- tailRecM go n
+  log $ "Done: " <> show res
+  where
+  go 0 = pure (Right 0)
+  go n | mod n 30000 == 0 = do
+    later' 10 (pure unit)
+    pure (Left (n - 1))
+  go n = pure (Left (n - 1))
 
 all :: forall eff. Int -> Aff (console :: CONSOLE, avar :: AVAR | eff) Unit
 all n = do
@@ -195,11 +211,14 @@ main = runAff throwException (const (pure unit)) $ do
   log "Testing cancel of Par (<|>)"
   test_cancelPar
 
+  log "Testing synchronous tailRecM"
+  test_syncTailRecM
+
   log "pre-delay"
   delay 1000
-
   log "post-delay"
-  loop 1000000
+
+  loopAndBounce 1000000
 
   all 100000