1 files changed, 70 insertions, 0 deletions

diff --git a/lib/Postdelay/PrioMap.hs b/lib/Postdelay/PrioMap.hs
new file mode 100644
index 0000000..2b75984
--- /dev/null
+++ b/lib/Postdelay/PrioMap.hs
@@ -0,0 +1,70 @@
+{-# LANGUAGE TemplateHaskell, ViewPatterns, RankNTypes, GeneralizedNewtypeDeriving, DeriveTraversable, ScopedTypeVariables, FlexibleInstances, MultiParamTypeClasses #-}
+
+module Postdelay.PrioMap
+  ( PrioMap, prioMap, prioMap', _Endo
+  , prio, prios
+  , squash
+
+  , PrioEndo, prioEndo, prioEndo'
+  ) where
+
+import Data.IntMap.Strict (IntMap)
+import qualified Data.IntMap.Strict as IntMap
+
+import Control.Lens
+
+import Data.Foldable
+import Data.Semigroup
+  
+
+makePrisms ''Endo
+  
+newtype PrioMap p a = PrioMap (IntMap a)
+  deriving (Functor, Foldable, Traversable)
+makePrisms ''PrioMap
+
+instance Semigroup a => Semigroup (PrioMap p a) where
+  (PrioMap a) <> (PrioMap b) = PrioMap $ IntMap.unionWith (<>) a b
+
+instance Semigroup a => Monoid (PrioMap p a) where
+  mempty = PrioMap mempty
+  mappend = (<>)
+
+instance Enum p => FunctorWithIndex p (PrioMap p) where
+  imap f (PrioMap intMap) = PrioMap $ imap (f . toEnum) intMap
+
+instance Enum p => FoldableWithIndex p (PrioMap p) where
+  ifoldMap f (PrioMap intMap) = ifoldMap (f . toEnum) intMap
+
+instance Enum p => TraversableWithIndex p (PrioMap p) where
+  itraverse f (PrioMap intMap) = PrioMap <$> itraverse (f . toEnum) intMap
+
+
+prioMap :: (Enum p, Monoid a) => Iso' (PrioMap p a) a
+-- ^ `prioMap` squashes priority information into `0`
+prioMap = prioMap' $ toEnum 0
+
+prioMap' :: (Enum p, Monoid a) => p -> Iso' (PrioMap p a) a
+-- ^ `prioMap' p` squashes priority information into `p`
+prioMap' (fromEnum -> p) = _PrioMap . iso fold (IntMap.singleton p)
+
+prio :: Enum p => p -> Lens' (PrioMap p a) (Maybe a)
+prio (fromEnum -> p) = _PrioMap . at p
+
+squash :: Semigroup a => PrioMap p (PrioMap p a) -> PrioMap p a
+squash = unwrapMonoid . foldMap WrapMonoid
+
+prios :: forall f p arr a. (Bounded p, Enum p, Semigroup a, Applicative f, Indexable p arr) => arr (Maybe a) (f (Maybe a)) -> (PrioMap p a -> f (PrioMap p a))
+prios (indexed -> alter) pMap = foldr cons_f (pure $ PrioMap IntMap.empty) [minBound .. maxBound]
+  where
+    cons_f :: p -> f (PrioMap p a) -> f (PrioMap p a)
+    cons_f p x = (<>) <$> (maybe mempty (PrioMap . IntMap.singleton (fromEnum p)) <$> alter p (pMap ^. prio p)) <*> x
+
+
+type PrioEndo p a = PrioMap p (Endo a)
+
+prioEndo :: Enum p => Iso' (PrioEndo p a) (Endo a)
+prioEndo = prioMap
+
+prioEndo' :: Enum p => p -> Iso' (PrioEndo p a) (Endo a)
+prioEndo' = prioMap'