Revsums.hs

{-# LANGUAGE RecursiveDo #-}

module Revsums where

import Control.Monad.Fix
import Control.Monad

-- >>> revsums []
-- []
-- >>> revsums [4]
-- [4]
-- >>> revsums [1,2,3]
-- [6,5,3]
-- >>> length . take 3 $ revsums [1..]
-- 3
revsums :: Traversable t => t Int -> t Int
revsums xs = evalRevState (traverse step xs) 0 where
    step x = mdo
        let acc' = acc + x
        put acc'
        acc <- get
        return acc'
    {-# INLINE step #-}
{-# INLINE revsums #-}

--------------------

newtype RevState s a = RevState
    { runRevState :: s -> (a, s)
    }

evalRevState :: RevState s a -> s -> a
evalRevState a s = fst $ runRevState a s

get :: RevState s s
get = RevState $ \s -> (s, s)

put :: s -> RevState s ()
put s' = RevState $ \_ -> ((), s')

instance Functor (RevState s) where
    fmap = liftM

instance Applicative (RevState s) where
    pure  = return
    (<*>) = ap

instance Monad (RevState s) where
    return x = RevState $ \s -> (x, s)

    a >>= f = RevState $ \s'' ->
        let (x, s)  = runRevState a s'
            (y, s') = runRevState (f x) s''
        in (y, s)

instance MonadFix (RevState s) where
    mfix f = RevState $ \s ->
        fix $ \(~(x, _)) -> runRevState (f x) s