{-# LANGUAGE ExistentialQuantification, FlexibleContexts #-}
{-# OPTIONS_GHC -Wall #-}
----------------------------------------------------------------------
-- |
-- Module      :  Data.Zip.FoldL
-- Copyright   :  (c) Conal Elliott 2008
-- License     :  BSD3
-- 
-- Maintainer  :  conal@conal.net
-- Stability   :  experimental
-- 
-- Zipping of non-strict left folds.
-- 
-- See <http://conal.net/blog/tag/zip>.  Inspired by "Beautiful Folds" by
-- Max Rabkin <http://squing.blogspot.com/2008/11/beautiful-folding.html>
----------------------------------------------------------------------

module Data.Zip.FoldL
  ( Fold(..), cfoldl, cfoldl'
  , WithCont , FoldC , cfoldlc
  , WithCont', FoldC', cfoldlc'
  , Zip'(..), P
  ) where

import Prelude hiding (zip)

import Data.Monoid
import Control.Applicative
import Data.List (foldl')

import Data.Zip

-- | Data representation of a left fold
data Fold b a = F (a -> b -> a) a

-- | Interpretation of a 'Fold' as non-strict
cfoldl :: Fold b a -> [b] -> a
cfoldl (F op e) = foldl op e

-- Non-strict left-fold zipping
zipF :: Fold b a -> Fold b a' -> Fold b (a,a')
F op e `zipF` F op' e' = F op'' (e,e')
 where
   (a,a') `op''` b = (a `op` b, a' `op'` b)

instance Zip (Fold b) where zip = zipF

-- | Interpretation of a 'Fold' as non-strict
cfoldl' :: Fold b a -> [b] -> a
cfoldl' (F op e) = foldl' op e

-- Strict left-fold zipping
zipF' :: Fold b a -> Fold b a' -> Fold b (P a a')
F op e `zipF'` F op' e' = F op'' (P e e')
 where
   P a a' `op''` b = P (a `op` b) (a' `op'` b)


instance Zip' (Fold b) where zip' = zipF'


-- | Add a continuation.
data WithCont h b c = forall a. WC (h b a) (a -> c)

instance Functor (WithCont h b) where
  fmap g (WC f k) = WC f (fmap g k)

instance Zip (h b) => Applicative (WithCont h b) where
  pure a = WC (error "unneeded pre-cont") (pure a)
  WC hf hk <*> WC xf xk =
    WC (hf `zip` xf) (\ (a,a') -> (hk a) (xk a'))

-- | Non-strict left fold with continuation.
type FoldC = WithCont Fold

-- | Interpretation of a 'FoldC'
cfoldlc :: FoldC b a -> [b] -> a
cfoldlc (WC f k) = fmap k (cfoldl f)



-- | Like 'WithCont' but with pair-strict '(<*>)'
data WithCont' h b c = forall a. WC' (h b a) (a -> c)

instance Functor (WithCont' h b) where
  fmap g (WC' f k) = WC' f (fmap g k)

instance Zip' (h b) => Applicative (WithCont' h b) where
  pure a = WC' (error "unneeded pre-cont") (pure a)
  WC' hf hk <*> WC' xf xk =
    WC' (hf `zip'` xf) (\ (P a a') -> (hk a) (xk a'))

-- | Strict left fold with continuation.
type FoldC' = WithCont' Fold

-- | Interpretation of a 'FoldC'
cfoldlc' :: FoldC' b a -> [b] -> a
cfoldlc' (WC' f k) = fmap k (cfoldl' f)


----

-- | Strict pairs
data P c c' = P !c !c'

-- | Strict generalized zip
class Zip' f where
  zip' :: f a -> f b -> f (P a b)

instance             Zip' []       where zip' = liftA2 P
instance Monoid u => Zip' ((,)  u) where zip' = liftA2 P
instance             Zip' ((->) u) where zip' = liftA2 P
instance             Zip' IO       where zip' = liftA2 P