{-|
Module      : IRTS.Lang
Description : Internal representation of Idris' constructs.

License     : BSD3
Maintainer  : The Idris Community.
-}
{-# LANGUAGE DeriveDataTypeable, DeriveFunctor, DeriveGeneric, FlexibleContexts,
             PatternGuards #-}

module IRTS.Lang where

import Idris.Core.CaseTree
import Idris.Core.TT

import Control.Monad.State hiding (lift)
import Data.Data (Data)
import Data.List
import qualified Data.Map.Strict as Map
import Data.Typeable (Typeable)
import GHC.Generics (Generic)

data Endianness = Native | BE | LE deriving (Show, Eq)

data LVar = Loc Int | Glob Name
  deriving (Show, Eq)

-- ASSUMPTION: All variable bindings have unique names here
-- Constructors commented as lifted are not present in the LIR provided to the different backends.
data LExp = LV Name
          | LApp Bool LExp [LExp]    -- True = tail call
          | LLazyApp Name [LExp]     -- True = tail call
          | LLazyExp LExp            -- lifted out before compiling
          | LForce LExp              -- make sure Exp is evaluted
          | LLet Name LExp LExp      -- name just for pretty printing
          | LLam [Name] LExp         -- lambda, lifted out before compiling
          | LProj LExp Int           -- projection
          | LCon (Maybe Name)        -- Location to reallocate, if available
                 Int Name [LExp]
          | LCase CaseType LExp [LAlt]
          | LConst Const
          | LForeign FDesc           -- Function descriptor (usually name as string)
                     FDesc           -- Return type descriptor
                     [(FDesc, LExp)] -- first LExp is the FFI type description
          | LOp PrimFn [LExp]
          | LNothing
          | LError String
  deriving (Eq, Ord)

data FDesc = FCon Name
           | FStr String
           | FUnknown
           | FIO FDesc
           | FApp Name [FDesc]
  deriving (Show, Eq, Ord)

data Export = ExportData FDesc -- Exported data descriptor (usually string)
            | ExportFun Name -- Idris name
                        FDesc -- Exported function descriptor
                        FDesc -- Return type descriptor
                        [FDesc] -- Argument types
  deriving (Show, Eq, Ord)

data ExportIFace = Export Name -- FFI descriptor
                          String -- interface file
                          [Export]
  deriving (Show, Eq, Ord)

-- Primitive operators. Backends are not *required* to implement all
-- of these, but should report an error if they are unable

data PrimFn = LPlus ArithTy | LMinus ArithTy | LTimes ArithTy
            | LUDiv IntTy | LSDiv ArithTy | LURem IntTy | LSRem ArithTy
            | LAnd IntTy | LOr IntTy | LXOr IntTy | LCompl IntTy
            | LSHL IntTy | LLSHR IntTy | LASHR IntTy
            | LEq ArithTy | LLt IntTy | LLe IntTy | LGt IntTy | LGe IntTy
            | LSLt ArithTy | LSLe ArithTy | LSGt ArithTy | LSGe ArithTy
            | LSExt IntTy IntTy | LZExt IntTy IntTy | LTrunc IntTy IntTy
            | LStrConcat | LStrLt | LStrEq | LStrLen
            | LIntFloat IntTy | LFloatInt IntTy | LIntStr IntTy | LStrInt IntTy
            | LFloatStr | LStrFloat | LChInt IntTy | LIntCh IntTy
            | LBitCast ArithTy ArithTy -- Only for values of equal width

            | LFExp | LFLog | LFSin | LFCos | LFTan | LFASin | LFACos | LFATan
            | LFATan2 | LFSqrt | LFFloor | LFCeil | LFNegate

            | LStrHead | LStrTail | LStrCons | LStrIndex | LStrRev | LStrSubstr
            | LReadStr | LWriteStr

            -- system info
            | LSystemInfo

            | LFork
            | LPar -- evaluate argument anywhere, possibly on another
                   -- core or another machine. 'id' is a valid implementation
            | LExternal Name
            | LCrash

            | LNoOp
  deriving (Show, Eq, Ord, Generic)

-- Supported target languages for foreign calls

data FCallType = FStatic | FObject | FConstructor
  deriving (Show, Eq, Ord)

data FType = FArith ArithTy
           | FFunction
           | FFunctionIO
           | FString
           | FUnit
           | FPtr
           | FManagedPtr
           | FCData
           | FAny
  deriving (Show, Eq, Ord)

-- FIXME: Why not use this for all the IRs now?
data LAlt' e = LConCase Int Name [Name] e
             | LConstCase Const e
             | LDefaultCase e
  deriving (Show, Eq, Ord, Functor, Data, Typeable)

type LAlt = LAlt' LExp

data LDecl = LFun [LOpt] Name [Name] LExp -- options, name, arg names, def
           | LConstructor Name Int Int -- constructor name, tag, arity
  deriving (Show, Eq, Ord)

type LDefs = Ctxt LDecl

data LOpt = Inline | NoInline
  deriving (Show, Eq, Ord)

addTags :: Int -> [(Name, LDecl)] -> (Int, [(Name, LDecl)])
addTags i ds = tag i ds []
  where tag i ((n, LConstructor n' (-1) a) : as) acc
            = tag (i + 1) as ((n, LConstructor n' i a) : acc)
        tag i ((n, LConstructor n' t a) : as) acc
            = tag i as ((n, LConstructor n' t a) : acc)
        tag i (x : as) acc = tag i as (x : acc)
        tag i [] acc  = (i, reverse acc)

data LiftState = LS (Maybe Name) Int [(Name, LDecl)]
                    (Map.Map ([Name], LExp) Name) -- map from args/expressions
                          -- to names, so we don't create the same function
                          -- multiple times

setBaseName :: Name -> State LiftState ()
setBaseName n
    = do LS _ i ds done <- get
         put (LS (Just n) i ds done)

lname (NS n x) i = NS (lname n i) x
lname (UN n) i = MN i n
lname x i = sMN i (showCG x ++ "_lam")

getNextName :: State LiftState Name
getNextName
    = do LS mn i ds done <- get
         let newn = case mn of
                         Nothing -> lname (sUN "_") i
                         Just n -> lname n i
         put (LS mn (i + 1) ds done)
         return newn

renameArgs :: [Name] -> LExp -> ([Name], LExp)
renameArgs args e
   = let newargNames = map (\i -> sMN i "lift") [0..]
         newargs = zip args newargNames in
         (map snd newargs, rename newargs e)

addFn :: Name -> LDecl -> State LiftState ()
addFn fn d
    = do LS n i ds done <- get
         put (LS n i ((fn, d) : ds) done)

makeFn :: [Name] -> LExp -> State LiftState Name
makeFn args exp
    = do fn <- getNextName
         let (args', exp') = renameArgs args exp
         LS n i ds done <- get
         case Map.lookup (args', exp') done of
              Just fn -> return fn
              Nothing ->
                do addFn fn (LFun [Inline] fn args' exp')
                   LS n i ds done <- get
                   put (LS n i ds (Map.insert (args', exp') fn done))
                   return fn

liftAll :: [(Name, LDecl)] -> [(Name, LDecl)]
liftAll xs =
  let (LS _ _ decls _) = execState (mapM_ liftDef xs) (LS Nothing 0 [] Map.empty) in
      decls

liftDef :: (Name, LDecl) -> State LiftState ()
liftDef (n, LFun opts _ args e) =
    do setBaseName n
       e' <- lift args e
       addFn n (LFun opts n args e')
liftDef (n, x) = addFn n x

lift :: [Name] -> LExp -> State LiftState LExp
lift env (LV v) = return (LV v) -- Lifting happens before these can exist...
lift env (LApp tc (LV n) args) = do args' <- mapM (lift env) args
                                    return (LApp tc (LV n) args')
lift env (LApp tc f args) = do f' <- lift env f
                               fn <- makeFn env f'
                               args' <- mapM (lift env) args
                               return (LApp tc (LV fn) (map LV env ++ args'))
lift env (LLazyApp n args) = do args' <- mapM (lift env) args
                                return (LLazyApp n args')
lift env (LLazyExp (LConst c)) = return (LConst c)
-- lift env (LLazyExp (LApp tc (LV (Glob f)) args))
--                       = lift env (LLazyApp f args)
lift env (LLazyExp e) = do e' <- lift env e
                           let usedArgs = nub $ usedIn env e'
                           fn <- makeFn usedArgs e'
                           return (LLazyApp fn (map LV usedArgs))
lift env (LForce e) = do e' <- lift env e
                         return (LForce e')
lift env (LLet n v e) = do v' <- lift env v
                           e' <- lift (env ++ [n]) e
                           return (LLet n v' e')
lift env (LLam args (LLam args' e)) = lift env (LLam (args ++ args') e)
lift env (LLam args e) = do e' <- lift (env ++ args) e
                            let usedArgs = nub $ usedIn env e'
                            fn <- makeFn (usedArgs ++ args) e'
                            return (LApp False (LV fn) (map LV usedArgs))
lift env (LProj t i) = do t' <- lift env t
                          return (LProj t' i)
lift env (LCon loc i n args) = do args' <- mapM (lift env) args
                                  return (LCon loc i n args')
lift env (LCase up e alts) = do alts' <- mapM liftA alts
                                e' <- lift env e
                                return (LCase up e' alts')
  where
    liftA (LConCase i n args e) = do e' <- lift (env ++ args) e
                                     return (LConCase i n args e')
    liftA (LConstCase c e) = do e' <- lift env e
                                return (LConstCase c e')
    liftA (LDefaultCase e) = do e' <- lift env e
                                return (LDefaultCase e')
lift env (LConst c) = return (LConst c)
lift env (LForeign t s args) = do args' <- mapM (liftF env) args
                                  return (LForeign t s args')
  where
    liftF env (t, e) = do e' <- lift env e
                          return (t, e')
lift env (LOp f args) = do args' <- mapM (lift env) args
                           return (LOp f args')
lift env (LError str) = return $ LError str
lift env LNothing = return LNothing

allocUnique :: LDefs -> (Name, LDecl) -> (Name, LDecl)
allocUnique defs p@(n, LConstructor _ _ _) = p
allocUnique defs (n, LFun opts fn args e)
    = let e' = evalState (findUp e) [] in
          (n, LFun opts fn args e')
  where
    -- Keep track of 'updatable' names in the state, i.e. names whose heap
    -- entry may be reused, along with the arity which was there
    findUp :: LExp -> State [(Name, Int)] LExp
    findUp (LApp t (LV n) as)
       | Just (LConstructor _ i ar) <- lookupCtxtExact n defs,
         ar == length as
          = findUp (LCon Nothing i n as)
    findUp (LV n)
       | Just (LConstructor _ i 0) <- lookupCtxtExact n defs
          = return $ LCon Nothing i n [] -- nullary cons are global, no need to update
    findUp (LApp t f as) = LApp t <$> findUp f <*> mapM findUp as
    findUp (LLazyApp n as) = LLazyApp n <$> mapM findUp as
    findUp (LLazyExp e) = LLazyExp <$> findUp e
    findUp (LForce e) = LForce <$> findUp e
    -- use assumption that names are unique!
    findUp (LLet n val sc) = LLet n <$> findUp val <*> findUp sc
    findUp (LLam ns sc) = LLam ns <$> findUp sc
    findUp (LProj e i) = LProj <$> findUp e <*> return i
    findUp (LCon (Just l) i n es) = LCon (Just l) i n <$> mapM findUp es
    findUp (LCon Nothing i n es)
           = do avail <- get
                v <- findVar [] avail (length es)
                LCon v i n <$> mapM findUp es
    findUp (LForeign t s es)
           = LForeign t s <$> mapM (\ (t, e) -> do e' <- findUp e
                                                   return (t, e')) es
    findUp (LOp o es) = LOp o <$> mapM findUp es
    findUp (LCase Updatable e@(LV n) as)
           = LCase Updatable e <$> mapM (doUpAlt n) as
    findUp (LCase t e as)
           = LCase t <$> findUp e <*> mapM findUpAlt as
    findUp t = return t

    findUpAlt (LConCase i t args rhs) = do avail <- get
                                           rhs' <- findUp rhs
                                           put avail
                                           return $ LConCase i t args rhs'
    findUpAlt (LConstCase i rhs) = LConstCase i <$> findUp rhs
    findUpAlt (LDefaultCase rhs) = LDefaultCase <$> findUp rhs

    doUpAlt n (LConCase i t args rhs)
           = do avail <- get
                put ((n, length args) : avail)
                rhs' <- findUp rhs
                put avail
                return $ LConCase i t args rhs'
    doUpAlt n (LConstCase i rhs) = LConstCase i <$> findUp rhs
    doUpAlt n (LDefaultCase rhs) = LDefaultCase <$> findUp rhs

    findVar _ [] i = return Nothing
    findVar acc ((n, l) : ns) i | l == i = do put (reverse acc ++ ns)
                                              return (Just n)
    findVar acc (n : ns) i = findVar (n : acc) ns i


-- Return variables in list which are used in the expression

usedArg env n | n `elem` env = [n]
              | otherwise = []

usedIn :: [Name] -> LExp -> [Name]
usedIn env (LV n) = usedArg env n
usedIn env (LApp _ e args) = usedIn env e ++ concatMap (usedIn env) args
usedIn env (LLazyApp n args) = concatMap (usedIn env) args ++ usedArg env n
usedIn env (LLazyExp e) = usedIn env e
usedIn env (LForce e) = usedIn env e
usedIn env (LLet n v e) = usedIn env v ++ usedIn (env \\ [n]) e
usedIn env (LLam ns e) = usedIn (env \\ ns) e
usedIn env (LCon v i n args) = let rest = concatMap (usedIn env) args in
                                   case v of
                                      Nothing -> rest
                                      Just n -> usedArg env n ++ rest
usedIn env (LProj t i) = usedIn env t
usedIn env (LCase up e alts) = usedIn env e ++ concatMap (usedInA env) alts
  where usedInA env (LConCase i n ns e) = usedIn env e
        usedInA env (LConstCase c e) = usedIn env e
        usedInA env (LDefaultCase e) = usedIn env e
usedIn env (LForeign _ _ args) = concatMap (usedIn env) (map snd args)
usedIn env (LOp f args) = concatMap (usedIn env) args
usedIn env _ = []

lsubst :: Name -> LExp -> LExp -> LExp
lsubst n new (LV x) | n == x = new
lsubst n new (LApp t e args) = let e' = lsubst n new e
                                   args' = map (lsubst n new) args in
                                   LApp t e' args'
lsubst n new (LLazyApp fn args) = let args' = map (lsubst n new) args in
                                      LLazyApp fn args'
lsubst n new (LLazyExp e) = LLazyExp (lsubst n new e)
lsubst n new (LForce e) = LForce (lsubst n new e)
lsubst n new (LLet v val sc) = LLet v (lsubst n new val) (lsubst n new sc)
lsubst n new (LLam ns sc) = LLam ns (lsubst n new sc)
lsubst n new (LProj e i) = LProj (lsubst n new e) i
lsubst n new (LCon lv t cn args) = let args' = map (lsubst n new) args in
                                       LCon lv t cn args'
lsubst n new (LOp op args) = let args' = map (lsubst n new) args in
                                 LOp op args'
lsubst n new (LForeign fd rd args)
     = let args' = map (\(d, a) -> (d, lsubst n new a)) args in
           LForeign fd rd args'
lsubst n new (LCase t e alts) = let e' = lsubst n new e
                                    alts' = map (fmap (lsubst n new)) alts in
                                    LCase t e' alts'
lsubst n new tm = tm

rename :: [(Name, Name)] -> LExp -> LExp
rename ns tm@(LV x)
   = case lookup x ns of
          Just n -> LV n
          _ -> tm
rename ns (LApp t e args)
    = let e' = rename ns e
          args' = map (rename ns) args in
          LApp t e' args'
rename ns (LLazyApp fn args)
    = let args' = map (rename ns) args in
          LLazyApp fn args'
rename ns (LLazyExp e) = LLazyExp (rename ns e)
rename ns (LForce e) = LForce (rename ns e)
rename ns (LLet v val sc) = LLet v (rename ns val) (rename ns sc)
rename ns (LLam args sc) = LLam args (rename ns sc)
rename ns (LProj e i) = LProj (rename ns e) i
rename ns (LCon lv t cn args) = let args' = map (rename ns) args in
                                    LCon lv t cn args'
rename ns (LOp op args) = let args' = map (rename ns) args in
                              LOp op args'
rename ns (LForeign fd rd args)
     = let args' = map (\(d, a) -> (d, rename ns a)) args in
           LForeign fd rd args'
rename ns (LCase t e alts) = let e' = rename ns e
                                 alts' = map (fmap (rename ns)) alts in
                                 LCase t e' alts'
rename ns tm = tm

instance Show LExp where
   show e = show' [] "" e where
     show' env ind (LV n) = show n

     show' env ind (LLazyApp e args)
        = show e ++ "|(" ++ showSep ", " (map (show' env ind) args) ++")"

     show' env ind (LApp _ e args)
        = show' env ind e ++ "(" ++ showSep ", " (map (show' env ind) args) ++")"

     show' env ind (LLazyExp e) = "lazy{ "  ++ show' env ind e ++ " }"
     show' env ind (LForce   e) = "force{ " ++ show' env ind e ++ " }"
     show' env ind (LLet n v e)
        = "let " ++ show n ++ " = " ++ show' env ind v
            ++ " in " ++ show' (env ++ [show n]) ind e

     show' env ind (LLam args e)
        = "(\\ " ++ showSep "," (map show args)
            ++ " => " ++ show' (env ++ (map show args)) ind e ++ ") "

     show' env ind (LProj t i) = show t ++ "!" ++ show i

     show' env ind (LCon loc i n args)
        = atloc loc ++ show n ++ "(" ++ showSep ", " (map (show' env ind) args) ++ ")"
       where atloc Nothing = ""
             atloc (Just l) = "@" ++ show (LV l) ++ ":"

     show' env ind (LCase up e alts)
        = "case" ++ update ++ "(" ++ show' env ind e ++ ") of \n" ++ fmt alts
       where
         update = case up of
                       Shared -> " "
                       Updatable -> "! "
         fmt [] = ""
         fmt [alt]
            = "\t" ++ ind ++ "| " ++ showAlt env (ind ++ "    ") alt
         fmt (alt:as)
            = "\t" ++ ind ++ "| " ++ showAlt env (ind ++ ".   ") alt
                ++ "\n" ++ fmt as

     show' env ind (LConst c) = show c

     show' env ind (LForeign ty n args) = concat
            [ "foreign{ "
            ,       show n ++ "("
            ,           showSep ", " (map (\(ty,x) -> show' env ind x ++ " : " ++ show ty) args)
            ,       ") : "
            ,       show ty
            , " }"
            ]

     show' env ind (LOp f args)
        = show f ++ "(" ++ showSep ", " (map (show' env ind) args) ++ ")"

     show' env ind (LError str) = "error " ++ show str
     show' env ind LNothing = "____"

     showAlt env ind (LConCase _ n args e)
          = show n ++ "(" ++ showSep ", " (map show args) ++ ") => "
             ++ show' env ind e

     showAlt env ind (LConstCase c e) = show c ++ " => " ++ show' env ind e
     showAlt env ind (LDefaultCase e) = "_ => " ++ show' env ind e

occName :: Name -> LExp -> Int
occName n (LV x) = if n == x then 1 else 0
occName n (LApp t e es) = occName n e + sum (map (occName n) es)
occName n (LLazyApp x es)
    = if n == x then 1 + sum (map (occName n) es)
                else sum (map (occName n) es)
occName n (LForce e) = occName n e
occName n (LLet x v sc)
    = if n == x then occName n v
                else occName n v + occName n sc
occName n (LLam ns sc)
    = if n `elem` ns then 0 else occName n sc
occName n (LProj e i) = occName n e
occName n (LCon _ _ _ es) = sum (map (occName n) es)
occName n (LCase t e alts) = occName n e + maximum (map occAlt alts)
  where
    occAlt (LConCase _ _ ns e)
        = if n `elem` ns then 0 else occName n e
    occAlt (LConstCase _ e) = occName n e
    occAlt (LDefaultCase e) = occName n e
occName n (LForeign _ _ es) = sum (map (occName n . snd) es)
occName n (LOp _ es) = sum (map (occName n) es)
occName n _ = 0