{-# LANGUAGE DeriveGeneric, DeriveAnyClass, MonadComprehensions, DeriveLift,
      DeriveDataTypeable #-}
{-|
  Module      : Text.ANTLR.MultiMap
  Description : A one-to-many key value map
  Copyright   : (c) Karl Cronburg, 2018
  License     : BSD3
  Maintainer  : karl@cs.tufts.edu
  Stability   : experimental
  Portability : POSIX

-}
module Text.ANTLR.MultiMap where
import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe)
import Text.ANTLR.Set (Generic(..), Hashable(..), Set(..))
import qualified Text.ANTLR.Set as S
import Prelude hiding (lookup)
import Text.ANTLR.Pretty

import Data.Data (Data(..))
import Language.Haskell.TH.Syntax (Lift(..))

instance (Lift k, Lift v, Data k, Data v, Ord k, Ord v) => Lift (M.Map k v)

-- | A multi 'Map' is a mapping from keys @k@ to sets of values @v@. A nice
--   invariant to maintain while using a multi-map is to never have empty
--   sets mapped to by some key.
newtype Map k v = Map (M.Map k (Set v))
  deriving ((forall x. Map k v -> Rep (Map k v) x)
-> (forall x. Rep (Map k v) x -> Map k v) -> Generic (Map k v)
forall x. Rep (Map k v) x -> Map k v
forall x. Map k v -> Rep (Map k v) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall k v x. Rep (Map k v) x -> Map k v
forall k v x. Map k v -> Rep (Map k v) x
$cto :: forall k v x. Rep (Map k v) x -> Map k v
$cfrom :: forall k v x. Map k v -> Rep (Map k v) x
Generic, Int -> Map k v -> Int
Map k v -> Int
(Int -> Map k v -> Int) -> (Map k v -> Int) -> Hashable (Map k v)
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
forall k v. (Hashable k, Hashable v) => Int -> Map k v -> Int
forall k v. (Hashable k, Hashable v) => Map k v -> Int
hash :: Map k v -> Int
$chash :: forall k v. (Hashable k, Hashable v) => Map k v -> Int
hashWithSalt :: Int -> Map k v -> Int
$chashWithSalt :: forall k v. (Hashable k, Hashable v) => Int -> Map k v -> Int
Hashable, Map k v -> Map k v -> Bool
(Map k v -> Map k v -> Bool)
-> (Map k v -> Map k v -> Bool) -> Eq (Map k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => Map k v -> Map k v -> Bool
/= :: Map k v -> Map k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => Map k v -> Map k v -> Bool
== :: Map k v -> Map k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => Map k v -> Map k v -> Bool
Eq, Int -> Map k v -> ShowS
[Map k v] -> ShowS
Map k v -> String
(Int -> Map k v -> ShowS)
-> (Map k v -> String) -> ([Map k v] -> ShowS) -> Show (Map k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> Map k v -> ShowS
forall k v. (Show k, Show v) => [Map k v] -> ShowS
forall k v. (Show k, Show v) => Map k v -> String
showList :: [Map k v] -> ShowS
$cshowList :: forall k v. (Show k, Show v) => [Map k v] -> ShowS
show :: Map k v -> String
$cshow :: forall k v. (Show k, Show v) => Map k v -> String
showsPrec :: Int -> Map k v -> ShowS
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> Map k v -> ShowS
Show, Map k v -> Q Exp
(Map k v -> Q Exp) -> Lift (Map k v)
forall t. (t -> Q Exp) -> Lift t
forall k v.
(Hashable v, Lift k, Lift v, Data k, Data v, Ord k, Ord v) =>
Map k v -> Q Exp
lift :: Map k v -> Q Exp
$clift :: forall k v.
(Hashable v, Lift k, Lift v, Data k, Data v, Ord k, Ord v) =>
Map k v -> Q Exp
Lift)

instance (Prettify k, Prettify v, Hashable v, Eq v) => Prettify (Map k v) where
  prettify :: Map k v -> Pretty
prettify (Map m :: Map k (Set v)
m) = Map k (Set v) -> Pretty
forall t. Prettify t => t -> Pretty
prettify Map k (Set v)
m

-- | The singleton multimap, given a single key and a __single__ value.
singleton :: (Hashable v, Eq v) => k -> v -> Map k v
singleton :: k -> v -> Map k v
singleton k :: k
k v :: v
v = Map k (Set v) -> Map k v
forall k v. Map k (Set v) -> Map k v
Map (k -> Set v -> Map k (Set v)
forall k a. k -> a -> Map k a
M.singleton k
k (v -> Set v
forall a. Hashable a => a -> HashSet a
S.singleton v
v))

-- | Construct a multi 'Map' from a list of key-value pairs.
fromList :: (Hashable v, Ord k, Eq k, Eq v) => [(k, v)] -> Map k v
fromList :: [(k, v)] -> Map k v
fromList kvs :: [(k, v)]
kvs = Map k (Set v) -> Map k v
forall k v. Map k (Set v) -> Map k v
Map ([(k, Set v)] -> Map k (Set v)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
  [ (k
k1, [v] -> Set v
forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList [v
v2 | (k2 :: k
k2, v2 :: v
v2) <- [(k, v)]
kvs, k
k1 k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k2])
  | (k1 :: k
k1, _) <- [(k, v)]
kvs])

-- | Same as 'fromList' but where the values in the key-value tuples are already in sets.
fromList' :: (Ord k, Eq k, Hashable v, Eq v) => [(k, Set v)] -> Map k v
fromList' :: [(k, Set v)] -> Map k v
fromList' kvs :: [(k, Set v)]
kvs = [(k, v)] -> Map k v
forall v k. (Hashable v, Ord k, Eq k, Eq v) => [(k, v)] -> Map k v
fromList [(k
k, v
v) | (k :: k
k, vs :: Set v
vs) <- [(k, Set v)]
kvs, v
v <- Set v -> [v]
forall a. HashSet a -> [a]
S.toList Set v
vs]

-- | Inverse of 'fromList\''.
toList :: Map k v -> [(k, Set v)]
toList :: Map k v -> [(k, Set v)]
toList (Map m :: Map k (Set v)
m) = Map k (Set v) -> [(k, Set v)]
forall k a. Map k a -> [(k, a)]
M.toList Map k (Set v)
m

-- | Take the union of two maps.
union :: (Ord k, Eq k, Hashable v, Eq v) => Map k v -> Map k v -> Map k v
union :: Map k v -> Map k v -> Map k v
union m1 :: Map k v
m1 m2 :: Map k v
m2 = [(k, Set v)] -> Map k v
forall k v.
(Ord k, Eq k, Hashable v, Eq v) =>
[(k, Set v)] -> Map k v
fromList' (Map k v -> [(k, Set v)]
forall k v. Map k v -> [(k, Set v)]
toList Map k v
m1 [(k, Set v)] -> [(k, Set v)] -> [(k, Set v)]
forall a. [a] -> [a] -> [a]
++ Map k v -> [(k, Set v)]
forall k v. Map k v -> [(k, Set v)]
toList Map k v
m2)

-- | The empty multi-map.
empty :: Map k v
empty :: Map k v
empty = Map k (Set v) -> Map k v
forall k v. Map k (Set v) -> Map k v
Map Map k (Set v)
forall k a. Map k a
M.empty

-- | Get the set of values mapped to by some key @k@.
lookup :: (Ord k, Hashable v, Eq v) => k -> Map k v -> Set v
lookup :: k -> Map k v -> Set v
lookup k :: k
k (Map m :: Map k (Set v)
m) = Set v -> Maybe (Set v) -> Set v
forall a. a -> Maybe a -> a
fromMaybe Set v
forall a. HashSet a
S.empty (k -> Map k (Set v) -> Maybe (Set v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k Map k (Set v)
m)

-- | Number of keys in the multi-map.
size :: Map k v -> Int
size (Map m :: Map k (Set v)
m) = Map k (Set v) -> Int
forall k a. Map k a -> Int
M.size Map k (Set v)
m

-- | Map difference of two multi-maps, deleting individual key-value pairs
--   rather than deleting the entire key. Invariant maintained is that
--   input maps with non-null value sets will result in an output with
--   non-null value sets.
difference :: Map k v -> Map k v -> Map k v
difference (Map m1 :: Map k (Set v)
m1) m2 :: Map k v
m2 = Map k (Set v) -> Map k v
forall k v. Map k (Set v) -> Map k v
Map (Map k (Set v) -> Map k v) -> Map k (Set v) -> Map k v
forall a b. (a -> b) -> a -> b
$ [(k, Set v)] -> Map k (Set v)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
  [ (k
k1, Set v
vs)
  | (k1 :: k
k1, vs1 :: Set v
vs1) <- Map k (Set v) -> [(k, Set v)]
forall k a. Map k a -> [(k, a)]
M.toList Map k (Set v)
m1
  , let vs2 :: Set v
vs2 = k
k1 k -> Map k v -> Set v
forall k v. (Ord k, Hashable v, Eq v) => k -> Map k v -> Set v
`lookup` Map k v
m2
  , let vs :: Set v
vs  = Set v
vs1 Set v -> Set v -> Set v
forall a. (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a
`S.difference` Set v
vs2
  , (Bool -> Bool
not (Bool -> Bool) -> (Set v -> Bool) -> Set v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set v -> Bool
forall a. HashSet a -> Bool
S.null) Set v
vs
  ]