{-|
Module      : IRTS.Defunctionalise
Description : Defunctionalise Idris' IR.

License     : BSD3
Maintainer  : The Idris Community.

To defunctionalise:

1. Create a data constructor for each function
2. Create a data constructor for each underapplication of a function
3. Convert underapplications to their corresponding constructors
4. Create an EVAL function which calls the appropriate function for data constructors
   created as part of step 1
5. Create an APPLY function which adds an argument to each underapplication (or calls
   APPLY again for an exact application)
6. Wrap overapplications in chains of APPLY
7. Wrap unknown applications (i.e. applications of local variables) in chains of APPLY
8. Add explicit EVAL to case, primitives, and foreign calls

-}
{-# LANGUAGE FlexibleContexts, PatternGuards #-}
module IRTS.Defunctionalise(module IRTS.Defunctionalise
                          , module IRTS.Lang
                          ) where

import Idris.Core.CaseTree
import Idris.Core.TT
import IRTS.Lang

import Control.Monad
import Control.Monad.State
import Data.List
import Data.Maybe

data DExp = DV Name
          | DApp Bool Name [DExp] -- True = tail call
          | DLet Name DExp DExp -- name just for pretty printing
          | DUpdate Name DExp -- eval expression, then update var with it
          | DProj DExp Int
          | DC (Maybe Name) Int Name [DExp]
          | DCase CaseType DExp [DAlt]
          | DChkCase DExp [DAlt] -- a case where the type is unknown (for EVAL/APPLY)
          | DConst Const
          | DForeign FDesc FDesc [(FDesc, DExp)]
          | DOp PrimFn [DExp]
          | DNothing -- erased value, can be compiled to anything since it'll never
                     -- be inspected
          | DError String
  deriving Eq

data DAlt = DConCase Int Name [Name] DExp
          | DConstCase Const DExp
          | DDefaultCase DExp
  deriving (Show, Eq)

data DDecl = DFun Name [Name] DExp -- name, arg names, definition
           | DConstructor Name Int Int -- constructor name, tag, arity
  deriving (Show, Eq)

type DDefs = Ctxt DDecl

defunctionalise :: Int -> LDefs -> DDefs
defunctionalise nexttag defs
     = let all = toAlist defs
           -- sort newcons so that EVAL and APPLY cons get sequential tags
           (allD, (enames, anames)) = runState (mapM (addApps defs) all) ([], [])
           anames' = sort (nub anames)
           enames' = nub enames
           newecons = sortBy conord $ concatMap (toCons enames') (getFn all)
           newacons = sortBy conord $ concatMap (toConsA anames') (getFn all)
           eval = mkEval newecons
           app = mkApply newacons
           app2 = mkApply2 newacons
           condecls = declare nexttag (newecons ++ newacons) in
           addAlist (eval : app : app2 : condecls ++ allD) emptyContext
   where conord (n, _, _) (n', _, _) = compare n n'

getFn :: [(Name, LDecl)] -> [(Name, Int)]
getFn xs = mapMaybe fnData xs
  where fnData (n, LFun _ _ args _) = Just (n, length args)
        fnData _ = Nothing

addApps :: LDefs -> (Name, LDecl) -> State ([Name], [(Name, Int)]) (Name, DDecl)
addApps defs o@(n, LConstructor _ t a)
    = return (n, DConstructor n t a)
addApps defs (n, LFun _ _ args e)
    = do e' <- aa args e
         return (n, DFun n args e')
  where
    aa :: [Name] -> LExp -> State ([Name], [(Name, Int)]) DExp
    aa env (LV n) | n `elem` env = return $ DV n
                         | otherwise = aa env (LApp False (LV n) [])
    aa env (LApp tc (LV n) args)
       = do args' <- mapM (aa env) args
            case lookupCtxtExact n defs of
                Just (LConstructor _ i ar) -> return $ DApp tc n args'
                Just (LFun _ _ as _) -> let arity = length as in
                                               fixApply tc n args' arity
                Nothing -> return $ chainAPPLY (DV n) args'
    aa env (LLazyApp n args)
       = do args' <- mapM (aa env) args
            case lookupCtxtExact n defs of
                Just (LConstructor _ i ar) -> return $ DApp False n args'
                Just (LFun _ _ as _) -> let arity = length as in
                                           fixLazyApply n args' arity
                Nothing -> return $ chainAPPLY (DV n) args'
    aa env (LForce (LLazyApp n args)) = aa env (LApp False (LV n) args)
    aa env (LForce e) = liftM eEVAL (aa env e)
    aa env (LLet n v sc) = liftM2 (DLet n) (aa env v) (aa (n : env) sc)
    aa env (LCon loc i n args) = liftM (DC loc i n) (mapM (aa env) args)
    aa env (LProj t@(LV n) i)
        | n `elem` env = do t' <- aa env t
                            return $ DProj (DUpdate n t') i
    aa env (LProj t i) = do t' <- aa env t
                            return $ DProj t' i
    aa env (LCase up e alts) = do e' <- aa env e
                                  alts' <- mapM (aaAlt env) alts
                                  return $ DCase up e' alts'
    aa env (LConst c) = return $ DConst c
    aa env (LForeign t n args)
        = do args' <- mapM (aaF env) args
             return $ DForeign t n args'
    aa env (LOp LFork args) = liftM (DOp LFork) (mapM (aa env) args)
    aa env (LOp f args) = do args' <- mapM (aa env) args
                             return $ DOp f args'
    aa env LNothing = return DNothing
    aa env (LError e) = return $ DError e

    aaF env (t, e) = do e' <- aa env e
                        return (t, e')

    aaAlt env (LConCase i n args e)
         = liftM (DConCase i n args) (aa (args ++ env) e)
    aaAlt env (LConstCase c e) = liftM (DConstCase c) (aa env e)
    aaAlt env (LDefaultCase e) = liftM DDefaultCase (aa env e)

    fixApply tc n args ar
        | length args == ar
             = return $ DApp tc n args
        | length args < ar
             = do (ens, ans) <- get
                  let alln = map (\x -> (n, x)) [length args .. ar]
                  put (ens, alln ++ ans)
                  return $ DApp tc (mkUnderCon n (ar - length args)) args
        | length args > ar
             = return $ chainAPPLY (DApp tc n (take ar args)) (drop ar args)

    fixLazyApply n args ar
        | length args == ar
             = do (ens, ans) <- get
                  put (n : ens, ans)
                  return $ DApp False (mkFnCon n) args
        | length args < ar
             = do (ens, ans) <- get
                  let alln = map (\x -> (n, x)) [length args .. ar]
                  put (ens, alln ++ ans)
                  return $ DApp False (mkUnderCon n (ar - length args)) args
        | length args > ar
             = return $ chainAPPLY (DApp False n (take ar args)) (drop ar args)

    chainAPPLY f [] = f
--     chainAPPLY f (a : b : as)
--          = chainAPPLY (DApp False (sMN 0 "APPLY2") [f, a, b]) as
    chainAPPLY f (a : as) = chainAPPLY (DApp False (sMN 0 "APPLY") [f, a]) as

eEVAL x = DApp False (sMN 0 "EVAL") [x]

data EvalApply a = EvalCase (Name -> a)
                 | ApplyCase a
                 | Apply2Case a

-- For a function name, generate a list of
-- data constuctors, and whether to handle them in EVAL or APPLY

toCons :: [Name] -> (Name, Int) -> [(Name, Int, EvalApply DAlt)]
toCons ns (n, i)
    | n `elem` ns
      = (mkFnCon n, i,
          EvalCase (\tlarg ->
            (DConCase (-1) (mkFnCon n) (take i (genArgs 0))
              (dupdate tlarg
                (DApp False n (map DV (take i (genArgs 0))))))))
          : [] -- mkApplyCase n 0 i
    | otherwise = []
  where dupdate tlarg x = DUpdate tlarg x

toConsA :: [(Name, Int)] -> (Name, Int) -> [(Name, Int, EvalApply DAlt)]
toConsA ns (n, i)
    | Just ar <- lookup n ns
--       = (mkFnCon n, i,
--           EvalCase (\tlarg ->
--             (DConCase (-1) (mkFnCon n) (take i (genArgs 0))
--               (dupdate tlarg
--                 (DApp False n (map DV (take i (genArgs 0))))))))
          = mkApplyCase n ar i
    | otherwise = []

mkApplyCase fname n ar | n == ar = []
mkApplyCase fname n ar
        = let nm = mkUnderCon fname (ar - n) in
              (nm, n, ApplyCase (DConCase (-1) nm (take n (genArgs 0))
                  (DApp False (mkUnderCon fname (ar - (n + 1)))
                       (map DV (take n (genArgs 0) ++
                         [sMN 0 "arg"])))))
                            :
              if (ar - (n + 2) >=0 )
                 then (nm, n, Apply2Case (DConCase (-1) nm (take n (genArgs 0))
                      (DApp False (mkUnderCon fname (ar - (n + 2)))
                       (map DV (take n (genArgs 0) ++
                         [sMN 0 "arg0", sMN 0 "arg1"])))))
                            :
                            mkApplyCase fname (n + 1) ar
                 else mkApplyCase fname (n + 1) ar

mkEval :: [(Name, Int, EvalApply DAlt)] -> (Name, DDecl)
mkEval xs = (sMN 0 "EVAL", DFun (sMN 0 "EVAL") [sMN 0 "arg"]
               (mkBigCase (sMN 0 "EVAL") 256 (DV (sMN 0 "arg"))
                  (mapMaybe evalCase xs ++
                      [DDefaultCase (DV (sMN 0 "arg"))])))
  where
    evalCase (n, t, EvalCase x) = Just (x (sMN 0 "arg"))
    evalCase _ = Nothing

mkApply :: [(Name, Int, EvalApply DAlt)] -> (Name, DDecl)
mkApply xs = (sMN 0 "APPLY", DFun (sMN 0 "APPLY") [sMN 0 "fn", sMN 0 "arg"]
                             (case mapMaybe applyCase xs of
                                [] -> DNothing
                                cases ->
                                    mkBigCase (sMN 0 "APPLY") 256
                                               (DV (sMN 0 "fn"))
                                              (cases ++
                                    [DDefaultCase DNothing])))
  where
    applyCase (n, t, ApplyCase x) = Just x
    applyCase _ = Nothing

mkApply2 :: [(Name, Int, EvalApply DAlt)] -> (Name, DDecl)
mkApply2 xs = (sMN 0 "APPLY2", DFun (sMN 0 "APPLY2") [sMN 0 "fn", sMN 0 "arg0", sMN 0 "arg1"]
                             (case mapMaybe applyCase xs of
                                [] -> DNothing
                                cases ->
                                    mkBigCase (sMN 0 "APPLY") 256
                                               (DV (sMN 0 "fn"))
                                              (cases ++
                                    [DDefaultCase
                                       (DApp False (sMN 0 "APPLY")
                                       [DApp False (sMN 0 "APPLY")
                                              [DV (sMN 0 "fn"),
                                               DV (sMN 0 "arg0")],
                                               DV (sMN 0 "arg1")])
                                               ])))
  where
    applyCase (n, t, Apply2Case x) = Just x
    applyCase _ = Nothing


declare :: Int -> [(Name, Int, EvalApply DAlt)] -> [(Name, DDecl)]
declare t xs = dec' t xs [] where
   dec' t [] acc = reverse acc
   dec' t ((n, ar, _) : xs) acc = dec' (t + 1) xs ((n, DConstructor n t ar) : acc)


genArgs i = sMN i "P_c" : genArgs (i + 1)

mkFnCon    n = sMN 0 ("P_" ++ show n)
mkUnderCon n 0       = n
mkUnderCon n missing = sMN missing ("U_" ++ show n)

instance Show DExp where
   show e = show' [] e where
     show' env (DV n) = show n
     show' env (DApp _ e args) = show e ++ "(" ++
                                   showSep ", " (map (show' env) args) ++")"
     show' env (DLet n v e) = "let " ++ show n ++ " = " ++ show' env v ++ " in " ++
                               show' (env ++ [show n]) e
     show' env (DUpdate n e) = "!update " ++ show n ++ "(" ++ show' env e ++ ")"
     show' env (DC loc i n args) = atloc loc ++ "CON " ++ show n ++ "(" ++ showSep ", " (map (show' env) args) ++ ")"
       where atloc Nothing = ""
             atloc (Just l) = "@" ++ show (LV l) ++ ":"
     show' env (DProj t i) = show t ++ "!" ++ show i
     show' env (DCase up e alts) = "case" ++ update ++ show' env e ++ " of {\n\t" ++
                                    showSep "\n\t| " (map (showAlt env) alts)
         where update = case up of
                           Shared -> " "
                           Updatable -> "! "
     show' env (DChkCase e alts) = "case' " ++ show' env e ++ " of {\n\t" ++
                                    showSep "\n\t| " (map (showAlt env) alts)
     show' env (DConst c) = show c
     show' env (DForeign ty n args)
           = "foreign " ++ show n ++ "(" ++ showSep ", " (map (show' env) (map snd args)) ++ ")"
     show' env (DOp f args) = show f ++ "(" ++ showSep ", " (map (show' env) args) ++ ")"
     show' env (DError str) = "error " ++ show str
     show' env DNothing = "____"

     showAlt env (DConCase _ n args e)
          = show n ++ "(" ++ showSep ", " (map show args) ++ ") => "
             ++ show' env e
     showAlt env (DConstCase c e) = show c ++ " => " ++ show' env e
     showAlt env (DDefaultCase e) = "_ => " ++ show' env e

-- | Divide up a large case expression so that each has a maximum of
-- 'max' branches
mkBigCase cn max arg branches
   | length branches <= max = DChkCase arg branches
   | otherwise = DChkCase arg branches

groupsOf :: Int -> [DAlt] -> [[DAlt]]
groupsOf x [] = []
groupsOf x xs = let (batch, rest) = span (tagLT (x + tagHead xs)) xs in
                    batch : groupsOf x rest
  where tagHead (DConstCase (I i) _ : _) = i
        tagHead (DConCase t _ _ _ : _) = t
        tagHead (DDefaultCase _ : _) = -1 -- must be the end

        tagLT i (DConstCase (I j) _) = i < j
        tagLT i (DConCase j _ _ _) = i < j
        tagLT i (DDefaultCase _) = False

dumpDefuns :: DDefs -> String
dumpDefuns ds = showSep "\n" $ map showDef (toAlist ds)
  where showDef (x, DFun fn args exp)
            = show fn ++ "(" ++ showSep ", " (map show args) ++ ") = \n\t" ++
              show exp ++ "\n"
        showDef (x, DConstructor n t a) = "Constructor " ++ show n ++ " " ++ show t