git.iko.soy/tix
on be81c19 help sign in
tix/src/Data/Map/MultiKey/Strict.hs
105 lines 2.9 kB history raw
1module Data.Map.MultiKey.Strict
2 ( Keyable (..),
3 Map (..),
4 singleton,
5 delete,
6 lookup,
7 lookupEmpty,
8 insert,
9 partition,
10 fromList,
11 toList,
12 elems,
13 keys,
14 map,
15 toMap,
16 )
17where
19import Control.Monad
20import Data.Bifunctor
21import Data.Foldable (fold)
22import Data.Kind
23import qualified Data.Map.Strict as M
24import Data.Maybe
25import Data.Monoid
26import qualified Data.Set as S
27import Prelude hiding (lookup, map)
29class (Ord (Key t), Ord t) => Keyable t where
30 type Key t :: Type
31 getKeys :: t -> [Key t]
33data Map t = Map (M.Map (Key t) (S.Set t)) (S.Set t)
35instance (Ord (Key t), Ord t) => Semigroup (Map t) where
36 (Map x x') <> (Map y y') = Map (M.unionWith S.union x y) (x' <> y')
38instance (Ord (Key t), Ord t) => Monoid (Map t) where
39 mempty = Map M.empty S.empty
41map :: (Keyable h, Ord t) => (t -> h) -> Map t -> Map h
42map f = fromList . fmap f . toList
44deriving stock instance (Show (Key t), Show t) => Show (Map t)
46singleton :: Keyable t => t -> Map t
47singleton t = case getKeys t of
48 [] -> Map M.empty (S.singleton t)
49 ks@(_ : _) -> Map (M.fromList $ fmap (,s) ks) S.empty
50 where
51 s = S.singleton t
53insert :: Keyable t => t -> Map t -> Map t
54insert t m = m <> singleton t
56lookup :: Keyable t => Key t -> Map t -> S.Set t
57lookup k (Map m _) = fromMaybe S.empty $ M.lookup k m
59lookupEmpty :: Map t -> S.Set t
60lookupEmpty (Map _ s) = s
62fromList :: Keyable t => [t] -> Map t
63fromList = foldMap singleton
65elems :: Ord t => Map t -> S.Set t
66elems (Map m s) = (fold . M.elems $ m) <> s
68toList :: Ord t => Map t -> [t]
69toList = S.toList . elems
71keys :: Map t -> S.Set (Key t)
72keys (Map m _) = M.keysSet m
74toMap :: Ord (Key t) => Map t -> M.Map (Maybe (Key t)) (S.Set t)
75toMap (Map m s) = M.insert Nothing s $ M.mapKeysMonotonic Just m
77delete :: Keyable t => t -> Map t -> Map t
78delete t (Map m empties) = Map (f m) (S.delete t empties)
79 where
80 s = S.singleton t
81 f =
82 appEndo $
83 foldMap
84 ( Endo
85 . M.update
86 (\s' -> if s == s then Nothing else Just (S.delete t s'))
87 )
88 (getKeys t)
90partition :: Keyable t => (Key t -> Bool) -> Map t -> (Map t, Map t)
91partition predicate (Map m empties) = (Map keep' empties, Map toss S.empty <> Map toss' S.empty)
92 where
93 (keep, toss) = M.partitionWithKey (\k _ -> predicate k) m
94 (keep', toss') =
95 join bimap M.fromAscList
96 . p (join bimap nothingNull . S.partition (all predicate . getKeys))
97 . M.toAscList
98 $ keep
99 p :: (v -> (Maybe v, Maybe v)) -> [(k, v)] -> ([(k, v)], [(k, v)])
100 p _ [] = ([], [])
101 p f ((k, x) : xs) = join bimap (maybe id ((:) . (k,))) (f x) `apBoth` p f xs
102 apBoth :: (a -> b, x -> y) -> (a, x) -> (b, y)
103 apBoth (f, g) (a, x) = (f a, g x)
104 nothingNull :: S.Set x -> Maybe (S.Set x)
105 nothingNull s = if S.null s then Nothing else Just s