content: {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE Arrows #-}
  {-# LANGUAGE TypeFamilies #-}
  module Control.Arrow.Transformer.Concrete.Environment where
  
  import           Prelude hiding ((.),read)
  
  import           Control.Category
  
  import           Control.Arrow
  import           Control.Arrow.Const
  import           Control.Arrow.Environment as Env
  import           Control.Arrow.Closure as Closure
  import           Control.Arrow.Except
  import           Control.Arrow.Fail
  import           Control.Arrow.Fix
  import           Control.Arrow.Frame
  import           Control.Arrow.Reader as Reader
  import           Control.Arrow.State
  import           Control.Arrow.Store
  import           Control.Arrow.Trans
  import           Control.Arrow.LetRec
  
  import           Control.Arrow.Transformer.Reader
  
  import           Data.Concrete.Closure
  import           Data.Identifiable
  import           Data.HashMap.Lazy (HashMap)
  import qualified Data.HashMap.Lazy as M
  import           Data.Profunctor
  import           Data.Profunctor.Unsafe
  import           Data.Coerce
  
  -- | Arrow transformer that adds an environment to a computation.
  newtype EnvT env c x y = EnvT (ReaderT env c x y)
    deriving (Profunctor,Category,Arrow,ArrowChoice,ArrowLift,ArrowTrans,
              ArrowFail e,ArrowExcept e,ArrowState s,ArrowConst r,
              ArrowStore var' val', ArrowRun)
  
  instance (Arrow c, Profunctor c) => ArrowFrame env (EnvT env c) where
    newFrame (EnvT f) = EnvT $ Reader.local f
    askFrame = EnvT Reader.ask
  
  instance (Identifiable var, ArrowChoice c, Profunctor c) => ArrowEnv var val (EnvT (HashMap var val) c) where
    type Join y (EnvT (HashMap var val) c) = ()
    lookup (EnvT f) (EnvT g) = EnvT $ proc (var,x) -> do
      env <- Reader.ask -< ()
      case M.lookup var env of
        Just val -> f -< (val,x)
        Nothing -> g -< x
    extend (EnvT f) = EnvT $ proc (var,val,x) -> do
      env <- Reader.ask -< ()
      Reader.local f -< (M.insert var val env, x)
  
  instance (ArrowChoice c, Profunctor c) => ArrowClosure expr (Closure expr env) (EnvT env c) where
    type Join y (Closure expr env) (EnvT env c) = ()
    closure = EnvT $ proc expr -> do
      env <- Reader.ask -< ()
      returnA -< Closure expr env
    apply (EnvT f) = EnvT $ proc (Closure expr env, x) ->
      Reader.local f -< (env,(expr,x))
  
  instance (Identifiable var, IsClosure val (HashMap var val), ArrowChoice c, Profunctor c)
    => ArrowLetRec var val (EnvT (HashMap var val) c) where
    letRec (EnvT f) = EnvT $ proc (bindings,x) -> do
      env <- Reader.ask -< ()
      let env' = foldr (\(var,val) -> M.insert var (setEnvironment env' val)) env bindings
      Reader.local f -< (env',x)
  
  instance (ArrowApply c,Profunctor c) => ArrowApply (EnvT env c) where
    app = EnvT $ app .# first coerce
  
  instance ArrowReader r c => ArrowReader r (EnvT env c) where
    ask = lift' Reader.ask
    local (EnvT (ReaderT f)) = EnvT (ReaderT (lmap (\(env,(r,x)) -> (r,(env,x))) (Reader.local f)))
  
  instance ArrowFix (Underlying (EnvT env c) x y) => ArrowFix (EnvT env c x y) where
    type Fix (EnvT env c x y) = Fix (Underlying (EnvT env c) x y)
  