Add MonadUnique instance for LlvmM

Reviewers: erikd, austin

Reviewed By: erikd

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D1994

2 files changed, 16 insertions, 9 deletions

diff --git a/compiler/llvmGen/LlvmCodeGen/Base.hs b/compiler/llvmGen/LlvmCodeGen/Base.hs
index 3a60891810..3e2b795650 100644
--- a/compiler/llvmGen/LlvmCodeGen/Base.hs
+++ b/compiler/llvmGen/LlvmCodeGen/Base.hs
@@ -18,7 +18,7 @@ module LlvmCodeGen.Base (
         runLlvm, liftStream, withClearVars, varLookup, varInsert,
         markStackReg, checkStackReg,
         funLookup, funInsert, getLlvmVer, getDynFlags, getDynFlag, getLlvmPlatform,
-        dumpIfSetLlvm, renderLlvm, runUs, markUsedVar, getUsedVars,
+        dumpIfSetLlvm, renderLlvm, markUsedVar, getUsedVars,
         ghcInternalFunctions,
 
         getMetaUniqueId,
@@ -224,6 +224,19 @@ instance Monad LlvmM where
 instance HasDynFlags LlvmM where
     getDynFlags = LlvmM $ \env -> return (envDynFlags env, env)
 
+instance MonadUnique LlvmM where
+    getUniqueSupplyM = do
+        us <- getEnv envUniq
+        let (us1, us2) = splitUniqSupply us
+        modifyEnv (\s -> s { envUniq = us2 })
+        return us1
+
+    getUniqueM = do
+        us <- getEnv envUniq
+        let (u,us') = takeUniqFromSupply us
+        modifyEnv (\s -> s { envUniq = us' })
+        return u
+
 -- | Lifting of IO actions. Not exported, as we want to encapsulate IO.
 liftIO :: IO a -> LlvmM a
 liftIO m = LlvmM $ \env -> do x <- m
@@ -323,12 +336,6 @@ renderLlvm sdoc = do
     dumpIfSetLlvm Opt_D_dump_llvm "LLVM Code" sdoc
     return ()
 
--- | 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 }
diff --git a/compiler/llvmGen/LlvmCodeGen/CodeGen.hs b/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
index 4bc6248470..fa47d6ada3 100644
--- a/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
+++ b/compiler/llvmGen/LlvmCodeGen/CodeGen.hs
@@ -144,7 +144,7 @@ getInstrinct2 fname fty@(LMFunction funSig) = do
         return []
       Nothing -> do
         funInsert fname fty
-        un <- runUs getUniqueM
+        un <- getUniqueM
         let lbl = mkAsmTempLabel un
         return [CmmData (Section Data lbl) [([],[fty])]]
 
@@ -1783,7 +1783,7 @@ getHsFunc' name fty
 -- | Create a new local var
 mkLocalVar :: LlvmType -> LlvmM LlvmVar
 mkLocalVar ty = do
-    un <- runUs getUniqueM
+    un <- getUniqueM
     return $ LMLocalVar un ty