content: {-# LANGUAGE DeriveFunctor #-}
  {-# LANGUAGE DeriveTraversable #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TupleSections #-}
  module Data.Abstract.Map (
    Map,singleton,empty,lookup,unsafeLookup, insert,
    insertWith,unsafeInsertWith,unsafeInsertWithLookup,
    delete,delete',union,unsafeUnion,adjust,toList,
    fromList,mapMaybe,map,compose,widening,fromThereList
  ) where
  
  import           Prelude hiding (lookup,map,Either(..),(**))
  
  import           Control.Arrow hiding ((<+>))
  
  import           Data.HashMap.Lazy (HashMap)
  import qualified Data.HashMap.Lazy as H
  import           Data.Hashable
  import           Data.Order
  import           Data.Identifiable
  import           Data.Maybe (fromJust)
  import           Data.Empty
  import           Prettyprinter
  
  import qualified Data.Abstract.Maybe as M
  import           Data.Abstract.There (There(..))
  import qualified Data.Abstract.There as T
  import           Data.Abstract.Widening
  import           Data.Abstract.Stable
  
  import           GHC.Exts
  
  import           Text.Printf
  import           Control.DeepSeq
  
  -- | Abstract hashmap
  newtype Map a b = Map (HashMap a (There,b)) deriving (Eq,Functor,Foldable,Traversable,Hashable,NFData)
  
  instance (Show a,Show b) => Show (Map a b) where
    show (Map h)
      | H.null h = "[]"
      | otherwise = "[" ++ init (unwords [ printf "%s%s -> %s," (show k) (show t) (show v) | (k,(t,v)) <- H.toList h]) ++ "]"
  
  instance (Pretty k, Pretty v) => Pretty (Map k v) where
    pretty (Map m) = list [ pretty k <> pretty t <+> " -> " <> pretty v | (k, (t,v)) <- H.toList m ]
  
  
  instance (Identifiable a, PreOrd b) => PreOrd (Map a b) where
    m1 ⊑ Map m2 = all (\(k,(t,v)) -> lookup k m1 ⊑ (case t of Must -> M.Just v; May -> M.JustNothing v))
                      (H.toList m2)
  
  instance (Identifiable a, Complete b) => Complete (Map a b) where
    (⊔) = toJoin1 widening (⊔)
  
  widening :: Identifiable a => Widening b -> Widening (Map a b)
  widening w (Map m1) (Map m2)
    | H.keysSet m1 == H.keysSet m2 = second Map $ sequence (H.intersectionWith (T.widening ** w) m1 m2)
    | otherwise = (Unstable, Map $
        -- this would be much more efficient with a Data.Map.merge operation.
        H.intersectionWith (\(t1,b1) (t2,b2) -> (t1 ⊔ t2, snd (w b1 b2))) m1 m2
          `H.union`
          H.map (\(_,b) -> (May,b)) (m1 `H.difference` m2)
          `H.union`
          H.map (\(_,b) -> (May,b)) (m2 `H.difference` m1)
        )
  {-# INLINE widening #-}
  
  instance (Identifiable a, PreOrd b) => LowerBounded (Map a b) where
    bottom = empty
  
  instance IsEmpty (Map a b) where
    empty = Map H.empty
  
  singleton :: Identifiable a => a -> b -> Map a b
  singleton a b = Map (H.singleton a (Must,b))
  {-# INLINE singleton #-}
  
  map :: (b -> b') -> Map a b -> Map a b'
  map f (Map h) = Map (H.map (second f) h)
  {-# INLINE map #-}
  
  mapMaybe :: (Identifiable a', Complete b') => ((a,b) -> Maybe (a',b')) -> Map a b -> Map a' b'
  mapMaybe f (Map h) = Map (H.fromListWith (⊔) [ (a',(here,b')) | (a,(here,b)) <- H.toList h, Just (a',b') <- return (f (a,b))])
  {-# INLINE mapMaybe #-}
  
  lookup :: (Identifiable a) => a -> Map a b -> M.Maybe b
  lookup a (Map m) = case H.lookup a m of
    Just (Must,b) -> M.Just b
    Just (May,b) -> M.JustNothing b
    Nothing -> M.Nothing
  {-# INLINE lookup #-}
  
  unsafeLookup :: (Identifiable a) => a -> Map a b -> Maybe b
  unsafeLookup a (Map m) = case H.lookup a m of
    Just (_,b) -> Just b
    Nothing -> Nothing
  {-# INLINE unsafeLookup #-}
  
  insert :: Identifiable a => a -> b -> Map a b -> Map a b
  insert a b (Map m) = Map (H.insert a (Must,b) m)
  {-# INLINE insert #-}
  
  insertWith :: (Identifiable a,Complete b) => (b -> b -> b) -> a -> b -> Map a b -> Map a b
  insertWith f a b (Map m) = Map (H.insertWith (\(_,new) (here,old) -> case here of
      Must -> (Must,f new old)
      May -> (Must,new ⊔ f new old)
    ) a (Must,b) m)
  {-# INLINE insertWith #-}
  
  unsafeInsertWith :: (Identifiable a) => (b -> b -> b) -> a -> b -> Map a b -> Map a b
  unsafeInsertWith f a b (Map m) = Map (H.insertWith (\(_,new) (_,old) -> (Must,f new old)) a (Must,b) m)
  {-# INLINE unsafeInsertWith #-}
  
  unsafeInsertWithLookup :: (Identifiable a) => (b -> b -> b) -> a -> b -> Map a b -> (b,Map a b)
  unsafeInsertWithLookup f a b m =
    let m' = unsafeInsertWith f a b m
    in (fromJust (unsafeLookup a m'),m')
  {-# INLINE unsafeInsertWithLookup #-}
  
  delete :: Identifiable a => a -> Map a b -> Map a b
  delete a (Map m) = Map (H.delete a m)
  {-# INLINE delete #-}
  
  delete' :: (Foldable f, Identifiable a) => f a -> Map a b -> Map a b
  delete' l m = foldl (flip delete) m l
  {-# INLINE delete' #-}
  
  union :: (Identifiable a, Complete b) => Map a b -> Map a b -> Map a b
  union (Map m1) (Map m2) = Map (H.unionWith (\(here,l) (_,r) -> case here of
      Must -> (Must,l)
      May -> (May,l ⊔ r)
    ) m1 m2)
  {-# INLINE union #-}
  
  unsafeUnion :: (Identifiable a) => Map a b -> Map a b -> Map a b
  unsafeUnion (Map m1) (Map m2) = Map (H.union m1 m2)
  {-# INLINE unsafeUnion #-}
  
  adjust :: Identifiable a => (b -> b) -> a -> Map a b -> Map a b
  adjust f a (Map m) = Map (H.adjust (second f) a m)
  {-# INLINE adjust #-}
  
  compose :: (Identifiable a, Identifiable b, Complete c) => [(a,b)] -> Map b c -> Map a c
  compose f (Map g) = Map $ H.fromListWith (⊔) [ (a,c) | (a,b) <- f, Just c <- return $ H.lookup b g ]
  {-# INLINE compose #-}
  
  instance Identifiable a => IsList (Map a b) where
    type Item (Map a b) = (a,b)
    toList (Map m) = H.toList (H.map snd m)
    fromList l = Map (H.fromList (fmap (second (Must,)) l))
  
  fromThereList :: Identifiable a => [(a,(There,b))] -> Map a b
  fromThereList = Map . H.fromList
  {-# INLINE fromThereList #-}
  