{-# LANGUAGE TypeFamilies  #-}
{-# LANGUAGE TypeOperators #-}

import Generics.Regular.Rewriting


-----------------------------------------------------------------------------
-- Types and conversions
-----------------------------------------------------------------------------

infixl 7 :**:
infixl 6 :++:

data Expr = Const Int | Expr :++: Expr | Expr :**: Expr deriving Show

type instance PF Expr = K Int :+: I :*: I :+: I :*: I
instance Regular Expr where
  from (Const n)    = L (K n)
  from (e1 :++: e2) = R (L  $ (I e1) :*: (I e2))
  from (e1 :**: e2) = R (R  $ (I e1) :*: (I e2))

  to (L (K n))                     = Const n
  to (R (L ((I r1) :*: (I r2)))) = r1 :++: r2
  to (R (R ((I r1) :*: (I r2)))) = r1 :**: r2

{-
 -- with Con type constructors to specify constructor names.
instance Regular Expr where
  type PF Expr = Con (K Int) :+: Con (I :*: I) :+: Con (I :*: I)

  from (Const n)    = L (Con "Const" (K n))
  from (e1 :++: e2) = R (L (Con "(:++:)" $ (I e1) :*: (I e2)))
  from (e1 :**: e2) = R (R (Con "(:**:)" $ (I e1) :*: (I e2)))

  to (L (Con _ (K n)))                     = Const n
  to (R (L (Con _ ((I r1) :*: (I r2))))) = r1 :++: r2
  to (R (R (Con _ ((I r1) :*: (I r2))))) = r1 :**: r2
-}

instance Rewrite Expr

-----------------------------------------------------------------------------
-- Example rules
-----------------------------------------------------------------------------

rule1 :: Rule Expr
rule1 = 
  rule $ \x -> x :++: Const 0 :~>
               x

rule2 :: Rule Expr
rule2 = 
  rule $ \x -> x :++: x :~> 
               Const 2 :**: x

rule3 :: Rule Expr
rule3 =
  rule $ \x y -> x :++: y :~> 
                 y :++: x

rule4 :: Rule Expr
rule4 =  
  rule $ \x y -> Const 2 :**: (x :++: y) :~>
                 (Const 2 :**: x) :++: (Const 2 :**: y)

rule5 :: Rule Expr
rule5 = 
  rule $ \x y z -> x :**: (y :++: z) :~>  
                   (x :**: y) :++: (x :**: z) 

rule6 :: Rule Expr
rule6 = 
  rule $ Const 1 :++: Const 1 :~>
         Const 2


-----------------------------------------------------------------------------
-- Tests
-----------------------------------------------------------------------------

test1 :: Maybe Expr
test1 = rewriteM rule1 (Const 2 :++: Const 0)

test2 :: Maybe Expr
test2 = rewriteM rule1 (Const 2 :++: Const 3)

test3 :: Maybe Expr
test3 = rewriteM rule2 (Const 4 :++: Const 3)

test4 :: Maybe Expr
test4 = rewriteM rule2 (Const 4 :++: Const 4)

test5 :: Maybe Expr
test5 = one (rewriteM rule1) ((Const 4 :++: Const 0) :**: Const 2)

-- This does not work because the optimisation target is not
-- an immediate child.
test6 :: Maybe Expr
test6 = one (rewriteM rule1) (((Const 4 :++: Const 0) :**: Const 2) :++: Const 7)

-- This works well, because once applies the rule to the optimisation
-- target exactly once.
test7 :: Maybe Expr
test7 = once (rewriteM rule1) (((Const 4 :++: Const 0) :**: Const 2) :++: Const 7)

test8 :: Maybe Expr
test8 = rewriteM rule3 ((Const 1) :++: (Const 2))

test9 :: Maybe Expr
test9 = rewriteM rule4 ((Const 2) :**: ((Const 3) :++: (Const 4)))

test10 :: Maybe Expr
test10 = rewriteM rule5 ((Const 1) :**: ((Const 2) :++: (Const 3)))

test11 :: Maybe Expr
test11 = rewriteM rule6 (Const 1 :++: Const 1)

allTests :: [Maybe Expr]
allTests = [ test1
           , test2
           , test3
           , test4
           , test5
           , test6
           , test7
           , test8
           , test9
           , test10
           , test11
           ]


-----------------------------------------------------------------------------
-- Running all the tests
-----------------------------------------------------------------------------

-- This main function is defined to solve a bug in GHC
main :: IO ()
main = do let resultsPP         = zipWith resultPP [1..] allTests
              resultPP n result = "test" ++ show n ++ ": " ++ show result
          putStr (unlines resultsPP)