{-# LANGUAGE TupleSections #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveFoldable #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Comp.Mapping
-- Copyright   :  (c) 2014 Patrick Bahr
-- License     :  BSD3
-- Maintainer  :  Patrick Bahr <paba@diku.dk>
-- Stability   :  experimental
-- Portability :  non-portable (GHC Extensions)
--
-- This module provides functionality to construct mappings from
-- positions in a functorial value.
--
--------------------------------------------------------------------------------

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)


-- | This type is used for numbering components of a functorial value.
data Numbered a = Numbered Int a

unNumbered :: Numbered a -> a
unNumbered :: Numbered a -> a
unNumbered (Numbered Int
_ a
x) = a
x


-- | This function numbers the components of the given functorial
-- value with consecutive integers starting at 0.
number :: Traversable f => f a -> f (Numbered a)
number :: f a -> f (Numbered a)
number f a
x = State Int (f (Numbered a)) -> Int -> f (Numbered a)
forall s a. State s a -> s -> a
evalState ((a -> StateT Int Identity (Numbered a))
-> f a -> State Int (f (Numbered a))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM a -> StateT Int Identity (Numbered a)
forall (m :: * -> *) a. MonadState Int m => a -> m (Numbered a)
run f a
x) Int
0 where
  run :: a -> m (Numbered a)
run a
b = do Int
n <- m Int
forall s (m :: * -> *). MonadState s m => m s
get
             Int -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
             Numbered a -> m (Numbered a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Numbered a -> m (Numbered a)) -> Numbered a -> m (Numbered a)
forall a b. (a -> b) -> a -> b
$ Int -> a -> Numbered a
forall a. Int -> a -> Numbered a
Numbered Int
n a
b


infix 1 |->
infixr 0 &


class Functor m => Mapping m k | m -> k where
    -- | left-biased union of two mappings.
    (&) :: m v -> m v -> m v

    -- | This operator constructs a singleton mapping.
    (|->) :: k -> v -> m v

    -- | This is the empty mapping.
    empty :: m v

    -- | This function constructs the pointwise product of two maps each
    -- with a default value.
    prodMapWith :: (v1 -> v2 -> v) -> v1 -> v2 -> m v1 -> m v2 -> m v

    -- | Returns the value at the given key or returns the given
    -- default when the key is not an element of the map.
    findWithDefault :: a -> k -> m a -> a

-- | This function constructs the pointwise product of two maps each
-- with a default value.
prodMap :: Mapping m k => v1 -> v2 -> m v1 -> m v2 -> m (v1, v2)
prodMap :: v1 -> v2 -> m v1 -> m v2 -> m (v1, v2)
prodMap = (v1 -> v2 -> (v1, v2)) -> v1 -> v2 -> m v1 -> m v2 -> m (v1, v2)
forall (m :: * -> *) k v1 v2 v.
Mapping m k =>
(v1 -> v2 -> v) -> v1 -> v2 -> m v1 -> m v2 -> m v
prodMapWith (,)

newtype NumMap k v = NumMap (IntMap v) deriving (a -> NumMap k b -> NumMap k a
(a -> b) -> NumMap k a -> NumMap k b
(forall a b. (a -> b) -> NumMap k a -> NumMap k b)
-> (forall a b. a -> NumMap k b -> NumMap k a)
-> Functor (NumMap k)
forall a b. a -> NumMap k b -> NumMap k a
forall a b. (a -> b) -> NumMap k a -> NumMap k b
forall k a b. a -> NumMap k b -> NumMap k a
forall k a b. (a -> b) -> NumMap k a -> NumMap k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> NumMap k b -> NumMap k a
$c<$ :: forall k a b. a -> NumMap k b -> NumMap k a
fmap :: (a -> b) -> NumMap k a -> NumMap k b
$cfmap :: forall k a b. (a -> b) -> NumMap k a -> NumMap k b
Functor,a -> NumMap k a -> Bool
NumMap k m -> m
NumMap k a -> [a]
NumMap k a -> Bool
NumMap k a -> Int
NumMap k a -> a
NumMap k a -> a
NumMap k a -> a
NumMap k a -> a
(a -> m) -> NumMap k a -> m
(a -> m) -> NumMap k a -> m
(a -> b -> b) -> b -> NumMap k a -> b
(a -> b -> b) -> b -> NumMap k a -> b
(b -> a -> b) -> b -> NumMap k a -> b
(b -> a -> b) -> b -> NumMap k a -> b
(a -> a -> a) -> NumMap k a -> a
(a -> a -> a) -> NumMap k a -> a
(forall m. Monoid m => NumMap k m -> m)
-> (forall m a. Monoid m => (a -> m) -> NumMap k a -> m)
-> (forall m a. Monoid m => (a -> m) -> NumMap k a -> m)
-> (forall a b. (a -> b -> b) -> b -> NumMap k a -> b)
-> (forall a b. (a -> b -> b) -> b -> NumMap k a -> b)
-> (forall b a. (b -> a -> b) -> b -> NumMap k a -> b)
-> (forall b a. (b -> a -> b) -> b -> NumMap k a -> b)
-> (forall a. (a -> a -> a) -> NumMap k a -> a)
-> (forall a. (a -> a -> a) -> NumMap k a -> a)
-> (forall a. NumMap k a -> [a])
-> (forall a. NumMap k a -> Bool)
-> (forall a. NumMap k a -> Int)
-> (forall a. Eq a => a -> NumMap k a -> Bool)
-> (forall a. Ord a => NumMap k a -> a)
-> (forall a. Ord a => NumMap k a -> a)
-> (forall a. Num a => NumMap k a -> a)
-> (forall a. Num a => NumMap k a -> a)
-> Foldable (NumMap k)
forall a. Eq a => a -> NumMap k a -> Bool
forall a. Num a => NumMap k a -> a
forall a. Ord a => NumMap k a -> a
forall m. Monoid m => NumMap k m -> m
forall a. NumMap k a -> Bool
forall a. NumMap k a -> Int
forall a. NumMap k a -> [a]
forall a. (a -> a -> a) -> NumMap k a -> a
forall k a. Eq a => a -> NumMap k a -> Bool
forall k a. Num a => NumMap k a -> a
forall k a. Ord a => NumMap k a -> a
forall m a. Monoid m => (a -> m) -> NumMap k a -> m
forall k m. Monoid m => NumMap k m -> m
forall k a. NumMap k a -> Bool
forall k a. NumMap k a -> Int
forall k a. NumMap k a -> [a]
forall b a. (b -> a -> b) -> b -> NumMap k a -> b
forall a b. (a -> b -> b) -> b -> NumMap k a -> b
forall k a. (a -> a -> a) -> NumMap k a -> a
forall b a. (b -> a -> b) -> b -> NumMap k a -> b
forall a b. (a -> b -> b) -> b -> NumMap k a -> b
forall k m a. Monoid m => (a -> m) -> NumMap k a -> m
forall k b a. (b -> a -> b) -> b -> NumMap k a -> b
forall k a b. (a -> b -> b) -> b -> NumMap k a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: NumMap k a -> a
$cproduct :: forall k a. Num a => NumMap k a -> a
sum :: NumMap k a -> a
$csum :: forall k a. Num a => NumMap k a -> a
minimum :: NumMap k a -> a
$cminimum :: forall k a. Ord a => NumMap k a -> a
maximum :: NumMap k a -> a
$cmaximum :: forall k a. Ord a => NumMap k a -> a
elem :: a -> NumMap k a -> Bool
$celem :: forall k a. Eq a => a -> NumMap k a -> Bool
length :: NumMap k a -> Int
$clength :: forall k a. NumMap k a -> Int
null :: NumMap k a -> Bool
$cnull :: forall k a. NumMap k a -> Bool
toList :: NumMap k a -> [a]
$ctoList :: forall k a. NumMap k a -> [a]
foldl1 :: (a -> a -> a) -> NumMap k a -> a
$cfoldl1 :: forall k a. (a -> a -> a) -> NumMap k a -> a
foldr1 :: (a -> a -> a) -> NumMap k a -> a
$cfoldr1 :: forall k a. (a -> a -> a) -> NumMap k a -> a
foldl' :: (b -> a -> b) -> b -> NumMap k a -> b
$cfoldl' :: forall k b a. (b -> a -> b) -> b -> NumMap k a -> b
foldl :: (b -> a -> b) -> b -> NumMap k a -> b
$cfoldl :: forall k b a. (b -> a -> b) -> b -> NumMap k a -> b
foldr' :: (a -> b -> b) -> b -> NumMap k a -> b
$cfoldr' :: forall k a b. (a -> b -> b) -> b -> NumMap k a -> b
foldr :: (a -> b -> b) -> b -> NumMap k a -> b
$cfoldr :: forall k a b. (a -> b -> b) -> b -> NumMap k a -> b
foldMap' :: (a -> m) -> NumMap k a -> m
$cfoldMap' :: forall k m a. Monoid m => (a -> m) -> NumMap k a -> m
foldMap :: (a -> m) -> NumMap k a -> m
$cfoldMap :: forall k m a. Monoid m => (a -> m) -> NumMap k a -> m
fold :: NumMap k m -> m
$cfold :: forall k m. Monoid m => NumMap k m -> m
Foldable,Functor (NumMap k)
Foldable (NumMap k)
Functor (NumMap k)
-> Foldable (NumMap k)
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> NumMap k a -> f (NumMap k b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    NumMap k (f a) -> f (NumMap k a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> NumMap k a -> m (NumMap k b))
-> (forall (m :: * -> *) a.
    Monad m =>
    NumMap k (m a) -> m (NumMap k a))
-> Traversable (NumMap k)
(a -> f b) -> NumMap k a -> f (NumMap k b)
forall k. Functor (NumMap k)
forall k. Foldable (NumMap k)
forall k (m :: * -> *) a.
Monad m =>
NumMap k (m a) -> m (NumMap k a)
forall k (f :: * -> *) a.
Applicative f =>
NumMap k (f a) -> f (NumMap k a)
forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> NumMap k a -> m (NumMap k b)
forall k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> NumMap k a -> f (NumMap k b)
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => NumMap k (m a) -> m (NumMap k a)
forall (f :: * -> *) a.
Applicative f =>
NumMap k (f a) -> f (NumMap k a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> NumMap k a -> m (NumMap k b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> NumMap k a -> f (NumMap k b)
sequence :: NumMap k (m a) -> m (NumMap k a)
$csequence :: forall k (m :: * -> *) a.
Monad m =>
NumMap k (m a) -> m (NumMap k a)
mapM :: (a -> m b) -> NumMap k a -> m (NumMap k b)
$cmapM :: forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> NumMap k a -> m (NumMap k b)
sequenceA :: NumMap k (f a) -> f (NumMap k a)
$csequenceA :: forall k (f :: * -> *) a.
Applicative f =>
NumMap k (f a) -> f (NumMap k a)
traverse :: (a -> f b) -> NumMap k a -> f (NumMap k b)
$ctraverse :: forall k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> NumMap k a -> f (NumMap k b)
$cp2Traversable :: forall k. Foldable (NumMap k)
$cp1Traversable :: forall k. Functor (NumMap k)
Traversable)

lookupNumMap :: a -> Int -> NumMap t a -> a
lookupNumMap :: a -> Int -> NumMap t a -> a
lookupNumMap a
d Int
k (NumMap IntMap a
m) = a -> Int -> IntMap a -> a
forall a. a -> Int -> IntMap a -> a
IntMap.findWithDefault a
d Int
k IntMap a
m

lookupNumMap' :: Int -> NumMap t a -> Maybe a
lookupNumMap' :: Int -> NumMap t a -> Maybe a
lookupNumMap' Int
k (NumMap IntMap a
m) = Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
k IntMap a
m

instance Mapping (NumMap k) (Numbered k) where
    NumMap IntMap v
m1 & :: NumMap k v -> NumMap k v -> NumMap k v
& NumMap IntMap v
m2 = IntMap v -> NumMap k v
forall k v. IntMap v -> NumMap k v
NumMap (IntMap v -> IntMap v -> IntMap v
forall a. IntMap a -> IntMap a -> IntMap a
IntMap.union IntMap v
m1 IntMap v
m2)
    Numbered Int
k k
_ |-> :: Numbered k -> v -> NumMap k v
|-> v
v = IntMap v -> NumMap k v
forall k v. IntMap v -> NumMap k v
NumMap (IntMap v -> NumMap k v) -> IntMap v -> NumMap k v
forall a b. (a -> b) -> a -> b
$ Int -> v -> IntMap v
forall a. Int -> a -> IntMap a
IntMap.singleton Int
k v
v
    empty :: NumMap k v
empty = IntMap v -> NumMap k v
forall k v. IntMap v -> NumMap k v
NumMap IntMap v
forall a. IntMap a
IntMap.empty

    findWithDefault :: a -> Numbered k -> NumMap k a -> a
findWithDefault a
d (Numbered Int
i k
_) NumMap k a
m = a -> Int -> NumMap k a -> a
forall a t. a -> Int -> NumMap t a -> a
lookupNumMap a
d Int
i NumMap k a
m

    prodMapWith :: (v1 -> v2 -> v)
-> v1 -> v2 -> NumMap k v1 -> NumMap k v2 -> NumMap k v
prodMapWith v1 -> v2 -> v
f v1
p v2
q (NumMap IntMap v1
mp) (NumMap IntMap v2
mq) = IntMap v -> NumMap k v
forall k v. IntMap v -> NumMap k v
NumMap (IntMap v -> NumMap k v) -> IntMap v -> NumMap k v
forall a b. (a -> b) -> a -> b
$ (Int -> v1 -> v2 -> Maybe v)
-> (IntMap v1 -> IntMap v)
-> (IntMap v2 -> IntMap v)
-> IntMap v1
-> IntMap v2
-> IntMap v
forall a b c.
(Int -> a -> b -> Maybe c)
-> (IntMap a -> IntMap c)
-> (IntMap b -> IntMap c)
-> IntMap a
-> IntMap b
-> IntMap c
IntMap.mergeWithKey Int -> v1 -> v2 -> Maybe v
forall p. p -> v1 -> v2 -> Maybe v
merge 
                                          ((v1 -> v) -> IntMap v1 -> IntMap v
forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map (v1 -> v2 -> v
`f` v2
q)) ((v2 -> v) -> IntMap v2 -> IntMap v
forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map (v1
p v1 -> v2 -> v
`f`)) IntMap v1
mp IntMap v2
mq
      where merge :: p -> v1 -> v2 -> Maybe v
merge p
_ v1
p v2
q = v -> Maybe v
forall a. a -> Maybe a
Just (v1
p v1 -> v2 -> v
`f` v2
q)