{-# LANGUAGE TupleSections #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveFoldable #-}
module Data.Comp.Mapping
( Numbered (..)
, unNumbered
, number
, Traversable ()
, Mapping (..)
, prodMap
, lookupNumMap
, lookupNumMap'
, NumMap) where
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.Traversable
import Data.Foldable
import Control.Monad.State hiding (mapM)
import Prelude hiding (mapM)
data Numbered a = Numbered Int a
unNumbered :: Numbered a -> a
unNumbered (Numbered _ x) = x
number :: Traversable f => f a -> f (Numbered a)
number x = evalState (mapM run x) 0 where
run b = do n <- get
put (n+1)
return $ Numbered n b
infix 1 |->
infixr 0 &
class Functor m => Mapping m k | m -> k where
(&) :: m v -> m v -> m v
(|->) :: k -> v -> m v
empty :: m v
prodMapWith :: (v1 -> v2 -> v) -> v1 -> v2 -> m v1 -> m v2 -> m v
findWithDefault :: a -> k -> m a -> a
prodMap :: Mapping m k => v1 -> v2 -> m v1 -> m v2 -> m (v1, v2)
prodMap = prodMapWith (,)
newtype NumMap k v = NumMap (IntMap v) deriving (Functor,Foldable,Traversable)
lookupNumMap :: a -> Int -> NumMap t a -> a
lookupNumMap d k (NumMap m) = IntMap.findWithDefault d k m
lookupNumMap' :: Int -> NumMap t a -> Maybe a
lookupNumMap' k (NumMap m) = IntMap.lookup k m
instance Mapping (NumMap k) (Numbered k) where
NumMap m1 & NumMap m2 = NumMap (IntMap.union m1 m2)
Numbered k _ |-> v = NumMap $ IntMap.singleton k v
empty = NumMap IntMap.empty
findWithDefault d (Numbered i _) m = lookupNumMap d i m
prodMapWith f p q (NumMap mp) (NumMap mq) = NumMap $ IntMap.mergeWithKey merge
(IntMap.map (`f` q)) (IntMap.map (p `f`)) mp mq
where merge _ p q = Just (p `f` q)