    module Vectorise.Builtins.Initialise (
    	-- * Initialisation
    	initBuiltins, initBuiltinVars, initBuiltinTyCons, initBuiltinDataCons,
    	initBuiltinPAs, initBuiltinPRs,
    	initBuiltinBoxedTyCons, initBuiltinScalars,
    ) where
    import Vectorise.Builtins.Base
    import Vectorise.Builtins.Modules
    import Vectorise.Builtins.Prelude

    import BasicTypes
    import PrelNames
    import TysPrim
    import DsMonad
    import IfaceEnv
    import InstEnv
    import TysWiredIn
    import DataCon
    import TyCon
    import Class
    import CoreSyn
    import Type
    import Name
    import Module
    import Var
    import Id
    import FastString
    import Outputable

    import Control.Monad
    import Data.Array
    import Data.List

    -- | Create the initial map of builtin types and functions.
    initBuiltins
    	:: PackageId 	-- ^ package id the builtins are in, eg dph-common
    	-> DsM Builtins

    initBuiltins pkg
     = do mapM_ load dph_Orphans

          -- From dph-common:Data.Array.Parallel.Lifted.PArray
          parrayTyCon	<- externalTyCon	dph_PArray	(fsLit "PArray")
          let [parrayDataCon] = tyConDataCons parrayTyCon

          pdataTyCon	<- externalTyCon	dph_PArray	(fsLit "PData")
          paTyCon		<- externalClassTyCon	dph_PArray	(fsLit "PA")
          let [paDataCon]	= tyConDataCons paTyCon

          preprTyCon	<- externalTyCon 	dph_PArray	(fsLit "PRepr")
          prTyCon		<- externalClassTyCon	dph_PArray	(fsLit "PR")
          let [prDataCon]	= tyConDataCons prTyCon

          closureTyCon	<- externalTyCon dph_Closure		(fsLit ":->")

          -- From dph-common:Data.Array.Parallel.Lifted.Repr
          voidTyCon		<- externalTyCon	dph_Repr	(fsLit "Void")
          wrapTyCon		<- externalTyCon	dph_Repr	(fsLit "Wrap")

          -- From dph-common:Data.Array.Parallel.Lifted.Unboxed
          sel_tys		<- mapM (externalType dph_Unboxed)
                               	(numbered "Sel" 2 mAX_DPH_SUM)

          sel_replicates	<- mapM (externalFun dph_Unboxed)
    				(numbered_hash "replicateSel" 2 mAX_DPH_SUM)

          sel_picks 	<- mapM (externalFun dph_Unboxed)
    				(numbered_hash "pickSel" 2 mAX_DPH_SUM)

          sel_tags		<- mapM (externalFun dph_Unboxed)
    				(numbered "tagsSel" 2 mAX_DPH_SUM)

          sel_els		<- mapM mk_elements
    				[(i,j) | i <- [2..mAX_DPH_SUM], j <- [0..i-1]]

          sum_tcs		<- mapM (externalTyCon dph_Repr)
    				(numbered "Sum" 2 mAX_DPH_SUM)

          let selTys        = listArray (2, mAX_DPH_SUM) sel_tys
              selReplicates = listArray (2, mAX_DPH_SUM) sel_replicates
              selPicks      = listArray (2, mAX_DPH_SUM) sel_picks
              selTagss      = listArray (2, mAX_DPH_SUM) sel_tags
              selEls        = array     ((2,0), (mAX_DPH_SUM, mAX_DPH_SUM)) sel_els
              sumTyCons     = listArray (2, mAX_DPH_SUM) sum_tcs

          voidVar          <- externalVar dph_Repr		(fsLit "void")
          pvoidVar         <- externalVar dph_Repr		(fsLit "pvoid")
          fromVoidVar      <- externalVar dph_Repr		(fsLit "fromVoid")
          punitVar         <- externalVar dph_Repr		(fsLit "punit")
          closureVar       <- externalVar dph_Closure	(fsLit "closure")
          applyVar         <- externalVar dph_Closure	(fsLit "$:")
          liftedClosureVar <- externalVar dph_Closure	(fsLit "liftedClosure")
          liftedApplyVar   <- externalVar dph_Closure	(fsLit "liftedApply")
          replicatePDVar   <- externalVar dph_PArray	(fsLit "replicatePD")
          emptyPDVar       <- externalVar dph_PArray	(fsLit "emptyPD")
          packByTagPDVar   <- externalVar dph_PArray	(fsLit "packByTagPD")

          combines 		<- mapM (externalVar dph_PArray)
                           		[mkFastString ("combine" ++ show i ++ "PD")
                              		| i <- [2..mAX_DPH_COMBINE]]
          let combinePDVars = listArray (2, mAX_DPH_COMBINE) combines

          scalarClass 	<- externalClass dph_PArray	(fsLit "Scalar")
          scalar_map	<- externalVar	dph_Scalar	(fsLit "scalar_map")
          scalar_zip2	<- externalVar	dph_Scalar	(fsLit "scalar_zipWith")
          scalar_zips	<- mapM (externalVar dph_Scalar)
                              	(numbered "scalar_zipWith" 3 mAX_DPH_SCALAR_ARGS)

          let scalarZips 	= listArray (1, mAX_DPH_SCALAR_ARGS)
                                     (scalar_map : scalar_zip2 : scalar_zips)

          closures 		<- mapM (externalVar dph_Closure)
                           		(numbered "closure" 1 mAX_DPH_SCALAR_ARGS)

          let closureCtrFuns = listArray (1, mAX_DPH_COMBINE) closures

          liftingContext	<- liftM (\u -> mkSysLocal (fsLit "lc") u intPrimTy)
    				newUnique

          return   $ Builtins
                   { dphModules       = mods
                   , parrayTyCon      = parrayTyCon
                   , parrayDataCon    = parrayDataCon
                   , pdataTyCon       = pdataTyCon
                   , paTyCon          = paTyCon
                   , paDataCon        = paDataCon
                   , preprTyCon       = preprTyCon
                   , prTyCon          = prTyCon
                   , prDataCon        = prDataCon
                   , voidTyCon        = voidTyCon
                   , wrapTyCon        = wrapTyCon
                   , selTys           = selTys
                   , selReplicates    = selReplicates
                   , selPicks         = selPicks
                   , selTagss         = selTagss
                   , selEls           = selEls
                   , sumTyCons        = sumTyCons
                   , closureTyCon     = closureTyCon
                   , voidVar          = voidVar
                   , pvoidVar         = pvoidVar
                   , fromVoidVar      = fromVoidVar
                   , punitVar         = punitVar
                   , closureVar       = closureVar
                   , applyVar         = applyVar
                   , liftedClosureVar = liftedClosureVar
                   , liftedApplyVar   = liftedApplyVar
                   , replicatePDVar   = replicatePDVar
                   , emptyPDVar       = emptyPDVar
                   , packByTagPDVar   = packByTagPDVar
                   , combinePDVars    = combinePDVars
                   , scalarClass      = scalarClass
                   , scalarZips       = scalarZips
                   , closureCtrFuns   = closureCtrFuns
                   , liftingContext   = liftingContext
                   }
      where
        mods@(Modules {
                   dph_PArray         = dph_PArray
                 , dph_Repr           = dph_Repr
                 , dph_Closure        = dph_Closure
                 , dph_Scalar         = dph_Scalar
                 , dph_Unboxed        = dph_Unboxed
                 })
          = dph_Modules pkg

        load get_mod = dsLoadModule doc mod
          where
            mod = get_mod mods
            doc = ppr mod <+> ptext (sLit "is a DPH module")

        -- Make a list of numbered strings in some range, eg foo3, foo4, foo5
        numbered :: String -> Int -> Int -> [FastString]
        numbered pfx m n = [mkFastString (pfx ++ show i) | i <- [m..n]]

        numbered_hash :: String -> Int -> Int -> [FastString]
        numbered_hash pfx m n = [mkFastString (pfx ++ show i ++ "#") | i <- [m..n]]

        mk_elements :: (Int, Int) -> DsM ((Int, Int), CoreExpr)
        mk_elements (i,j)
          = do
              v <- externalVar dph_Unboxed
                 $ mkFastString ("elementsSel" ++ show i ++ "_" ++ show j ++ "#")
              return ((i,j), Var v)

    -- | Get the mapping of names in the Prelude to names in the DPH library.
    initBuiltinVars :: Builtins -> DsM [(Var, Var)]
    initBuiltinVars (Builtins { dphModules = mods })
      = do
          uvars <- zipWithM externalVar umods ufs
          vvars <- zipWithM externalVar vmods vfs
          cvars <- zipWithM externalVar cmods cfs
          return $ [(v,v) | v <- map dataConWorkId defaultDataConWorkers]
                   ++ zip (map dataConWorkId cons) cvars
                   ++ zip uvars vvars
      where
        (umods, ufs, vmods, vfs) = unzip4 (preludeVars mods)
        (cons, cmods, cfs)       = unzip3 (preludeDataCons mods)

        defaultDataConWorkers :: [DataCon]
        defaultDataConWorkers = [trueDataCon, falseDataCon, unitDataCon]

    preludeDataCons :: Modules -> [(DataCon, Module, FastString)]
    preludeDataCons (Modules { dph_Prelude_Tuple = dph_Prelude_Tuple })
      = [mk_tup n dph_Prelude_Tuple (mkFastString $ "tup" ++ show n) | n <- [2..3]]
      where
        mk_tup n mod name = (tupleCon Boxed n, mod, name)

    -- | Get a list of names to `TyCon`s in the mock prelude.
    initBuiltinTyCons :: Builtins -> DsM [(Name, TyCon)]
    initBuiltinTyCons bi
      = do
          -- parr <- externalTyCon dph_Prelude_PArr (fsLit "PArr")
          dft_tcs <- defaultTyCons
          return $ (tyConName funTyCon, closureTyCon bi)
                 : (parrTyConName,      parrayTyCon bi)

                 -- FIXME: temporary
                 : (tyConName $ parrayTyCon bi, parrayTyCon bi)

                 : [(tyConName tc, tc) | tc <- dft_tcs]

      where	defaultTyCons :: DsM [TyCon]
    	defaultTyCons
      	 = do	word8 <- dsLookupTyCon word8TyConName
    		return [intTyCon, boolTyCon, doubleTyCon, word8]

    -- | Get a list of names to `DataCon`s in the mock prelude.
    initBuiltinDataCons :: Builtins -> [(Name, DataCon)]
    initBuiltinDataCons _
      = [(dataConName dc, dc)| dc <- defaultDataCons]
      where	defaultDataCons :: [DataCon]
    	defaultDataCons = [trueDataCon, falseDataCon, unitDataCon]

    -- | Get the names of all buildin instance functions for the PA class.
    initBuiltinPAs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
    initBuiltinPAs (Builtins { dphModules = mods }) insts
      = liftM (initBuiltinDicts insts) (externalClass (dph_PArray mods) (fsLit "PA"))

    -- | Get the names of all builtin instance functions for the PR class.
    initBuiltinPRs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
    initBuiltinPRs (Builtins { dphModules = mods }) insts
      = liftM (initBuiltinDicts insts) (externalClass (dph_PArray mods) (fsLit "PR"))

    -- | Get the names of all DPH instance functions for this class.
    initBuiltinDicts :: (InstEnv, InstEnv) -> Class -> [(Name, Var)]
    initBuiltinDicts insts cls = map find $ classInstances insts cls
      where
        find i | [Just tc] <- instanceRoughTcs i 	= (tc, instanceDFunId i)
               | otherwise				= pprPanic "Invalid DPH instance" (ppr i)

    -- | Get a list of boxed `TyCons` in the mock prelude. This is Int only.
    initBuiltinBoxedTyCons :: Builtins -> DsM [(Name, TyCon)]
    initBuiltinBoxedTyCons
      = return . builtinBoxedTyCons
      where	builtinBoxedTyCons :: Builtins -> [(Name, TyCon)]
    	builtinBoxedTyCons _
      		= [(tyConName intPrimTyCon, intTyCon)]

    -- | Get a list of all scalar functions in the mock prelude.
    initBuiltinScalars :: Builtins -> DsM [Var]
    initBuiltinScalars bi
      = mapM (uncurry externalVar) (preludeScalars $ dphModules bi)

    -- | Lookup some variable given its name and the module that contains it.
    externalVar :: Module -> FastString -> DsM Var
    externalVar mod fs
      = dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)

    -- | Like `externalVar` but wrap the `Var` in a `CoreExpr`
    externalFun :: Module -> FastString -> DsM CoreExpr
    externalFun mod fs
     = do var <- externalVar mod fs
          return $ Var var

    -- | Lookup some `TyCon` given its name and the module that contains it.
    externalTyCon :: Module -> FastString -> DsM TyCon
    externalTyCon mod fs
      = dsLookupTyCon =<< lookupOrig mod (mkTcOccFS fs)

    -- | Lookup some `Type` given its name and the module that contains it.
    externalType :: Module -> FastString -> DsM Type
    externalType mod fs
     = do  tycon <- externalTyCon mod fs
           return $ mkTyConApp tycon []

    -- | Lookup some `Class` given its name and the module that contains it.
    externalClass :: Module -> FastString -> DsM Class
    externalClass mod fs
      = dsLookupClass =<< lookupOrig mod (mkClsOccFS fs)

    -- | Like `externalClass`, but get the TyCon of of the class.
    externalClassTyCon :: Module -> FastString -> DsM TyCon
    externalClassTyCon mod fs = liftM classTyCon (externalClass mod fs)