map.hs

{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE DeriveFoldable    #-}
{-# LANGUAGE DeriveGeneric     #-}
{-# LANGUAGE DeriveTraversable #-}

module Data.Diet.Map where

import qualified Data.Sequence as Seq
import Data.Sequence (Seq, ViewL(..), (<|), (><))

import qualified Data.Interval as I
import Data.Interval (Interval(I))
import Data.Relation
import Data.Semiring

import Prelude hiding (null, (+), foldl, foldr)
import GHC.Generics
import qualified Prelude (null)
import Control.Applicative ((<$>))
import Control.Monad.Fix
import Data.Foldable (Foldable(foldMap))
import Data.Functor.Classes
import Data.Monoid
import Data.Traversable (Traversable(traverse))
import Data.Semigroup
import Data.Diet.Node
import qualified Data.Foldable as Foldable

data Map v a
  = Empty
  | LF !(Interval v) a
  | BR !(Interval v) a (Map v a) (Map v a)

instance (Enum v, Ord v, Semiring v, Eq a) => Eq (Map v a) where
  m1 == m2 = toList m1 == toList m2 
  
instance (Enum v, Ord v, Show v, Show a) => Show (Map v a) where
  showsPrec d m = showParen (d > 10) $
    showString "fromList ". shows (toList m)

instance (Enum v, Ord v) => Foldable.Foldable (Map v) where
  foldr = foldr

toList :: (Enum k, Ord k) => Map k v -> [(Interval k, v)]
toList = foldrWithKey (\k x xs -> (k,x):xs) []

null :: Map k v -> Bool
null Empty = True
null _     = False

empty :: Map k v
empty = Empty

--singleton :: Interval v -> a -> Map v a
--singleton i x = Map $ Seq.singleton $ Node i x

atNode :: Map k v -> Interval k
atNode Empty        = I.Empty
atNode (LF i _)     = i
atNode (BR i _ _ _) = i

insert
  :: (Enum v, Ord v, Semiring v, Semiring a)
  => Interval v
  -> a
  -> Map v a
  -> Map v a
insert i x lf@(LF i' x') = if I.valid i
  then if I.mergeable i i'
    then LF (i <> i') (x + x')
    else lf -- wrong
  else lf -- wrong
insert i@(I a b) x br@(BR i' x' l r) = if I.invalid i
  then br
  else if I.mergeable i i'
    then if I.mergeable i li
      then insert i x l
      else if I.mergeable i ri
        then insert i x l
        else br
    else br
  where
    li = atNode l
    ri = atNode r

-- map from 'Interval Int' to 'Int'
mapTest :: Map Int Int
mapTest = insert (I.interval 3 10) 3 empty

foldr :: (Enum k, Ord k) => (v -> b -> b) -> b -> Map k v -> b
foldr _ z (LF _ _) = z
foldr f z (BR _ v l r) = foldr f (f v (foldr f z r)) l

foldl :: (Enum k, Ord k) => (b -> v -> b) -> b -> Map k v -> b
foldl _ z (LF _ _) = z
foldl f z (BR _ v l r) = foldl f (f (foldl f z l) v) r

foldrWithKey :: (Enum k, Ord k) => (Interval k -> v -> a -> a) -> a -> Map k v -> a
foldrWithKey _ z (LF _ _) = z
foldrWithKey f z (BR k v l r) = foldrWithKey f (f k v (foldrWithKey f z r)) l

foldlWithKey :: (Enum k, Ord k) => (a -> Interval k -> v -> a) -> a -> Map k v -> a
foldlWithKey _ z (LF _ _) = z
foldlWithKey f z (BR k v l r) = foldlWithKey f (f (foldlWithKey f z l) k v) r

foldr' :: (Enum k, Ord k) => (a -> b -> b) -> b -> Map k a -> b
foldr' _ !z (LF _ _) = z
foldr' f !z (BR _ v l r) = foldr' f (f v (foldr' f z r)) l

foldl' :: (Enum k, Ord k) => (b -> a -> b) -> b -> Map k a -> b
foldl' _ !z (LF _ _) = z
foldl' f !z (BR k v l r) = foldl' f (f (foldl' f z l) v) r

foldrWithKey' :: (Enum k, Ord k) => (Interval k -> v -> a -> a) -> a -> Map k v -> a
foldrWithKey' _ !z (LF _ _) = z
foldrWithKey' f !z (BR k v l r) = foldrWithKey' f (f k v (foldrWithKey' f z r)) l

foldlWithKey' :: (Enum k, Ord k) => (a -> Interval k -> v -> a) -> a -> Map k v -> a
foldlWithKey' _ !z (LF _ _) = z
foldlWithKey' f !z (BR k v l r) = foldlWithKey' f (f (foldlWithKey' f z l) k v) r

foldMapWithKey :: (Enum k, Ord k, Semiring k, Monoid m, Eq v, Semiring v) => (Interval k -> v -> m) -> Map k v -> m
foldMapWithKey f = go
  where
    go (LF _ _) = mempty
    go br@(BR k v l r) = if null br
      then f k v
      else go l `mappend` (f k v `mappend` go r)