-----------------------------------------------------------------------------
-- |
-- Module      :  Transform.Rewriting
-- Copyright   :  (c) 2010 University of Minho
-- License     :  BSD3
--
-- Maintainer  :  hpacheco@di.uminho.pt
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Pointless Rewrite:
-- automatic transformation system for point-free programs
-- 
-- Point-free type-preserving rewriting.
--
-----------------------------------------------------------------------------

module Transform.Rewriting where

import Data.Type
import Data.Pf
import Data.Spine
import Data.Equal hiding (replace)

import Data.List
import Control.Monad
import Control.Monad.RWS
import Control.Monad.State
import System.IO

import Generics.Pointless.Combinators

-- Generic queries

gmapQ :: GenericQ r -> GenericQ [r]
gmapQ q t x = aux q (toSpine t x)
    where aux :: GenericQ r -> (forall a. Spine a -> [r])
          aux q (_ `As` _)      = []
          aux q (Ap f (t :| x)) = aux q f ++ [q t x]

type Query r = forall a . Typed a -> r

everything :: (r -> r -> r) -> GenericQ r -> Query r
everything op q (t :| x) = aux op q t x
    where aux :: (r -> r -> r) -> GenericQ r -> GenericQ r
	  aux op q t x = foldl1 op ([q t x] ++ gmapQ (aux op q) t x)
	       
-- Locations and contexts

type Location = [Int]

down :: Location -> Location
down = (0:)

next :: Location -> Location
next (h:t) = (h+1):t

replace :: MonadPlus m => Location -> Typed b -> Typed a -> m a
replace [0] (b :| x) (a :| y) = coerce b a x
replace l   (b :| x) (a :| y) = do s <- aux (last l) (init l) (toSpine a y)
                                   return $ fromSpine s
    where aux :: MonadPlus m => Int -> Location -> Spine a -> m (Spine a)
          aux 0     l (Ap f (a :| y)) = do
              z <- replace l (b :| x) (a :| y)
              return $ Ap f (a :| z)
          aux n l (Ap f (a :| y)) = do
              g <- aux (pred n) l f
              return $ Ap g (a :| y)
          aux _ _ _ = mzero

hole :: Type a -> a
hole (Pf _) = BOT

puthole :: Location -> Dynamic -> Dynamic
puthole l (Dyn t x) = Dyn t (xua l (t :| x))
    where xua :: Location -> Typed a -> a
          xua [0] (a :| y) = hole a
	  xua l   (a :| y) = fromSpine $ aux (last l) (init l) (toSpine a y)
	  aux :: Int -> Location -> Spine a -> Spine a
          aux 0     l (Ap f (a :| y)) = Ap f (a :| xua l (a :| y))
          aux n l (Ap f (a :| y)) = Ap (aux (pred n) l f) (a :| y)

-- The basic type of rules
type GenericM m = forall a . Type a -> Pf a -> m (Pf a)

-- Generic one-level map with location update

gmapMo :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
gmapMo h t y = aux h (toSpine (Pf t) y)
    where aux :: (MonadReader Location m, MonadPlus m) => GenericM m -> Spine a -> m a
          aux h (c `As` _) = mzero
          aux h (Ap f (Pf t :| x)) = (do
              g <- local next $ aux h f
              return $ g x)
              `mplus` (do
              let g = fromSpine f
              y <- local down $ h t x
              return $ g y)
          aux h _ = mzero

gmapMo' :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
gmapMo' h t y = aux h (toSpine (Pf t) y)
    where aux :: (MonadReader Location m, MonadPlus m) => GenericM m -> Spine a -> m a
          aux h (c `As` _) = mzero
          aux h (Ap f (Pf t :| x)) = (do
              let g = fromSpine f
              y <- local down $ h t x
              return $ g y)
              `mplus` (do
              g <- local next $ aux h f
              return $ g x)
          aux h _ = mzero

gmapM :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
gmapM h t y = aux h (toSpine (Pf t) y)
    where aux :: (MonadReader Location m, MonadPlus m) => GenericM m -> Spine a -> m a
          aux h x@(c `As` _) = return c
          aux h (Ap f (Pf t :| x)) = do 
            g <- local next $ aux h f
            y <- local down $ h t x
            return $ g y
          aux h (Ap f (t :| x)) = do 
            g <- local next $ aux h f
            return $ g x

-- Promoting rules to traversals by updating context and propagating the type
top :: (MonadReader Location m, MonadState Dynamic m, MonadPlus m) => GenericM m -> GenericM m
top f t x = do
    y <- f t x
    (Dyn u c) <- get
    l <- ask
    let z = replace l (Pf t :| y) (u :| c)
    case z of
        Just z' -> put (Dyn u z')
        Nothing -> put (Dyn One _L)
    return y

gtop :: (MonadReader Location m, MonadState Dynamic m, MonadPlus m) => Rule -> GenericM m -> GenericM m
gtop r f t x = do
    y <- f t x
    (Dyn (Pf u) c) <- get
    l <- ask
    let Just z = replace l (Pf t :| y) ((Pf u) :| c)
        c' = reducePf r u c
        z' = reducePf r u z
    guard (geq (Pf u) c' z')
    put (Dyn (Pf u) z)
    return y

-- Strategy combinators
once :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
once f = f ||| gmapMo (once f)

once' :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
once' f = f ||| gmapMo' (once' f)

everywhere :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
everywhere f = f >>> gmapM (everywhere f)

everywhere' :: (MonadReader Location m, MonadPlus m) => GenericM m -> GenericM m
everywhere' f = gmapM (everywhere f) >>> f

innermost :: Rule -> Rule
innermost r = gmapM (innermost r) >>> try (r >>> innermost r)

outermost :: Rule -> Rule
outermost r = try (many1 (once r))

outermost1 :: Rule -> Rule
outermost1 r = many1 (once r)

(>>>) :: Monad m => GenericM m -> GenericM m -> GenericM m
(f >>> g) t x = f t x >>= g t

(|||) :: MonadPlus m => GenericM m -> GenericM m -> GenericM m
(f ||| g) t x = f t x `mplus` g t x

many :: MonadPlus m => GenericM m -> GenericM m
many r = (r >>> many r) ||| nop

many1 :: MonadPlus m => GenericM m -> GenericM m
many1 r = r >>> many r

many2 :: MonadPlus m => GenericM m -> GenericM m
many2 r = r >>> r >>> many r

nop :: Monad m => GenericM m
nop t = return

try :: MonadPlus m => GenericM m -> GenericM m
try x = x ||| nop

-- Rewriting monad

type Log = [(String,String)]

type Rewrite = RWST Location Log Dynamic Maybe

type Rule = GenericM Rewrite

printRule :: String -> Type a -> a -> Rewrite ()
printRule n t v = tell [(n,gshow Dynamic (Dyn t v))] 

debug :: String -> Type a -> a -> Rewrite ()
debug n t v = --trace ("entering " ++ n ++ ": " ++ gshow t v) $
	return ()

debugT :: String -> Type a -> Rewrite ()
debugT n t = --trace ("entering " ++ n ++ ": " ++ show t) $
	return ()

success :: String -> a -> Rewrite a
success n x =
    do z@(Dyn t v) <- get
       --trace n $ printRule n t v
       return x

context :: Rewrite (Typed Dynamic)
context =
    do l <- ask
       y <- get
       return (Dynamic :| puthole l y)

-- Simplification wrapers

simplify :: Typeable a => Bool -> Rule -> Pf a -> IO (Pf a)
simplify = rewrite typeof

rewrite :: Type a -> Bool -> Rule -> Pf a -> IO (Pf a)
rewrite t b s e = do
  Just (x,l) <- return $ evalRWST (s t e) [0] (Dyn (Pf t) e)
  when b $ sequence_ (map aux l)
  putStrLn ("  "++(gshow (Pf t) x))
  return x
 where aux (n,y) =putStrLn ("  "++y) >> putStrLn ("=   { "++n++" }") 

reduce :: Rule -> Type a -> Pf a -> (Pf a,[String])
reduce s t x = maybe (x,[]) (id >< map fst) (evalRWST (s t x) [0] (Dyn (Pf t) x))

reduceIO :: Rule -> Type a -> Pf a -> IO (Pf a)
reduceIO s t x = do
	putStrLn "Running optimizations..."
        let (l,r) = maybe (x,[]) id (evalRWST (s t x) [0] (Dyn (Pf t) x))        
        hPutRuleTree stdout r
        putStrLn ""
--        putStrLn $ gshow (Pf t) l
        return l

writeIO :: FilePath -> Rule -> Type a -> Pf a -> IO (Pf a)
writeIO f s t x = do
        h <- openFile f WriteMode
        let (l,r) = maybe (x,[]) id (evalRWST (s t x) [0] (Dyn (Pf t) x))
        putStr $ gshow (Pf t) l
        putStrLn ""
        hPutRuleTree h r
        return l

reducePfMb :: Rule -> Type a -> Pf a -> Maybe (Pf a)
reducePfMb s t x = liftM fst (evalRWST (s t x) [0] (Dyn (Pf t) x))

reducePf :: Rule -> Type a -> Pf a -> Pf a
reducePf s t x = maybe x fst (evalRWST (s t x) [0] (Dyn (Pf t) x))

reduceCount :: Rule -> Type a -> Pf a -> (Pf a,Int)
reduceCount s t x = maybe (x,0) (id >< length) (evalRWST (s t x) [0] (Dyn (Pf t) x))

hPutRuleTree :: Handle -> Log -> IO ()
hPutRuleTree h l = evalStateT (hPutRuleTree' h l) 0

hPutRuleTree' :: Handle -> Log -> StateT Int IO ()
hPutRuleTree' _ [] = return ()
hPutRuleTree' h ((r,e):xs) = do
        if (isSuffixOf ":" r)
            then do
                i <- get
                liftIO $ hPutStrLn h $ printRuleNode i True (init r)
                modify succ
            else if (isPrefixOf ":" r)
                then do
                modify pred
                i <- get
                liftIO $ hPutStrLn h $ printRuleNode i False (tail r)
                else do
                i <- get
                liftIO $ hPutStrLn h $ printRuleNode i True r
        hPutRuleTree' h xs
                
printRuleNode :: Int -> Bool -> String -> String
printRuleNode n True s = replicate n ' ' ++ "|- " ++ s
printRuleNode n False s = replicate n ' ' ++ "/- " ++ s

proof_strat :: Rule -> Type a -> Pf a -> Pf a -> Rewrite ()
proof_strat r t f g = do
    eq1 <- r t f
    eq2 <- r t g
    debug "proof1: " (Pf t) eq1
    debug "proof2: " (Pf t) eq2
    guard $ (geq (Pf t) eq1 eq2)

proof_strat' :: Rule -> Type a -> Pf a -> Type b -> Pf b -> Rewrite ()
proof_strat' r a f b g = do
    eq1 <- r a f
    eq2 <- r b g
    guard $ (geq' (Pf a) eq1 (Pf b) eq2)

proof :: Rule -> Type a -> Pf a -> Pf a -> Bool
proof r t f g = maybe False (const True) $ evalRWST (proof_strat r t f g) [0] (Dyn (Pf t) f)

proof' :: Rule -> Type a -> Pf a -> Type b -> Pf b -> Bool
proof' r a f b g = maybe False (const True) $ evalRWST (proof_strat' r a f b g) [0] (Dyn (Pf a) f)