{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE KindSignatures        #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Generics.MultiRec.Rewriting.HZip
-- Copyright   :  (c) 2009 Universiteit Utrecht
-- License     :  BSD3
--
-- Maintainer  :  generics@haskell.org
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Summary: Monadic zip on functorial values. 
-----------------------------------------------------------------------------

module Generics.MultiRec.Rewriting.HZip (
    HZip, hzipM
  , hzip, hzip'
  , combine, eq'
  ) where

import Control.Monad

import Generics.MultiRec.Base


-- * Generic zip
class HZip (s :: * -> *) (f :: (* -> *) -> * -> *) where
  hzipM :: Monad m =>
           (forall ix'. s ix' -> r1 ix' -> r2 ix' -> m (r3 ix')) ->
           f r1 ix -> f r2 ix -> m (f r3 ix)

instance HZip s U where
  hzipM _ U U = return U

instance (El s xi) => HZip s (I xi) where
  hzipM f (I x) (I y) = liftM I (f proof x y)

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

instance (HZip s a, HZip s b) => HZip s (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 _ _     _     = fail "zip failed"

instance (HZip s a) => HZip s (C c a) where
  hzipM f (C x) (C y) = liftM C (hzipM f x y)

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

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

-- | Monadic zip but argument is not monadic
hzip :: (HZip s f, Monad m) =>
        (forall ix'. s ix' -> r1 ix' -> r2 ix' -> r3 ix') ->
        s ix -> f r1 ix -> f r2 ix -> m (f r3 ix)
hzip f _ = hzipM (\w x y -> return (f w x y))

-- | Unsafe zip
hzip'  :: (HZip s f) =>
          (forall ix'. s ix' -> r1 ix' -> r2 ix' -> r3 ix') ->
          s ix -> f r1 ix -> f r2 ix -> f r3 ix
hzip' f _ a b = case hzipM (\w x -> return . f w x) a b of
                  Nothing  -> error "generic zip failed"
                  Just res -> res

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

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