-- Initial preparations of Core before conversion to Erlang starts.

module Yhc.Core.BackEnd.Erlang.CorePrep where

import Yhc.Core.Extra
import Yhc.Core.BackEnd.Erlang.EGMonad
import Data.Maybe
import qualified Data.Map as M

-- Name of the function which forces evaluation of its argument and returns
-- the evaluated form.

fnforce = "_f_"

-- Transform Haskell function name to conform with Erlang Core requirements.
-- Single quotes are replaced with dollar signs. Additionally, if function's
-- name is not on the exports list, it is replaced with somewhat shorter,
-- and the mapping is stored in the functions' map.

sq2d '\'' = '~'
sq2d '-' = '_'
sq2d ':' = '@'
sq2d z = z

transFname :: CoreFuncName -> EGM CoreFuncName

transFname cfn = do
  let cfn' = map sq2d cfn
  rr <- gets expFunc
  if (cfn `notElem` rr)
    then do
      m <- gets funcMap
      let cfn'' = fromMaybe "" $ M.lookup cfn m
      if null cfn''
        then do
          c <- getCnt
          let cft = dropModule cfn
              commas = if (length cft < 2) then "" else tail $ reverse $ tail cft
              tupcon = not (null commas) && all (== ',') commas
              tuplen = if tupcon then length commas else 0
              nn = "." ++ if tupcon then "TUP" ++ show tuplen else show c
              m' = M.insert cfn nn m
          st <- get
          put st {funcMap = m'}
          return nn
        else return cfn''
    else return $ dropModule cfn'

-- Transform a Haskell constructor tag name. Most of tags are preserved,
-- while few are converted (like Prelude;True).

transCName :: CoreCtorName -> EGM CoreCtorName

transCName "Prelude;True" = return "true"
transCName "Prelude;False" = return "false"
transCName "Prelude;:" = return ".CONS"
transCName "Prelude;[]" = return ".EOL"
transCName "Prelude;EQ" = return ".EQ"
transCName "Prelude;LT" = return ".LT"
transCName "Prelude;GT" = return ".GT"
transCName c = return $ map sq2d c

-- Based on the customized strictness analysis function, insert
-- additional code to force function's arguments it is strict on.

strictify :: (Core -> (CoreFuncName -> [Bool])) -> Core -> Core

strictify sf core = core {coreFuncs = map f' (coreFuncs core)} where
  f' f@(CoreFunc {}) =
    let args = coreFuncArgs f
        args_f = map (++ "_f") args
        strct = sf core (coreFuncName f)
        argzf = zip3 strct args_f args
        arglet = concat $ map (\(a, b, c) -> if a 
          then [(c, CoreApp (CoreFun fnforce) [CoreVar b])]
          else []) argzf
        newargs = map (\(a, b, c) -> if a then b else c) argzf
        newbody = (if null arglet then id else CoreLet arglet) (coreFuncBody f)
    in f {coreFuncArgs = newargs, coreFuncBody = newbody}
  f' p@(CorePrim {}) = p

-- Map variable names within an expression using the given map.

mapVarsInExpr :: M.Map String String -> CoreExpr -> CoreExpr

mapVarsInExpr vm (CoreVar s) = CoreVar (mapVariable vm s)
mapVarsInExpr vm (CoreLet bnds e) = 
  CoreLet (map (\(s, c) -> (mapVariable vm s, c)) bnds) e
mapVarsInExpr vm (CoreCase expr alts) = 
  CoreCase (mapVarsInExpr vm expr) (map (mapalt vm) alts) where
    mapalt vm (altpat, altexpr) = (mappat vm altpat, mapVarsInExpr vm altexpr) where
      mappat vm (PatCon s cvs) = PatCon s (map (mapVariable vm) cvs)
      mappat _ x = x
mapVarsInExpr vm (CoreApp fn args) = 
  let fn':args' = map (mapVarsInExpr vm) (fn:args)
  in  CoreApp fn' args'
mapVarsInExpr vm z = z

mapVariable varmap s = let ns = M.lookup s varmap in fromMaybe s ns