Merge pull request #5499 from unisonweb/cp/inline-ec

Add missing INLINE's on EnumContainers
unisonweb · Dec 12, 2024 · f9cc70e · f9cc70e
2 parents deefb3d + 5adb267
commit f9cc70e
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 58 deletions.
diff --git a/parser-typechecker/src/Unison/Util/EnumContainers.hs b/parser-typechecker/src/Unison/Util/EnumContainers.hs
@@ -43,11 +43,15 @@ class EnumKey k where
   intToKey :: Int -> k
 
 instance EnumKey Word64 where
+  {-# INLINE keyToInt #-}
   keyToInt e = fromIntegral e
+  {-# INLINE intToKey #-}
   intToKey i = fromIntegral i
 
 instance EnumKey Word16 where
+  {-# INLINE keyToInt #-}
   keyToInt e = fromIntegral e
+  {-# INLINE intToKey #-}
   intToKey i = fromIntegral i
 
 newtype EnumMap k a = EM (IM.IntMap a)
@@ -77,24 +81,31 @@ newtype EnumSet k = ES IS.IntSet
       Semigroup
     )
 
+{-# INLINE mapFromList #-}
 mapFromList :: (EnumKey k) => [(k, a)] -> EnumMap k a
 mapFromList = EM . IM.fromList . fmap (first keyToInt)
 
+{-# INLINE setFromList #-}
 setFromList :: (EnumKey k) => [k] -> EnumSet k
 setFromList = ES . IS.fromList . fmap keyToInt
 
+{-# INLINE setToList #-}
 setToList :: (EnumKey k) => EnumSet k -> [k]
 setToList (ES s) = intToKey <$> IS.toList s
 
+{-# INLINE mapSingleton #-}
 mapSingleton :: (EnumKey k) => k -> a -> EnumMap k a
 mapSingleton e a = EM $ IM.singleton (keyToInt e) a
 
+{-# INLINE setSingleton #-}
 setSingleton :: (EnumKey k) => k -> EnumSet k
 setSingleton e = ES . IS.singleton $ keyToInt e
 
+{-# INLINE mapInsert #-}
 mapInsert :: (EnumKey k) => k -> a -> EnumMap k a -> EnumMap k a
 mapInsert e x (EM m) = EM $ IM.insert (keyToInt e) x m
 
+{-# INLINE unionWith #-}
 unionWith ::
   (EnumKey k) =>
   (a -> a -> a) ->
@@ -103,60 +114,77 @@ unionWith ::
   EnumMap k a
 unionWith f (EM l) (EM r) = EM $ IM.unionWith f l r
 
+{-# INLINE intersectionWith #-}
 intersectionWith ::
   (a -> b -> c) ->
   EnumMap k a ->
   EnumMap k b ->
   EnumMap k c
 intersectionWith f (EM l) (EM r) = EM $ IM.intersectionWith f l r
 
+{-# INLINE keys #-}
 keys :: (EnumKey k) => EnumMap k a -> [k]
 keys (EM m) = fmap intToKey . IM.keys $ m
 
+{-# INLINE keysSet #-}
 keysSet :: (EnumKey k) => EnumMap k a -> EnumSet k
 keysSet (EM m) = ES (IM.keysSet m)
 
+{-# INLINE restrictKeys #-}
 restrictKeys :: (EnumKey k) => EnumMap k a -> EnumSet k -> EnumMap k a
 restrictKeys (EM m) (ES s) = EM $ IM.restrictKeys m s
 
+{-# INLINE withoutKeys #-}
 withoutKeys :: (EnumKey k) => EnumMap k a -> EnumSet k -> EnumMap k a
 withoutKeys (EM m) (ES s) = EM $ IM.withoutKeys m s
 
+{-# INLINE mapDifference #-}
 mapDifference :: (EnumKey k) => EnumMap k a -> EnumMap k b -> EnumMap k a
 mapDifference (EM l) (EM r) = EM $ IM.difference l r
 
+{-# INLINE member #-}
 member :: (EnumKey k) => k -> EnumSet k -> Bool
 member e (ES s) = IS.member (keyToInt e) s
 
+{-# INLINE hasKey #-}
 hasKey :: (EnumKey k) => k -> EnumMap k a -> Bool
 hasKey k (EM m) = IM.member (keyToInt k) m
 
+{-# INLINE lookup #-}
 lookup :: (EnumKey k) => k -> EnumMap k a -> Maybe a
 lookup e (EM m) = IM.lookup (keyToInt e) m
 
+{-# INLINE lookupWithDefault #-}
 lookupWithDefault :: (EnumKey k) => a -> k -> EnumMap k a -> a
 lookupWithDefault d e (EM m) = IM.findWithDefault d (keyToInt e) m
 
+{-# INLINE mapWithKey #-}
 mapWithKey :: (EnumKey k) => (k -> a -> b) -> EnumMap k a -> EnumMap k b
 mapWithKey f (EM m) = EM $ IM.mapWithKey (f . intToKey) m
 
+{-# INLINE foldMapWithKey #-}
 foldMapWithKey :: (EnumKey k) => (Monoid m) => (k -> a -> m) -> EnumMap k a -> m
 foldMapWithKey f (EM m) = IM.foldMapWithKey (f . intToKey) m
 
+{-# INLINE mapToList #-}
 mapToList :: (EnumKey k) => EnumMap k a -> [(k, a)]
 mapToList (EM m) = first intToKey <$> IM.toList m
 
+{-# INLINE (!) #-}
 (!) :: (EnumKey k) => EnumMap k a -> k -> a
 (!) (EM m) e = m IM.! keyToInt e
 
+{-# INLINE findMin #-}
 findMin :: (EnumKey k) => EnumSet k -> k
 findMin (ES s) = intToKey $ IS.findMin s
 
+{-# INLINE traverseSet_ #-}
 traverseSet_ ::
   (Applicative f) => (EnumKey k) => (k -> f ()) -> EnumSet k -> f ()
 traverseSet_ f (ES s) =
   IS.foldr (\i r -> f (intToKey i) *> r) (pure ()) s
 
+{-# INLINE interverse #-}
 interverse ::
   (Applicative f) =>
   (a -> b -> f c) ->
@@ -166,6 +194,7 @@ interverse ::
 interverse f (EM l) (EM r) =
   fmap EM . traverse id $ IM.intersectionWith f l r
 
+{-# INLINE traverseWithKey #-}
 traverseWithKey ::
   (Applicative f) =>
   (EnumKey k) =>
@@ -174,5 +203,6 @@ traverseWithKey ::
   f (EnumMap k b)
 traverseWithKey f (EM m) = EM <$> IM.traverseWithKey (f . intToKey) m
 
+{-# INLINE setSize #-}
 setSize :: EnumSet k -> Int
 setSize (ES s) = IS.size s