{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Comp.Multi.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.Multi.Mapping
    ( Numbered (..)
    , unNumbered
    , number
    , HTraversable ()
    , Mapping (..)
    , lookupNumMap) where

import Data.Comp.Multi.HFunctor
import Data.Comp.Multi.HTraversable

import Control.Monad.State

import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap


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

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


-- | This function numbers the components of the given functorial
-- value with consecutive integers starting at 0.
number :: HTraversable f => f a :-> f (Numbered a)
number :: f a :-> f (Numbered a)
number x :: f a i
x = State Int (f (Numbered a) i) -> Int -> f (Numbered a) i
forall s a. State s a -> s -> a
evalState (NatM (StateT Int Identity) a (Numbered a)
-> f a i -> State Int (f (Numbered a) i)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) (a :: * -> *)
       (b :: * -> *).
(HTraversable t, Monad m) =>
NatM m a b -> NatM m (t a) (t b)
hmapM NatM (StateT Int Identity) a (Numbered a)
forall (m :: * -> *) (a :: * -> *) i.
MonadState Int m =>
a i -> m (Numbered a i)
run f a i
x) 0 where
  run :: a i -> m (Numbered a i)
run b :: a i
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
+1)
             Numbered a i -> m (Numbered a i)
forall (m :: * -> *) a. Monad m => a -> m a
return (Numbered a i -> m (Numbered a i))
-> Numbered a i -> m (Numbered a i)
forall a b. (a -> b) -> a -> b
$ Int -> a i -> Numbered a i
forall (a :: * -> *) i. Int -> a i -> Numbered a i
Numbered Int
n a i
b



infix 1 |->
infixr 0 &


class 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 i -> 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.
    prodMap :: v1 -> v2 -> m v1 -> m v2 -> m (v1, v2)

    -- | 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 i -> m a -> a


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 (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (k :: * -> *) a b. a -> NumMap k b -> NumMap k a
forall (k :: * -> *) a b. (a -> b) -> NumMap k a -> NumMap k b
<$ :: 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

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

instance Mapping (NumMap k) (Numbered k) where
    NumMap m1 :: IntMap v
m1 & :: NumMap k v -> NumMap k v -> NumMap k v
& NumMap m2 :: 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 k :: Int
k _ |-> :: Numbered k i -> v -> NumMap k v
|-> 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 i -> NumMap k a -> a
findWithDefault d :: a
d (Numbered i :: Int
i _) m :: 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

    prodMap :: v1 -> v2 -> NumMap k v1 -> NumMap k v2 -> NumMap k (v1, v2)
prodMap p :: v1
p q :: v2
q (NumMap mp :: IntMap v1
mp) (NumMap mq :: IntMap v2
mq) = IntMap (v1, v2) -> NumMap k (v1, v2)
forall (k :: * -> *) v. IntMap v -> NumMap k v
NumMap (IntMap (v1, v2) -> NumMap k (v1, v2))
-> IntMap (v1, v2) -> NumMap k (v1, v2)
forall a b. (a -> b) -> a -> b
$ (Int -> v1 -> v2 -> Maybe (v1, v2))
-> (IntMap v1 -> IntMap (v1, v2))
-> (IntMap v2 -> IntMap (v1, v2))
-> IntMap v1
-> IntMap v2
-> IntMap (v1, v2)
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 (v1, v2)
forall p a b. p -> a -> b -> Maybe (a, b)
merge 
                                          ((v1 -> (v1, v2)) -> IntMap v1 -> IntMap (v1, v2)
forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map (,v2
q)) ((v2 -> (v1, v2)) -> IntMap v2 -> IntMap (v1, v2)
forall a b. (a -> b) -> IntMap a -> IntMap b
IntMap.map (v1
p,)) IntMap v1
mp IntMap v2
mq
      where merge :: p -> a -> b -> Maybe (a, b)
merge _ p :: a
p q :: b
q = (a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just (a
p,b
q)