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

module Generics.MultiRec.Rewriting.Rules where

import Generics.MultiRec
import Generics.MultiRec.LR
import Generics.MultiRec.HZip

-----------------------------------------------------------------------------
-- Rule specification.
-----------------------------------------------------------------------------

-- | Specifies a rule as a value of a datatype.
infix 5 :~>
data RuleSpec a = a :~> a

-- | Returns the left-hand side of a rule.
lhsR :: RuleSpec a -> a
lhsR (x :~> _) = x

-- | Returns the right-hand side of a rule.
rhsR :: RuleSpec a -> a
rhsR (_ :~> y) = y

-----------------------------------------------------------------------------
-- Representation of a rule.
-----------------------------------------------------------------------------
-- | Extends a pattern functor with a case for a metavariable.
type Ext phi  = K Metavar :+: PF phi
type Metavar  = Int

-- | Recursively extends a type with a case for a metavariable.
type Scheme phi = HFix (Ext phi)

-- | Allows metavariables on either side of a rule.
data Rule phi a where 
  Rule :: phi ix -> RuleSpec (Scheme phi ix) -> Rule phi ix

-- | Constructs a metavariable.
metavar :: phi ix -> Metavar -> Scheme phi ix
metavar _ = HIn . L . K

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

-----------------------------------------------------------------------------
-- Builder for transforming a rule specification to a rule.
-----------------------------------------------------------------------------

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

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

instance (Builder phi a, Fam phi, LR phi (PF phi), El phi b)
         => Builder phi (b -> a) where
  type Target (b -> a) = Target a
  base ix f            = base ix (f (left  (proof :: phi b)))
  diag ix f            = base ix (f (right (proof :: phi b))) :
                         diag ix (f (left  (proof :: phi b)))

rule :: forall phi r. (Fam phi, Builder phi r, HZip phi (PF phi), 
                       El phi (Target r), EqS phi, HFunctor phi (PF phi))
        => r -> Rule phi (Target r)
rule f = Rule ix $ foldr1 mergeRules rules
  where
    ix = proof :: phi (Target r)
    mergeRules x y = 
      mergeSchemes ix (lhsR x) (lhsR y) :~>
      mergeSchemes ix (rhsR x) (rhsR y)
    rules          = zipWith (ins (base ix f)) (diag ix f) [0..]   
    ins x y v      = 
      insertMVar v ix (I0 (lhsR x)) (I0 (lhsR y)) :~>
      insertMVar v ix (I0 (rhsR x)) (I0 (rhsR y))

mergeSchemes :: HZip phi (PF phi)
                => phi ix -> Scheme phi ix -> Scheme phi ix -> Scheme phi ix
mergeSchemes p a@(HIn x) b@(HIn y) = case (x,y) of
  (L _,_) -> a
  (_,L _) -> b
  _       -> HIn (hzip' mergeSchemes p x y)

insertMVar :: forall phi ix. (Fam phi, HZip phi (PF phi), El phi ix)
              => Metavar -> phi ix -> I0 ix -> I0 ix -> Scheme phi ix
insertMVar name p (I0 x) (I0 y) =
  case hzip (insertMVar name) p (from p x) (from p y) of
    Just struc -> pf p struc
    Nothing    -> metavar p name