{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE Arrows #-}
{-# LANGUAGE TypeFamilies #-}
module Control.Arrow.Transformer.FreeVars where

import Prelude hiding ((.),read,Maybe(..))

import           Control.Category
import           Control.Arrow
import           Control.Arrow.Const
import           Control.Arrow.Environment as Env
import           Control.Arrow.Except
import           Control.Arrow.Fail
import           Control.Arrow.Fix
import           Control.Arrow.Order
import           Control.Arrow.State
import           Control.Arrow.Store
import           Control.Arrow.Trans
import           Control.Arrow.Writer as Writer

import           Control.Arrow.Transformer.Writer

import           Data.Identifiable
import           Data.Profunctor
import           Data.Profunctor.Unsafe((.#))
import           Data.Coerce
import           Data.HashSet(HashSet)
import qualified Data.HashSet as H

newtype FreeVarsT var c x y = FreeVarsT (WriterT (HashSet var) c x y)
  deriving (Profunctor, Category, Arrow, ArrowChoice, ArrowLift, ArrowTrans, ArrowLowerBounded a,
            ArrowState s, ArrowFail e, ArrowExcept e, ArrowStore var' val', ArrowConst k)

runFreeVarsT :: Profunctor c => FreeVarsT var c x y -> c x y
runFreeVarsT (FreeVarsT f) = rmap snd (runWriterT f)
{-# INLINE runFreeVarsT #-}

instance (Identifiable var, ArrowRun c) => ArrowRun (FreeVarsT var c) where
  type Run (FreeVarsT var c) x y = Run c x y
  run = run . runFreeVarsT
  {-# INLINE run #-}

instance (Identifiable var, ArrowEnv var val c) => ArrowEnv var val (FreeVarsT var c) where
  type Join y (FreeVarsT var c) = Env.Join (HashSet var,y) c
  lookup (FreeVarsT f) (FreeVarsT g) = FreeVarsT $ proc (var,x) -> do
    tell -< H.singleton var
    Env.lookup f g -< (var,x)
  extend (FreeVarsT f) = FreeVarsT $ proc (var,val,x) ->
    censor (\(var,_,_) -> H.delete var) (Env.extend f) -< (var,val,x)

instance (Identifiable var,ArrowApply c, Profunctor c) => ArrowApply (FreeVarsT var c) where
  app = FreeVarsT (app .# first coerce)

instance (ArrowFix (Underlying (FreeVarsT var c) x y)) => ArrowFix (FreeVarsT var c x y) where
  type Fix (FreeVarsT var c x y) = Fix (Underlying (FreeVarsT var c) x y)

deriving instance (Identifiable var,ArrowComplete (HashSet var,y) c) => ArrowComplete y (FreeVarsT var c)
