Binding on unfoldr

Unfold.hs

module Data.Unfold where

import Control.Applicative
import Data.Foldable
import Data.Maybe (fromMaybe)
import Prelude hiding (foldl, foldr, repeat)

-- This is a useful abstraction that I haven't found on Hackage, but that
-- doesn't mean it's not there by some other name.

-- | The 'Unfold' constructor just binds the arguments to 'unfoldr'. Now we
-- can do any kind of fold without allocating intermediate data. This is
-- like stream fusion for lists, but unlike lists, it won't let you do
-- anything that doesn't fuse! I've been looking at this as a replacement
-- for lists-as-control-structures in tight places where I don't want to
-- trust inlining and rewrite rules to get fusion.
data Unfold b a = Unfold (b -> Maybe (a, b)) b

instance Foldable (Unfold a) where
    foldr f acc (Unfold unfold start) =
      fromMaybe acc $ do
        (x, next) <- unfold start
        return $ f x $ foldr f acc $ Unfold unfold next

    foldl f acc (Unfold unfold start) =
      fromMaybe acc $ do
        (x, next) <- unfold start
        return $ foldl f (f acc x) $ Unfold unfold next

instance Functor (Unfold a) where
    fmap f (Unfold unfold start) =
      Unfold (fmap (\(a, b) -> (f a, b)) . unfold) start

repeat :: a -> Unfold b a
repeat x = Unfold (\y -> Just (x, y)) undefined

instance Applicative (Unfold b) where
    pure = repeat

    (Unfold unfoldB _) <*> (Unfold unfoldA start) = Unfold unfoldC start
      where
        unfoldC next = do
          (a, next') <- unfoldA next
          (f, next'') <- unfoldB next'
          return (f a, next'')
        

iterate :: (a -> a) -> a -> Unfold a a
iterate f = Unfold (\next -> let x = f next in Just (x, x))

replicate :: Int -> a -> Unfold Int a
replicate n x =
    Unfold (\next -> if next < n then Just (x, pred next) else Nothing) 0

unfoldr :: (b -> Maybe (a, b)) -> b -> Unfold b a
unfoldr = Unfold

scanr :: (a -> b -> b) -> b -> Unfold c a -> Unfold (c, b) b
scanr f startB (Unfold unfoldC startC) = Unfold unfoldCB (startC, startB)
  where
    unfoldCB (c, b) =
      fmap (\(a, c') -> let b' = f a b in (b', (c', b'))) (unfoldC c)