ChangingFolds.hs

{-# LANGUAGE ExistentialQuantification #-}


module MyLib where


import Data.Foldable (toList)

import Control.Applicative

import Prelude hiding (sum, length, drop)

data Fold a b = forall x. Fold (x -> a -> Either x (Fold a b)) x (x -> b)



drop :: Int -> Fold a b -> Fold a b
drop n (Fold step0 x0 done0) = Fold step n done where
    done _ = done0 x0
    step 0 a = case step0 x0 a of
        Left x'                     -> Right (Fold step0 x' done0)
        Right (Fold step1 x1 done1) -> Right (Fold step1 x1 done1)
    step n _ = Left (n-1)

fold :: Foldable f => Fold a b -> f a -> b
fold (Fold step0 x0 done0) t = go step0 x0 done0 (toList t) where

    go :: forall x a b. (x -> a -> Either x (Fold a b)) -> x -> (x -> b) -> [a] -> b
    go st x done [] = done x
    go st x done (a:as) = case st x a of
        Right (Fold st' x' done') -> go st' x' done' as
        Left x'                   -> go st  x' done  as


instance Functor (Fold a) where
    fmap f (Fold step x done) = Fold (\x a -> fmap (fmap f) $ step x a) x (f . done)

instance Applicative (Fold a) where
    {-# INLINE pure #-}
    pure a = Fold (\x a -> Left x) () (const a)
    {-# INLINE (<*>) #-}
    (Fold stepL0 xL0 doneL0) <*> (Fold stepR0 xR0 doneR0) = Fold step x' done
        where
            step (xL, xR) a = case (stepL0 xL a, stepR0 xR a) of
                (Left xL'    , Left xR'    ) -> Left (xL', xR')
                (Right foldL', Left xR'    ) -> Right (foldL' <*> Fold stepR0 xR' doneR0)
                (Left xL'    , Right foldR') -> Right (Fold stepL0 xL' doneL0 <*> foldR')
                (Right foldL', Right foldR') -> Right (foldL' <*> foldR')
            done (xL,xR) = doneL0 xL (doneR0 xR)
            x' = (xL0, xR0)


sum :: Num a => Fold a a
sum = Fold (\n a -> Left (n+a)) 0 id

length :: Fold a Int
length = Fold (\n a -> Left (n+1)) 0 id

genericLength :: Num b => Fold a b
genericLength = Fold (\n a -> Left (n+1)) 0 id


instance Num b => Num (Fold a b) where
    (+) = liftA2 (+)
    (*) = liftA2 (*)
    abs = fmap abs
    signum = fmap signum
    negate = fmap negate
    fromInteger = pure . fromInteger

instance Fractional b => Fractional (Fold a b) where
    (/) = liftA2 (/)
    recip = fmap recip
    fromRational = pure . fromRational


average = sum / drop 1 genericLength