-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Equal
-- 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
-- 
-- Implementation of type and function equality at the value-level.
--
-----------------------------------------------------------------------------

module Data.Equal where

import Data.Type
import Data.Pf
import Data.Spine

import Control.Monad hiding (Functor(..))
import Unsafe.Coerce
import Control.Monad.State as ST hiding (Functor(..))
import Control.Monad.Reader hiding (Functor(..))
import Data.Map as Map
import Data.List as List
import Prelude hiding (Functor(..))

import Generics.Pointless.Functors hiding (rep)

data Equal a b where
    Eq :: Equal a a

teqBool :: Type a -> Type b -> Bool
teqBool a b = maybe False (const True) (teq a b)

type Vars = Map String DynType

-- type equality where the left-side type may have unbounded variables, representing pattern-matching
teqvar :: MonadPlus m => Type a -> Type b -> StateT Vars m (Equal a b)
teqvar (Var n) a = do
    vars <- ST.get
    case (Map.lookup n vars) of
    { Just (DynT t) -> do
        Eq <- teq t a
        return (unsafeCoerce Eq)
    ; otherwise -> do
        ST.put (Map.insert n (DynT a) vars)
        return (unsafeCoerce Eq)
    }
teqvar Any _ = return (unsafeCoerce Eq)
teqvar _ Any = return (unsafeCoerce Eq)
teqvar (Id a) (Id b) = teqvar a b
teqvar One One = return Eq
teqvar Int Int = return Eq
teqvar Bool Bool = return Eq
teqvar Char Char = return Eq
teqvar (Prod a b) (Prod c d) = do
	Eq <- teqvar a c
	Eq <- teqvar b d
	return Eq
teqvar (Either a b) (Either c d) = do
	Eq <- teqvar a c
	Eq  <- teqvar b d
	return Eq
teqvar (Data s fx) (Data s' fy) = do
    guard (sameName s s')
    Eq <- feqvar fx fy
    return (unsafeCoerce Eq)
teqvar (NewData s fx) (NewData s' fy) = do
    guard (sameName s s')
    Eq <- feqvar fx fy
    return (unsafeCoerce Eq)
teqvar (List a) (List b) = do
    Eq <- teqvar a b
    return Eq
teqvar (Fun a b) (Fun c d) = do
    Eq <- teqvar a c
    Eq <- teqvar b d
    return Eq
teqvar (Lns a b) (Lns c d) = do
    Eq <- teqvar a c
    Eq <- teqvar b d
    return Eq
teqvar (Pf a) (Pf b) = do
    Eq <- teqvar a b
    return Eq
teqvar Dynamic Dynamic = return Eq
teqvar TP TP = return Eq
teqvar (TU a) (TU b) = do
    Eq <- teqvar a b
    return Eq
teqvar _ _ = mzero

teqvars :: MonadPlus m => Type a -> Type b -> m (Equal a b,Vars)
teqvars a b = runStateT (teqvar a b) Map.empty

-- regular type equality
teq :: MonadPlus m => Type a -> Type b -> m (Equal a b)
teq a b = evalStateT (teqvar a b) Map.empty

feqvar :: MonadPlus m => Fctr f -> Fctr g -> StateT Vars m (Equal (Fix f) (Fix g))
feqvar I I = return Eq
feqvar (K a) (K b) = do
    Eq <- teqvar a b
    return Eq
feqvar L L = return Eq
feqvar (f :*!: g) (h :*!: i) = do
    Eq <- feqvar f h
    Eq <- feqvar g i
    return Eq
feqvar (f :+!: g) (h :+!: i) = do
    Eq <- feqvar f h
    Eq <- feqvar g i
    return Eq
feqvar (f :@!: g) (h :@!: i) = do
    Eq <- feqvar f h
    Eq <- feqvar g i
    return Eq
feqvar AnyF f = return (unsafeCoerce Eq)
feqvar f AnyF = return (unsafeCoerce Eq)
feqvar _ _ = mzero

feq :: MonadPlus m => Fctr f -> Fctr g -> m (Equal (Fix f) (Fix g))
feq f g = evalStateT (feqvar f g) Map.empty

-- | Tests if a functor is recursive or not, by applying it to two distinct types.
isRec :: Fctr f -> Bool
isRec fctr = case teq (rep fctr Int) (rep fctr One) of { Just Eq -> False ; otherwise -> True }

-- | Syntactic equality, with the exception of protected values.
geq :: Type a -> a -> a -> Bool
geq (Pf t) (PROTECT x) y = geq (Pf t) x y
geq (Pf t) x (PROTECT y) = geq (Pf t) x y
geq (Pf t) (PROTECT_LNS x) y = geq (Pf t) x y
geq (Pf t) x (PROTECT_LNS y) = geq (Pf t) x y
geq t x y = geq' t x t y

geq' :: Type a -> a -> Type b -> b -> Bool
geq' a x b y = aux a b x y
    where aux :: Type a -> Type b -> a -> b -> Bool
          aux t1 t2 x y = aux' t1 (toSpine t1 x) t2 (toSpine t2 y)
          aux' :: Type a -> Spine a -> Type b -> Spine b -> Bool
          aux' _ (_ `As` c1) _ (_ `As` c2) = name c1 == name c2
          aux' t (f1 `Ap` (t1 :| a1)) t' (f2 `Ap` (t2 :| a2)) = aux' (Fun t1 t) f1 (Fun t2 t') f2 && aux t1 t2 a1 a2
          aux' _ _ _ _ = False

-- | Clone of |geq| with a specific case for top.
geqt :: Type a -> a -> a -> Bool
geqt (Pf t) (PROTECT x) y = geqt (Pf t) x y
geqt (Pf t) x (PROTECT y) = geqt (Pf t) x y
geqt (Pf t) (PROTECT_LNS x) y = geqt (Pf t) x y
geqt (Pf t) x (PROTECT_LNS y) = geqt (Pf t) x y
geqt t x y = geqt' t x t y

geqt' :: Type a -> a -> Type b -> b -> Bool
geqt' a x b y = aux a b x y
    where aux :: Type a -> Type b -> a -> b -> Bool
          aux t1 t2 x y = aux' t1 (toSpine t1 x) t2 (toSpine t2 y)
          aux' :: Type a -> Spine a -> Type b -> Spine b -> Bool
          aux' _ _ (Pf _) (TOP `As` _) = True
          aux' (Pf _) (As TOP _) _ _ = True
          aux' _ (_ `As` c1) _ (_ `As` c2) = name c1 == name c2
          aux' t (f1 `Ap` (t1 :| a1)) t' (f2 `Ap` (t2 :| a2)) = aux' (Fun t1 t) f1 (Fun t2 t') f2 && aux t1 t2 a1 a2
          aux' _ _ _ _ = False

-- | Explicitly coerce a value of a given type to another given type.
coerce :: MonadPlus m => Type a -> Type b -> a -> m b
coerce a b x = do Eq <- teq a b
                  return x

collectDyn :: MonadPlus m => Type a -> a -> m DynType
collectDyn a v = case collectDyn' a v of { Just d -> return d; otherwise -> mzero }
collectDyn' :: Type a -> a -> Maybe DynType
collectDyn' = collect q plus
    where q :: MonadPlus m => GenericQ (m DynType)
          q (Pf _) (MKDYN a) = return $ DynT a
          q _ _ = mzero
          plus :: Maybe DynType -> Maybe DynType -> Maybe DynType
          plus (Just (DynT a)) (Just (DynT b)) = teq a b >> return (DynT a)
          plus m Nothing = m
          plus Nothing n = n

collectNewNames :: Type a -> [String]
collectNewNames = List.map fst . Map.toList . collectNewDatas

collectNewDatas :: Type a -> Map String DynFctr
collectNewDatas = maybe Map.empty id . collect q mcat TypeRep
    where q :: MonadPlus m => GenericQ (m (Map String DynFctr))
          q TypeRep (NewData s f) = return $ Map.singleton s (DynF f)
          q _ _ = return Map.empty
          mcat m n = do { x <- m; y <- n; return (x `Map.union` y) }

{-
showDatas :: Type a -> String
showDatas = maybe [] id . collect q mcat TypeRep
    where q :: MonadPlus m => GenericQ (m String)
          q TypeRep d@(isData -> True) = return (showData d ++ "\n")
          q _ _ = return []
          mcat m n = do { x <- m; y <- n; return (x ++ y) }
-}
collectVars :: Type a -> [String]
collectVars = maybe [] id . collect q mcat TypeRep
    where q :: MonadPlus m => GenericQ (m [String])
          q TypeRep (Var s) = return [s]
          q _ _ = return []
          mcat m n = do { x <- m; y <- n; return (x ++ y) }

collect :: MonadPlus m => GenericQ (m r) -> (m r -> m r -> m r) -> Type a -> a -> m r
collect (q :: GenericQ (m r)) plus a x = collectSpine a (toSpine a x)
    where collectSpine :: MonadPlus m => Type a -> Spine a -> m r
          collectSpine t s@(As _ _) = q t (fromSpine s)
          collectSpine t s@(Ap f (a :| v)) = q t (fromSpine s)
            `plus` (collectSpine (Fun a t) f)
            `plus` (collectSpine a (toSpine a v))

-- | Find a value of type b inside a value of type a
find :: Type b -> b -> Type a -> a -> Bool
find b y a x = findSpine a (toSpine a x)
    where findSpine :: Type a -> Spine a -> Bool
          findSpine t (As v con) = case teq t b of {
              Just Eq   -> geqt t v y;
              otherwise -> False
              }
          findSpine t s@(Ap f (a :| v)) = (case teq t b of {
              Just Eq   -> geqt b y (fromSpine s);
              otherwise -> False
              })
              || findSpine (Fun a t) f
              || findSpine a (toSpine a v)

removeIds :: Type a -> a -> a
removeIds t x = fromSpine $ removeIdSpine t $ toSpine t x

removeIdSpine :: Type a -> Spine a -> Spine a
removeIdSpine TypeRep s@(fromSpine -> (Id a)) = removeIdSpine TypeRep (toSpine TypeRep a)
removeIdSpine t (As v con) = As v con
removeIdSpine t s@(Ap f (a :| v)) = Ap (removeIdSpine (Fun a t) f) (a :| fromSpine (removeIdSpine a (toSpine a v)))

unDyn :: Type a -> Dynamic -> a
unDyn t (Dyn a x) = case teq t a of { Just Eq -> x; otherwise -> error "unDyn failed"}

cast :: Type a -> Type b -> b -> a
cast a Dynamic (Dyn b x) = cast a b x
cast a b@(Data s f) x | isBasic a = case teq (rep f b) a of { Just Eq -> out x; otherwise -> error "type cast failed"}
cast a b@(NewData s f) x | isBasic a = case teq (rep f b) a of { Just Eq -> out x; otherwise -> error "type cast failed"}
cast a b x = case teq a b of { Just Eq -> x; otherwise -> error "type cast failed"}

isInt :: Type a -> Maybe (Equal a Int)
isInt a = teq a Int
isList :: Type a -> Maybe (Equal a [b])
isList a = teq a (List Any)
isNat :: Type a -> Maybe (Equal a Nat)
isNat a = teq a nat

-- infers a new functor for newly created recursive types
reshape :: MonadPlus m => Type a -> m DynType
reshape (NewData s f) = do
	let mark = Id Any
	DynF g <- reshapeF f
	FRep h <- inferFctr mark (rep g mark)
	return $ DynT $ NewData s h
reshape (Prod a b) = do
	DynT c <- reshape a	
	DynT d <- reshape b
	return $ DynT $ Prod c d
reshape (Either a b) = do
	DynT c <- reshape a	
	DynT d <- reshape b
	return $ DynT $ Either c d
reshape (List a) = do
	DynT b <- reshape a
	return $ DynT $ List b
reshape a = return $ DynT a
	
reshapeF :: MonadPlus m => Fctr f -> m DynFctr
reshapeF I = return $ DynF I
reshapeF (K a) = do
	DynT b <- reshape a
	return $ DynF $ K b
reshapeF L = return $ DynF L
reshapeF (f :*!: g) = do
	DynF h <- reshapeF f
	DynF i <- reshapeF g
	return $ DynF $ h :*!: i
reshapeF (f :+!: g) = do
	DynF h <- reshapeF f
	DynF i <- reshapeF g
	return $ DynF $ h :+!: i
reshapeF (f :@!: g) = do
	DynF h <- reshapeF f
	DynF i <- reshapeF g
	return $ DynF $ h :@!: i

data FctrRep a b where
    FRep :: (Functor f,Rep f a ~ b) => Fctr f -> FctrRep a b

-- Infers a new functor from a base type and an identity marker
inferFctr :: MonadPlus m => Type a -> Type b -> m (FctrRep a b)
inferFctr a (Prod x y) = do
    FRep f <- inferFctr a x
    FRep g <- inferFctr a y
    return $ FRep $ f :*!: g
inferFctr a (Either x y) = do
    FRep f <- inferFctr a x
    FRep g <- inferFctr a y
    return $ FRep $ f :+!: g
inferFctr a (List x) = do
    FRep f <- inferFctr a x
    return $ FRep $ L :@!: f
inferFctr a x = (do
    Eq <- teq a x
    return $ FRep I)
        `mplus` (do
    return $ FRep (K x))

-- Infers a new constant functor from a base type
-- The functor is always constant, i.e., forall a,b. Rep f a ~ Rep f b, altough this escapes the type-checker.
inferKFctr :: MonadPlus m => Type b -> m (FctrRep Dynamic b)
inferKFctr (Prod x y) = do
    FRep f <- inferKFctr x
    FRep g <- inferKFctr y
    return $ FRep $ f :*!: g
inferKFctr (Either x y) = do
    FRep f <- inferKFctr x
    FRep g <- inferKFctr y
    return $ FRep $ f :+!: g
inferKFctr (List x) = do
    FRep f <- inferKFctr x
    return $ FRep $ L :@!: f
inferKFctr x = return $ FRep (K x)

type TypeRule s = MonadPlus m => forall a. Type a -> StateT s m (Type a)
type FctrRule s = MonadPlus m => forall f. Fctr f -> StateT s m (Fctr f)

-- replaces the variables in an argument type with the concrete instantiations in the context.
replacevar :: MonadPlus m => Type a -> Vars -> m (Type a)
replacevar t vars = evalStateT (replace var none t) vars
	where
	var :: TypeRule Vars
	var (Var s) = do
		ctx <- ST.get
		case (Map.lookup s ctx) of
    			{ Just (DynT a) -> return (unsafeCoerce a)
    			; otherwise -> mzero }
        var _ = mzero
        none :: FctrRule Vars
        none f = mzero

replacedyn :: Type a -> Type a
replacedyn t = maybe t id $ evalStateT (replace dyn kdyn t) ()
	where dyn :: TypeRule ()
	      dyn Dynamic = return Any
	      dyn _ = mzero
	      kdyn :: FctrRule ()
	      kdyn (K Dynamic) = return AnyF
	      kdyn _ = mzero

replace,replace' :: TypeRule s -> FctrRule s -> TypeRule s
replace tr fr t = tr t `mplus` replace' tr fr t
replace' tr fr (Var s) = return $ Var s
replace' tr fr (Id a) = do
	x <- replace tr fr a
	return (Id x)
replace' tr fr Int = return Int
replace' tr fr Bool = return Bool
replace' tr fr Char = return Char
replace' tr fr One = return One
replace' tr fr (Either a b) = do
	x <- replace tr fr a
	y <- replace tr fr b
	return (Either x y)
replace' tr fr (Prod a b) = do
	x <- replace tr fr a
	y <- replace tr fr b
	return (Prod x y)
replace' tr fr (Fun a b) = do
	x <- replace tr fr a
	y <- replace tr fr b
	return (Fun x y)
replace' tr fr (Lns a b) = do
	x <- replace tr fr a
	y <- replace tr fr b
	return (Lns x y)
replace' tr fr (List a) = do
	x <- replace tr fr a
	return (List x)
replace' tr fr (Data s f) = do
	g <- replaceF tr fr f
	Eq <- feq f g
	return (Data s g)
replace' tr fr (NewData s f) = do
	g <- replaceF tr fr f
	return (NewData s g)
replace' tr fr (Pf a) = do
	x <- replace tr fr a
	return (Pf x)
replace' tr fr TP = return TP
replace' tr fr (TU a) = do
	x <- replace tr fr a
	return $ TU a
replace' tr fr Any = return Any
replace' tr fr Dynamic = return Dynamic

replaceF,replaceF' :: TypeRule s -> FctrRule s -> FctrRule s
replaceF tr fr f = fr f `mplus` replaceF' tr fr f
replaceF' tr fr I = return I
replaceF' tr fr (K a) = do
	x <- replace tr fr a 
	return (K x)
replaceF' tr fr L = return L
replaceF' tr fr (f :*!: g) = do
	x <- replaceF tr fr f
	y <- replaceF tr fr g
	return (x :*!: y)
replaceF' tr fr (f :+!: g) = do
	x <- replaceF tr fr f
	y <- replaceF tr fr g
	return (x :+!: y)
replaceF' tr fr (f :@!: g) = do
	x <- replaceF tr fr f
	y <- replaceF tr fr g
	return (x :@!: y)