ATs

{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, TypeFamilies #-}
module AT where

{- --- ORIGINAL --- -}

data Step s a = Done
              | Skip !s
              | Yield !a !s

{- --- WITH ATs --- -}

class AStep s a where
    data Stepper s a
    done :: Stepper s a
    skip :: s -> Stepper s a
    yield :: a -> s -> Stepper s a

instance (AStep s) Char where
    data Stepper s Char = CDone
                        | CSkip !s
                        | CYield {-# UNPACK #-} !Char !s
    done = CDone
    {-# INLINE done #-}
    skip = CSkip
    {-# INLINE skip #-}
    yield = CYield
    {-# INLINE yield #-}

{- --- CRAZINESS --- -}

-- Use the type parameter a as a phantom type.
data Step s a = Done
              | Skip !s
              | CYield {-# UNPACK #-} !Char !s
              | WYield {-# UNPACK #-} !Word8 !s

eq (Stream next1 s1 _) (Stream next2 s2 _) = loop (next1 s1) (next2 s2)
where
loop Done Done = True
loop (Skip s1') (Skip s2') = loop (next1 s1') (next2 s2')
loop (Skip s1') x2 = loop (next1 s1') x2
loop x1 (Skip s2') = loop x1 (next2 s2')
loop Done _ = False
loop _ Done = False
loop (Yield x1 s1') (Yield x2 s2') = x1 == x2 &&
loop (next1 s1') (next2 s2')