wait using MVars

Author: vmchale

clash-lib/src/Clash/Core/VarEnv.hs

@@ -38,6 +38,8 @@ module Clash.Core.VarEnv
     -- ** Conversions
     -- *** Lists
   , eltsVarEnv
+  , toListVarEnv
+  , listToVarEnv
     -- * Sets of variables
   , VarSet
     -- ** Construction
@@ -261,6 +263,15 @@ eltsVarEnv
   -> [a]
 eltsVarEnv = eltsUniqMap
 
+toListVarEnv :: VarEnv a -> [(Unique, a)]
+toListVarEnv = toListUniqMap
+
+listToVarEnv
+  :: Uniquable a
+  => [(a, b)]
+  -> VarEnv b
+listToVarEnv = listToUniqMap
+
 -- | Does the variable exist in the environment
 elemVarEnv
   :: Var a

clash-lib/src/Clash/Normalize.hs

@@ -70,8 +70,9 @@ import           Clash.Core.Type                  (isPolyTy)
 import           Clash.Core.Var                   (Id, varName, varType)
 import           Clash.Core.VarEnv
   (VarEnv, VarSet, elemVarSet, eltsVarEnv, emptyInScopeSet, emptyVarEnv, emptyVarSet,
-   extendVarEnv, extendVarSet, lookupVarEnv, mapVarEnv, mapMaybeVarEnv,
-   mkVarEnv, mkVarSet, notElemVarEnv, notElemVarSet, nullVarEnv)
+   extendVarEnv, extendVarSet, lookupVarEnv, mapMaybeVarEnv,
+   mkVarEnv, mkVarSet, notElemVarEnv, notElemVarSet, nullVarEnv,
+   listToVarEnv, toListVarEnv)
 import           Clash.Debug                      (traceIf)
 import           Clash.Driver.Types
   (BindingMap, Binding(..), DebugOpts(..), ClashEnv(..))
@@ -234,11 +235,11 @@ normalize' nm q = do
             normV <- Lens.use (extra.normalized)
 
             toNormalize <-
-              MVar.withMVar normV $ \norm ->
-                let prevNorm = mapVarEnv bindingId norm
-                    toNormalize = filter (`notElemVarSet` topEnts)
+              MVar.withMVar normV $ \norm -> do
+                prevNorm <- listToVarEnv <$> traverse (\(k, v) -> (k,) . bindingId <$> MVar.readMVar v) (toListVarEnv norm)
+                let toNormalize = filter (`notElemVarSet` topEnts)
                                 $ filter (`notElemVarEnv` extendVarEnv nm nm prevNorm) usedBndrs
-                 in pure toNormalize
+                  in pure toNormalize
 
             traverse_ (Monad.liftIO . MS.pushL q) toNormalize
             pure (nm, tmNorm)

clash-lib/src/Clash/Normalize/Types.hs

@@ -23,13 +23,13 @@ import Clash.Core.Term        (Term)
 import Clash.Core.Type        (Type)
 import Clash.Core.Var         (Id)
 import Clash.Core.VarEnv      (VarEnv)
-import Clash.Driver.Types     (BindingMap)
+import Clash.Driver.Types     (Binding)
 import Clash.Rewrite.Types    (Rewrite, RewriteMonad)
 
 -- | State of the 'NormalizeMonad'
 data NormalizeState
   = NormalizeState
-  { _normalized          :: MVar BindingMap
+  { _normalized          :: MVar (VarEnv (MVar (Binding Term)))
   -- ^ Global binders
   , _specialisationCache :: MVar (Map (Id,Int,Either Term Type) Id)
   -- ^ Cache of previously specialized functions:

clash-lib/src/Clash/Normalize/Util.hs

@@ -424,16 +424,19 @@ normalizeTopLvlBndr
 normalizeTopLvlBndr isTop nm (Binding nm' sp inl pr tm _) = do
   normalizedV <- Lens.use (extra.normalized)
 
-  mValue <- MVar.withMVar normalizedV (pure . lookupVarEnv nm)
-
   -- TODO This was a call to makeCachedU, but since there was no variation
   -- for MVar, I unrolled everything. Maybe there should be MVar versions of
   -- the makeCachedX functions needed in normalization.
-  case mValue of
-    Just value ->
-      pure value
 
+  cache <- MVar.takeMVar normalizedV
+  case lookupVarEnv nm cache of
+    Just vMVar -> do
+      MVar.putMVar normalizedV cache
+      MVar.readMVar vMVar
     Nothing -> do
+      tmp <- MVar.newEmptyMVar
+      MVar.putMVar normalizedV (extendVarEnv nm tmp cache)
+
       tcm <- Lens.view tcCache
       let nmS = showPpr (varName nm)
       -- We deshadow the term because sometimes GHC gives us
@@ -454,8 +457,8 @@ normalizeTopLvlBndr isTop nm (Binding nm' sp inl pr tm _) = do
       let r' = nm' `globalIdOccursIn` tm3
       let value = Binding nm'{varType = ty'} sp inl pr tm3 r'
 
-      MVar.modifyMVar normalizedV $ \binders ->
-        pure (extendVarEnv nm value binders, value)
+      MVar.putMVar tmp value
+      pure value
 
 -- | Turn type equality constraints into substitutions and apply them.
 --