content: {-# LANGUAGE Arrows #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE ImplicitParams #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE UnboxedTuples #-}
  {-# LANGUAGE UndecidableInstances #-}
  module Control.Arrow.Transformer.Abstract.Fix.Cache.Immutable where
  
  import           Prelude hiding (pred,lookup,map,head,iterate,(.),id,truncate,elem,product,(**))
  
  import           Control.Category
  import           Control.Arrow hiding ((<+>))
  import           Control.Arrow.Primitive
  import           Control.Arrow.Strict
  import           Control.Arrow.Trans
  import           Control.Arrow.State
  import           Control.Arrow.Fix.Reuse(ArrowReuse(..))
  import           Control.Arrow.Fix.CallCount as CallCount
  import           Control.Arrow.Fix.ControlFlow as ControlFlow
  import           Control.Arrow.Fix.Context as Context
  import           Control.Arrow.Fix.Cache as Cache
  import           Control.Arrow.Order (ArrowLowerBounded)
  import qualified Control.Arrow.Order as Order
  import           Control.Arrow.Transformer.State
  
  import           Data.Profunctor.Unsafe
  import           Data.Empty
  import           Data.Order hiding (lub)
  import           Data.Coerce
  import           Data.Identifiable
  import           Data.Hashable(hash)
  import           Data.HashMap.Strict (HashMap)
  import qualified Data.HashMap.Strict as M
  import           Data.Monoidal
  import           Data.Maybe
  import           Prettyprinter
  
  import           Data.Abstract.Stable
  import qualified Data.Abstract.Widening as W
  
  import           GHC.Exts
  
  newtype CacheT cache a b c x y = CacheT { unCacheT :: StateT (cache a b) c x y}
    deriving (Profunctor,Category,Arrow,ArrowChoice,ArrowStrict,ArrowTrans,
              ArrowState (cache a b),ArrowControlFlow stmt, ArrowPrimitive, ArrowCFG graph,
              ArrowCallCount a', ArrowCallSite ctx, ArrowContext ctx a')
  
  instance (IsEmpty (cache a b), ArrowRun c) => ArrowRun (CacheT cache a b c) where
    type Run (CacheT cache a b c) x y = Run c x (cache a b,y)
    run f = run (lmap (\x -> (empty, x)) (unlift f))
    {-# INLINE run #-}
  
  instance ArrowLift (CacheT cache a b c) where
    type Underlying (CacheT cache a b c) x y = c (cache a b, x) (cache a b, y)
  
  instance (Arrow c, Profunctor c) => ArrowGetCache (cache a b) (CacheT cache a b c) where
    getCache = CacheT get
  
  instance (Profunctor c,ArrowApply c) => ArrowApply (CacheT cache a b c) where
    app = CacheT (app .# first coerce)
    {-# INLINE app #-}
  
  ----- Basic Cache -----
  newtype Cache a b = Cache { getMap :: HashMap a (Stable,b)}
  
  instance IsEmpty (Cache a b) where
    empty = Cache M.empty
  
  instance (Show a, Show b) => Show (Cache a b) where
    show (Cache m) = show (M.toList m)
  
  instance (Identifiable a, Pretty a, Identifiable b, Pretty b) => Pretty (Cache a b) where
    pretty (Cache m) = list [ pretty k<> "@" <> pretty (hash k) <+>
                              prettyStable s <+>
                              pretty v <> "@" <> pretty (hash v)
                            | (k,(s,v)) <- M.toList m]
      where
        prettyStable Stable = "->"
        prettyStable Unstable = "~>"
  
  instance (Identifiable a, LowerBounded b, ArrowChoice c, Profunctor c)
    => ArrowCache a b (CacheT Cache a b c) where
    type Widening (CacheT Cache a b c) = W.Widening b
  
    initialize = CacheT $ modify' $ \(a,Cache cache) ->
      let cache' = M.insertWith (\_ _old -> _old) a (Unstable,bottom) cache
          ~(_,b) = M.lookupDefault (Unstable,bottom) a cache
      in (b,Cache cache')
    lookup = CacheT $ proc a -> do
      Cache cache <- get -< ()
      returnA -< M.lookup a cache
    update = CacheT $ proc (st,a,b) -> do
      Cache cache <- get -< ()
      case M.lookup a cache of
        Just (_,b') -> do
          let (st',b'') = ?cacheWidening b' b
          put -< Cache (M.insert a (st ⊔ st',b'') cache)
          returnA -< (st',a,b'')
        Nothing -> do
          put -< Cache (M.insert a (Unstable,b) cache)
          returnA -< (Unstable,a,b)
    write = CacheT $ modify' (\((a,b,s),Cache cache) -> ((),Cache (M.insert a (s,b) cache)))
    setStable = CacheT $ modify' $ \((s,a),Cache cache) -> ((),Cache (M.adjust (first (const s)) a cache))
    {-# INLINE initialize #-}
    {-# INLINE lookup #-}
    {-# INLINE write #-}
    {-# INLINE update #-}
    {-# INLINE setStable #-}
    {-# SCC initialize #-}
    {-# SCC lookup #-}
    {-# SCC write #-}
    {-# SCC update #-}
    {-# SCC setStable #-}
  
  instance (Arrow c, Profunctor c) => ArrowIterateCache a b (CacheT Cache a b c) where
    nextIteration = CacheT $ proc x -> do
      put -< empty
      returnA -< x
    {-# INLINE nextIteration #-}
  
  instance (PreOrd a, Arrow c, Profunctor c) => ArrowReuse a b (CacheT Cache a b c) where
    reuse f = CacheT $ proc (a,s) -> do
      Cache cache <- get -< ()
      returnA -< M.foldlWithKey' (\m a' (s',b') -> if s' ⊑ s && a ⊑ a' then m <> f a a' s' b' else m) mempty cache
    {-# INLINE reuse #-}
  
  instance (LowerBounded b, Profunctor c, Arrow c) => ArrowLowerBounded b (CacheT Cache a b c) where
    bottom = proc _ -> returnA -< bottom
    {-# INLINE bottom #-}
  
  instance Identifiable a => IsList (Cache a b) where
    type Item (Cache a b) = (a,b,Stable)
    toList (Cache m) = [ (a,b,s) | (a,(s,b)) <- M.toList m]
    fromList l = Cache $ M.fromList [ (a,(s,b)) | (a,b,s) <- l]
  
  ------ Group Cache ------
  data Group cache a b where
    Groups :: HashMap k (cache a b) -> Group cache (k,a) b
  
  instance (Identifiable k, Pretty k, Pretty (cache a b)) => Pretty (Group cache (k, a) b) where
    pretty (Groups m) = vsep [pretty k <> "@" <> pretty (hash k) <+> "->" <> line <> indent 4 (pretty cache) | (k, cache) <- M.toList m]
  
  instance IsEmpty (Group cache (k,a) b) where
    empty = Groups empty
    {-# INLINE empty #-}
  
  instance (Identifiable k, Arrow c, Profunctor c, ArrowCache a b (CacheT cache a b c), IsEmpty (cache a b)) => ArrowCache (k,a) b (CacheT (Group cache) (k,a) b c) where
    type Widening (CacheT (Group cache) (k,a) b c) = Cache.Widening (CacheT cache a b c)
    initialize = withGroup Cache.initialize
    lookup = withGroup Cache.lookup
    update = proc (st,(k,a),b) -> do
      (st',a',b') <- withGroup Cache.update -< (k,(st,a,b))
      returnA -< (st',(k,a'),b')
    write = lmap (\((k,a),b,s) -> (k,(a,b,s))) (withGroup Cache.write)
    setStable = lmap shuffle1 (withGroup Cache.setStable)
    {-# INLINE initialize #-}
    {-# INLINE lookup #-}
    {-# INLINE update #-}
    {-# INLINE write #-}
    {-# INLINE setStable #-}
    {-# SCC initialize #-}
    {-# SCC lookup #-}
    {-# SCC write #-}
    {-# SCC update #-}
    {-# SCC setStable #-}
  
  withGroup :: (Identifiable k, IsEmpty (cache a b),
                Profunctor c, Arrow c)
            => CacheT cache a b c x y -> CacheT (Group cache) (k,a) b c (k,x) y
  withGroup f = lift $
    dimap (\(Groups groups,(k,x)) -> ((groups,k),(fromMaybe empty (M.lookup k groups),x)))
          (\((groups,k),(cache,y)) -> (Groups (M.insert k cache groups), y))
          (second (unlift f))
  {-# INLINE withGroup #-}
  
  
  instance (Identifiable k, IsEmpty (cache a b), ArrowApply c, Profunctor c, ArrowReuse a b (CacheT cache a b c)) => ArrowReuse (k,a) b (CacheT (Group cache) (k,a) b c) where
    reuse f = proc ((k,a0),s) -> withGroup (reuse (\a a' -> f (k,a) (k,a'))) -<< (k,(a0,s))
    {-# INLINE reuse #-}
  
  instance Identifiable k => IsList (Group cache (k,a) b) where
    type Item (Group cache (k,a) b) = (k,cache a b)
    toList (Groups m) = M.toList m
    fromList l = Groups $ M.fromList l
  
  instance (Show k, Show (cache a b)) => Show (Group cache (k,a) b) where
    show (Groups m) = show (M.toList m)
  
  ------ Parallel Cache ------
  data Parallel cache a b =
    Parallel { old :: cache a b
             , new :: cache a b
             , stable :: !Stable
             }
  
  instance Pretty (cache a b) => Show (Parallel cache a b)   where show = show . pretty
  instance Pretty (cache a b) => Pretty (Parallel cache a b) where
    pretty (Parallel o n s) =
      vsep ["Parallel",
            "Old:" <+> align (pretty o),
            "New:" <+> align (pretty n),
            "Stable:" <> viaShow s
           ]
  
  instance IsEmpty (cache a b) => IsEmpty (Parallel cache a b) where
    empty = Parallel { old = empty, new = empty, stable = Stable }
  
  instance (Profunctor c, Arrow c,
            ArrowLowerBounded b (CacheT cache a b c))
      => ArrowLowerBounded b (CacheT (Parallel cache) a b c) where
    bottom = newCache Order.bottom
    {-# INLINE bottom #-}
  
  instance (Profunctor c, ArrowChoice c,
            ArrowIterateCache a b (CacheT cache a b c),
            ArrowCache a b (CacheT cache a b c))
      => ArrowParallelCache a b (CacheT (Parallel cache) a b c) where
    lookupOldCache = oldCache $ proc a -> do
      m <- lookup -< a; case m of
        Just (_,b) -> returnA -< b
        Nothing    -> initialize -< a
    lookupNewCache = newCache (rmap (fmap snd) lookup)
    updateNewCache = dimap (\(a,b) -> (Stable,a,b)) (\(_,_,b) -> b) update
    isStable = modify' (\(_,p) -> (stable p,p))
    {-# INLINE lookupOldCache #-}
    {-# INLINE lookupNewCache #-}
    {-# INLINE updateNewCache #-}
    {-# INLINE isStable #-}
    {-# SCC lookupOldCache #-}
    {-# SCC lookupNewCache #-}
    {-# SCC updateNewCache #-}
    {-# SCC isStable #-}
  
  instance (Profunctor c, ArrowChoice c,
            ArrowIterateCache a b (CacheT cache a b c))
      => ArrowIterateCache a b (CacheT (Parallel cache) a b c) where
    nextIteration = proc x -> do
      modify' (\(_,p) -> ((),p { old = new p, stable = Stable })) -< ()
      newCache nextIteration -< x
    {-# INLINE nextIteration #-}
  
  instance (Profunctor c, ArrowChoice c,
            ArrowCache a b (CacheT cache a b c))
      => ArrowCache a b (CacheT (Parallel cache) a b c) where
    type Widening (CacheT (Parallel cache) a b c) = Cache.Widening (CacheT cache a b c)
    initialize = proc a -> do
      m <- oldCache lookup -< a; case m of
        Just (s,b) -> do
          newCache write -< (a,b,s)
          returnA -< b
        Nothing -> do
          modify' (\((),cache) -> ((),cache { stable = Unstable })) -< ()
          newCache initialize -< a
    lookup = newCache lookup
    update = proc (st,a,b) -> do
      (st',a',b') <- newCache update -< (st,a,b)
      st'' <- modify' (\(st',cache) -> let st'' = stable cache ⊔ st' in (st'',cache { stable = st'' })) -< st'
      returnA -< (st'',a',b')
    write = newCache write
    setStable = newCache setStable
    {-# INLINE initialize #-}
    {-# INLINE lookup #-}
    {-# INLINE write #-}
    {-# INLINE update #-}
    {-# INLINE setStable #-}
  
  -- Note: All operations on the old cache are read-only.
  oldCache :: (Arrow c, Profunctor c) => CacheT cache a b c x y -> CacheT (Parallel cache) a b c x y
  oldCache f = lift $ dimap (\(p,x) -> (p,(old p,x))) (\(p,(_,y)) -> (p, y)) (second (unlift f))
  {-# INLINE oldCache #-}
  
  newCache :: (Arrow c, Profunctor c) => CacheT cache a b c x y -> CacheT (Parallel cache) a b c x y
  newCache f = lift $
    dimap (\(p,x) -> (p,(new p,x)))
          (\(p,(n,y)) -> (p { new = n }, y))
          (second (unlift f))
  {-# INLINE newCache #-}
  
  ------ Monotone Cache ------
  data Monotone a b where
    Monotone :: s -> HashMap a (Stable,s,b) -> Monotone (s,a) (s,b)
  
  instance IsEmpty s => IsEmpty (Monotone (s,a) (s,b)) where
    empty = Monotone empty empty
  
  instance (Show s, Show a, Show b) => Show (Monotone (s,a) (s,b)) where
    show (Monotone s m) = show (s,m)
  
  instance (Pretty s, Pretty a, Pretty b) => Pretty (Monotone (s,a) (s,b)) where
    pretty (Monotone _ m) =
      align (list [ pretty s <+> "|" <+> pretty a <+> showArrow st <+> pretty b | (a,(st,s,b)) <- M.toList m])
  
  instance (Identifiable a, PreOrd s, LowerBounded b, ArrowChoice c, Profunctor c)
      => ArrowCache (s,a) (s,b) (CacheT Monotone (s,a) (s,b) c) where
    type Widening (CacheT Monotone (s,a) (s,b) c) = (W.Widening s,W.Widening b)
    initialize = CacheT $ modify' $ \((sNew,a),Monotone s cache) ->
      case M.lookup a cache of
        Just (_,sOld,b)
          | sNew ⊑ sOld -> ((sNew,b),      Monotone s cache)
          | otherwise   -> ((sNew,b),      Monotone s (M.insert a (Unstable,sNew,b) cache))
        Nothing         -> ((sNew,bottom), Monotone s (M.insert a (Unstable,sNew,bottom) cache))
    lookup = CacheT $ modify' $ \((s,a),m@(Monotone _ cache)) ->
      case M.lookup a cache of
        Just (st,s',b) | s ⊑ s' -> (Just (st,(s,b)), m)
        _ -> (Nothing, m)
    update = CacheT $ modify' $ \((stable0,(sNew,a),(sNew',bNew)),Monotone s cache) ->
      let (widenS,widenB) = ?cacheWidening
          (_,sWiden) = widenS s sNew'
      in case M.lookup a cache of
        Just (_,sOld,bOld) ->
            let stable1 = if sNew ⊑ sOld then Stable else Unstable
                (stable2,bWiden) = widenB bOld bNew
            in ((stable1 ⊔ stable2, (sWiden,a), (sWiden,bWiden)),
                Monotone sWiden (M.insert a (stable0 ⊔ stable1 ⊔ stable2,sNew,bWiden) cache))
        Nothing -> ((Unstable,(sWiden,a),(sWiden,bNew)), Monotone sWiden (M.insert a (Unstable,sNew,bNew) cache))
    write = CacheT $ modify' $ \(((sNew, a), (_,b), st), Monotone s cache) ->
      ((), Monotone s (M.insert a (st, sNew, b) cache))
    setStable = CacheT $ proc _ -> returnA -< ()
    {-# INLINE initialize #-}
    {-# INLINE lookup #-}
    {-# INLINE write #-}
    {-# INLINE update #-}
    {-# INLINE setStable #-}
    {-# SCC initialize #-}
    {-# SCC lookup #-}
    {-# SCC write #-}
    {-# SCC update #-}
    {-# SCC setStable #-}
  
  instance (Arrow c, Profunctor c) => ArrowIterateCache (s,a) (s,b) (CacheT Monotone (s,a) (s,b) c) where
    nextIteration = CacheT $ proc ((_,a),(sNew,b)) -> do
      put -< Monotone sNew empty
      returnA -< ((sNew,a),(sNew,b))
    {-# INLINE nextIteration #-}
  
  -- | Cache for an analysis with abstract garbage collection. The type variable
  -- @s@ stands for the store.
  data GarbageCollect a b where
    GarbageCollect :: HashMap a (Stable,s,s,b) -> GarbageCollect (s,a) (s,b)
  
  instance IsEmpty (GarbageCollect (s,a) (s,b)) where
    empty = GarbageCollect empty
  
  instance (Show s, Show a, Show b) => Show (GarbageCollect (s,a) (s,b)) where
    show (GarbageCollect m) = show m
  
  instance (Pretty s, Pretty a, Pretty b) => Pretty (GarbageCollect (s,a) (s,b)) where
    pretty (GarbageCollect m) =
      align (list [ pretty s1 <+> "×" <+> pretty a <+> showArrow st <+> pretty s2 <+> "×" <+> pretty b | (a,(st,s1,s2,b)) <- M.toList m])
  
  instance (Identifiable a, PreOrd s, LowerBounded b, ArrowChoice c, Profunctor c)
      => ArrowCache (s,a) (s,b) (CacheT GarbageCollect (s,a) (s,b) c) where
    type Widening (CacheT GarbageCollect (s,a) (s,b) c) = (W.Widening s,W.Widening b)
    initialize = CacheT $ modify' $ \((sNew1,a),GarbageCollect cache) ->
      let (widenS,_) = ?cacheWidening
      in case M.lookup a cache of
        Just (_,sOld1,sOld2,b)
          | sNew1 ⊑ sOld1 -> ((sOld2,b), GarbageCollect cache)
          | otherwise     -> let (_,sWiden1) = widenS sOld1 sNew1
                                 (_,sWiden2) = widenS sOld2 sNew1
                             in ((sWiden2,b), GarbageCollect (M.insert a (Unstable,sWiden1,sOld2,b) cache))
        Nothing           -> ((sNew1,bottom), GarbageCollect (M.insert a (Unstable,sNew1,sNew1,bottom) cache))
    lookup = CacheT $ modify' $ \((sNew1,a),m@(GarbageCollect cache)) ->
      case M.lookup a cache of
        Just (st,sOld1,sOld2,b)
          | sNew1 ⊑ sOld1 -> (Just (st,(sOld2,b)),       m)
          | otherwise     -> (Just (Unstable,(sOld2,b)), m)
        _ -> (Nothing, m)
    update = CacheT $ modify' $ \((stable0,(sNew1,a),(sNew2,bNew)),GarbageCollect cache) ->
      let (widenS,widenB) = ?cacheWidening
      in case M.lookup a cache of
        Just (_,sOld1,sOld2,bOld) ->
            let (stable1,sWiden1) = widenS sOld1 sNew1
                (stable2,sWiden2) = widenS sOld2 sNew2
                (stable3,bWiden)  = widenB bOld bNew
            in ((stable1 ⊔ stable2 ⊔ stable3, (sWiden1,a), (sWiden2,bWiden)),
                GarbageCollect (M.insert a (stable0 ⊔ stable1 ⊔ stable2 ⊔ stable3, sWiden1, sWiden2, bWiden) cache))
        Nothing -> ((Unstable,(sNew1,a),(sNew2,bNew)), GarbageCollect (M.insert a (Unstable,sNew1,sNew2,bNew) cache))
    write = CacheT $ modify' $ \(((sNew1, a), (sNew2,b), st), GarbageCollect cache) ->
      ((), GarbageCollect (M.insert a (st, sNew1, sNew2, b) cache))
    setStable = CacheT $ proc _ -> returnA -< ()
    {-# INLINE initialize #-}
    {-# INLINE lookup #-}
    {-# INLINE write #-}
    {-# INLINE update #-}
    {-# INLINE setStable #-}
    {-# SCC initialize #-}
    {-# SCC lookup #-}
    {-# SCC write #-}
    {-# SCC update #-}
    {-# SCC setStable #-}
  
  instance (Arrow c, Profunctor c) => ArrowIterateCache (s,a) (s,b) (CacheT GarbageCollect (s,a) (s,b) c) where
    nextIteration = CacheT $ proc ((_,a),(sNew2,b)) -> do
      put -< GarbageCollect empty
      returnA -< ((sNew2,a),(sNew2,b))
    {-# INLINE nextIteration #-}
  
  ------ Monotone Cache that factors out the monotone element ------
  data MonotoneFactor a b where
    MonotoneFactor :: s -> HashMap a b -> MonotoneFactor (s,a) (s,b)
  
  instance IsEmpty s => IsEmpty (MonotoneFactor (s,a) (s,b)) where
    empty = MonotoneFactor empty empty
  
  instance (Show s, Show a, Show b) => Show (MonotoneFactor (s,a) (s,b)) where
    show (MonotoneFactor s m) = show (s,m)
  
  instance (Pretty s, Pretty a, Pretty b) => Pretty (MonotoneFactor (s,a) (s,b)) where
    pretty (MonotoneFactor s m) =
      vsep [ "Monotone:" <+> pretty s
           , "NonMonotone:" <+> align (list [ pretty k <+> "->" <+> pretty v | (k,v) <- M.toList m])
           ]
  
  instance (Identifiable a, LowerBounded b,
            ArrowChoice c, Profunctor c)
      => ArrowCache (s,a) (s,b) (CacheT MonotoneFactor (s,a) (s,b) c) where
    type Widening (CacheT MonotoneFactor (s,a) (s,b) c) = (W.Widening s,W.Widening b)
    initialize = CacheT $ modify' $ \((s,a),MonotoneFactor s' cache) ->
      let cache' = M.insertWith (\_ _old -> _old) a bottom cache
          b = M.lookupDefault bottom a cache'
      in ((s,b),MonotoneFactor s' cache')
    lookup = CacheT $ modify' $ \((s,a),m@(MonotoneFactor _ cache)) ->
      ((\b -> (Unstable,(s,b))) <$> M.lookup a cache, m)
    update = CacheT $ modify' $ \((_,(_,a),(sNew,b)),MonotoneFactor sOld cache) ->
      let (widenS,widenB) = ?cacheWidening
          (stable1,sWiden) = widenS sOld sNew
      in case M.lookup a cache of
          Just b' ->
            let (stable2,b'') = widenB b' b
            in ((stable1 ⊔ stable2, (sWiden,a), (sWiden,b'')),
                MonotoneFactor sWiden (M.insert a b'' cache))
          Nothing -> ((Unstable,(sWiden,a), (sWiden,b)),MonotoneFactor sWiden (M.insert a b cache))
    write = CacheT $ modify' $ \(((_, a), (_, b), _),MonotoneFactor s cache) -> ((),MonotoneFactor s (M.insert a b cache))
    setStable = CacheT $ proc _ -> returnA -< ()
    {-# INLINE initialize #-}
    {-# INLINE lookup #-}
    {-# INLINE write #-}
    {-# INLINE update #-}
    {-# INLINE setStable #-}
    {-# SCC initialize #-}
    {-# SCC lookup #-}
    {-# SCC write #-}
    {-# SCC update #-}
    {-# SCC setStable #-}
  
  instance (Arrow c, Profunctor c) => ArrowIterateCache (s,a) (s,b) (CacheT MonotoneFactor (s,a) (s,b) c) where
    nextIteration = CacheT $ proc ((_,a),(sNew,b)) -> do
      put -< MonotoneFactor sNew empty
      returnA -< ((sNew,a),(sNew,b))
    {-# INLINE nextIteration #-}
  