{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs            #-}
{-# LANGUAGE TypeFamilies     #-}

module ASTExamples where

import Data.Maybe (fromJust)
import Control.Arrow ((>>>))
import Control.Monad ((>=>))

import Base
import AST
import Compos
import Fold
import Zipper
import HZip
import Rewriting
import Rules

-- | Example expression

example = Let ("x" := Mul (Const 6) (Const 9))
              (Add (EVar "x") (EVar "y"))

-- | Renaming variables using 'compos'

renameVar :: Expr -> Expr
renameVar = renameVar' Expr
  where
    renameVar' :: Ix AST a => AST a -> a -> a
    renameVar' Var x = x ++ "_"
    renameVar' _   x = compos renameVar' x

-- | Test for 'renameVar'

testRename :: Expr
testRename = renameVar example

-- | Result of evaluating an expression

data family Value aT :: *
data instance Value Expr  =  EV  (Env -> Int)
data instance Value Decl  =  DV  (Env -> Env)
data instance Value Var   =  VV  Var

type Env = [(Var, Int)]

-- | Algebra for evaluating an expression

evalAlgebra :: Algebra AST Value
evalAlgebra _ =  
     tag (\ (K x)                   -> EV (const x))
  &  tag (\ (I (EV x) :*: I (EV y)) -> EV (\ env -> x env  +  y env))
  &  tag (\ (I (EV x) :*: I (EV y)) -> EV (\ env -> x env  *  y env))
  &  tag (\ (I (VV x))              -> EV (fromJust . lookup x))
  &  tag (\ (I (DV e) :*: I (EV x)) -> EV (\ env -> x (e env)))
  &  tag (\ (I (VV x) :*: I (EV v)) -> DV (\ env -> (x, v env) : env ))
  &  tag (\ (I (DV f) :*: I (DV g)) -> DV (g . f))
  &  tag (\ (K x)                   -> VV x)

-- | Evaluator

eval :: Expr -> Env -> Int
eval x = let (EV f) = fold evalAlgebra x in f

-- | Test for 'eval'

testEval :: Int
testEval = eval example [("y", -12)] 

-- | Test for the generic zipper

testZipper :: Maybe Expr
testZipper =
    enter Expr   >>>
    down         >=>
    down         >=>
    right        >=>
    update solve >>>
    leave        >>>
    return        $  example
  where
    solve :: AST ix -> ix -> ix
    solve Expr _ = Const 42
    solve _    x = x

-- | Testing generic equality of expressions.

testEqual = (geq Expr example example,                -- should be True
             geq Expr example (fromJust testZipper))  -- should be False

instance Rewrite AST -- That's a bit nasty, should probably remove it

-- | A rewrite rule for expressions.

inlineRule var expr = Let (var := expr) (EVar var) :~> expr

-- | An example expression that matches the rule.

example2 = Let ("x" := Mul (Const 6) (Const 9))
               (EVar "x")

-- | Testing generic rewriting

testRewriting = rewrite2 Expr inlineRule example2

-- This more flexible version does not type in GHC 6.8.3 .
-- But it does type under the HEAD (and in mine!).
--
-- testRewriting = rewrite (rule Expr inlineRule) example2