Major Llvm refactoring

This combined patch reworks the LLVM backend in a number of ways:

1. Most prominently, we introduce a LlvmM monad carrying the contents of
   the old LlvmEnv around. This patch completely removes LlvmEnv and
   refactors towards standard library monad combinators wherever possible.

2. Support for streaming - we can now generate chunks of Llvm for Cmm as
   it comes in. This might improve our speed.

3. To allow streaming, we need a more flexible way to handle forward
   references. The solution (getGlobalPtr) unifies LlvmCodeGen.Data
   and getHsFunc as well.

4. Skip alloca-allocation for registers that are actually never written.
   LLVM will automatically eliminate these, but output is smaller and
   friendlier to human eyes this way.

5. We use LlvmM to collect references for llvm.used. This allows places
   other than cmmProcLlvmGens to generate entries.

1 files changed, 295 insertions, 98 deletions

diff --git a/compiler/llvmGen/LlvmCodeGen/Base.hs b/compiler/llvmGen/LlvmCodeGen/Base.hs
index 8de52eb0ba..95d3abdc27 100644
--- a/compiler/llvmGen/LlvmCodeGen/Base.hs
+++ b/compiler/llvmGen/LlvmCodeGen/Base.hs
@@ -13,15 +13,23 @@ module LlvmCodeGen.Base (
         LlvmVersion, defaultLlvmVersion, minSupportLlvmVersion,
         maxSupportLlvmVersion,
 
-        LlvmEnv, initLlvmEnv, clearVars, varLookup, varInsert,
-        funLookup, funInsert, getLlvmVer, setLlvmVer, getLlvmPlatform,
-        getDflags, ghcInternalFunctions,
+        LlvmM,
+        runLlvm, liftStream, withClearVars, varLookup, varInsert,
+        markStackReg, checkStackReg,
+        funLookup, funInsert, getLlvmVer, getDynFlags, getDynFlag, getLlvmPlatform,
+        renderLlvm, runUs, markUsedVar, getUsedVars,
+        ghcInternalFunctions,
+
+        getMetaUniqueId,
+        setUniqMeta, getUniqMeta,
+        freshSectionId,
 
         cmmToLlvmType, widthToLlvmFloat, widthToLlvmInt, llvmFunTy,
         llvmFunSig, llvmStdFunAttrs, llvmFunAlign, llvmInfAlign,
         llvmPtrBits, mkLlvmFunc, tysToParams,
 
-        strCLabel_llvm, genCmmLabelRef, genStringLabelRef
+        strCLabel_llvm, strDisplayName_llvm, strProcedureName_llvm,
+        getGlobalPtr, generateAliases,
 
     ) where
 
@@ -36,9 +44,16 @@ import DynFlags
 import FastString
 import Cmm
 import qualified Outputable as Outp
+import qualified Pretty as Prt
 import Platform
 import UniqFM
 import Unique
+import MonadUtils ( MonadIO(..) )
+import BufWrite   ( BufHandle )
+import UniqSet
+import UniqSupply
+import ErrUtils
+import qualified Stream
 
 -- ----------------------------------------------------------------------------
 -- * Some Data Types
@@ -93,30 +108,32 @@ llvmGhcCC dflags
  | otherwise                                      = CC_Ncc 10
 
 -- | Llvm Function type for Cmm function
-llvmFunTy :: DynFlags -> LiveGlobalRegs -> LlvmType
-llvmFunTy dflags live = LMFunction $ llvmFunSig' dflags live (fsLit "a") ExternallyVisible
+llvmFunTy :: LiveGlobalRegs -> LlvmM LlvmType
+llvmFunTy live = return . LMFunction =<< llvmFunSig' live (fsLit "a") ExternallyVisible
 
 -- | Llvm Function signature
-llvmFunSig :: LlvmEnv -> LiveGlobalRegs -> CLabel -> LlvmLinkageType -> LlvmFunctionDecl
-llvmFunSig env live lbl link
-    = llvmFunSig' (getDflags env) live (strCLabel_llvm env lbl) link
-
-llvmFunSig' :: DynFlags -> LiveGlobalRegs -> LMString -> LlvmLinkageType -> LlvmFunctionDecl
-llvmFunSig' dflags live lbl link
-  = let toParams x | isPointer x = (x, [NoAlias, NoCapture])
-                   | otherwise   = (x, [])
-    in LlvmFunctionDecl lbl link (llvmGhcCC dflags) LMVoid FixedArgs
-                        (map (toParams . getVarType) (llvmFunArgs dflags live))
-                        (llvmFunAlign dflags)
+llvmFunSig :: LiveGlobalRegs ->  CLabel -> LlvmLinkageType -> LlvmM LlvmFunctionDecl
+llvmFunSig live lbl link = do
+  lbl' <- strCLabel_llvm lbl
+  llvmFunSig' live lbl' link
+
+llvmFunSig' :: LiveGlobalRegs -> LMString -> LlvmLinkageType -> LlvmM LlvmFunctionDecl
+llvmFunSig' live lbl link
+  = do let toParams x | isPointer x = (x, [NoAlias, NoCapture])
+                      | otherwise   = (x, [])
+       dflags <- getDynFlags
+       return $ LlvmFunctionDecl lbl link (llvmGhcCC dflags) LMVoid FixedArgs
+                                 (map (toParams . getVarType) (llvmFunArgs dflags live))
+                                 (llvmFunAlign dflags)
 
 -- | Create a Haskell function in LLVM.
-mkLlvmFunc :: LlvmEnv -> LiveGlobalRegs -> CLabel -> LlvmLinkageType -> LMSection -> LlvmBlocks
-           -> LlvmFunction
-mkLlvmFunc env live lbl link sec blks
-  = let dflags = getDflags env
-        funDec = llvmFunSig env live lbl link
-        funArgs = map (fsLit . Outp.showSDoc dflags . ppPlainName) (llvmFunArgs dflags live)
-    in LlvmFunction funDec funArgs llvmStdFunAttrs sec blks
+mkLlvmFunc :: LiveGlobalRegs -> CLabel -> LlvmLinkageType -> LMSection -> LlvmBlocks
+           -> LlvmM LlvmFunction
+mkLlvmFunc live lbl link sec blks
+  = do funDec <- llvmFunSig live lbl link
+       dflags <- getDynFlags
+       let funArgs = map (fsLit . Outp.showSDoc dflags . ppPlainName) (llvmFunArgs dflags live)
+       return $ LlvmFunction funDec funArgs llvmStdFunAttrs sec blks
 
 -- | Alignment to use for functions
 llvmFunAlign :: DynFlags -> LMAlign
@@ -172,96 +189,276 @@ maxSupportLlvmVersion = 33
 -- * Environment Handling
 --
 
--- two maps, one for functions and one for local vars.
-newtype LlvmEnv = LlvmEnv (LlvmEnvMap, LlvmEnvMap, LlvmVersion, DynFlags)
+data LlvmEnv = LlvmEnv
+  { envFunMap :: LlvmEnvMap
+  , envVarMap :: LlvmEnvMap
+  , envStackRegs :: [GlobalReg]
+  , envUsedVars :: [LlvmVar]
+  , envAliases :: UniqSet LMString
+  , envLabelMap :: [(CLabel, CLabel)]
+  , envVersion :: LlvmVersion
+  , envDynFlags :: DynFlags
+  , envOutput :: BufHandle
+  , envUniq :: UniqSupply
+  , envFreshMeta :: Int
+  , envUniqMeta :: UniqFM Int
+  , envNextSection :: Int
+  }
 
 type LlvmEnvMap = UniqFM LlvmType
 
--- | Get initial Llvm environment.
-initLlvmEnv :: DynFlags -> LlvmEnv
-initLlvmEnv dflags = LlvmEnv (initFuncs, emptyUFM, defaultLlvmVersion, dflags)
-    where initFuncs = listToUFM $ [ (n, LMFunction ty) | (n, ty) <- ghcInternalFunctions dflags ]
+-- | The Llvm monad. Wraps @LlvmEnv@ state as well as the @IO@ monad
+newtype LlvmM a = LlvmM { runLlvmM :: LlvmEnv -> IO (a, LlvmEnv) }
+instance Monad LlvmM where
+    return x = LlvmM $ \env -> return (x, env)
+    m >>= f  = LlvmM $ \env -> do (x, env') <- runLlvmM m env
+                                  runLlvmM (f x) env'
+instance Functor LlvmM where
+    fmap f m = LlvmM $ \env -> do (x, env') <- runLlvmM m env
+                                  return (f x, env')
+instance MonadIO LlvmM where
+    liftIO m = LlvmM $ \env -> do x <- m
+                                  return (x, env)
+
+instance HasDynFlags LlvmM where
+    getDynFlags = LlvmM $ \env -> return (envDynFlags env, env)
 
--- | Here we pre-initialise some functions that are used internally by GHC
--- so as to make sure they have the most general type in the case that
--- user code also uses these functions but with a different type than GHC
--- internally. (Main offender is treating return type as 'void' instead of
--- 'void *'. Fixes trac #5486.
-ghcInternalFunctions :: DynFlags -> [(LMString, LlvmFunctionDecl)]
-ghcInternalFunctions dflags =
-    [ mk "memcpy" i8Ptr [i8Ptr, i8Ptr, llvmWord dflags]
-    , mk "memmove" i8Ptr [i8Ptr, i8Ptr, llvmWord dflags]
-    , mk "memset" i8Ptr [i8Ptr, llvmWord dflags, llvmWord dflags]
-    , mk "newSpark" (llvmWord dflags) [i8Ptr, i8Ptr]
-    ]
-  where
-    mk n ret args =
-        let n' = fsLit n
-        in (n', LlvmFunctionDecl n' ExternallyVisible CC_Ccc ret
-                                 FixedArgs (tysToParams args) Nothing)
-
--- | Clear variables from the environment.
-clearVars :: LlvmEnv -> LlvmEnv
-clearVars (LlvmEnv (e1, _, n, p)) = {-# SCC "llvm_env_clear" #-}
-    LlvmEnv (e1, emptyUFM, n, p)
-
--- | Insert local variables into the environment.
-varInsert :: Uniquable key => key -> LlvmType -> LlvmEnv -> LlvmEnv
-varInsert s t (LlvmEnv (e1, e2, n, p)) = {-# SCC "llvm_env_vinsert" #-}
-    LlvmEnv (e1, addToUFM e2 s t, n, p)
-
--- | Insert functions into the environment.
-funInsert :: Uniquable key => key -> LlvmType -> LlvmEnv -> LlvmEnv
-funInsert s t (LlvmEnv (e1, e2, n, p)) = {-# SCC "llvm_env_finsert" #-}
-    LlvmEnv (addToUFM e1 s t, e2, n, p)
-
--- | Lookup local variables in the environment.
-varLookup :: Uniquable key => key -> LlvmEnv -> Maybe LlvmType
-varLookup s (LlvmEnv (_, e2, _, _)) = {-# SCC "llvm_env_vlookup" #-}
-    lookupUFM e2 s
-
--- | Lookup functions in the environment.
-funLookup :: Uniquable key => key -> LlvmEnv -> Maybe LlvmType
-funLookup s (LlvmEnv (e1, _, _, _)) = {-# SCC "llvm_env_flookup" #-}
-    lookupUFM e1 s
+-- | Get initial Llvm environment.
+runLlvm :: DynFlags -> LlvmVersion -> BufHandle -> UniqSupply -> LlvmM () -> IO ()
+runLlvm dflags ver out us m = do
+    _ <- runLlvmM m env
+    return ()
+  where env = LlvmEnv { envFunMap = emptyUFM
+                      , envVarMap = emptyUFM
+                      , envStackRegs = []
+                      , envUsedVars = []
+                      , envAliases = emptyUniqSet
+                      , envLabelMap = []
+                      , envVersion = ver
+                      , envDynFlags = dflags
+                      , envOutput = out
+                      , envUniq = us
+                      , envFreshMeta = 0
+                      , envUniqMeta = emptyUFM
+                      , envNextSection = 1
+                      }
+
+-- | Get environment (internal)
+getEnv :: (LlvmEnv -> a) -> LlvmM a
+getEnv f = LlvmM (\env -> return (f env, env))
+
+-- | Modify environment (internal)
+modifyEnv :: (LlvmEnv -> LlvmEnv) -> LlvmM ()
+modifyEnv f = LlvmM (\env -> return ((), f env))
+
+-- | Lift a stream into the LlvmM monad
+liftStream :: Stream.Stream IO a x -> Stream.Stream LlvmM a x
+liftStream s = Stream.Stream $ do
+  r <- liftIO $ Stream.runStream s
+  case r of
+    Left b        -> return (Left b)
+    Right (a, r2) -> return (Right (a, liftStream r2))
+
+-- | Clear variables from the environment for a subcomputation
+withClearVars :: LlvmM a -> LlvmM a
+withClearVars m = LlvmM $ \env -> do
+    (x, env') <- runLlvmM m env { envVarMap = emptyUFM, envStackRegs = [] }
+    return (x, env' { envVarMap = emptyUFM, envStackRegs = [] })
+
+-- | Insert variables or functions into the environment.
+varInsert, funInsert :: Uniquable key => key -> LlvmType -> LlvmM ()
+varInsert s t = modifyEnv $ \env -> env { envVarMap = addToUFM (envVarMap env) s t }
+funInsert s t = modifyEnv $ \env -> env { envFunMap = addToUFM (envFunMap env) s t }
+
+-- | Lookup variables or functions in the environment.
+varLookup, funLookup :: Uniquable key => key -> LlvmM (Maybe LlvmType)
+varLookup s = getEnv (flip lookupUFM s . envVarMap)
+funLookup s = getEnv (flip lookupUFM s . envFunMap)
+
+-- | Set a register as allocated on the stack
+markStackReg :: GlobalReg -> LlvmM ()
+markStackReg r = modifyEnv $ \env -> env { envStackRegs = r : envStackRegs env }
+
+-- | Check whether a register is allocated on the stack
+checkStackReg :: GlobalReg -> LlvmM Bool
+checkStackReg r = getEnv ((elem r) . envStackRegs)
+
+-- | Allocate a new global unnamed metadata identifier
+getMetaUniqueId :: LlvmM Int
+getMetaUniqueId = LlvmM $ \env -> return (envFreshMeta env, env { envFreshMeta = envFreshMeta env + 1})
 
 -- | Get the LLVM version we are generating code for
-getLlvmVer :: LlvmEnv -> LlvmVersion
-getLlvmVer (LlvmEnv (_, _, n, _)) = n
+getLlvmVer :: LlvmM LlvmVersion
+getLlvmVer = getEnv envVersion
 
--- | Set the LLVM version we are generating code for
-setLlvmVer :: LlvmVersion -> LlvmEnv -> LlvmEnv
-setLlvmVer n (LlvmEnv (e1, e2, _, p)) = LlvmEnv (e1, e2, n, p)
+-- | Get the platform we are generating code for
+getDynFlag :: (DynFlags -> a) -> LlvmM a
+getDynFlag f = getEnv (f . envDynFlags)
 
 -- | Get the platform we are generating code for
-getLlvmPlatform :: LlvmEnv -> Platform
-getLlvmPlatform (LlvmEnv (_, _, _, d)) = targetPlatform d
+getLlvmPlatform :: LlvmM Platform
+getLlvmPlatform = getDynFlag targetPlatform
+
+-- | Prints the given contents to the output handle
+renderLlvm :: Outp.SDoc -> LlvmM ()
+renderLlvm sdoc = LlvmM $ \env -> do
+
+    -- Write to output
+    let doc = Outp.withPprStyleDoc (envDynFlags env) (Outp.mkCodeStyle Outp.CStyle) sdoc
+    Prt.bufLeftRender (envOutput env) doc
+
+    -- Dump, if requested
+    dumpIfSet_dyn (envDynFlags env) Opt_D_dump_llvm "LLVM Code" sdoc
+    return ((), env)
+
+-- | Run a @UniqSM@ action with our unique supply
+runUs :: UniqSM a -> LlvmM a
+runUs m = LlvmM $ \env -> do
+    let (x, us') = initUs (envUniq env) m
+    return (x, env { envUniq = us' })
+
+-- | Marks a variable as "used"
+markUsedVar :: LlvmVar -> LlvmM ()
+markUsedVar v = modifyEnv $ \env -> env { envUsedVars = v : envUsedVars env }
+
+-- | Return all variables marked as "used" so far
+getUsedVars :: LlvmM [LlvmVar]
+getUsedVars = getEnv envUsedVars
+
+-- | Saves that at some point we didn't know the type of the label and
+-- generated a reference to a type variable instead
+saveAlias :: LMString -> LlvmM ()
+saveAlias lbl = modifyEnv $ \env -> env { envAliases = addOneToUniqSet (envAliases env) lbl }
+
+-- | Sets metadata node for a given unique
+setUniqMeta :: Unique -> Int -> LlvmM ()
+setUniqMeta f m = modifyEnv $ \env -> env { envUniqMeta = addToUFM (envUniqMeta env) f m }
+-- | Gets metadata node for given unique
+getUniqMeta :: Unique -> LlvmM (Maybe Int)
+getUniqMeta s = getEnv (flip lookupUFM s . envUniqMeta)
+
+-- | Returns a fresh section ID
+freshSectionId :: LlvmM Int
+freshSectionId = LlvmM $ \env -> return (envNextSection env, env { envNextSection = envNextSection env + 1})
+
+-- ----------------------------------------------------------------------------
+-- * Internal functions
+--
 
--- | Get the DynFlags for this compilation pass
-getDflags :: LlvmEnv -> DynFlags
-getDflags (LlvmEnv (_, _, _, d)) = d
+-- | Here we pre-initialise some functions that are used internally by GHC
+-- so as to make sure they have the most general type in the case that
+-- user code also uses these functions but with a different type than GHC
+-- internally. (Main offender is treating return type as 'void' instead of
+-- 'void *'). Fixes trac #5486.
+ghcInternalFunctions :: LlvmM ()
+ghcInternalFunctions = do
+    dflags <- getDynFlags
+    mk "memcpy" i8Ptr [i8Ptr, i8Ptr, llvmWord dflags]
+    mk "memmove" i8Ptr [i8Ptr, i8Ptr, llvmWord dflags]
+    mk "memset" i8Ptr [i8Ptr, llvmWord dflags, llvmWord dflags]
+    mk "newSpark" (llvmWord dflags) [i8Ptr, i8Ptr]
+  where
+    mk n ret args = do
+      let n' = fsLit n
+          decl = LlvmFunctionDecl n' ExternallyVisible CC_Ccc ret
+                                 FixedArgs (tysToParams args) Nothing
+      renderLlvm $ ppLlvmFunctionDecl decl
+      funInsert n' (LMFunction decl)
 
 -- ----------------------------------------------------------------------------
 -- * Label handling
 --
 
 -- | Pretty print a 'CLabel'.
-strCLabel_llvm :: LlvmEnv -> CLabel -> LMString
-strCLabel_llvm env l = {-# SCC "llvm_strCLabel" #-}
-    (fsLit . toString . pprCLabel (getLlvmPlatform env)) l
-    where dflags = getDflags env
-          style = Outp.mkCodeStyle Outp.CStyle
-          toString doc = Outp.renderWithStyle dflags doc style
-
--- | Create an external definition for a 'CLabel' defined in another module.
-genCmmLabelRef :: LlvmEnv -> CLabel -> LMGlobal
-genCmmLabelRef env = genStringLabelRef (getDflags env) . strCLabel_llvm env
-
--- | As above ('genCmmLabelRef') but taking a 'LMString', not 'CLabel'.
-genStringLabelRef :: DynFlags -> LMString -> LMGlobal
-genStringLabelRef dflags cl
-  = let ty = LMPointer $ LMArray 0 (llvmWord dflags)
-    in LMGlobal (LMGlobalVar cl ty External Nothing Nothing Global) Nothing
+strCLabel_llvm :: CLabel -> LlvmM LMString
+strCLabel_llvm lbl = do
+    platform <- getLlvmPlatform
+    dflags <- getDynFlags
+    let sdoc = pprCLabel platform lbl
+        str = Outp.renderWithStyle dflags sdoc (Outp.mkCodeStyle Outp.CStyle)
+    return (fsLit str)
+
+strDisplayName_llvm :: CLabel -> LlvmM LMString
+strDisplayName_llvm lbl = do
+    platform <- getLlvmPlatform
+    dflags <- getDynFlags
+    let sdoc = pprCLabel platform lbl
+        depth = Outp.PartWay 1
+        style = Outp.mkUserStyle (const Outp.NameNotInScope2, const True) depth
+        str = Outp.renderWithStyle dflags sdoc style
+    return (fsLit (dropInfoSuffix str))
+
+dropInfoSuffix :: String -> String
+dropInfoSuffix = go
+  where go "_info"        = []
+        go "_static_info" = []
+        go "_con_info"    = []
+        go (x:xs)         = x:go xs
+        go []             = []
+
+strProcedureName_llvm :: CLabel -> LlvmM LMString
+strProcedureName_llvm lbl = do
+    platform <- getLlvmPlatform
+    dflags <- getDynFlags
+    let sdoc = pprCLabel platform lbl
+        depth = Outp.PartWay 1
+        style = Outp.mkUserStyle (const Outp.NameUnqual, const False) depth
+        str = Outp.renderWithStyle dflags sdoc style
+    return (fsLit str)
+
+-- ----------------------------------------------------------------------------
+-- * Global variables / forward references
+--
+
+-- | Create/get a pointer to a global value. Might return an alias if
+-- the value in question hasn't been defined yet. We especially make
+-- no guarantees on the type of the returned pointer.
+getGlobalPtr :: LMString -> LlvmM LlvmVar
+getGlobalPtr llvmLbl = do
+  m_ty <- funLookup llvmLbl
+  let mkGlbVar lbl ty = LMGlobalVar lbl (LMPointer ty) Private Nothing Nothing
+  case m_ty of
+    -- Directly reference if we have seen it already
+    Just ty -> return $ mkGlbVar llvmLbl ty Global
+    -- Otherwise use a forward alias of it
+    Nothing -> do
+      saveAlias llvmLbl
+      return $ mkGlbVar (llvmLbl `appendFS` fsLit "$alias") i8 Alias
+
+-- | Generate definitions for aliases forward-referenced by @getGlobalPtr@.
+--
+-- Must be called at a point where we are sure that no new global definitions
+-- will be generated anymore!
+generateAliases :: LlvmM ([LMGlobal], [LlvmType])
+generateAliases = do
+  delayed <- fmap uniqSetToList $ getEnv envAliases
+  defss <- flip mapM delayed $ \lbl -> do
+    let var      ty = LMGlobalVar lbl (LMPointer ty) External Nothing Nothing Global
+        aliasLbl    = lbl `appendFS` fsLit "$alias"
+        aliasVar    = LMGlobalVar aliasLbl i8Ptr Private Nothing Nothing Alias
+    -- If we have a definition, set the alias value using a
+    -- cost. Otherwise, declare it as an undefined external symbol.
+    m_ty <- funLookup lbl
+    case m_ty of
+      Just ty -> return [LMGlobal aliasVar $ Just $ LMBitc (LMStaticPointer (var ty)) i8Ptr]
+      Nothing -> return [LMGlobal (var i8) Nothing,
+                         LMGlobal aliasVar $ Just $ LMStaticPointer (var i8) ]
+  -- Reset forward list
+  modifyEnv $ \env -> env { envAliases = emptyUniqSet }
+  return (concat defss, [])
+
+-- Note [Llvm Forward References]
+--
+-- The issue here is that LLVM insists on being strongly typed at
+-- every corner, so the first time we mention something, we have to
+-- settle what type we assign to it. That makes things awkward, as Cmm
+-- will often reference things before their definition, and we have no
+-- idea what (LLVM) type it is going to be before that point.
+--
+-- Our work-around is to define "aliases" of a standard type (i8 *) in
+-- these kind of situations, which we later tell LLVM to be either
+-- references to their actual local definitions (involving a cast) or
+-- an external reference. This obviously only works for pointers.
 
 -- ----------------------------------------------------------------------------
 -- * Misc