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

module Rules where

import Prelude hiding (zip)

import Base
import HFunctor
import HZip
import LeftRight

-- * Rule types

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

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

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

type Rule   s a = RuleSpec (Scheme s a)

-- * Schemes

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

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

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

toScheme :: (HFunctor (PF s), Ix s ix) => ix -> Scheme s ix
toScheme = pf . hmap (\ _ -> toScheme . unI0) . from

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

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

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

-- * Making a rule

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

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

instance (Builder a s, LR (PF s) s, Ix s b) => Builder (b -> a) s 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 r s, HZip (PF s),Ix s (Target r)) =>
        s (Target r) -> r -> Rule s (Target r)
rule ix f = foldr1 mergeRules rules
  where
    mergeRules x y = 
      mergeSchemes (lhsR x) (lhsR y) :~>
      mergeSchemes (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))

-- | Merge two schemes by accepting a
-- meta-variable case in one of the 
-- schemes, and otherwise zip the 
-- remaining parts.

mergeSchemes :: HZip (PF s) => Scheme s ix -> Scheme s ix -> Scheme s ix
mergeSchemes p@(HIn x) q@(HIn y)
  = case viewScheme p of
      MVar i -> p
      _     -> case viewScheme q of
                 MVar j -> q
                 _     -> HIn (hzip' mergeSchemes 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 :: (HZip (PF s),Ix s ix) => MVar -> s ix -> I0 ix -> I0 ix -> Scheme s ix
insertMVar name ix (I0 x) (I0 y)
  = case hzip (insertMVar name) (from'' x) (from'' y) of
      Just struc -> pf struc
      Nothing    -> mvar name