{-# 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 :: forall v. MultiMap v
empty = forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap forall a. IntMap a
I.empty

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

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

-- | O(1)
singleton :: FilePath -> v -> MultiMap v
singleton :: forall v. FilePath -> v -> MultiMap v
singleton FilePath
path v
v = forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap forall a b. (a -> b) -> a -> b
$ forall a. Key -> a -> IntMap a
I.singleton (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 :: forall v. FilePath -> MultiMap v -> Maybe v
lookup FilePath
path (MultiMap IntMap [(FilePath, v)]
mm) = case forall a. Key -> IntMap a -> Maybe a
I.lookup (forall a. Hashable a => a -> Key
hash FilePath
path) IntMap [(FilePath, v)]
mm of
    Maybe [(FilePath, v)]
Nothing -> forall a. Maybe a
Nothing
    Just [(FilePath, v)]
s  -> 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 :: forall v. FilePath -> v -> MultiMap v -> MultiMap v
insert FilePath
path v
v (MultiMap IntMap [(FilePath, v)]
mm) = forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap
  forall a b. (a -> b) -> a -> b
$ forall a. (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a
I.insertWith forall a. Semigroup a => a -> a -> a
(<>) (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 :: forall v. MultiMap v -> [(FilePath, v)]
toList (MultiMap IntMap [(FilePath, v)]
mm) = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ 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 :: forall v.
MultiMap v -> ((FilePath, v) -> IO Bool) -> IO (MultiMap v)
pruneWith (MultiMap IntMap [(FilePath, v)]
mm) (FilePath, v) -> IO Bool
action
  = forall a b. (Key -> a -> b -> b) -> b -> IntMap a -> b
I.foldrWithKey forall {b}.
Key
-> [(FilePath, v)]
-> (IntMap [(FilePath, v)] -> IO b)
-> IntMap [(FilePath, v)]
-> IO b
go (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap) IntMap [(FilePath, v)]
mm 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 <- 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 forall a b. (a -> b) -> a -> b
$! 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 :: forall v. MultiMap v -> MultiMap v -> MultiMap v
merge (MultiMap IntMap [(FilePath, v)]
m1) (MultiMap IntMap [(FilePath, v)]
m2) = forall v. IntMap [(FilePath, v)] -> MultiMap v
MultiMap forall a b. (a -> b) -> a -> b
$ forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
I.unionWith forall a. Semigroup a => a -> a -> a
(<>) IntMap [(FilePath, v)]
m1 IntMap [(FilePath, v)]
m2