content: {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE Arrows #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE PolyKinds #-}
  {-# OPTIONS_GHC -fno-warn-redundant-constraints #-}
  module Control.Arrow.Transformer.Abstract.FiniteEnvironment where
  
  import Prelude hiding ((.),read,Maybe(..),lookup,map)
  import qualified Prelude as P
  
  import Control.Category
  import Control.Arrow
  import Control.Arrow.Cont
  import Control.Arrow.Const
  import Control.Arrow.Transformer.Const
  import Control.Arrow.Transformer.Reader
  import Control.Arrow.Reader as Reader
  import Control.Arrow.Transformer.State
  import Control.Arrow.State as State
  import Control.Arrow.Store
  import Control.Arrow.Fail
  import Control.Arrow.Except
  import Control.Arrow.Trans
  import Control.Arrow.Fix.Context hiding (Widening)
  import Control.Arrow.Environment as Env
  import Control.Arrow.Closure
  import Control.Arrow.Fix
  import Control.Arrow.Order
  import Control.Arrow.Utils
  import Control.Arrow.LetRec
  
  import Data.Abstract.Widening (Widening)
  import Data.Abstract.Closure (Closure)
  import qualified Data.Abstract.Closure as Cls
  
  import Data.HashSet(HashSet)
  import qualified Data.HashSet as Set
  import Data.HashMap.Lazy (HashMap)
  import qualified Data.HashMap.Lazy as Map
  import Data.Order
  import Data.Identifiable
  import Data.Profunctor
  import Data.Profunctor.Unsafe((.#))
  import Data.Coerce
  
  type Alloc var addr val c = EnvT var addr val c (var,val) addr
  newtype EnvT var addr val c x y = EnvT (ConstT (Alloc var addr val c, Widening val) (ReaderT (HashMap var addr) (StateT (HashMap addr val) c)) x y)
    deriving (Profunctor,Category,Arrow,ArrowChoice,ArrowLift, ArrowLowerBounded a,
              ArrowFail e, ArrowExcept e, ArrowStore var' val', ArrowRun, ArrowCont,
              ArrowCallSite ctx)
  
  instance (Identifiable var, Identifiable addr, Complete val, ArrowChoice c, Profunctor c) => ArrowEnv var val (EnvT var addr val c) where
    type Join y (EnvT var addr val c) = ()
    lookup (EnvT f) (EnvT g) = EnvT $ proc (var,x) -> do
      env <- Reader.ask -< ()
      store <- State.get -< ()
      case do { addr <- Map.lookup var env; Map.lookup addr store } of
        P.Just val -> f -< (val,x)
        P.Nothing   -> g -< x
    extend (EnvT f) = EnvT $ askConst $ \(EnvT alloc,widening) -> proc (var,val,x) -> do
      env <- Reader.ask -< ()
      store <- State.get -< ()
      case Map.lookup var env of
        P.Just addr -> do
          State.put -< Map.insertWith (\old new -> snd (widening old new)) addr val store
          f -< x
        P.Nothing -> do
          addr <- alloc -< (var,val)
          State.put -< Map.insertWith (\old new -> snd (widening old new)) addr val store
          Reader.local f -< (Map.insert var addr env, x)
    {-# INLINE lookup #-}
    {-# INLINE extend #-}
  
  instance (Identifiable var, Identifiable addr, Identifiable expr, ArrowChoice c, Profunctor c) =>
    ArrowClosure expr (Closure expr (HashSet (HashMap var addr))) (EnvT var addr val c) where
    type Join y (Closure expr (HashSet (HashMap var addr))) (EnvT var addr val c) = Complete y
    closure = EnvT $ proc expr -> do
      env <- Reader.ask -< ()
      returnA -< Cls.closure expr (Set.singleton env)
    apply (EnvT f) = Cls.apply $ proc ((expr,envs),x) ->
      (| joinList (error "encountered an empty set of environments" -< ())
                  (\env -> EnvT (Reader.local f) -< (env,(expr,x))) |) (Set.toList envs)
    {-# INLINE closure #-}
    {-# INLINE apply #-}
  
  instance (Identifiable var, Identifiable addr, Complete val, IsClosure val (HashSet (HashMap var addr)), ArrowChoice c, Profunctor c) => ArrowLetRec var val (EnvT var addr val c) where
    letRec (EnvT f) = EnvT $ askConst $ \(EnvT alloc,widening) -> proc (bindings,x) -> do
      env <- Reader.ask -< ()
      addrs <- map alloc -< bindings
      let env' = Map.fromList [ (var,addr) | ((var,_), addr) <- zip bindings addrs ] `Map.union` env
          vals = Map.fromList [ (addr, setEnvironment (Set.singleton env') val) | (addr, (_,val)) <- zip addrs bindings ]
      State.modify' (\(vals,store) -> ((), Map.unionWith (\old new -> snd (widening old new)) store vals)) -< vals
      Reader.local f -< (env',x)
    {-# INLINE letRec #-}
  
  instance (ArrowApply c, Profunctor c) => ArrowApply (EnvT var addr val c) where
    app = EnvT (app .# first coerce)
    {-# INLINE app #-}
  
  instance ArrowTrans (EnvT var addr val) where
    lift' = EnvT . lift' . lift' . lift'
    {-# INLINE lift' #-}
  
  instance (Complete y, Arrow c, Profunctor c) => ArrowComplete y (EnvT var addr val c) where
    EnvT f <⊔> EnvT g = EnvT (rmap (uncurry (⊔)) (f &&& g))
    {-# INLINE (<⊔>) #-}
  
  instance (Profunctor c, Arrow c, ArrowFix (ReaderT (HashMap var addr) (StateT (HashMap addr val) c) x y)) => ArrowFix (EnvT var addr val c x y) where
    type Fix (EnvT var addr val c x y) = Fix (ReaderT (HashMap var addr) (StateT (HashMap addr val) c) x y)
    fix f = EnvT (fix (coerce f))
    {-# INLINE fix #-}
  