-- GRIN-like backend for Yhc Core.
-- Translation from Yhc Core to GRIN.

module Yhc.Core.GRIN.TransCore  where

import Yhc.Core.GRIN.Type
import Yhc.Core.GRIN.GGMonad
import Yhc.Core.GRIN.SubstVars
import Yhc.Core.GRIN.HeapSeed
import Yhc.Core.GRIN.HeapPointsTo
import Yhc.Core.GRIN.OptFetch
import Yhc.Core.GRIN.OptEvals
import Yhc.Core.GRIN.OptApply
import Yhc.Core.GRIN.OptCases
import Yhc.Core.GRIN.ElimDead
import Yhc.Core.GRIN.FlatInline
import Yhc.Core.GRIN.Sharing
import Yhc.Core.Extra
import Data.Maybe
import qualified Data.Map as M
import qualified Data.Set as S

import Debug.Trace

-- |Translate a linked Yhc Core to GRIN.

core2GRIN :: [GName] -> Core -> CoreAnnotations -> GRIN

core2GRIN roots core canno = 
  let seed = coreHeapSeed core canno
      env = GG {stateCnt = 1
               ,currFun = CoreFunc {coreFuncName = ""
                                   ,coreFuncArgs = []
                                   ,coreFuncBody = CoreLit (CoreInt 0)}
               ,funStack = []
               ,coreRef = core
               ,varMap = M.empty
               ,funMap = toCoreFuncMap core
               ,liftSet = S.empty
               ,autoFuncs = []
               ,heapSeed = seed
               ,coreAnno = canno}
  in  flip evalState env $ do
        core <- gets coreRef
        seed <- gets heapSeed
        gfuncs <- mapM cf2gf (coreFuncs core)
        auto <- gets autoFuncs
        grin <- addEvalApply $ GRIN (concat gfuncs ++ auto)
        iter roots grin 


iter roots g = do
  seed <- gets heapSeed
  let hm = let m = heapPointsTo g `M.union` seed in trace ("* " ++ (show $ M.size m)) m
      post = elimDeadFuncs roots seed .
             elimDictFields seed .
             optFetch seed .
             propVal seed . 
             flattenInline . 
             inlineCalls isLinear . 
             optCases
      replEA gg = rmEvalApply gg >>= addEvalApply >>= return . post
      (g1, f1) = elimWhnfEvalsLoc hm g
      (g3, f3) = optStrict seed g
  if f1 
    then replEA g1 >>= iter roots
    else if f3
      then replEA g3 >>= iter roots
      else do
        g' <- replEA g
        let hm' = let m' = heapPointsTo g' `M.union` seed in trace ("+ " ++ (show $ M.size m')) m'
        if M.size hm' < M.size hm 
          then iter roots g'
          else return g'
  

rmEvalApply g = do
  let flt (GFunc {gFuncName = "GRIN;eval"}) = False
      flt (GFunc {gFuncName = "GRIN;apply"}) = False
      flt _ = True
  return $ g {gFuncs = filter flt (gFuncs g)}

addEvalApply g = do
  s <- gets heapSeed
  let hm = heapPointsTo g `M.union` s
  apply <- buildApply hm
  let g' = g {gFuncs = gFuncs g ++ [apply]}
      hm' = heapPointsTo g' `M.union` s
  eval <- buildEval hm'
  return $ g' {gFuncs = gFuncs g' ++ [eval]}


-- Translate a single function.

cf2gf :: CoreFunc -> GGM [GFunc]

cf2gf cf@CoreFunc {} = do
  core <- gets coreRef
  st <- get
  pushFun cf
  args <- mapM mapVarName (coreFuncArgs cf)
  gb <- schR (coreFuncBody cf)
  let gf = [GFunc {
      gFuncName = coreFuncName cf
     ,gFuncArgs = args
     ,gFuncBody = gb}]
  popFun
  return gf

cf2gf _ = return []

-- Translate a function body into a block of bindings.

-- Three "schemes" as described in the paper.

-- The "R" scheme to translate whe whole function body.
-- Translates the top-level expression, binds it to a variable
-- and finishes with return of that variable evaluated.

schR :: CoreExpr -> GGM GBlock

schR (CoreLet lts ex) = do
  ltexs <- mapM oneLet lts
  r <- schR ex
  return $ concat ltexs ++ r

schR (CoreCase e cpxs) = do
  (ev, ex) <- schE e
  pats <- mapM (schP (unGVar ev)) cpxs
  return $ ex ++ [GCase ev pats]

schR ex = do
  (rv, exc) <- schE ex
  let excr = reverse exc
  case head excr of
    GBind (GUnit r) _ -> return $ reverse (GSimple (GUnit r) : tail excr)
    _                 -> return $ exc ++ [GSimple (GUnit rv)]


-- The "C" scheme to translate an expression.
-- Returns a pair of the variable where result is bound,
-- and the block of bindings.

schC :: CoreExpr -> GGM (GVal, GBlock)

schC (CoreCase e cpxs) = do
  nv <- newGVar
  (ev, ex) <- schE e
  pats <- mapM (schP (unGVar ev)) cpxs
  let gbnd = GBind (GInline [GCase ev pats]) nv
  return (nv, ex ++ [gbnd])


-- Special case: fetch a numbered element (n) from a node (v)
-- n is expected to be a constant and is not evaluated.

schC (CoreApp (CoreFun "SEL_ELEM") [t, v, (CoreLit (CoreInt n))]) =
  schC (CoreApp (CoreFun "SEL_ELEM") [t, v, (CoreLit (CoreInteger $ fromIntegral n))])

schC (CoreApp (CoreFun "SEL_ELEM") [(CoreCon c), v, (CoreLit (CoreInteger n))]) = do
  (GVal (GVar fv), fblk) <- schE v
  nv <- newGVar
  let fbnd = GBind (GFetch fv (mkTagName c) (Just $ fromIntegral n)) nv
  return (nv, fblk ++ [fbnd])

-- Special case: strict application. Evaluate the argument first, then
-- apply the function to it.

-- General application.

schC ap@(CoreApp ex args) = do
  ai <- getAppInfo ap
  nv <- newGVar
  xargs <- mapM schC args
  let blks = concat $ map snd xargs
      apply f x = GCall "GRIN;apply" [f, x]
  case (appType ai, appSat ai) of
    (AppPrim, EQ) -> do                           -- saturated application of a primitive
      plift <- mkLift ai                          -- autocreate a lifting function
      let ftag = GTagged ("F@" ++ plift) (map (val2sval . fst) xargs)
          gbnd = GBind (GStore ftag) nv
      return (nv, blks ++ [gbnd])
    (AppFunc, EQ) -> do                           -- saturated application of a function
      let (CoreFun f) = appOf ai
          ftag = case length args of
            0 -> GTag ("K@" ++ f)
            _ -> GTagged ("F@" ++ f) (map (val2sval . fst) xargs)
          gbnd = GBind (GStore ftag) nv
      return (nv, blks ++ [gbnd])
    (aptype, LT) | aptype `elem` [AppCon, AppPrim] -> do -- partial application of a ctor/primitive
      plift <- mkLift ai                          -- autocreate a lifting function
      let missarg = appArity ai - appNargs ai
          ftag = GTagged ("P@" ++ show missarg ++ ":" ++ plift) (map (val2sval . fst) xargs)
          gbnd = GBind (GStore ftag) nv
      return (nv, blks ++ [gbnd])
    (AppFunc, LT) -> do                           -- partial application of a function
      let (CoreFun f) = appOf ai
          missarg = appArity ai - appNargs ai
          ftag = GTagged ("P@" ++ show missarg ++ ":" ++ f) (map (val2sval . fst) xargs)
          gbnd = GBind (GStore ftag) nv
      return (nv, blks ++ [gbnd])
    (aptype, GT) | aptype `elem` [AppFunc, AppPrim] -> do -- oversaturated application, known function
      let (fargs, extra) = splitAt (appArity ai) args
      schC (CoreApp (CoreApp (appOf ai) fargs) extra)
    (AppUnk, GT) -> do                            -- oversaturated application, unknown function
      (av, blk) <- schC (appOf ai)
      bvs <- replicateM (length xargs) newGVar
      chain <- mkChain apply (av : map fst xargs) (tail bvs ++ [nv]) []
      return (nv, blk ++ blks ++ chain)
    (AppUnk, EQ) -> schC (appOf ai)               -- unnecessarily nested
    (AppCon, EQ) -> do                            -- saturated application of a constructor
      let (CoreCon c) = appOf ai
          ctag = GTagged (mkTagName c) (map (val2sval . fst) xargs)
          gbnd = GBind (GStore ctag) nv
      return (nv, blks ++ [gbnd])
    (x, y) -> return (gVar "_", [GError $ "schC ap " ++ show x ++ " " ++ show y ++ " " ++ showRaw ap])


-- Function encodes as a partial application unless it is a CAF itself;
-- in the latter case it is stored as a special tag.

schC (CoreFun f) = do
  cmap <- gets funMap
  let mbcf = coreFuncMapMaybe cmap f
  case mbcf of
    Nothing -> return (gVar "_", [GError $ "undefined function " ++ f])
    Just cf | coreFuncArity cf > 0 ->               -- partial application
      schC (CoreApp (CoreFun f) [])
    _ -> do                                         -- CAF
      nv <- newGVar
      let ktag = GTag ("K@" ++ f)
      return (nv, [GBind (GStore ktag) nv])


-- Single constructor encodes as its tag alone if nullary, or in an incomplete
-- call to a lifting function otherwise.

schC (CoreCon c) = do
  nv <- newGVar
  core <- gets coreRef
  let mbct = coreCtorMaybe core c
  case mbct of
    Nothing -> return (gVar "_", [GError $ "undefined constructor " ++ c])
    Just ct | (not . null) (coreCtorFields ct) ->
      schC (CoreApp (CoreCon c) [])
    _ -> do
      let ftag = GTag (mkTagName c)
      return (nv, [GBind (GStore ftag) nv])

schC (CoreLit (CoreInteger i)) = do
  let lit = GVal $ GLitVal (GInteger i)
  nv <- newGVar
  let gbnd = GBind (GStore lit) nv
  return (nv, [gbnd])

schC (CoreLit (CoreInt i)) = do
  let lit = GVal $ GLitVal (GInt i)
  nv <- newGVar
  let gbnd = GBind (GStore lit) nv
  return (nv, [gbnd])

schC (CoreLit (CoreDouble d)) = do
  let lit = GVal $ GLitVal (GDouble d)
  nv <- newGVar
  let gbnd = GBind (GStore lit) nv
  return (nv, [gbnd])

schC (CoreLit (CoreFloat d)) = do
  let lit = GVal $ GLitVal (GFloat d)
  nv <- newGVar
  let gbnd = GBind (GStore lit) nv
  return (nv, [gbnd])

schC l@(CoreLit (CoreChr c)) = do
  let lit = GVal $ GLitVal (GChar c)
  nv <- newGVar
  let gbnd = GBind (GStore lit) nv
  return (nv, [gbnd])

-- String literals are represented as explicit lists of characters.

schC l@(CoreLit (CoreStr s)) = schC (s2l s) where
  s2l [] = CoreCon "Prelude;[]"
  s2l (c:cs) = CoreApp (CoreCon "Prelude;:") [CoreLit (CoreChr c), s2l cs]


schC (CoreVar v) = do
  nv <- mapVarName v >>= return . gVar
  return (nv, [])
  

schC z = return (gVar "_", [GError $ show z])

-- The "E" scheme to translate expressions in strict context.

schE :: CoreExpr -> GGM (GVal, GBlock)

schE (CoreCase e cpxs) = schC (CoreCase e cpxs)

-- Special case of the SEL_ELEM primitive: same as schC

schE ap@(CoreApp (CoreFun "SEL_ELEM") _) = schC ap

-- General application of a function. If saturated, evaluate all arguments,
-- and call the function.

schE ap@(CoreApp ex args) = do
  ai <- getAppInfo ap
  nv <- newGVar
  let argeval = case (appType ai, appSat ai) of
        (AppPrim, EQ) -> schE                    -- saturated app of primitive: evaluate args
        (_, _) -> schC                           -- just store them in other cases
  xargs <- mapM argeval args
  let blks = concat $ map snd xargs
      apply f x = GCall "GRIN;apply" [f, x]
  case (appType ai, appSat ai) of
    (aptype, EQ) | aptype `elem` [AppFunc, AppPrim] -> do -- saturated application of a function
      let (CoreFun f) = appOf ai                          -- or a primitive
      case f of
        "STRICT_APP" -> do                        -- special primitive to apply argument strictly
          let [a1, a2] = map fst xargs
              gbnd = GBind (apply a1 a2) nv
          env <- newGVar
          let ebnd = GBind (GCall "GRIN;eval" [nv]) env
          return (env, blks ++ [gbnd, ebnd])
        _ -> do
          bc <- bldCall aptype f (map fst xargs) nv
          return (fst bc, blks ++ snd bc)
    (aptype, LT) | aptype `elem` [AppCon, AppPrim] -> 
      schC ap                                     -- partial application of a ctor/primitive
    (AppFunc, LT) -> do                           -- partial application of a function
      let (CoreFun f) = appOf ai
          missarg = appArity ai - appNargs ai
          ftag = GTagged ("P@" ++ show missarg ++ ":" ++ f) (map (val2sval . fst) xargs)
          gbnd = GBind (GUnit ftag) nv
      return (nv, blks ++ [gbnd])
    (aptype, GT) | aptype `elem` [AppFunc, AppPrim] -> do -- oversaturated application, known function
      let (fargs, extra) = splitAt (appArity ai) args
      schE (CoreApp (CoreApp (appOf ai) fargs) extra)
    (AppUnk, GT) -> do                            -- oversaturated application, unknown function
      (av, blk) <- schE (appOf ai)
      bvs <- replicateM (length xargs) newGVar
      chain <- mkChain apply (av : map fst xargs) (tail bvs ++ [nv]) []
      env <- newGVar
      let ebnd = GBind (GCall "GRIN;eval" [nv]) env
      return (env, blk ++ blks ++ chain ++ [ebnd])
    (AppUnk, EQ) -> schE (appOf ai)               -- unnecessarily nested
    (AppCon, EQ) -> do                            -- saturated application of a constructor
      let (CoreCon c) = appOf ai
          ctag = GTagged (mkTagName c) (map (val2sval . fst) xargs)
          gbnd = GBind (GUnit ctag) nv
      return (nv, blks ++ [gbnd])
    (x, y) -> return (gVar "_", [GError $ "schC ap " ++ show x ++ " " ++ show y ++ " " ++ showRaw ap])

-- Same as in the non-strict scheme, but with GUnit instead of GStore.

schE (CoreCon c) = do
  nv <- newGVar
  core <- gets coreRef
  let mbct = coreCtorMaybe core c
  case mbct of
    Nothing -> return (gVar "_", [GError $ "undefined constructor " ++ c])
    Just ct | (not . null) (coreCtorFields ct) ->
      schC (CoreApp (CoreCon c) [])
    _ -> do
      let ftag = GTag (mkTagName c)
      return (nv, [GBind (GUnit ftag) nv])

schE ex = do
  (v, b) <- schC ex
  case v of
    GVal (GVar ev) -> do
      evprim <- newGVar
      let ebnd = GBind (GCall "GRIN;eval" [v]) evprim
      return (evprim, b ++ [ebnd])
    _ -> return (v, b)

-- The "P" scheme to translate case patterns.

schP :: GName -> (CorePat, CoreExpr) -> GGM (GPat, GBlock)

schP _ (PatDefault, ex) = do
  cex <- schR ex
  return (GPatDefault, cex)

schP _ (PatLit l, ex) = do
  let (GLitVal p) = lit2lit (CoreLit l)
  cex <- schR ex
  return (GPatLit p, cex)

schP cv (PatCon con args, ex) = do
  flds <- mapM mapVarName args
  cex <- schR ex
  let tag = mkTagName con
      ftchs = zipWith3 (\a b c -> GBind (GFetch a tag (Just b)) c) 
                       (repeat cv) [1 .. ] 
                       (map gVar flds)
      p = GPatTag tag
  return (p, ftchs ++ cex)

-- Build a saturated call to a function or a primitive.

bldCall :: AppType -> GName -> [GVal] -> GVal -> GGM (GVal, GBlock)

bldCall AppFunc f args nv = return (nv, [GBind (GCall f args) nv])

bldCall AppPrim f args nv = do
  let sig (HasType t) = [t]
      sig _ = []
  ts <- gets heapSeed >>= return . concat . map sig . M.findWithDefault [] (HFunc f)
  case ts of
    [] -> bldCall AppFunc f args nv
    tt -> do
      let call = GBind (GCall f args) nv
      return $ (nv, [call])
  
unBoxArg gv _ = return (gv, [])

boxRet gv _ = return (gv, [])

-- Create a lifting function for partially applied primitive or a constructor.
-- A lifting function can be applied partially while a primitive or a constructor
-- cannot. For a primitive, all its arguments will be evaluated once the function
-- is called. Name of the new function will be returned. The function is added to the
-- autocreated functions list.

mkLift :: AppInfo -> GGM GName

mkLift ai' = do
  let ai | (appType ai') `notElem` [AppPrim, AppCon] = error "mkLift: need a ctor or prim app"
         | appSat ai' == GT = error "mkLift: need a partial or saturated app"
         | otherwise = ai'
      fname = case appOf ai of
        CoreFun f -> f
        CoreCon c -> c
  lset <- gets liftSet
  cnt <- getCnt
  let lname = "l@" ++ fname --  ++ show cnt
  if lname `S.member` lset
    then return lname
    else do
      let lset' = S.insert lname lset
      st0 <- get
      put (st0 {liftSet = lset'})
      let fargs = take (appArity ai) (map (('_':) . show) [1 .. ])
          lift = CoreFunc {
             coreFuncName = lname
            ,coreFuncArgs = fargs
            ,coreFuncBody = CoreApp (appOf ai) (map CoreVar fargs)}
      gfun <- cf2gf lift
      st <- get
      put (st {autoFuncs = autoFuncs st ++ gfun})
      return lname

-- Gather information about an application of a known function,
-- unknown function, or a constructor. Crashes if a function
-- or a primitive is not defined.

getAppInfo :: CoreExpr -> GGM AppInfo

getAppInfo (CoreApp ex args) = do
  core <- gets coreRef
  cmap <- gets funMap
  let (atype, arity) = case ex of
        CoreFun f -> case fromJust (coreFuncMapMaybe cmap f) of
          cf@CoreFunc {} -> (AppFunc, coreFuncArity cf)
          cp@CorePrim {} -> (AppPrim, coreFuncArity cp)
        CoreCon c -> let ct = coreCtor core c
                     in  (AppCon, length $ coreCtorFields ct)
        _ -> (AppUnk, 0)
      nargs = length args
  return $ AppInfo {
    appSat = compare nargs arity 
   ,appType = atype
   ,appOf = ex
   ,appArity = arity
   ,appNargs = nargs
   ,appArgs = args}

getAppInfo z = error $ "getAppInfo: not an application: " ++ show z

-- Short-hand functions

oneLet (v, e) = do
  (lv, ex) <- schC e
  nv <- mapVarName v >>= return . gVar
  let bnd = GBind (GUnit lv) nv
  return $ ex ++ [bnd]

lit2lit :: CoreExpr -> GSval

lit2lit (CoreLit (CoreInteger i)) = GLitVal (GInteger i)
lit2lit (CoreLit (CoreInt i)) = GLitVal (GInt i)
lit2lit (CoreLit (CoreChr c)) = GLitVal (GChar c)
lit2lit (CoreLit (CoreDouble d)) = GLitVal (GDouble d)
lit2lit (CoreLit (CoreFloat d)) = GLitVal (GFloat d)

-- Make a chain of F-nodes or apply calls for oversaturated applications.

mkChain :: (GVal -> GVal -> GSexpr) -> [GVal] -> [GVal] -> [GExpr] -> GGM [GExpr]

mkChain f [e] [] exs = return exs

mkChain f (e:e':es) (b:bs) exs = do
  bb <- newGVar
  mkChain f (b:es) bs (exs ++ [GBind (GCall "GRIN;eval" [e]) bb, GBind (f bb e') b])