{-# LANGUAGE DeriveDataTypeable, PatternGuards #-}

module Type(module Type, Literal(..)) where

import Language.Haskell.Exts(Literal(..))

import Data.Data
import Data.Maybe
import Data.List
import Data.Generics.PlateData
import Control.Arrow
import Control.Monad
import Control.Monad.State
import Control.Monad.Identity
import qualified Data.Map as Map

type Var = String
type Con = String
type Fun = String


type Prog = [Func]

data Func = Func {funcName :: Fun, funcArgs :: [Var], funcBody :: Expr}
            deriving (Data,Typeable,Show)

data Expr = ECon  Con
          | EFun  Fun
          | ELit  Literal
          | EVar  Var
          | EApp  I Expr Expr
          | ECase I  Expr [(Patt,Expr)]
          | ELet  I [(Var,Expr)] Expr
          | ELam  I Var Expr -- not supported in many places!
            deriving (Data,Typeable,Show,Eq,Ord)


data Patt = Patt Con [Var]
          | PattAny
          | PattLit Literal
            deriving (Data,Typeable,Show,Eq,Ord)

pattVars (Patt _ vs) = vs
pattVars _ = []
isPattAny PattAny{} = True ; isPattAny _ = False


getFunc :: Prog -> (Var -> Func)
getFunc prog = \name -> fromMaybe (error $ "Couldn't find function " ++ name) $ func name
    where func = getFuncMaybe prog

getFuncMaybe :: Prog -> (Var -> Maybe Func)
getFuncMaybe prog = \name -> Map.lookup name mp
    where mp = Map.fromList $ map (funcName &&& id) prog


-- the Expr is varNames is always undefined, just for a fast subst test
type I = Map.Map Var Int

getI :: Expr -> Map.Map Var Int
getI (ECon _) = Map.empty
getI (ELit _) = Map.empty
getI (EFun _) = Map.empty
getI (EVar x) = Map.singleton x 1
getI (EApp i _ _) = i
getI (ECase i _ _) = i
getI (ELet i _ _) = i
getI (ELam i _ _) = i

unionI = Map.unionWith (+)
unionsI = Map.unionsWith (+)
minusI x [] = x
minusI x (y:ys) = minusI (Map.delete y x) ys

eCon x = ECon x
eFun x = EFun x
eLit x = ELit x
eVar x = EVar x
eApp x y = EApp (getI x `unionI` getI y) x y
eCase x y = ECase (unionI (getI x) alts) x y
    where alts = Map.unionsWith max [minusI (getI x) (pattVars p) | (p,x) <- y] 
eLet [] y = y
eLet x y = ELet (unionI binds body) x y
    where
        binds = unionsI $ map (getI . snd) x
        body = minusI (getI y) $ map fst x

eLam v x = ELam (minusI (getI x) [v]) v x

eApps x [] = x
eApps x y = eApp (eApps x (init y)) (last y)

eLams [] x = x
eLams (v:vs) x = eLam v (eLams vs x)

fromELams (ELam _ v x) = (v:a, b)
    where (a,b) = fromELams x
fromELams x = ([], x)

fromEApps (EApp _ x y) = (a, b ++ [y])
    where (a,b) = fromEApps x
fromEApps x = (x, [])


subst :: Var -> Expr -> Expr -> Expr
subst v b = runIdentity . substsWith return (Map.singleton v b)


substs :: [(Var,Expr)] -> Expr -> Expr
substs bs = runIdentity . substsWith return (Map.fromList bs)

substsWith :: Monad m => (Expr -> m Expr) -> Map.Map Var Expr -> Expr -> m Expr
substsWith op mp x | Map.null mp2 = return x
                   | otherwise = op =<< case x of
        EVar v -> return $ Map.findWithDefault (error "substs logic error") v mp2 -- must be there, or would have gone the empty route
        EApp _ x y -> liftM2 eApp (f mp2 x) (f mp2 y)
        ECase _ x y -> liftM2 eCase (f mp2 x) $ sequence [liftM ((,) a) $ f (minusI mp2 $ pattVars a) b | (a,b) <- y]
        ELet _ xy z -> liftM2 eLet (mapM (\(x,y) -> liftM ((,) x) $ f mp2 y) xy) (f (minusI mp2 $ map fst xy) z)
        _ -> error "substs logic error (2)"
    where mp2 = mp `Map.intersection` getI x
          f = substsWith op


-- is the variable used at most once down any evaluation path
linear :: Var -> Expr -> Bool
linear v x = Map.findWithDefault 0 v (getI x) <= 1


freeVars :: Expr -> [Var]
freeVars = Map.keys . getI


disjoint x y = null $ x `intersect` y


-- normalise bound variables and make them all unique
normalise :: Func -> Unique Func
normalise (Func name args bod) = do
    vs <- freshN (length args)
    fmap (Func name vs) $ f (Map.fromList $ zip args vs) bod
    where
        f mp (EVar v) = return $ eVar $ Map.findWithDefault v v mp
        f mp (EApp _ x y) = liftM2 eApp (f mp x) (f mp y)
        f mp (ECase _ x y) = do
            x <- f mp x
            y <- mapM (g mp) y
            return $ eCase x y
        f mp (ELet _ xy z) = do
            let (xs,ys) = unzip xy
            vs <- freshN (length xy)
            xy <- fmap (zip vs) $ mapM (f mp) ys
            z <- f (Map.fromList (zip xs vs) `Map.union` mp) z
            return $ eLet xy z
        f mp (ELam _ x y) = do
            v <- fresh
            y <- f (Map.singleton x v `Map.union` mp) y
            return $ eLam v y
        f mp x = return x

        g mp (Patt c vs, x) = do
            vs2 <- freshN (length vs)
            fmap ((,) (Patt c vs2)) $ f (Map.fromList (zip vs vs2) `Map.union` mp) x
        g mp (p,x) = fmap ((,) p) $ f mp x


normaliseExpr :: Expr -> Unique Expr
normaliseExpr = fmap funcBody . normalise . Func "" []


type Unique a = State Int a


runUnique :: Unique a -> Int -> (a, Int)
runUnique = runState

unique :: Unique a -> a
unique = flip evalState 1

fresh :: Unique String
fresh = do
    i <- get
    put $ i+1
    return $ "v" ++ show i

freshN :: Int -> Unique [String]
freshN n = replicateM n fresh

fromUnique :: Unique a -> (Int,a)
fromUnique = (\(x,y) -> (y,x)) . flip runState 1

toUnique :: Int -> Unique a -> Unique a
toUnique i x = do
    put i
    x


---------------------------------------------------------------------
-- REDEX

data Redex = RLet [(Var, Redex)] Redex
           | RVar Var
           | RFun Fun
           | RApp Var Var
           | RCon Con [Var]
           | RCase Var [(Patt, Redex)]
           | RLit Literal
             deriving (Data, Typeable, Show, Eq, Ord)


fromRedex :: Redex -> Expr
fromRedex (RLet xs y) = eLet (map (second fromRedex) xs) (fromRedex y)
fromRedex (RApp x y) = eApp (EVar x) (EVar y)
fromRedex (RVar x) = EVar x
fromRedex (RFun x) = EFun x
fromRedex (RCase x xs) = eCase (EVar x) (map (second fromRedex) xs)
fromRedex (RCon x xs) = eApps (ECon x) (map EVar xs)
fromRedex (RLit x) = ELit x

fromRedexShort :: Redex -> Expr
fromRedexShort (RCase x xs) = fromRedex $ RCase x $ map (second $ const $ RVar "...") xs
fromRedexShort x = fromRedex x

rLet xs x = if null xs then x else RLet xs x

toRedex :: Func -> Unique ([Var], Redex)
toRedex func = do
    Func _ args bod <- normalise func
    fmap ((,) args) $ f bod
    where
        f (EApp _ x y) = bind [x,y] $ \[x,y] -> RApp x y
        f (ECon x) = return $ RCon x []
        f (EFun x) = return $ RFun x
        f (EVar x) = return $ RVar x
        f (ELit x) = return $ RLit x
        f (ECase _ x xs) = do
            let (x1,x2) = unzip xs
            x2 <- mapM f x2
            bind [x] $ \[x] -> RCase x $ zip x1 x2
        f (ELet _ bind x) = do
            let (x1,x2) = unzip bind
            x2 <- mapM f x2
            x <- f x
            return $ RLet (zip x1 x2) x

        bind xs op = do
            (new,vs) <- mapAndUnzipM g xs
            return $ rLet (concat new) $ op vs

        g (EVar v) = return ([], v)
        g x = do
            v <- fresh
            y <- f x
            return ([(v,y)],v)


-- | Note: Currently approximate
freeVarsR :: Redex -> [Var]
freeVarsR (RLet xs x) = concatMap (freeVarsR . snd) xs ++ freeVarsR x
freeVarsR (RApp x y) = [x,y]
freeVarsR (RCase x xs) = x : concatMap (freeVarsR . snd) xs
freeVarsR (RCon x xs) = xs
freeVarsR (RVar x) = [x]
freeVarsR _ = []