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

module Generics.MultiRec.Rewriting.Rules where

import Prelude hiding (zip)

import Generics.MultiRec.Base
import Generics.MultiRec.HFunctor
import Generics.MultiRec.Rewriting.HZip
import Generics.MultiRec.Rewriting.LeftRight
import Generics.MultiRec.HFix

import Control.Monad (join)

-- * Rule types

infix 5 :~>
data RuleSpec a = a :~> a

lhsR :: RuleSpec a -> a
lhsR (lhs :~> _) = lhs

rhsR :: RuleSpec a -> a
rhsR (_ :~> rhs) = rhs

type Rule   s a = RuleSpec (Scheme s a)

-- * Schemes

type MVar = Int
type Scheme s = HFix (K MVar :+: PF s)

mvar :: s ix -> MVar -> Scheme s ix
mvar _ = HIn . L . K

pf :: s ix -> (PF s) (Scheme s) ix -> Scheme s ix
pf _ = HIn . R

toScheme :: (HFunctor s (PF s), Fam s) => s ix -> ix -> Scheme s ix
toScheme p x = pf p $ hmap (\p' -> toScheme p' . unI0) p $ from p x

-- This can probably be written in a cleaner way
fromScheme :: (HFunctor phi (PF phi), Fam phi)
           => phi ix -> Scheme phi ix -> Maybe ix
fromScheme p = fmap (to p) . join
             . fmap (hmapM (\p' -> fmap I0 . fromScheme p') p) . unpf p where
  unpf :: phi ix -> Scheme phi ix -> Maybe ((PF phi) (Scheme phi) ix)
  unpf _ (HIn (R x)) = Just x
  unpf _ _           = Nothing

data ViewScheme s ix = MVar MVar | PF (PF s (Scheme s) ix)

viewScheme :: s ix -> Scheme s ix -> ViewScheme s ix
viewScheme _ (HIn (L (K x))) = MVar x
viewScheme _ (HIn (R r))     = PF r

foldScheme :: (HFunctor s (PF s))
           => s ix
           -> (forall ix'. MVar -> r ix')
           -> (forall ix'. () => PF s r ix' -> r ix')
           -> Scheme s ix -> r ix
foldScheme p f g scheme =
  case viewScheme p scheme of
    MVar x -> f x
    PF r   -> g (hmap (\p' -> foldScheme p' f g) p r)

-- * Making a rule

class Builder s a where
  type Target a :: *
  base          :: s (Target a) -> a -> RuleSpec (Target a)
  diag          :: s (Target a) -> a -> [RuleSpec (Target a)]

instance Builder s (RuleSpec a) where
  type Target (RuleSpec a) = a
  base _ x                 = x
  diag _ x                 = [x]

instance (Builder s a, Single s (PF s), El s b, Fam s) => Builder s (b -> a) where
  type Target (b -> a) = Target a
  base ix f            = base ix (f (left  (index::s b)))
  diag ix f            = base ix (f (right (index::s b))) :
                         diag ix (f (left  (index::s b)))

rule :: (Builder s r, HZip s (PF s), Fam s) =>
        s (Target r) -> r -> Rule s (Target r)
rule p f = foldr1 mergeRules rules
  where
    mergeRules x y = 
      mergeSchemes p (lhsR x) (lhsR y) :~>
      mergeSchemes p (rhsR x) (rhsR y)
    rules          = zipWith (ins (base p f)) (diag p f) [0..]
    ins x y v      = 
      insertMVar v p (I0 (lhsR x)) (I0 (lhsR y)) :~>
      insertMVar v p (I0 (rhsR x)) (I0 (rhsR y))

-- | Merge two schemes by accepting a meta-variable case in one of the 
-- schemes, and otherwise zip the remaining parts.
mergeSchemes :: HZip s (PF s) 
             => s ix -> Scheme s ix -> Scheme s ix -> Scheme s ix
mergeSchemes w p@(HIn x) q@(HIn y)
  = case viewScheme w p of
      MVar _ -> p
      _      -> case viewScheme w q of
                  MVar _ -> q
                  _      -> HIn (hzip' mergeSchemes w x y)

-- | Extend a value with a meta-variable case. Two structures are traversed 
-- in parallel, whenever the structures mismatch, insert a meta-variable.
-- Otherwise, lift the matching parts.
insertMVar :: forall s ix .(HZip s (PF s), Fam s) 
           => MVar -> s ix -> I0 ix -> I0 ix -> Scheme s ix
insertMVar name w (I0 x) (I0 y)
  = case hzip (insertMVar name) w (from w x) (from w y) of
      Just struc -> pf w struc
      Nothing    -> mvar w name