{-# 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 (Generic, Hashable, Eq, Show, Lift)

instance (Prettify k, Prettify v, Hashable v, Eq v) => Prettify (Map k v) where
  prettify (Map m) = prettify 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 (M.singleton k (S.singleton 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 kvs = Map (M.fromList
  [ (k1, S.fromList [v2 | (k2, v2) <- kvs, k1 == k2])
  | (k1, _) <- 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' kvs = fromList [(k, v) | (k, vs) <- kvs, v <- S.toList vs]

-- | Inverse of 'fromList\''.
toList :: Map k v -> [(k, Set v)]
toList (Map m) = M.toList 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 m1 m2 = fromList' (toList m1 ++ toList m2)

-- | The empty multi-map.
empty :: Map k v
empty = Map 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 m) = fromMaybe S.empty (M.lookup k m)

-- | Number of keys in the multi-map.
size (Map m) = M.size 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 m1) m2 = Map $ M.fromList
  [ (k1, vs)
  | (k1, vs1) <- M.toList m1
  , let vs2 = k1 `lookup` m2
  , let vs  = vs1 `S.difference` vs2
  , (not . S.null) vs
  ]