-----------------------------------------------------------------------------
-- |
-- 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.Pf
import Data.Lens
import Data.Equal
import Transform.Rewriting

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

-- ** Combinators

comp :: Rule -> Rule
comp r (Fun d a) (COMP b f (COMP c g h)) = do
    fg <- r (Fun c a) (COMP b f g)
    return $ COMP c fg h
comp r (Fun d a) (COMP c (COMP b f g) h) = do
    gh <- r (Fun d b) (COMP c g h)
    return $ COMP b f gh
comp r t f = r t f

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 t@(Fun (Prod a b) (Prod c d)) v@(PROD ID g) = do
    COMP y g' g'' <- rightmost' (Fun b d) g
    return $ COMP (Prod a y) (ID ><= g') (ID ><= g'')
rightmost_prod t@(Fun (Prod a b) (Prod c d)) v@(PROD f ID) = do
    COMP x f' f'' <- rightmost' (Fun a c) f
    return $ COMP (Prod x b) (f' ><= ID) (f'' ><= ID)
rightmost_prod t@(Fun (Prod a b) (Prod c d)) v@(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' t@(Fun a One) v@(COMP b BANG f) = do
    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' _ _ = 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' _ _ = 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_wunfusion :: Rule
prod_wunfusion t@(Fun a _) (COMP x f g `SPLIT` COMP y h g') = do
    Eq <- teq x y
    guard $ geq (Pf $ Fun a x) g g'
    success "prod-Unfusion" $ COMP x (f `SPLIT` h) g
prod_wunfusion _ _ = 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_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

dyn_cancel, dyn_cancel' :: Rule
dyn_cancel = comp dyn_cancel'
dyn_cancel' _ (COMP _ (UNDYN a) (MKDYN b)) = do
    Eq <- teq a b
    success "dyn-Cancel" ID
dyn_cancel' _ _ = mzero

cast_cancel, cast_cancel' :: Rule
cast_cancel = comp cast_cancel'
cast_cancel' _ (COMP _ (CAST a) (MKDYN b)) = do
	cast_cancel' (Fun b a) (CAST a)
cast_cancel' (Fun b@(Data s f) _) (CAST a) | isBasic a = do
	Eq <- teq (rep f b) a
	success "cast-Cancel" OUT
cast_cancel' (Fun b@(NewData s f) _) (CAST a) | isBasic a = do
	Eq <- teq (rep f b) a
	success "cast-Cancel" OUT
cast_cancel' (Fun b _) (CAST a) = do
	Eq <- teq a b
	success "cast-Cancel" ID
cast_cancel' _ _ = mzero

primitives :: Rule
primitives = top comp_assocr ||| top nat_id ||| top dyn_cancel ||| top cast_cancel ||| top top_fusion

bangs :: Rule
bangs = top bang_reflex ||| top bang_fusion ||| top bang_uniq

-- ** Relating sums with products

abides = abides'
abides' :: Rule
abides' (Fun _ _) ((f `SPLIT` g) `EITHER` (h `SPLIT` i)) =
    success "abides" $ (f \/= h) /\= (g \/= i)
abides' _ _ = mzero

unabides = unabides'
unabides' :: Rule
unabides' (Fun _ _) ((f `EITHER` h) `SPLIT` (g `EITHER` i)) =
    success "unabides" $ (f /\= g) \/= (h /\= i)
unabides' _ _ = mzero