{-# LANGUAGE ConstraintKinds, ExistentialQuantification, FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, UndecidableInstances, ScopedTypeVariables, DeriveDataTypeable, TypeFamilies, StandaloneDeriving, Trustworthy #-}

-- | PolyTraversals that add cycle detection and fixed point construction to any R/W strategy.
module Data.Columbia.CycleDetection where

import Data.Map
import Data.Set (Set)
import Data.Int
import Data.Word
import Data.IORef
import Data.Dynamic hiding (Proxy)
import Data.Generics.SYB.WithClass.Basics
import Data.Columbia.CompoundData
import Data.Columbia.FRecord
import Data.Columbia.Orphans
import Generics.Pointless.Functors
import Generics.Pointless.Combinators
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Fix
import Prelude hiding (lookup)

import Data.Array
import qualified Data.Array.Unboxed as U

class (Monad m) => StateM m where
	type StateOf m :: *
	mfy' :: (StateOf m -> (StateOf m, a)) -> m a

instance (Monad m) => StateM(StateT s m) where
	type StateOf(StateT s m) = s
	mfy' f = get>>= \s -> let(s', x) = f s in (put$!s')>>return x
instance StateM(ReaderT(IORef s) IO) where
	type StateOf(ReaderT(IORef s) IO) = s
	mfy' f = ask>>=lift.(`atomicModifyIORef'` f)

get' :: (StateM m) => m(StateOf m)
get' = mfy'(\s -> (s, s))
put' s = mfy'(\_ -> (s, ()))

-- | 'isKeyed' may always return false, but if it returns true ever, 'keyCompare'
--   must be well-defined and a valid equivalence relation on values for which
--   'isKeyed' returns true (i.e. where all values concerned have /isKeyed x=true/).
--The default is to have 'isKeyed' return false on all values.
class KeyComparable t where
	isKeyed :: t->Bool
	isKeyed _ = False
	keyCompare :: t -> t->Ordering
	keyCompare _ _ = error"KeyComparable.keyCompare: is not a keyed data type"

data KeyCtx t = (KeyComparable t) => KeyCtx

instance (KeyComparable t) => Sat(KeyCtx t) where dict = KeyCtx

data DynamicWithCtx ctx = forall t. (ctx t, Typeable t) => DynamicWithCtx !t

dynamicWithCtx :: (ctx t, Typeable t) => t -> DynamicWithCtx ctx
dynamicWithCtx = DynamicWithCtx

instance Eq(DynamicWithCtx Eq) where
	DynamicWithCtx x == DynamicWithCtx x2 = maybe False(==x2) (cast x)
instance KeyComparable(DynamicWithCtx KeyComparable) where
	isKeyed(DynamicWithCtx x) = isKeyed x
	keyCompare(DynamicWithCtx x) (DynamicWithCtx x2) = maybe(compare(typeOf x) (typeOf x2)) (`keyCompare` x2) (cast x)
instance Ord(DynamicWithCtx KeyComparable) where
	-- Here I'm going to assume that all data being compared have keys.
	compare = keyCompare
instance Eq(DynamicWithCtx KeyComparable) where
	(==) d = (==EQ).keyCompare d

type CycleDetectionR m = StateT(Map Word32 Dynamic) m
type CycleDetectionW m = StateT(Map(DynamicWithCtx KeyComparable) Word32) m
type CycleDetectionRIO = ReaderT(IORef(Map Word32 Dynamic)) IO
type CycleDetectionWIO = ReaderT(IORef(Map(DynamicWithCtx KeyComparable) Word32)) IO

cycleDetect :: forall ctx m t. (MonadFix m, StateM m, StateOf m ~ Map Word32 Dynamic, HasField ctx RWCtx, Data ctx t)
	=> Proxy ctx -> PolyTraversal ctx m t
cycleDetect proxy m =
	getPosition>>= \n->
	lift get'>>=
	maybe
		(mfix(\x -> do
			lift$mfy'(\mp -> (insert n(toDyn x) mp, ()))
			m))
		(maybe
			(fail$"cycleDetect: type mismatch")
			return
			.fromDynamic)
		.lookup n

cycleDetectW :: forall ctx m t. (StateM m, StateOf m ~ Map(DynamicWithCtx KeyComparable) Word32, HasField ctx RWCtx, HasField ctx KeyCtx, Data ctx t)
	=> Proxy ctx
	-> PolyTraversalW ctx m t
cycleDetectW proxy f x =
	case hasField(dict :: ctx t) :: KeyCtx t of
		KeyCtx ->
			let d = dynamicWithCtx x in
			lift get'>>= \mp->
			maybe
				(do
				addr <- getWriterPosition
				when(isKeyed d)$lift$put'$insert d addr mp
				f x)
				return
				$lookup d mp

runCycleDetectionR :: (Monad m)
	=> ReaderT(SeekableStream(CycleDetectionR m) Word8) (CycleDetectionR m) t
	-> ReaderT(SeekableStream m Word8) m t
runCycleDetectionR m = do
	s <- ask
	lift$evalStateT(runReaderT m(hoistStream lift s)) empty

runCycleDetectionW :: (Monad m)
	=> ReaderT(SeekableWriter(CycleDetectionW m) Word8) (CycleDetectionW m) t
	-> ReaderT(SeekableWriter m Word8) m t
runCycleDetectionW m = do
	sw <- ask
	lift$evalStateT(runReaderT m(hoistWriter lift sw)) empty

-----------------------------------

-- | A 'Pair' is a good data structures for associating a key and a value.
data Pair k v = Pair k v deriving (Read, Show, Typeable, Eq, Ord)

instance (Ord k) => KeyComparable(Pair k v) where
	isKeyed _ = True
	keyCompare(Pair k _) (Pair k2 _) = compare k k2

pairCtor = Constr(AlgConstr 1) "Pair" [] Prefix pairDataType
pairDataType = DataType "Data.Columbia.CycleDetection"(AlgRep[pairCtor])
instance (Sat(ctx(Pair k v)), Data ctx k, Data ctx v) => Data ctx(Pair k v) where
	gfoldl _ o f (Pair k v) = f Pair `o` k `o` v
	gunfold _ k f _ = k(k(f Pair))
	dataTypeOf _ _ = pairDataType
	toConstr _ _ = pairCtor
	dataCast2 _ f = gcast2 f

instance (Typeable k, Typeable v) => RW(Pair k v)

type instance Rep(Pair k) v = Pair k v

instance ToRep(Pair k) where
	rep = id
	unrep _ _ = id
	val _ = ann
	fun _ = ann

-- Some boilerplate instances
instance KeyComparable(t1,t2)
instance KeyComparable(t1,t2,t3)
instance KeyComparable(t1,t2,t3,t4)
instance KeyComparable(t1,t2,t3,t4,t5)
instance KeyComparable(Either t t2)
instance KeyComparable(Maybe t)
instance KeyComparable[t]
instance KeyComparable(Map t u)
instance KeyComparable(Set t)
instance KeyComparable(Array i e)
instance KeyComparable(U.UArray i e)
instance KeyComparable Int
instance KeyComparable Word
instance KeyComparable Int8
instance KeyComparable Word8
instance KeyComparable Int16
instance KeyComparable Word16
instance KeyComparable Int32
instance KeyComparable Word32
instance KeyComparable Int64
instance KeyComparable Word64
instance KeyComparable Float
instance KeyComparable Char
instance KeyComparable Ordering
instance KeyComparable Bool
instance KeyComparable ()
instance KeyComparable ((:+:) f g x)
instance KeyComparable ((:*:) f g x)
instance KeyComparable ((:@:) f g x)
instance KeyComparable (Fix f)
instance KeyComparable (LazyFix f)
instance KeyComparable (Const x x2)
instance KeyComparable (Id x)