Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add missing INLINE's on EnumContainers #5499

Merged
merged 2 commits into from
Dec 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions parser-typechecker/src/Unison/Util/EnumContainers.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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) ->
Expand All @@ -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) ->
Expand All @@ -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) =>
Expand All @@ -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
Loading
Loading