-- Linker utility functions for Hugs Core

module Yhc.Core.FrontEnd.Hugs.LinkUtil where

import Control.Monad
import Data.Maybe
import Yhc.Core.Extra
import qualified Data.Set as S
import qualified Data.Map as M
import qualified Data.List as L
import Yhc.Core.FrontEnd.Hugs.ParseUtil (nameGen)
import Yhc.Core.FrontEnd.Hugs.PrimTable (hugsPrimTable)

import Debug.Trace

-- Taken from Yhc Linker

mergeCores :: String -> [Core] -> Core
mergeCores modname cores = Core modname [] (concat datas) (concat funcs)
    where (datas,funcs) = unzip $ map (\x -> (coreDatas x, coreFuncs x)) cores

-- Fix core by deriving some missing objects such as tuple constructors,
-- or some other known constructors and functions.

fixCore :: Core -> Core

fixCore c = 
  let urd = coreUnreachableDatas c 
      c' = fixDatas urd c
      urf = coreUnreachableFuncs c'
  in  fixFuncs urf c' where
  fixFuncs _ c = c
  fixDatas mcs c = foldl fixOneData c mcs where
    isTuple ('H':'u':'g':'s':'.':'P':'r':'e':'l':'u':'d':'e':';':'(' : rtpl)
      | all (== ',') ((reverse . tail . reverse) rtpl) = length rtpl
    isTuple _ = 0
    fixOneData c tpl | isTuple tpl > 0 =
      let ntpl = isTuple tpl
          cargs = take ntpl nameGen
          tpldata = CoreData {
            coreDataName = tpl,
            coreDataTypes = [],
            coreDataCtors = [CoreCtor {
              coreCtorName = tpl,
              coreCtorFields = zip cargs (repeat Nothing)}]}
      in c {coreDatas = tpldata : coreDatas c,
            coreFuncs = coreFuncs c}
    fixOneData c _ = c

-- Convert cases with calls to PrimPmXXX to case with constant patterns.
-- Hugs generates:
-- case (Prelude;primPmInt Prelude;v30 0 _1) of
--   Prelude;True -> Prelude;v1489 0
--   _ -> foo
--
-- This has to be transformed to:
-- case _1 of
--   0 -> Prelude;v1489 0
--   _ -> foo
-- In this case, the first argument of primPmInt is a dictionary 
-- that would not be used under the new scheme. The second is
-- a constant to match, and the third is a variable to become
-- the case scrutinee variable under the new scheme.

pmfuns = map ("Prelude;primPm" ++) ["Int", "Integer", "Flt"]

fixCasePats :: Core -> Core

fixCasePats core = transformExpr (f core) core where
  f core (CoreCase (CoreApp (CoreFun pmfun) [_, CoreLit clit, cvar]) 
            [(PatCon "Prelude;True" [], exmat), d@(PatDefault, exdfl)]) 
              | pmfun `elem` pmfuns =
    CoreCase cvar [(PatLit clit, exmat), d]
  f _ x = x

-- Remap function names in the case Hugs defines them in an undesired way.
-- Basically, everything defined in Hugs.Prelude goes to Prelude.

mapFuns :: Core -> Core

mapFuns core = mapFunNames funmap core where
  funmap = M.fromList $ zip hpfuns pfuns
  hpfuns = filter toRemap allfuns
  allfuns = map coreFuncName $ coreFuncs core
  pfuns = map (("Prelude;" ++ ) . dropModule) hpfuns


-- Remap constructor names in the case Hugs defines them in an undesired way.
-- Basically, everything defined in Hugs.Prelude goes to Prelude.

mapCons :: Core -> Core

mapCons core = mapConNames ctormap core where
  ctormap = M.fromList $ zip hpctors pctors
  hpctors = filter toRemap allctors
  allctors = map coreCtorName $ concat $ map coreDataCtors $ coreDatas core
  pctors = map (("Prelude;" ++ ) . dropModule) hpctors

-- Remap data objects (LHS of data XXX) in the case Hugs defines them in an undesired way.
-- Basically, everything defined in Hugs.Prelude goes to Prelude.

mapDatas :: Core -> Core

mapDatas core = mapDataNames dtmap core where
  dtmap = M.fromList $ zip hpdatas pdatas
  hpdatas = filter toRemap alldatas
  alldatas = map coreDataName $ coreDatas core
  pdatas = map (("Prelude;" ++ ) . dropModule) hpdatas

-- Remap Hugs primitives into Yhc Core normal primitives.
-- The mapping table is in the Yhc.Core.FrontEnd.Hugs.PrimTable module.

mapPrims :: Core -> Core

mapPrims = mapFunNames (M.fromList hugsPrimTable)

-- Filter for function names to remap by mapFuns.

toRemap ('H':'u':'g':'s':'.':'P':'r':'e':'l':'u':'d':'e':';': _) = True
toRemap _ = False

-- Transform all calls to constructor lifting functions into
-- calls to conctructors themselves provided that call to the lifting
-- function was saturated.

unLiftCtors :: Core -> Core
unLiftCtors core = transformExpr (fx core) core where
  fx core ap@(CoreApp (CoreFun f) args) | coreSaturated core ap =
    fromMaybe ap $ do
      fdef <- coreFuncMaybe core f
      when (isCorePrim fdef) $ fail $ "unLiftCtors: primitive " ++ f
      let fbdy = coreFuncBody fdef
          fargs = coreFuncArgs fdef
      case fbdy of
        CoreApp cc@(CoreCon c) cargs
          | length fargs == length cargs &&
            all isCoreVar cargs &&
            all id (zipWith (==) (map CoreVar fargs) cargs) -> return $ CoreApp cc args
        _ -> return ap
  fx _ x = x