Merge pull request #5591 from IntersectMBO/qnikst/map-extras

Extras for working with maps

Author: lehins

libs/cardano-data/CHANGELOG.md

@@ -5,6 +5,8 @@
 * Add `Data.Map.NonEmpty`
 * Add `Data.Set.NonEmpty`
 * Add `lookupInternMap`
+* Add `Data.MapExtras.boundedEnumMap`
+* Add `Data.MapExtras.lookupMapFail`
 * Replace `okeyL` with `toOKey`
 
 ## 1.2.4.1

libs/cardano-data/src/Data/MapExtras.hs

@@ -24,6 +24,8 @@ module Data.MapExtras (
   intersectDomPLeft,
   intersectMapSetFold,
   disjointMapSetFold,
+  boundedEnumMap,
+  lookupMapFail,
   extractKeys,
   extractKeysSmallSet,
   fromKeys,
@@ -284,3 +286,13 @@ lookupInternMap k = go
         LT -> go l
         GT -> go r
         EQ -> Just (kx, v)
+
+boundedEnumMap :: (Ord k, Bounded a, Enum a) => (a -> k) -> Map k a
+boundedEnumMap toTxt = Map.fromList [(toTxt t, t) | t <- [minBound .. maxBound]]
+
+lookupMapFail ::
+  (Ord k, Show k, MonadFail m) => String -> Map k v -> k -> m v
+lookupMapFail name m key =
+  case Map.lookup key m of
+    Just t -> pure t
+    Nothing -> fail $ "Unrecognized " <> name <> ": " ++ show key

libs/cardano-data/test/Test/Cardano/Data/MapExtrasSpec.hs

@@ -6,6 +6,7 @@ import Control.Monad
 import Data.Map (Map)
 import qualified Data.Map.Strict as Map
 import Data.MapExtras
+import Data.Word (Word8)
 import Test.Cardano.Data
 import Test.Hspec
 import Test.Hspec.QuickCheck
@@ -39,3 +40,9 @@ mapExtrasSpec =
           ma = intersectDomPLeft f m1 m2
       expectValidMap ma
       ma `shouldBe` Map.mapMaybeWithKey (\k v -> v <$ (guard . f k =<< Map.lookup k m2)) m1
+    prop "boundedEnumMap" $ do
+      let convert = fromIntegral :: Word8 -> Int
+          m = boundedEnumMap convert :: Map Int Word8
+      expectValidMap m
+      forM_ [minBound .. maxBound] $ \t -> do
+        Map.lookup (convert t) m `shouldBe` Just t