Refactoring RWST to an ADT

Improves performance to put it on par with StateT.

Author: ethul

docs/Control/Monad/RWS/Trans.md

@@ -5,13 +5,7 @@ This module defines the reader-writer-state monad transformer, `RWST`.
 #### `See`
 
 ``` purescript
-type See s a w = { state :: s, result :: a, log :: w }
-```
-
-#### `mkSee`
-
-``` purescript
-mkSee :: forall s a w. (Monoid w) => s -> a -> w -> See s a w
+data See state result writer
 ```
 
 #### `RWST`
@@ -26,7 +20,7 @@ of `ReaderT`, `WriterT` and `StateT` into a single monad transformer.
 
 ##### Instances
 ``` purescript
-instance functorRWST :: (Functor m) => Functor (RWST r w s m)
+instance functorRWST :: (Functor m, Monoid w) => Functor (RWST r w s m)
 instance applyRWST :: (Bind m, Monoid w) => Apply (RWST r w s m)
 instance bindRWST :: (Bind m, Monoid w) => Bind (RWST r w s m)
 instance applicativeRWST :: (Monad m, Monoid w) => Applicative (RWST r w s m)

src/Control/Monad/RWS/Trans.purs

@@ -1,7 +1,7 @@
 -- | This module defines the reader-writer-state monad transformer, `RWST`.
 
 module Control.Monad.RWS.Trans
-  ( See(), mkSee
+  ( See()
   , RWST(..), runRWST, evalRWST, execRWST, mapRWST, withRWST
   , module Control.Monad.Trans
   , module Control.Monad.RWS.Class
@@ -22,14 +22,7 @@ import Control.Monad.State.Class
 import Control.Monad.Trans
 import Control.Monad.Writer.Class
 
-type See s a w =
-  { state  :: s
-  , result :: a
-  , log    :: w
-  }
-
-mkSee :: forall s a w. (Monoid w) => s -> a -> w -> See s a w
-mkSee s a w = { state: s, result: a, log: w }
+data See state result writer = See state result writer
 
 -- | The reader-writer-state monad transformer, which combines the operations
 -- | of `ReaderT`, `WriterT` and `StateT` into a single monad transformer.
@@ -41,11 +34,11 @@ runRWST (RWST x) = x
 
 -- | Run a computation in the `RWST` monad, discarding the final state.
 evalRWST :: forall r w s m a. (Monad m) => RWST r w s m a -> r -> s -> m (Tuple a w)
-evalRWST m r s = runRWST m r s >>= \see -> return (Tuple see.result see.log)
+evalRWST m r s = runRWST m r s >>= \(See _ result writer) -> return (Tuple result writer)
 
 -- | Run a computation in the `RWST` monad, discarding the result.
 execRWST :: forall r w s m a. (Monad m) => RWST r w s m a -> r -> s -> m (Tuple s w)
-execRWST m r s = runRWST m r s >>= \see -> return (Tuple see.state see.log)
+execRWST m r s = runRWST m r s >>= \(See state _ writer) -> return (Tuple state writer)
 
 -- | Change the result and accumulator types in a `RWST` monad action.
 mapRWST :: forall r w1 w2 s m1 m2 a1 a2. (m1 (See s a1 w1) -> m2 (See s a2 w2)) -> RWST r w1 s m1 a1 -> RWST r w2 s m2 a2
@@ -55,43 +48,43 @@ mapRWST f m = RWST \r s -> f $ runRWST m r s
 withRWST :: forall r1 r2 w s m a. (r2 -> s -> Tuple r1 s) -> RWST r1 w s m a -> RWST r2 w s m a
 withRWST f m = RWST \r s -> uncurry (runRWST m) (f r s)
 
-instance functorRWST :: (Functor m) => Functor (RWST r w s m) where
-  map f m = RWST \r s -> (\see -> see{result = f see.result}) <$> runRWST m r s
+instance functorRWST :: (Functor m, Monoid w) => Functor (RWST r w s m) where
+  map f m = RWST \r s -> (\(See state result writer) -> See state (f result) writer) <$> runRWST m r s
 
 instance applyRWST :: (Bind m, Monoid w) => Apply (RWST r w s m) where
   apply f m = RWST \r s ->
-    runRWST f r s  >>= \{state = s',  result = f',  log = w'}  ->
-    runRWST m r s' <#> \{state = s'', result = a'', log = w''} ->
-    mkSee s'' (f' a'') (w' ++ w'')
+    runRWST f r s  >>= \(See s' f' w') ->
+    runRWST m r s' <#> \(See s'' a'' w'') ->
+    See s'' (f' a'') (w' ++ w'')
 
 instance bindRWST :: (Bind m, Monoid w) => Bind (RWST r w s m) where
   bind m f = RWST \r s ->
-    runRWST m     r s  >>= \{result = a, state = s', log = l} ->
-    runRWST (f a) r s' <#> \see' ->
-    see' { log = l ++ see'.log }
+    runRWST m     r s  >>= \(See s' a w) ->
+    runRWST (f a) r s' <#> \(See state result writer) ->
+    See state result (w ++ writer)
 
 instance applicativeRWST :: (Monad m, Monoid w) => Applicative (RWST r w s m) where
-  pure a = RWST \_ s -> pure $ mkSee s a mempty
+  pure a = RWST \_ s -> pure $ See s a mempty
 
 instance monadRWST :: (Monad m, Monoid w) => Monad (RWST r w s m)
 
 instance monadTransRWST :: (Monoid w) => MonadTrans (RWST r w s) where
-  lift m = RWST \_ s -> m >>= \a -> return $ mkSee s a mempty
+  lift m = RWST \_ s -> m >>= \a -> return $ See s a mempty
 
 instance monadEffRWS :: (Monad m, Monoid w, MonadEff eff m) => MonadEff eff (RWST r w s m) where
   liftEff = lift <<< liftEff
 
 instance monadReaderRWST :: (Monad m, Monoid w) => MonadReader r (RWST r w s m) where
-  ask = RWST \r s -> pure $ mkSee s r mempty
+  ask = RWST \r s -> pure $ See s r mempty
   local f m = RWST \r s -> runRWST m (f r) s
 
 instance monadStateRWST :: (Monad m, Monoid w) => MonadState s (RWST r w s m) where
-  state f = RWST \_ s -> case f s of Tuple a s' -> pure $ mkSee s' a mempty
+  state f = RWST \_ s -> case f s of Tuple a s' -> pure $ See s' a mempty
 
 instance monadWriterRWST :: (Monad m, Monoid w) => MonadWriter w (RWST r w s m) where
-  writer (Tuple a w) = RWST \_ s -> pure $ { state: s, result: a, log: w }
-  listen m = RWST \r s -> runRWST m r s >>= \{ state: s', result: a, log: w} -> pure { state: s', result: Tuple a w, log: w }
-  pass m = RWST \r s -> runRWST m r s >>= \{ result: Tuple a f, state: s', log: w} -> pure { state: s', result: a, log: f w }
+  writer (Tuple a w) = RWST \_ s -> pure $ See s a w
+  listen m = RWST \r s -> runRWST m r s >>= \(See s' a w) -> pure $ See s' (Tuple a w) w
+  pass m = RWST \r s -> runRWST m r s >>= \(See s' (Tuple a f) w) -> pure $ See s' a (f w)
 
 instance monadRWSRWST :: (Monad m, Monoid w) => MonadRWS r w s (RWST r w s m)
 
@@ -100,10 +93,10 @@ instance monadErrorRWST :: (MonadError e m, Monoid w) => MonadError e (RWST r w
   catchError m h = RWST $ \r s -> catchError (runRWST m r s) (\e -> runRWST (h e) r s)
 
 instance monadRecRWST :: (Monoid w, MonadRec m) => MonadRec (RWST r w s m) where
-  tailRecM k a = RWST \r s -> tailRecM (k' r) { writer: mempty, state: s, result: a }
+  tailRecM k a = RWST \r s -> tailRecM (k' r) (See s a mempty)
     where
-    k' r o = do
-      see <- runRWST (k o.result) r o.state
-      return case see.result of
-                  Left a -> Left { state: see.state, result: a, writer: o.writer <> see.log }
-                  Right b -> Right (mkSee see.state b (o.writer <> see.log))
+    k' r (See state result writer) = do
+      See state' result' writer' <- runRWST (k result) r state
+      return case result' of
+                  Left a -> Left (See state' a (writer <> writer'))
+                  Right b -> Right (See state' b (writer <> writer'))

test/Example/RWS.purs

@@ -8,6 +8,7 @@ import Control.Monad.RWS
 import Control.Monad.RWS.Trans
 import Control.Monad.Rec.Class
 import Control.Monad.State
+import Control.Monad.State.Trans
 import Control.Monad.Writer
 
 import Data.Either
@@ -24,12 +25,24 @@ loop n = tailRecM go n
     put (x + 1)
     return (Left (n - 1))
 
+loopState :: Int -> StateT Int Identity Unit
+loopState n = tailRecM go n
+  where
+  go 0 = do
+    return (Right unit)
+  go n = do
+    x <- get
+    put (x + 1)
+    return (Left (n - 1))
+
 main = do
   t1 <- t
-  res <- pure $ runIdentity (runRWST (loop 10000) "" 0)
+  res1 <- pure $ runIdentity (runRWST (loop 1000000) "" 0)
   t2 <- t
-  print $ "RWST.state: " ++ show res.state
-  print $ "RWST.log: " ++ show res.log
-  print $ "t2 - t1 = " ++ show (t2 - t1)
+  print $ "RWST: " ++ show (t2 - t1)
+  t3 <- t
+  res2 <- pure $ execState (loopState 1000000) 0
+  t4 <- t
+  print $ "StateT: " ++ show (t4 - t3)
 
 foreign import t :: forall eff. Eff eff Number