{-|
Module      : IRTS.Simplified
Description : Simplified expressions, where functions/constructors can only be applied to variables.

License     : BSD3
Maintainer  : The Idris Community.
-}
{-# LANGUAGE FlexibleContexts #-}
module IRTS.Simplified(simplifyDefs, SDecl(..), SExp(..), SAlt(..)) where

import Idris.Core.CaseTree
import Idris.Core.TT
import IRTS.Defunctionalise

import Control.Monad.State

data SExp = SV LVar
          | SApp Bool Name [LVar]
          | SLet LVar SExp SExp
          | SUpdate LVar SExp
          | SCon (Maybe LVar) -- location to reallocate, if available
                 Int Name [LVar]
          | SCase CaseType LVar [SAlt]
          | SChkCase LVar [SAlt]
          | SProj LVar Int
          | SConst Const
          -- Keep DExps for describing foreign things, because they get
          -- translated differently
          | SForeign FDesc FDesc [(FDesc, LVar)]
          | SOp PrimFn [LVar]
          | SNothing -- erased value, will never be inspected
          | SError String
  deriving Show

data SAlt = SConCase Int Int Name [Name] SExp
          | SConstCase Const SExp
          | SDefaultCase SExp
  deriving Show

data SDecl = SFun Name [Name] Int SExp
  deriving Show


ldefs :: State (DDefs, Int) DDefs
ldefs = do (l, h) <- get
           return l

-- | Simplify an expression by let-binding argument expressions that
-- are not variables
-- The boolean parameter indicates whether the expression is at tail
-- call position.
simplify :: Bool -> DExp -> State (DDefs, Int) SExp
simplify tl (DV x)
    = do ctxt <- ldefs
         case lookupCtxtExact x ctxt of
              Just (DConstructor _ t 0) -> return $ SCon Nothing t x []
              _ -> return $ SV (Glob x)
simplify tl (DApp tc n args) = bindExprs args (SApp (tl || tc) n)
simplify tl (DForeign ty fn args)
    = let (fdescs, exprs) = unzip args
      in bindExprs exprs (\vars -> SForeign ty fn (zip fdescs vars))
simplify tl (DLet n v e) = do v' <- simplify False v
                              e' <- simplify tl e
                              return (SLet (Glob n) v' e')
simplify tl (DUpdate n e) = do e' <- simplify False e
                               return (SUpdate (Glob n) e')
simplify tl (DC loc i n args) = bindExprs args (SCon (Glob <$> loc) i n)
simplify tl (DProj t i) = bindExpr t (\var -> SProj var i)
simplify tl (DCase up e alts)
    = do alts' <- mapM (sAlt tl) alts
         bindExpr e (\var -> SCase up var alts')
simplify tl (DChkCase e alts)
    = do alts' <- mapM (sAlt tl) alts
         bindExpr e (\var -> SChkCase var alts')
simplify tl (DConst c) = return (SConst c)
simplify tl (DOp p args) = bindExprs args (SOp p)
simplify tl DNothing = return SNothing
simplify tl (DError str) = return $ SError str


-- | Let-bind a list of expressions to variables and construct the
-- inner expression with the bound variables.
-- If an expression in the list is already a variable we don’t bind it
-- again.
bindExprs :: [DExp] -> ([LVar] -> SExp) -> State (DDefs, Int) SExp
bindExprs es f = bindExprs' es f [] where
    bindExprs' [] f vars = return $ f (reverse vars)
    bindExprs' (e:es) f vars =
        bindExprM e (\var -> bindExprs' es f (var:vars))

-- | Special case of 'bindExprs' for just one expression
bindExpr :: DExp -> (LVar -> SExp) -> State (DDefs, Int) SExp
bindExpr e f = bindExprM e (return . f)

bindExprM :: DExp -> (LVar -> State (DDefs, Int) SExp) -> State (DDefs, Int) SExp
bindExprM (DV x) f
    = do ctxt <- ldefs
         case lookupCtxtExact x ctxt of
              Just (DConstructor _ t 0) -> bindExprM (DC Nothing t x []) f
              _ -> f (Glob x)
bindExprM e f =
    do e' <- simplify False e
       var <- freshVar
       f' <- f var
       return $ SLet var e' f'
    where
        freshVar = do (defs, i) <- get
                      put (defs, i + 1)
                      return (Glob (sMN i "R"))


sAlt :: Bool -> DAlt -> State (DDefs, Int) SAlt
sAlt tl (DConCase i n args e) = do e' <- simplify tl e
                                   return (SConCase (-1) i n args e')
sAlt tl (DConstCase c e) = do e' <- simplify tl e
                              return (SConstCase c e')
sAlt tl (DDefaultCase e) = do e' <- simplify tl e
                              return (SDefaultCase e')

simplifyDefs :: DDefs -> [(Name, DDecl)] -> TC [(Name, SDecl)]
simplifyDefs ctxt [] = return []
simplifyDefs ctxt (con@(n, DConstructor _ _ _) : xs)
    = do xs' <- simplifyDefs ctxt xs
         return xs'
simplifyDefs ctxt ((n, DFun n' args exp) : xs)
    = do let sexp = evalState (simplify True exp) (ctxt, 0)
         (exp', locs) <- runStateT (scopecheck n ctxt (zip args [0..]) sexp) (length args)
         xs' <- simplifyDefs ctxt xs
         return ((n, SFun n' args ((locs + 1) - length args) exp') : xs')

lvar v = do i <- get
            put (max i v)

scopecheck :: Name -> DDefs -> [(Name, Int)] -> SExp -> StateT Int TC SExp
scopecheck fn ctxt envTop tm = sc envTop tm where
    failsc err = fail $ "Codegen error in " ++ show fn ++ ":" ++ err

    sc env (SV (Glob n)) =
       case lookup n (reverse env) of -- most recent first
              Just i -> do lvar i; return (SV (Loc i))
              Nothing -> case lookupCtxtExact n ctxt of
                              Just (DConstructor _ i ar) ->
                                  if True -- ar == 0
                                     then return (SCon Nothing i n [])
                                     else failsc $ "Constructor " ++ show n ++
                                                 " has arity " ++ show ar
                              Just _ -> return (SV (Glob n))
                              Nothing -> failsc $ "No such variable " ++ show n
    sc env (SApp tc f args)
       = do args' <- mapM (scVar env) args
            case lookupCtxtExact f ctxt of
                Just (DConstructor n tag ar) ->
                    if True -- (ar == length args)
                       then return $ SCon Nothing tag n args'
                       else failsc $ "Constructor " ++ show f ++
                                   " has arity " ++ show ar
                Just _ -> return $ SApp tc f args'
                Nothing -> failsc $ "No such variable " ++ show f
    sc env (SForeign ty f args)
       = do args' <- mapM (\ (t, a) -> do a' <- scVar env a
                                          return (t, a')) args
            return $ SForeign ty f args'
    sc env (SCon loc tag f args)
       = do loc' <- case loc of
                         Nothing -> return Nothing
                         Just l -> do l' <- scVar env l
                                      return (Just l')
            args' <- mapM (scVar env) args
            case lookupCtxtExact f ctxt of
                Just (DConstructor n tag ar) ->
                    if True -- (ar == length args)
                       then return $ SCon loc' tag n args'
                       else failsc $ "Constructor " ++ show f ++
                                     " has arity " ++ show ar
                _ -> failsc $ "No such constructor " ++ show f
    sc env (SProj e i)
       = do e' <- scVar env e
            return (SProj e' i)
    sc env (SCase up e alts)
       = do e' <- scVar env e
            alts' <- mapM (scalt env) alts
            return (SCase up e' alts')
    sc env (SChkCase e alts)
       = do e' <- scVar env e
            alts' <- mapM (scalt env) alts
            return (SChkCase e' alts')
    sc env (SLet (Glob n) v e)
       = do let env' = env ++ [(n, length env)]
            v' <- sc env v
            n' <- scVar env' (Glob n)
            e' <- sc env' e
            return (SLet n' v' e')
    sc env (SUpdate (Glob n) e)
       = do -- n already in env
            e' <- sc env e
            n' <- scVar env (Glob n)
            return (SUpdate n' e')
    sc env (SOp prim args)
       = do args' <- mapM (scVar env) args
            return (SOp prim args')
    sc env x = return x

    scVar env (Glob n) =
       case lookup n (reverse env) of -- most recent first
              Just i -> do lvar i; return (Loc i)
              Nothing -> case lookupCtxtExact n ctxt of
                              Just (DConstructor _ i ar) ->
                                  failsc "can't pass constructor here"
                              Just _ -> return (Glob n)
                              Nothing -> failsc $ "No such variable " ++ show n ++
                                               " in " ++ show tm ++ " " ++ show envTop
    scVar _ x = return x

    scalt env (SConCase _ i n args e)
       = do let env' = env ++ zip args [length env..]
            tag <- case lookupCtxtExact n ctxt of
                        Just (DConstructor _ i ar) ->
                             if True -- (length args == ar)
                                then return i
                                else failsc $ "Constructor " ++ show n ++
                                            " has arity " ++ show ar
                        _ -> failsc $ "No constructor " ++ show n
            e' <- sc env' e
            return (SConCase (length env) tag n args e')
    scalt env (SConstCase c e) = do e' <- sc env e
                                    return (SConstCase c e')
    scalt env (SDefaultCase e) = do e' <- sc env e
                                    return (SDefaultCase e')