Eliminate the separate implementation of NameMap, as it's just NameMapE applied to LiftE.

Author: axch

src/lib/CheckType.hs

@@ -56,12 +56,13 @@ liftTyperM cont =
 affineUsed :: AtomName r o -> TyperM r i o ()
 affineUsed name = TyperM $ do
   affines <- get
-  case lookupNameMap name affines of
-    Just n -> if n > 0 then
-                throw TypeErr $ "Affine name " ++ pprint name ++ " used " ++ show (n + 1) ++ " times."
-              else
-                put $ insertNameMap name (n + 1) affines
-    Nothing -> put $ insertNameMap name 1 affines
+  case lookupNameMapE name affines of
+    Just (LiftE n) ->
+        if n > 0 then
+          throw TypeErr $ "Affine name " ++ pprint name ++ " used " ++ show (n + 1) ++ " times."
+        else
+          put $ insertNameMapE name (LiftE $ n + 1) affines
+    Nothing -> put $ insertNameMapE name (LiftE 1) affines
 
 parallelAffines :: [TyperM r i o a] -> TyperM r i o [a]
 parallelAffines actions = TyperM $ do
@@ -77,7 +78,7 @@ parallelAffines actions = TyperM $ do
     result <- runTyperT' act
     (result,) <$> get
   put affines
-  forM_ (toListNameMap $ unionsWithNameMap max isolateds) \(name, ct) ->
+  forM_ (toListNameMapE $ unionsWithNameMapE max isolateds) \(name, (LiftE ct)) ->
     case ct of
       0 -> return ()
       1 -> runTyperT' $ affineUsed name

src/lib/Core.hs

@@ -482,5 +482,3 @@ freshNameM hint = do
   Distinct <- getDistinct
   return $ withFresh hint scope \b -> Abs b (binderName b)
 {-# INLINE freshNameM #-}
-
-type AtomNameMap r = NameMap (AtomNameC r)

src/lib/Lower.hs

@@ -217,17 +217,18 @@ lowerCase maybeDest scrut alts resultTy = do
 -- so that it never allocates scratch space for its result, but will put it directly in
 -- the corresponding slice of the full 2D buffer.
 
-type DestAssignment (i'::S) (o::S) = AtomNameMap SimpIR (ProjDest o) i'
+type DestAssignment (i'::S) (o::S) = NameMap (AtomNameC SimpIR) (ProjDest o) i'
 
 data ProjDest o
   = FullDest (Dest SimpIR o)
   | ProjDest (NE.NonEmpty Projection) (Dest SimpIR o)  -- dest corresponds to the projection applied to name
+  deriving (Show)
 
 instance SinkableE ProjDest where
   sinkingProofE = todoSinkableProof
 
 lookupDest :: DestAssignment i' o -> SAtomName i' -> Maybe (ProjDest o)
-lookupDest = flip lookupNameMap
+lookupDest dests = fmap fromLiftE . flip lookupNameMapE dests
 
 -- Matches up the free variables of the atom, with the given dest. For example, if the
 -- atom is a pair of two variables, the dest might be split into per-component dests,
@@ -238,10 +239,10 @@ lookupDest = flip lookupNameMap
 -- XXX: When adding more cases, be careful about potentially repeated vars in the output!
 decomposeDest :: Emits o => Dest SimpIR o -> SAtom i' -> LowerM i o (Maybe (DestAssignment i' o))
 decomposeDest dest = \case
-  Var v -> return $ Just $ singletonNameMap (atomVarName v) $ FullDest dest
+  Var v -> return $ Just $ singletonNameMapE (atomVarName v) $ LiftE $ FullDest dest
   ProjectElt _ p x -> do
     (ps, v) <- return $ asNaryProj p x
-    return $ Just $ singletonNameMap (atomVarName v) $ ProjDest ps dest
+    return $ Just $ singletonNameMapE (atomVarName v) $ LiftE $ ProjDest ps dest
   _ -> return Nothing
 
 lowerBlockWithDest :: Emits o => Dest SimpIR o -> SBlock i -> LowerM i o (SAtom o)
@@ -258,7 +259,7 @@ lowerBlockWithDest dest (Abs decls ans) = do
         Just DistinctBetween -> do
           s' <- traverseDeclNestWithDestS destMap s decls
           -- But we have to emit explicit writes, for all the vars that are not defined in decls!
-          forM_ (toListNameMap $ hoistFilterNameMap decls destMap) \(n, d) -> do
+          forM_ (toListNameMapE $ hoistNameMap decls destMap) \(n, (LiftE d)) -> do
             x <- case s ! n of
               Rename v -> Var <$> toAtomVar v
               SubstVal a -> return a

src/lib/MTL1.hs

@@ -223,8 +223,8 @@ instance HoistableState UnitE where
   hoistState _ _ UnitE = UnitE
   {-# INLINE hoistState #-}
 
-instance HoistableState (NameMap c a) where
-  hoistState _ b m = hoistFilterNameMap b m
+instance Show a => HoistableState (NameMap c a) where
+  hoistState _ b m = hoistNameMap b m
   {-# INLINE hoistState #-}
 
 -------------------- ScopedT1 --------------------

src/lib/Name.hs

@@ -481,7 +481,7 @@ newtype NonEmptyListE (e::E) (n::S) = NonEmptyListE { fromNonEmptyListE :: NonEm
         deriving (Show, Eq, Generic)
 
 newtype LiftE (a:: *) (n::S) = LiftE { fromLiftE :: a }
-        deriving (Show, Eq, Generic, Monoid, Semigroup)
+        deriving (Show, Eq, Ord, Generic, Monoid, Semigroup)
 
 newtype ComposeE (f :: * -> *) (e::E) (n::S) =
   ComposeE { fromComposeE :: (f (e n)) }
@@ -3262,60 +3262,6 @@ instance HoistableB b => HoistableB (WithAttrB a b) where
 
 -- Hoisting the map removes entries that are no longer in scope.
 
-newtype NameMap (c::C) (a:: *) (n::S) = UnsafeNameMap (RawNameMap a)
-                                 deriving (Eq, Semigroup, Monoid, Store)
-
-hoistFilterNameMap :: BindsNames b => b n l -> NameMap c a l -> NameMap c a n
-hoistFilterNameMap b (UnsafeNameMap raw) =
-  UnsafeNameMap $ raw `R.difference` frag
-  where UnsafeMakeScopeFrag frag = toScopeFrag b
-{-# INLINE hoistFilterNameMap #-}
-
-insertNameMap :: Name c n -> a -> NameMap c a n -> NameMap c a n
-insertNameMap (UnsafeMakeName n) x (UnsafeNameMap raw) = UnsafeNameMap $ R.insert n x raw
-{-# INLINE insertNameMap #-}
-
-lookupNameMap :: Name c n -> NameMap c a n -> Maybe a
-lookupNameMap (UnsafeMakeName n) (UnsafeNameMap raw) = R.lookup n raw
-{-# INLINE lookupNameMap #-}
-
-singletonNameMap :: Name c n -> a -> NameMap c a n
-singletonNameMap (UnsafeMakeName n) x = UnsafeNameMap $ R.singleton n x
-{-# INLINE singletonNameMap #-}
-
-toListNameMap :: NameMap c a n -> [(Name c n, a)]
-toListNameMap (UnsafeNameMap raw) = R.toList raw <&> \(r, x) -> (UnsafeMakeName r, x)
-{-# INLINE toListNameMap #-}
-
-unionWithNameMap :: (a -> a -> a) -> NameMap c a n -> NameMap c a n -> NameMap c a n
-unionWithNameMap f (UnsafeNameMap raw1) (UnsafeNameMap raw2) =
-  UnsafeNameMap $ R.unionWith f raw1 raw2
-{-# INLINE unionWithNameMap #-}
-
-unionsWithNameMap :: (Foldable f) => (a -> a -> a) -> f (NameMap c a n) -> NameMap c a n
-unionsWithNameMap func maps =
-  foldl' (unionWithNameMap func) mempty maps
-{-# INLINE unionsWithNameMap #-}
-
-traverseNameMap :: (Applicative f) => (a -> f b)
-                 -> NameMap c a n -> f (NameMap c b n)
-traverseNameMap f (UnsafeNameMap raw) = UnsafeNameMap <$> traverse f raw
-{-# INLINE traverseNameMap #-}
-
-mapNameMap :: (a -> b) -> NameMap c a n -> (NameMap c b n)
-mapNameMap f (UnsafeNameMap raw) = UnsafeNameMap $ fmap f raw
-{-# INLINE mapNameMap #-}
-
-keysNameMap :: NameMap c a n -> [Name c n]
-keysNameMap = map fst . toListNameMap
-{-# INLINE keysNameMap #-}
-
-keySetNameMap :: (Color c) => NameMap c a n -> NameSet n
-keySetNameMap nmap = freeVarsE $ ListE $ keysNameMap nmap
-
-instance SinkableE (NameMap c a) where
-  sinkingProofE = undefined
-
 newtype NameMapE (c::C) (e:: E) (n::S) = UnsafeNameMapE (RawNameMap (e n))
   deriving (Eq, Semigroup, Monoid, Store)
 
@@ -3353,6 +3299,11 @@ unionWithNameMapE f (UnsafeNameMapE raw1) (UnsafeNameMapE raw2) =
   UnsafeNameMapE $ R.unionWith f raw1 raw2
 {-# INLINE unionWithNameMapE #-}
 
+unionsWithNameMapE :: (Foldable f) => (e n -> e n -> e n) -> f (NameMapE c e n) -> NameMapE c e n
+unionsWithNameMapE func maps =
+  foldl' (unionWithNameMapE func) mempty maps
+{-# INLINE unionsWithNameMapE #-}
+
 traverseNameMapE :: (Applicative f) => (e1 n -> f (e2 n))
                  -> NameMapE c e1 n -> f (NameMapE c e2 n)
 traverseNameMapE f (UnsafeNameMapE raw) = UnsafeNameMapE <$> traverse f raw
@@ -3379,6 +3330,13 @@ instance RenameE e => RenameE (NameMapE c e) where
 instance HoistableE e => HoistableE (NameMapE c e) where
   freeVarsE = undefined
 
+type NameMap (c::C) (a:: *) = NameMapE c (LiftE a)
+
+hoistNameMap :: (BindsNames b, Show a)
+  => b n l -> NameMap c a l -> (NameMap c a n)
+hoistNameMap b = ignoreHoistFailure . hoistNameMapE b
+{-# INLINE hoistNameMap #-}
+
 -- === E-kinded IR coercions ===
 
 -- XXX: the intention is that we won't have to use this much

src/lib/Occurrence.hs

@@ -88,11 +88,6 @@ class MaxPlus a where
   max :: a -> a -> a
   plus :: a -> a -> a
 
-instance (MaxPlus a) => MaxPlus (NameMap c a n) where
-  zero = mempty
-  max  = unionWithNameMap max
-  plus = unionWithNameMap plus
-
 instance (MaxPlus (e n)) => MaxPlus (NameMapE c e n) where
   zero = mempty
   max  = unionWithNameMapE max

src/lib/Vectorize.hs

@@ -131,7 +131,7 @@ askVectorByteWidth :: TopVectorizeM i o Word32
 askVectorByteWidth = TopVectorizeM $ SubstReaderT $ lift $ lift11 (fromLiftE <$> ask)
 
 extendCommuteMap :: AtomName SimpIR o -> MonoidCommutes -> TopVectorizeM i o a -> TopVectorizeM i o a
-extendCommuteMap name commutativity = local $ insertNameMap name commutativity
+extendCommuteMap name commutativity = local $ insertNameMapE name $ LiftE commutativity
 
 vectorizeLoopsDestBlock :: DestBlock i
   -> TopVectorizeM i o (DestBlock o)
@@ -309,9 +309,9 @@ vectorSafeEffect (EffectRow effs NoTail) = allM safe $ eSetToList effs where
   safe (RWSEffect Writer (Var h)) = do
     h' <- renameM $ atomVarName h
     commuteMap <- ask
-    case lookupNameMap h' commuteMap of
-      Just Commutes -> return True
-      Just DoesNotCommute -> return False
+    case lookupNameMapE h' commuteMap of
+      Just (LiftE Commutes) -> return True
+      Just (LiftE DoesNotCommute) -> return False
       Nothing -> error $ "Handle " ++ pprint h ++ " not present in commute map?"
   safe _ = return False