{-# 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.
-- 
-- Based on "Beautiful Folds" by Max Rabkin
-- <http://squing.blogspot.com/2008/11/beautiful-folding.html>
--
-- See also <http://conal.net/blog/posts/more-beautiful-fold-zipping>
----------------------------------------------------------------------

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

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

import Data.Pair

-- | 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 Pair (Fold b) where pair = 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 Pair' (Fold b) where pair' = 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 Pair (h b) => Applicative (WithCont h b) where
  pure a = WC (error "unneeded pre-cont") (pure a)
  WC f k <*> WC f' k' =
    WC (f `pair` f') (\ (a,a') -> (k a) (k' 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 Pair' (h b) => Applicative (WithCont' h b) where
  pure a = WC' (error "unneeded pre-cont") (pure a)
  WC' f k <*> WC' f' k' =
    WC' (f `pair'` f') (\ (P a a') -> (k a) (k' 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 pair
class Pair' f where
  pair' :: f a -> f b -> f (P a b)

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