-----------------------------------------------------------------------------
-- |
-- Module      :  Transform.Rules.PF.Combinators
-- 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
-- 
-- Combinators for the rewriting of point-free functions.
--
-----------------------------------------------------------------------------

module Transform.Rules.PF.Combinators where

import Data.Type
import Data.Lens
import Data.Equal
import Transform.Rewriting

import Prelude hiding (Functor(..))
import Control.Monad hiding (Functor(..))

-- ** Combinators

protect_lns :: Rule -> Rule
protect_lns r (Fun c a) (PROTECT f) =
    r (Fun c a) f
protect_lns r t f = r t f

unprotect :: Rule
unprotect (Fun c a) (PROTECT (CATA l1)) = mzero
unprotect (Fun c a) (PROTECT (ANA l1)) = mzero
unprotect (Fun c a) (PROTECT (COMP b l1 l2)) =
    return $ COMP b (PROTECT l1) (PROTECT l2)
unprotect (Fun (Prod c d) (Prod a b)) (PROTECT (PROD l1 l2)) =
    return $ PROD (PROTECT l1) (PROTECT l2)
unprotect (Fun (Either c d) (Either a b)) (PROTECT (SUM l1 l2)) = do
    return $ SUM (PROTECT l1) (PROTECT l2)
unprotect (Fun c a) (PROTECT l1) = do
    debug "safeUnprotect" (Pf $ Fun c a) l1
    return l1
unprotect _ _ = mzero

comp :: Rule -> Rule
comp r t@(Fun d a) e = r t e
    `mplus` (do
    COMP b f (COMP c g h) <- nop t e
    fg <- r (Fun c a) (COMP b f g)
    return $ COMP c fg h)
    `mplus` (do
    COMP c (COMP b f g) h <- nop t e
    gh <- r (Fun d b) (COMP c g h)
    return $ COMP b f gh)
comp _ _ _ = mzero

comp1 :: Rule -> Rule
comp1 r (Fun a c) (COMP b f g) = do
    f' <- r (Fun b c) f
    return $ COMP b f' g
comp1 _ _ _ = mzero

comp2 :: Rule -> Rule
comp2 r (Fun a c) (COMP b f g) = do
    g' <- r (Fun a b) g
    return $ COMP b f g'
comp2 _ _ _ = mzero

prod1 :: Rule -> Rule
prod1 r (Fun (Prod a b) (Prod c d)) (f `PROD` g) = do
    f' <- r (Fun a c) f
    return $ f' `PROD` g
prod1 _ _ _ = mzero

prod2 :: Rule -> Rule
prod2 r (Fun (Prod a b) (Prod c d)) (f `PROD` g) = do
    g' <- r (Fun b d) g
    return $ f `PROD` g'
prod2 _ _ _ = mzero

sum1 :: Rule -> Rule
sum1 r (Fun (Either a b) (Either c d)) (f `SUM` g) = do
    f' <- r (Fun a c) f
    return $ f' `SUM` g
sum1 _ _ _ = mzero

sum2 :: Rule -> Rule
sum2 r (Fun (Either a b) (Either c d)) (f `SUM` g) = do
    g' <- r (Fun b d) g
    return $ f `SUM` g'
sum2 _ _ _ = mzero

precomp :: Rule -> Rule -> Rule
precomp r1 r2 = comp $ r2 ||| (comp1 r1 >>> comp_assocr >>> comp2 r2)

postcomp :: Rule -> Rule -> Rule
postcomp r1 r2 = comp $ r2 ||| (comp2 r1 >>> comp_assocl >>> comp1 r2)

rightmost :: Rule
rightmost (Fun a c) (COMP b f g) = do
    g' <- rightmost' (Fun a b) g
    try comp_assocl (Fun a c) $ COMP b f g'
rightmost (Fun a c) f =
    return $ COMP c ID f
rightmost _ _ = mzero
rightmost' :: Rule
rightmost' (Fun a c) (COMP b f g) = do
    g' <- rightmost' (Fun a b) g
    try comp_assocl (Fun a c) $ COMP b f g'
rightmost' (Fun a c) f = return f
rightmost' _ _ = mzero

leftmost :: Rule
leftmost (Fun a c) (COMP b f g) = do
    f' <- leftmost' (Fun b c) f
    try comp_assocr (Fun a c) $ COMP b f' g
leftmost (Fun a c) f =
    return $ COMP a f ID
leftmost _ _ = mzero
leftmost' :: Rule
leftmost' (Fun a c) (COMP b f g) = do
    f' <- leftmost' (Fun b c) f
    try comp_assocr (Fun a c) $ COMP b f' g
leftmost' (Fun a c) f = return f
leftmost' _ _ = mzero

leftmost_sum :: Rule
leftmost_sum (Fun (Either a b) (Either c d)) (SUM ID ID) = mzero
leftmost_sum (Fun (Either a b) (Either c d)) (SUM ID g) = do
    COMP y g' g'' <- leftmost' (Fun b d) g
    return $ COMP (Either a y) (ID -|-= g') (ID -|-= g'')
leftmost_sum (Fun (Either a b) (Either c d)) (SUM f ID) = do
    COMP x f' f'' <- leftmost' (Fun a c) f
    return $ COMP (Either x b) (f' -|-= ID) (f'' -|-= ID)
leftmost_sum (Fun (Either a b) (Either c d)) (SUM f g) = do
    COMP x f' f'' <- leftmost' (Fun a c) f
    COMP y g' g'' <- leftmost' (Fun b d) g
    return $ COMP (Either x y) (f' -|-= g') (f'' -|-= g'')
leftmost_sum _ _ = mzero

rightmost_sum :: Rule
rightmost_sum (Fun (Either a b) (Either c d)) (SUM ID ID) = mzero
rightmost_sum (Fun (Either a b) (Either c d)) (SUM ID g) = do
    COMP y g' g'' <- rightmost' (Fun b d) g
    return $ COMP (Either a y) (ID -|-= g') (ID -|-= g'')
rightmost_sum (Fun (Either a b) (Either c d)) (SUM f ID) = do
    COMP x f' f'' <- rightmost' (Fun a c) f
    return $ COMP (Either x b) (f' -|-= ID) (f'' -|-= ID)
rightmost_sum (Fun (Either a b) (Either c d)) (SUM f g) = do
    COMP x f' f'' <- rightmost' (Fun a c) f
    COMP y g' g'' <- rightmost' (Fun b d) g
    return $ COMP (Either x y) (f' -|-= g') (f'' -|-= g'')
rightmost_sum _ _ = mzero

leftmost_prod :: Rule
leftmost_prod (Fun (Prod a b) (Prod c d)) (PROD ID ID) = mzero
leftmost_prod (Fun (Prod a b) (Prod c d)) (PROD ID g) = do
    COMP y g' g'' <- leftmost' (Fun b d) g
    return $ COMP (Prod a y) (ID ><= g') (ID ><= g'')
leftmost_prod (Fun (Prod a b) (Prod c d)) (PROD f ID) = do
    COMP x f' f'' <- leftmost' (Fun a c) f
    return $ COMP (Prod x b) (f' ><= ID) (f'' ><= ID)
leftmost_prod (Fun (Prod a b) (Prod c d)) (PROD f g) = do
    COMP x f' f'' <- leftmost' (Fun a c) f
    COMP y g' g'' <- leftmost' (Fun b d) g
    return $ COMP (Prod x y) (f' ><= g') (f'' ><= g'')
leftmost_prod _ _ = mzero

rightmost_prod :: Rule
rightmost_prod (Fun (Prod a b) (Prod c d)) (PROD ID ID) = mzero
rightmost_prod (Fun (Prod a b) (Prod c d)) (PROD ID g) = do
    COMP y g' g'' <- rightmost' (Fun b d) g
    return $ COMP (Prod a y) (ID ><= g') (ID ><= g'')
rightmost_prod (Fun (Prod a b) (Prod c d)) (PROD f ID) = do
    COMP x f' f'' <- rightmost' (Fun a c) f
    return $ COMP (Prod x b) (f' ><= ID) (f'' ><= ID)
rightmost_prod (Fun (Prod a b) (Prod c d)) (PROD f g) = do
    COMP x f' f'' <- rightmost' (Fun a c) f
    COMP y g' g'' <- rightmost' (Fun b d) g
    return $ COMP (Prod x y) (f' ><= g') (f'' ><= g'')
rightmost_prod _ _ = mzero


-- ** Identity and Composition

nat_id = comp nat_id'
nat_id' :: Rule
nat_id' _ (COMP _ ID f) = return f
nat_id' _ (COMP _ f ID) = return f
nat_id' _ _ = mzero

comp_assocr :: Rule
comp_assocr _ (COMP a (COMP b f g) h) =
    return $ (COMP b f (COMP a g h))
comp_assocr _ _ = mzero

comp_assocl :: Rule
comp_assocl _ (COMP a f (COMP b g h)) =
    return $ (COMP b (COMP a f g) h)
comp_assocl _ _ = mzero

-- ** Bangs

bang_reflex :: Rule
bang_reflex (Fun One One) BANG =
    success "bang-Reflex" ID
bang_reflex _ _ = mzero

bang_fusion = comp bang_fusion'
bang_fusion' :: Rule
bang_fusion' (Fun a One) (COMP b BANG f) =
    success "bang-Fusion" BANG
bang_fusion' _ _ = mzero

bang_uniq :: Rule
bang_uniq (Fun _ _) ID = mzero
bang_uniq (Fun _ _) BANG = mzero
bang_uniq (Fun a One) l1 =
    success "bang-Uniq" BANG
bang_uniq _ _ = mzero
    
-- ** Lens laws

create_def :: Rule
create_def t@(Fun a c) (CREATE l) = do
    success "create-Def" $ createof (Lns c a) l
create_def _ _ = mzero

get_def :: Rule
get_def (Fun c a) (GET l) = do
    success "get-Def" $ getof (Lns c a) l
get_def _ _ = mzero

put_def :: Rule
put_def (Fun (Prod a c) _) (PUT l) = do
    success "put-Def" $ putof (Lns c a) l
put_def _ _ = mzero

-- * Lens laws

create_get = comp create_get'
create_get' :: Rule
create_get' (Fun a a') (COMP c (GET f) (CREATE f')) = do
    Eq <- teq a a'
    guard $ geq (Pf $ Lns c a) f f'
    success "Create-Get" ID
--create_get' (Fun a a') (COMP c (GET f) g) = do
--    Eq <- teq a a'
--    proof_strat optimise_pf (Fun a c) (createof (Lns c a) f) g
--    success "Create-Get" ID
create_get' _ _ = mzero

put_get = comp put_get'
put_get' :: Rule
put_get' (Fun (Prod a c) a') (COMP _ (GET f) (PUT f')) = do
    Eq <- teq a a'
    guard $ geq (Pf $ Lns c a) f f'
    success "Put-Get" FST
--put_get' (Fun (Prod a c) a') (COMP c' (GET f) g) = do
--    Eq <- teq c c'
--    Eq <- teq a a'
--    proof_strat optimise_pf (Fun (Prod a c) c) (putof (Lns c a) f) g
--    success "Put-Get" FST
put_get' _ _ = mzero

get_put = comp get_put'
get_put' :: Rule
get_put' (Fun c c') (COMP (Prod a _) (PUT f) (GET f' `SPLIT` ID)) = do
    Eq <- teq c c'
    guard $ geq (Pf $ Lns c a) f f'
    success "Get-Put" ID
get_put' _ _ = mzero

create_put = comp create_put'
create_put' :: Rule
create_put' (Fun a c) (COMP (Prod _ c') (PUT f) (ID `SPLIT` CREATE f')) = do
    Eq <- teq c c'
    guard $ geq (Pf $ Lns c a) f f'
    success "Create-Put" $ CREATE f
create_put' _ _ = mzero

put_twice = comp put_twice'
put_twice' :: Rule
put_twice' (Fun (Prod a c) c') (COMP _ (PUT l) (FST `SPLIT` PUT l')) = do
    Eq <- teq c c'
    guard $ geq (Pf $ Lns c a) l l'
    success "Put-Twice" $ PUT l
put_twice' (Fun a c) (COMP (Prod b _) (PUT l) (f `SPLIT` (COMP (Prod b' c') (PUT l') (f' `SPLIT` g)))) = do
    Eq <- teq b b'
    Eq <- teq c c'
    guard $ geq (Pf $ Fun a b) f f'
    guard $ geq (Pf $ Lns c b) l l'
    success "Put-Twice" $ COMP (Prod b c) (PUT l) $ f /\= g 
put_twice' _ _ = mzero

-- ** Backtracking sums and products

prod_undef :: Rule
prod_undef t@(Fun a (Prod b c)) (f `SPLIT` g) = do
    COMP _ f' FST <- rightmost (Fun a b) f
    COMP _ g' SND <- rightmost (Fun a c) g
    success "prod-UnDef" $ f' ><= g'
prod_undef _ _ = mzero

prod_unfusion :: Rule
prod_unfusion _ (ID `SPLIT` ID) = mzero
prod_unfusion t@(Fun a (Prod b c)) w@(f `SPLIT` g) = do
    COMP x f' h  <- (sum_unfusion ||| rightmost) (Fun a b) f
    COMP y g' h' <- (sum_unfusion ||| rightmost) (Fun a c) g
    Eq <- teq x y
    guard $ geq (Pf $ Fun a x) h h'
    res <- try (comp1 prod_unfusion >>> comp_assocr) t (COMP x (f' /\= g') h)
    success "prod-UnFusion" res
prod_unfusion _ _ = mzero

sum_undef :: Rule
sum_undef t@(Fun (Either a b) c) (f `EITHER` g) = do
    COMP _ INL f' <- leftmost (Fun a c) f
    COMP _ INR g' <- leftmost (Fun b c) g
    success "sum-UnDef" $ f' -|-= g'
sum_undef _ _ = mzero

sum_unfusion :: Rule
sum_unfusion _ (ID `EITHER` ID) = mzero
sum_unfusion t@(Fun (Either a b) c) w@(f `EITHER` g) = do
    COMP x h  f' <- (prod_unfusion ||| leftmost) (Fun a c) f
    COMP y h' g' <- (prod_unfusion ||| leftmost) (Fun b c) g
    Eq <- teq x y
    guard $ geq (Pf $ Fun x c) h h'
    res <- try (comp2 sum_unfusion >>> comp_assocl) t (COMP x h (f' \/= g'))
    success "sum-UnFusion" res
sum_unfusion _ _ = mzero    

-- ** Tops and Bottoms

top_fusion = comp top_fusion'
top_fusion' :: Rule
top_fusion' (Fun _ _) (COMP _ TOP f) =
    success "top-Fusion" TOP
top_fusion' (Fun _ _) (COMP _ f TOP) =
    success "top-Fusion" TOP
top_fusion' _ _ = mzero