{-# LANGUAGE BangPatterns #-}

module Network.Wai.Handler.Warp.MultiMap (
    MultiMap
  , isEmpty
  , empty
  , singleton
  , insert
  , Network.Wai.Handler.Warp.MultiMap.lookup
  , pruneWith
  , toList
  , merge
  ) where

import Data.Hashable (hash)
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as I
import Data.Semigroup
import Prelude -- Silence redundant import warnings

----------------------------------------------------------------

-- | 'MultiMap' is used for cache of file descriptors.
--   Since multiple threads would open file descriptors for
--   the same file simultaneously, multiple entries must
--   be contained for the file.
--   Since hash values of file pathes are used as outer keys,
--   collison would happen for multiple file pathes.
--   Becase only positive entries are contained,
--   a bad guy cannot be cause the hash collision intentinally.
--   So, lists are good enough.
newtype MultiMap v = MultiMap (IntMap [(FilePath,v)])

----------------------------------------------------------------

-- | O(1)
empty :: MultiMap v
empty :: MultiMap v
empty = IntMap [(FilePath, v)] -> MultiMap v
forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap (IntMap [(FilePath, v)] -> MultiMap v)
-> IntMap [(FilePath, v)] -> MultiMap v
forall a b. (a -> b) -> a -> b
$ IntMap [(FilePath, v)]
forall a. IntMap a
I.empty

-- | O(1)
isEmpty :: MultiMap v -> Bool
isEmpty :: MultiMap v -> Bool
isEmpty (MultiMap IntMap [(FilePath, v)]
mm) = IntMap [(FilePath, v)] -> Bool
forall a. IntMap a -> Bool
I.null IntMap [(FilePath, v)]
mm

----------------------------------------------------------------

-- | O(1)
singleton :: FilePath -> v -> MultiMap v
singleton :: FilePath -> v -> MultiMap v
singleton FilePath
path v
v = IntMap [(FilePath, v)] -> MultiMap v
forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap IntMap [(FilePath, v)]
mm
  where
    !h :: Int
h = FilePath -> Int
forall a. Hashable a => a -> Int
hash FilePath
path
    !mm :: IntMap [(FilePath, v)]
mm = Int -> [(FilePath, v)] -> IntMap [(FilePath, v)]
forall a. Int -> a -> IntMap a
I.singleton Int
h [(FilePath
path,v
v)]

----------------------------------------------------------------

-- | O(N)
lookup :: FilePath -> MultiMap v -> Maybe v
lookup :: FilePath -> MultiMap v -> Maybe v
lookup FilePath
path (MultiMap IntMap [(FilePath, v)]
mm) = case Int -> IntMap [(FilePath, v)] -> Maybe [(FilePath, v)]
forall a. Int -> IntMap a -> Maybe a
I.lookup Int
h IntMap [(FilePath, v)]
mm of
    Maybe [(FilePath, v)]
Nothing -> Maybe v
forall a. Maybe a
Nothing
    Just [(FilePath, v)]
s  -> FilePath -> [(FilePath, v)] -> Maybe v
forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup FilePath
path [(FilePath, v)]
s
  where
    !h :: Int
h = FilePath -> Int
forall a. Hashable a => a -> Int
hash FilePath
path

----------------------------------------------------------------

-- | O(log n)
insert :: FilePath -> v -> MultiMap v -> MultiMap v
insert :: FilePath -> v -> MultiMap v -> MultiMap v
insert FilePath
path v
v (MultiMap IntMap [(FilePath, v)]
mm) = IntMap [(FilePath, v)] -> MultiMap v
forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap IntMap [(FilePath, v)]
mm'
  where
    !h :: Int
h = FilePath -> Int
forall a. Hashable a => a -> Int
hash FilePath
path
    !mm' :: IntMap [(FilePath, v)]
mm' = ([(FilePath, v)] -> [(FilePath, v)] -> [(FilePath, v)])
-> Int
-> [(FilePath, v)]
-> IntMap [(FilePath, v)]
-> IntMap [(FilePath, v)]
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
I.insertWith [(FilePath, v)] -> [(FilePath, v)] -> [(FilePath, v)]
forall a. Semigroup a => a -> a -> a
(<>) Int
h [(FilePath
path,v
v)] IntMap [(FilePath, v)]
mm

----------------------------------------------------------------

-- | O(n)
toList :: MultiMap v -> [(FilePath,v)]
toList :: MultiMap v -> [(FilePath, v)]
toList (MultiMap IntMap [(FilePath, v)]
mm) = ((Int, [(FilePath, v)]) -> [(FilePath, v)])
-> [(Int, [(FilePath, v)])] -> [(FilePath, v)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Int, [(FilePath, v)]) -> [(FilePath, v)]
forall a b. (a, b) -> b
snd ([(Int, [(FilePath, v)])] -> [(FilePath, v)])
-> [(Int, [(FilePath, v)])] -> [(FilePath, v)]
forall a b. (a -> b) -> a -> b
$ IntMap [(FilePath, v)] -> [(Int, [(FilePath, v)])]
forall a. IntMap a -> [(Int, a)]
I.toAscList IntMap [(FilePath, v)]
mm

----------------------------------------------------------------

-- | O(n)
pruneWith :: MultiMap v
          -> ((FilePath,v) -> IO Bool)
          -> IO (MultiMap v)
pruneWith :: MultiMap v -> ((FilePath, v) -> IO Bool) -> IO (MultiMap v)
pruneWith (MultiMap IntMap [(FilePath, v)]
mm) (FilePath, v) -> IO Bool
action = IntMap [(FilePath, v)] -> MultiMap v
forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap (IntMap [(FilePath, v)] -> MultiMap v)
-> IO (IntMap [(FilePath, v)]) -> IO (MultiMap v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (IntMap [(FilePath, v)])
mm'
  where
    !mm' :: IO (IntMap [(FilePath, v)])
mm' = [(Int, [(FilePath, v)])] -> IntMap [(FilePath, v)]
forall a. [(Int, a)] -> IntMap a
I.fromAscList ([(Int, [(FilePath, v)])] -> IntMap [(FilePath, v)])
-> IO [(Int, [(FilePath, v)])] -> IO (IntMap [(FilePath, v)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, [(FilePath, v)])]
-> [(Int, [(FilePath, v)])] -> IO [(Int, [(FilePath, v)])]
forall a.
[(a, [(FilePath, v)])]
-> [(a, [(FilePath, v)])] -> IO [(a, [(FilePath, v)])]
go (IntMap [(FilePath, v)] -> [(Int, [(FilePath, v)])]
forall a. IntMap a -> [(Int, a)]
I.toDescList IntMap [(FilePath, v)]
mm) []
    go :: [(a, [(FilePath, v)])]
-> [(a, [(FilePath, v)])] -> IO [(a, [(FilePath, v)])]
go []          ![(a, [(FilePath, v)])]
acc = [(a, [(FilePath, v)])] -> IO [(a, [(FilePath, v)])]
forall (m :: * -> *) a. Monad m => a -> m a
return [(a, [(FilePath, v)])]
acc
    go ((a
h,[(FilePath, v)]
s):[(a, [(FilePath, v)])]
kss) ![(a, [(FilePath, v)])]
acc = do
        [(FilePath, v)]
rs <- ((FilePath, v) -> IO Bool) -> [(FilePath, v)] -> IO [(FilePath, v)]
forall v.
((FilePath, v) -> IO Bool) -> [(FilePath, v)] -> IO [(FilePath, v)]
prune (FilePath, v) -> IO Bool
action [(FilePath, v)]
s
        case [(FilePath, v)]
rs of
            [] -> [(a, [(FilePath, v)])]
-> [(a, [(FilePath, v)])] -> IO [(a, [(FilePath, v)])]
go [(a, [(FilePath, v)])]
kss [(a, [(FilePath, v)])]
acc
            [(FilePath, v)]
_  -> [(a, [(FilePath, v)])]
-> [(a, [(FilePath, v)])] -> IO [(a, [(FilePath, v)])]
go [(a, [(FilePath, v)])]
kss ((a
h,[(FilePath, v)]
rs) (a, [(FilePath, v)])
-> [(a, [(FilePath, v)])] -> [(a, [(FilePath, v)])]
forall a. a -> [a] -> [a]
: [(a, [(FilePath, v)])]
acc)

----------------------------------------------------------------

-- O(n + m) where N is the size of the second argument
merge :: MultiMap v -> MultiMap v -> MultiMap v
merge :: MultiMap v -> MultiMap v -> MultiMap v
merge (MultiMap IntMap [(FilePath, v)]
m1) (MultiMap IntMap [(FilePath, v)]
m2) = IntMap [(FilePath, v)] -> MultiMap v
forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap IntMap [(FilePath, v)]
mm
  where
    !mm :: IntMap [(FilePath, v)]
mm = ([(FilePath, v)] -> [(FilePath, v)] -> [(FilePath, v)])
-> IntMap [(FilePath, v)]
-> IntMap [(FilePath, v)]
-> IntMap [(FilePath, v)]
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
I.unionWith [(FilePath, v)] -> [(FilePath, v)] -> [(FilePath, v)]
forall a. Semigroup a => a -> a -> a
(<>) IntMap [(FilePath, v)]
m1 IntMap [(FilePath, v)]
m2

----------------------------------------------------------------

prune :: ((FilePath,v) -> IO Bool) -> [(FilePath,v)] -> IO [(FilePath,v)]
prune :: ((FilePath, v) -> IO Bool) -> [(FilePath, v)] -> IO [(FilePath, v)]
prune (FilePath, v) -> IO Bool
action [(FilePath, v)]
xs0 = [(FilePath, v)] -> IO [(FilePath, v)]
go [(FilePath, v)]
xs0
  where
    go :: [(FilePath, v)] -> IO [(FilePath, v)]
go []     = [(FilePath, v)] -> IO [(FilePath, v)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
    go ((FilePath, v)
x:[(FilePath, v)]
xs) = do
        Bool
keep <- (FilePath, v) -> IO Bool
action (FilePath, v)
x
        [(FilePath, v)]
rs <- [(FilePath, v)] -> IO [(FilePath, v)]
go [(FilePath, v)]
xs
        [(FilePath, v)] -> IO [(FilePath, v)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(FilePath, v)] -> IO [(FilePath, v)])
-> [(FilePath, v)] -> IO [(FilePath, v)]
forall a b. (a -> b) -> a -> b
$ if Bool
keep then (FilePath, v)
x(FilePath, v) -> [(FilePath, v)] -> [(FilePath, v)]
forall a. a -> [a] -> [a]
:[(FilePath, v)]
rs else [(FilePath, v)]
rs