{-# OPTIONS -fglasgow-exts -fallow-undecidable-instances #-}

module Rewriting where

import Base
import Representations

import Control.Monad
import qualified Data.Map as M
import Data.Maybe
import Debug.Trace


-- --------------------------------------
-- Rule types
-- --------------------------------------
type Rule0 a = Rule a
type Rule1 a = a -> Rule a
type Rule2 a = a -> a -> Rule a
type Rule3 a = a -> a -> a -> Rule a

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

lhsR (lhs :~> rhs) = lhs
rhsR (lhs :~> rhs) = rhs


-- --------------------------------------
-- Patterns
-- --------------------------------------
type Var   = Int
type Pat f = Fix (Sum (K Var) f)

var :: Var -> Pat f
var = In . Inl . K

pat :: f (Pat f) -> Pat f
pat = In . Inr

toPat :: Functor f => Fix f -> Pat f
toPat = pat . fmap toPat . out

data ViewPat f = Var Var | Pat (f (Pat f))

viewPat :: Pat f -> ViewPat f
viewPat (In (Inl (K x))) = Var x
viewPat (In (Inr r))     = Pat r


-- --------------------------------------
-- Class synonyms
-- --------------------------------------
class (Functor f,Zip f,Flatten f,Left (Fix f),Right (Fix f),GShow f,GMap f) => Match f
instance (Eq a, Left (K a (Fix (K a))), Right (K a (Fix (K a))),
          Show a) => Match (K a)
instance (Zero (b (Fix (Sum a b))),
          Zero (a (Fix (Sum a b))),
          Flatten b,
          Flatten a,
          Zip b,
          Zip a,
          Functor b,
          Functor a,
          GShow a,
          GShow b,
          GMap a,
          GMap b)
         => Match (Sum a b)

class (View a, Match (PF a)) => Rewrite a


-- --------------------------------------
-- Strategies
-- --------------------------------------
-- Same monad to that in the SYB3 paper
data S m a = S a (m a)

instance MonadPlus m => Monad (S m) where
  return x = S x mzero
  (S x xs) >>= k
    = S r (rs1 `mplus` rs2)
      where
        S r rs1 = k x
        rs2 = do x <- xs
                 let S r _ = k x
                 return r

-- This is like Stratego's one.
-- It applies f to one of the immediate children of x.
one :: (Rewrite a, MonadPlus m, Functor m) => (a -> m a) -> a -> m a
one f x = fmap (to . In) rs
  where
    S _ rs = fmapM try_f $ out $ from x
    try_f x = S x (f_str x)
    f_str = fmap from . f . to

-- Apply f once to one of the nodes of x
-- Mirrors the corresponding definition in Stratego
once :: (Rewrite a, MonadPlus m, Functor m) => (a -> m a) -> a -> m a
once f x = f x `mplus` one (once f) x


-- --------------------------------------
-- Making a rule
-- --------------------------------------
class MkRule a where
  type MkRuleDep a :: *
  mkRule :: a -> Rule (Pat (PF (MkRuleDep a)))

instance Rewrite a => MkRule (Rule0 a) where
  type MkRuleDep (Rule0 a) = a
  mkRule = liftRule0

instance Rewrite a => MkRule (Rule1 a) where
  type MkRuleDep (Rule1 a) = a
  mkRule = liftRule1

instance Rewrite a => MkRule (Rule2 a) where
  type MkRuleDep (Rule2 a) = a
  mkRule = liftRule2
  
instance Rewrite a => MkRule (Rule3 a) where
  type MkRuleDep (Rule3 a) = a
  mkRule = liftRule3


-- --------------------------------------
-- Rewriting a term
-- --------------------------------------
rewrite :: (Rewrite a, Monad m, Functor m) => Rule (Pat (PF a)) -> a -> m a 
rewrite f = fmap to . rewriteF f . from

rewrite0 :: (Rewrite a, Monad m, Functor m) => Rule0 a -> a -> m a
rewrite0 f = fmap to . rewriteF (liftRule0 f) . from

rewrite1 :: (Rewrite a, Monad m, Functor m) => Rule1 a -> a -> m a
rewrite1 f = fmap to . rewriteF (liftRule1 f) . from

rewrite2 :: (Rewrite a, Monad m, Functor m) => Rule2 a -> a -> m a
rewrite2 f = fmap to . rewriteF (liftRule2 f) . from

rewrite3 :: (Rewrite a, Monad m, Functor m) => Rule3 a -> a -> m a
rewrite3 f = fmap to . rewriteF (liftRule3 f) . from

rewriteF :: (Match f, Monad m, Functor m) => Rule (Pat f) -> Fix f -> m (Fix f)
rewriteF (lhs :~> rhs) t = flip inst rhs `fmap` match lhs t


-- --------------------------------------
-- Lifting a rule
-- --------------------------------------
liftRule0 :: Rewrite a => Rule0 a -> Rule (Pat (PF a))
liftRule0 f = liftTerm0 (liftArrow0 (lhsR f)) :~>
              liftTerm0 (liftArrow0 (rhsR f))

liftRule1 :: Rewrite a => Rule1 a -> Rule (Pat (PF a))
liftRule1 f = liftTerm1 0 (liftArrow1 (lhsR . f)) :~>
              liftTerm1 0 (liftArrow1 (rhsR . f))

liftRule2 :: Rewrite a => Rule2 a -> Rule (Pat (PF a))
liftRule2 f = liftTerm2 0 1 (liftArrow2 (\x y -> lhsR (f x y))) :~>
              liftTerm2 0 1 (liftArrow2 (\x y -> rhsR (f x y)))

liftRule3 :: Rewrite a => Rule3 a -> Rule (Pat (PF a))
liftRule3 f = liftTerm3 0 1 2 (liftArrow3 (\x y z -> lhsR (f x y z))) :~>
              liftTerm3 0 1 2 (liftArrow3 (\x y z -> rhsR (f x y z)))


-- --------------------------------------
-- Lifting a term, filling its holes
-- with arbitrary different values
-- --------------------------------------
liftTerm0 :: Functor f => Fix f -> Pat f
liftTerm0 = toPat

liftTerm1 :: Match f => Var -> (Fix f -> Fix f) -> Pat f
liftTerm1 name f = diff name (f left) (f right)

liftTerm2 :: Match f => Var -> Var -> (Fix f -> Fix f -> Fix f) -> Pat f
liftTerm2 name1 name2 f = let pat1 = diff name1 (f left left) (f right left)
                              pat2 = diff name2 (f left left) (f left right)
                          in  pat1 `merge` pat2

liftTerm3 :: Match f => Var -> Var -> Var -> (Fix f -> Fix f -> Fix f -> Fix f) -> Pat f
liftTerm3 name1 name2 name3 f = let pat1 = diff name1 (f left left left) (f right left left)
                                    pat2 = diff name2 (f left left left) (f left right left)
                                    pat3 = diff name3 (f left left left) (f left left right)
                                in  pat1 `merge` pat2 `merge` pat3


-- --------------------------------------
-- Lifting a function expecting arguments
-- --------------------------------------
liftArrow0 :: View a => a -> Str a
liftArrow0 = from

liftArrow1 :: View a => (a -> a) -> (Str a -> Str a)
liftArrow1 f = from . f . to

liftArrow2 :: View a => (a -> a -> a) -> (Str a -> Str a -> Str a)
liftArrow2 f = \x y -> from (f (to x) (to y))

liftArrow3 :: View a => (a -> a -> a -> a) -> (Str a -> Str a -> Str a -> Str a)
liftArrow3 f = \x y z -> from (f (to x) (to y) (to z))


-- --------------------------------------
-- 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.
-- --------------------------------------
diff :: Match f => Var -> Fix f -> Fix f -> Pat f
diff name (In x) (In y)
  = case fzipM' (diff name) x y of
    Just struc -> pat struc
    Nothing    -> var name


-- --------------------------------------
-- Merge two patterns by accepting a
-- meta-variable case in one of the 
-- patterns, and otherwise zip the 
-- remaining parts.
-- --------------------------------------
merge :: Match f => Pat f -> Pat f -> Pat f
merge p@(In x) q@(In y)
  = case viewPat p of
      Var i -> p
      _     -> case viewPat q of
                 Var j -> q
                 _     -> In (fzip merge x y)


-- --------------------------------------
-- Matching a term
-- --------------------------------------
match :: (Match f,Monad m) => Pat f -> Fix f -> m (Subst f)
match p e@(In r2)
  = case viewPat p of
      Var x  -> return (M.singleton x e)
      Pat r1 -> do
                struTerms <- fzipM' (,) r1 r2
                let terms = flatten struTerms
                    matchAll = foldl matchOne (return M.empty) terms
                    matchOne msubst (t1,t2)
                      = do
                        s1 <- msubst
                        s2 <- match (apply s1 t1) t2
                          -- note that the range of s1 does not
                          -- contain meta-variables, therefore
                          -- union is enough to combine substitutions
                        return (M.union s1 s2)
                matchAll


-- --------------------------------------
-- Building a term
-- --------------------------------------
type Subst f = M.Map Var (Fix f)

-- | Apply a substition to a pattern
apply :: Functor f => Subst f -> Pat f -> Pat f
apply s p
  = case viewPat p of
      Var x -> maybe (var x) toPat (M.lookup x s)
      -- ah, mistmatch. Here comes a var
      Pat r -> In (Inr (fmap (apply s) r))

-- | Apply a substitution assuming that the result contains
-- no meta-variables. 
inst :: Functor f => Subst f -> Pat f -> Fix f
inst s p
  = case viewPat p of
      Var x -> fromJust (M.lookup x s)
      Pat r -> In (fmap (inst s) r)


-- --------------------------------------
-- Example instance
-- --------------------------------------
instance View [a] where
  type PF [a] = Sum Unit (Prod (K a) Id)

  from []       = In (Inl Unit)
  from (x : xs) = In (Inr (Prod (K x) (Id (from xs))))

  to (In (Inl Unit))                  = []
  to (In (Inr (Prod (K x) (Id (r))))) = x : to r


-- --------------------------------------
-- Trace and show functions
-- --------------------------------------
trPat pat = trace (showPat pat "") pat
trFix pat = trace (showFix pat "") pat

showRule1 :: Rewrite a => Rule1 a -> String
showRule1 f = (showPat (liftTerm1 0 $ liftArrow1 $ lhsR . f) .
               showString " :~> " .
               showPat (liftTerm1 0 $ liftArrow1 $ rhsR . f)
              ) ""

showPat :: forall f. GShow f => Pat f -> ShowS
showPat pat = case viewPat pat of
                Var _ -> showString "X" --fake meta var show
                Pat r -> sh r
  where
    sh :: f (Pat f) -> ShowS
    sh = gshow showPat

showFix :: forall f. GShow f => Fix f -> ShowS
showFix (In r) = gshow showFix r

showSubst s = foldr (\ (k,v) s -> showString "|->" . showFix v . s) id $ M.toList s