{-# LANGUAGE BangPatterns #-}

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

import Control.Monad (filterM)
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 as a cache of file descriptors.
--   Since multiple threads could open file descriptors for
--   the same file simultaneously, there could be multiple entries
--   for one file.
--   Since hash values of file paths are used as outer keys,
--   collison would happen for multiple file paths.
--   Because only positive entries are stored,
--   Malicious attack cannot cause the inner list to blow up.
--   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)]
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)] -> MultiMap v)
-> IntMap [(FilePath, v)] -> MultiMap v
forall a b. (a -> b) -> a -> b
$ Key -> [(FilePath, v)] -> IntMap [(FilePath, v)]
forall a. Key -> a -> IntMap a
I.singleton (FilePath -> Key
forall a. Hashable a => a -> Key
hash FilePath
path) [(FilePath
path,v
v)]

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

-- | O(M) where M is the number of entries per file
lookup :: FilePath -> MultiMap v -> Maybe v
lookup :: FilePath -> MultiMap v -> Maybe v
lookup FilePath
path (MultiMap IntMap [(FilePath, v)]
mm) = case Key -> IntMap [(FilePath, v)] -> Maybe [(FilePath, v)]
forall a. Key -> IntMap a -> Maybe a
I.lookup (FilePath -> Key
forall a. Hashable a => a -> Key
hash FilePath
path) 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

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

-- | 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)] -> MultiMap v)
-> IntMap [(FilePath, v)] -> MultiMap v
forall a b. (a -> b) -> a -> b
$ ([(FilePath, v)] -> [(FilePath, v)] -> [(FilePath, v)])
-> Key
-> [(FilePath, v)]
-> IntMap [(FilePath, v)]
-> IntMap [(FilePath, v)]
forall a. (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
I.insertWith [(FilePath, v)] -> [(FilePath, v)] -> [(FilePath, v)]
forall a. Semigroup a => a -> a -> a
(<>) (FilePath -> Key
forall a. Hashable a => a -> Key
hash FilePath
path) [(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) = ((Key, [(FilePath, v)]) -> [(FilePath, v)])
-> [(Key, [(FilePath, v)])] -> [(FilePath, v)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Key, [(FilePath, v)]) -> [(FilePath, v)]
forall a b. (a, b) -> b
snd ([(Key, [(FilePath, v)])] -> [(FilePath, v)])
-> [(Key, [(FilePath, v)])] -> [(FilePath, v)]
forall a b. (a -> b) -> a -> b
$ IntMap [(FilePath, v)] -> [(Key, [(FilePath, v)])]
forall a. IntMap a -> [(Key, 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
  = (Key
 -> [(FilePath, v)]
 -> (IntMap [(FilePath, v)] -> IO (MultiMap v))
 -> IntMap [(FilePath, v)]
 -> IO (MultiMap v))
-> (IntMap [(FilePath, v)] -> IO (MultiMap v))
-> IntMap [(FilePath, v)]
-> IntMap [(FilePath, v)]
-> IO (MultiMap v)
forall a b. (Key -> a -> b -> b) -> b -> IntMap a -> b
I.foldrWithKey Key
-> [(FilePath, v)]
-> (IntMap [(FilePath, v)] -> IO (MultiMap v))
-> IntMap [(FilePath, v)]
-> IO (MultiMap v)
forall b.
Key
-> [(FilePath, v)]
-> (IntMap [(FilePath, v)] -> IO b)
-> IntMap [(FilePath, v)]
-> IO b
go (MultiMap v -> IO (MultiMap v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MultiMap v -> IO (MultiMap v))
-> (IntMap [(FilePath, v)] -> MultiMap v)
-> IntMap [(FilePath, v)]
-> IO (MultiMap v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap [(FilePath, v)] -> MultiMap v
forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap) IntMap [(FilePath, v)]
mm IntMap [(FilePath, v)]
forall a. IntMap a
I.empty
  where
    go :: Key
-> [(FilePath, v)]
-> (IntMap [(FilePath, v)] -> IO b)
-> IntMap [(FilePath, v)]
-> IO b
go Key
h [(FilePath, v)]
s IntMap [(FilePath, v)] -> IO b
cont IntMap [(FilePath, v)]
acc = do
      [(FilePath, v)]
rs <- ((FilePath, v) -> IO Bool) -> [(FilePath, v)] -> IO [(FilePath, v)]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM (FilePath, v) -> IO Bool
action [(FilePath, v)]
s
      case [(FilePath, v)]
rs of
        [] -> IntMap [(FilePath, v)] -> IO b
cont IntMap [(FilePath, v)]
acc
        [(FilePath, v)]
_  -> IntMap [(FilePath, v)] -> IO b
cont (IntMap [(FilePath, v)] -> IO b) -> IntMap [(FilePath, v)] -> IO b
forall a b. (a -> b) -> a -> b
$! Key
-> [(FilePath, v)]
-> IntMap [(FilePath, v)]
-> IntMap [(FilePath, v)]
forall a. Key -> a -> IntMap a -> IntMap a
I.insert Key
h [(FilePath, v)]
rs IntMap [(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)] -> MultiMap v)
-> IntMap [(FilePath, v)] -> MultiMap v
forall a b. (a -> b) -> a -> b
$ ([(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