{-# LANGUAGE GADTs            #-}
{-# LANGUAGE RankNTypes       #-}
{-# LANGUAGE TypeOperators    #-}
{-# LANGUAGE FlexibleContexts #-}

module HZip where

import Control.Monad

import Base

-- * Generic zip

class HZip f where
  hzipM :: Monad m =>
           (forall ix. Ix s ix => s ix -> r ix -> r' ix -> m (r'' ix)) ->
           f s r ix -> f s r' ix -> m (f s r'' ix)

instance HZip (I xi) where
  hzipM f (I x) (I y) = liftM I (f index x y)

instance Eq a => HZip (K a) where
  hzipM f (K x) (K y) | x == y    = return (K x)
                      | otherwise = fail "zip failed in K"

instance (HZip a, HZip b) => HZip (a :+: b) where
  hzipM f (L x) (L y) = liftM L (hzipM f x y)
  hzipM f (R x) (R y) = liftM R (hzipM f x y)
  hzipM f _     _     = fail "zip failed"

instance (HZip a, HZip b) => HZip (a :*: b) where
  hzipM f (x1 :*: y1) (x2 :*: y2) = liftM2 (:*:) (hzipM f x1 x2) (hzipM f y1 y2)

instance HZip f => HZip (f :>: xi) where
  hzipM f (Tag x) (Tag y) = liftM Tag (hzipM f x y)

-- | Monadic zip but argument is not monadic
hzip :: (HZip f, Monad m) =>
        (forall ix. Ix l ix => l ix -> r ix -> s ix -> t ix) ->
        f l r ix -> f l s ix -> m (f l t ix)
hzip f = hzipM (\w x y -> return (f w x y))

-- | Unsafe zip
hzip'  :: (HZip f) =>
          (forall ix. r ix -> r' ix -> r'' ix) ->
          f s r ix -> f s r' ix -> f s r'' ix
hzip' f a b = case hzip (\ w x y -> f x y) a b of
                Nothing  -> error "generic zip failed"
                Just res -> res

-- | Combine two structures monadically only
combine :: (Monad m, HZip f) =>
           (forall ix. Ix s ix => s ix -> r ix -> r' ix -> m ()) ->
           s ix -> f s r ix -> f s r' ix -> m ()
combine f l x y = hzipM wrapf x y >> return ()
  where
    wrapf ix x y = f ix x y >> return (K0 ())

-- | Generic equality
geq :: (HZip (PF s), Ix s ix) => s ix -> ix -> ix -> Bool
geq ix x y = maybe False (const True) (geq' ix (I0 x) (I0 y))

-- | Monadic generic equality (just for the sake of the monad!)
geq' :: (Monad m, HZip (PF s), Ix s ix) => s ix -> I0 ix -> I0 ix -> m ()
geq' ix (I0 x) (I0 y) = combine geq' ix (from'' x) (from'' y)