{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}

-- Data.Generics with GPS (Generic Positioning System;-)
--
-- GPS employs Maps, to avoid getting lost in Data:
--
-- - for each traversed type, build a Map TypeKey TypeSet, mapping all
--   substructure types of the given type to their substructure types 
--
-- - traversals are short-circuited when the domain types of their queries 
--   or transformations cannot be found in the current substructure types
--
-- - domains of queries and transformations are computed on construction
--
-- GPS is inspired by Uniplate's PlateData direction finder (contains and
-- DataBox are copied from the Uniplate paper), generalised to tackle SYB's
-- more general queries and transformations (instead of oracles telling
-- whether to Stop, Follow, or Find in a search for type b in type a, we 
-- build IntSets of TypeRep keys, both for the domains of traversals and 
-- for substructure types; then several short-circuiting decisions can be 
-- based on fast intersection tests with the same IntSet).
--
-- TypeRep keys combined with IntSets or IntMaps can be used to speed up other
-- generic programming problems as well, including typecase and extensible
-- records libraries.

module Data.Generics.GPS(everything,everywhere,
                         mkQ,extQ,mkT,extT,
                         module Data.Generics) where
import Data.Generics hiding (everything,everywhere,mkQ,extQ,mkT,extT)
import qualified Data.Generics as DG(everything,everywhere,mkQ,extQ,mkT,extT)
import CompanyDatatypes
import Data.List hiding((\\))
import System.IO.Unsafe(unsafePerformIO)
import Data.IntSet as IS
import qualified Data.IntMap as Map
import Debug.Trace

type TypeSet = IntSet
type TypeKey = Int
type TypeMap = Map.IntMap

data GenericDomainQ r = 
  GenericDomainQ { queryDomain  :: TypeSet,
                   defaultValue :: r,
                   genericQuery :: GenericQ r }

mkQ ::  (Typeable a) => b -> (a -> b) -> GenericDomainQ b
def `mkQ` spec = 
  GenericDomainQ{ queryDomain  = singleton (getDomainKey spec),
                  defaultValue = def,
                  genericQuery = def `DG.mkQ` spec }

extQ :: (Typeable a) => GenericDomainQ b -> (a -> b) -> GenericDomainQ b
gen `extQ` spec = 
  gen{ queryDomain  = IS.insert (getDomainKey spec) (queryDomain gen),
       genericQuery = genericQuery gen `DG.extQ` spec }

data GenericDomainT = 
  GenericDomainT { transDomain  :: TypeSet,
                   defaultTrans :: GenericT,
                   genericTrans :: GenericT }

mkT ::  (Typeable b) => (b -> b) -> GenericDomainT
mkT spec = 
  GenericDomainT { transDomain  = singleton (getDomainKey spec),
                   defaultTrans = id,
                   genericTrans = DG.mkT spec }

extT :: forall a . Data a => GenericDomainT -> (a -> a) -> GenericDomainT
gen `extT` spec =
  gen { transDomain  = singleton (getDomainKey spec),
        genericTrans = genericTrans gen `DG.extT` spec }

everywhere :: forall a . Data a => GenericDomainT -> a -> a
everywhere = everywhereWithMap (getSubs subMap)
  where subMap = trace ("subMap{"++show (typeOf (undefined::a))++"}") 
               $ fromRoot (undefined::a) Map.empty

everywhereWithMap :: forall a . Data a 
                  => (forall a . Data a => a -> TypeSet) 
                  -> GenericDomainT -> a -> a
everywhereWithMap getSubs
                  gdt@(GenericDomainT{transDomain=domain,defaultTrans=dt,genericTrans=t})
                  x
  | not $ IS.null $ domain `intersection` getSubs x
  = t (gmapT (everywhereWithMap getSubs gdt) x)
  | otherwise
  = dt x

everything :: forall a r . Data a => (r -> r -> r) -> GenericDomainQ r -> a -> r
everything = everythingWithMap (getSubs subMap)
  where subMap = trace ("subMap{"++show (typeOf (undefined::a))++"}") 
               $ fromRoot (undefined::a) Map.empty

everythingWithMap :: forall a r . Data a 
                  => (forall a . Data a => a -> TypeSet) 
                  -> (r -> r -> r) -> GenericDomainQ r -> a -> r
everythingWithMap getSubs
                  k
                  gdq@(GenericDomainQ{queryDomain=domain,defaultValue=z,genericQuery=q})
                  x
  | not $ IS.null $ domain `intersection` getSubs x
  = foldl k (q x) (gmapQ (everythingWithMap getSubs k gdq) x)
  | otherwise
  = z

getSubs :: forall a . Data a => TypeMap TypeSet -> a -> TypeSet
getSubs subMap = \x->Map.findWithDefault (error ("missing key: "++show (typeOf x))) (key x) subMap

fromRoot :: forall a . Data a => a -> TypeMap TypeSet -> TypeMap TypeSet
fromRoot root map | Map.member (key root) map = map
fromRoot root map | otherwise                 = fromRoots (contains root) map'
  where map'             = Map.insert (key root) (allSubs root) map
        fromRoots rs map = foldl' (\m (DataBox x)->fromRoot x m) map rs

allSubs :: Data a => a -> TypeSet
allSubs x = IS.insert trx $ allSubs' x (singleton trx)
  where trx = typeRepKey' $ typeOf x

allSubs' :: Data a => a -> TypeSet -> TypeSet
allSubs' x but = subKeys `IS.union` (unions $ Prelude.map f subBoxes)
  where subBoxes = contains x \\ but
          where s \\ but = Prelude.filter (not . (`member` but) . (typeRepKey' . typeOfIn)) s
        subKeys = fromList (Prelude.map (typeRepKey' . typeOfIn) subBoxes) 
        f (DataBox x) = allSubs' x $ but `IS.union` subKeys

getDomainKey :: Typeable a => (a->b) -> TypeKey
getDomainKey f = key (getDomain f)
  where getDomain :: (a->b) -> a
        getDomain = undefined

typeOfIn (DataBox x) = typeOf x
typeRepKey' tr = unsafePerformIO (typeRepKey tr)
key x = typeRepKey' (typeOf x)

data DataBox = forall a . Data a => DataBox a 
contains :: Data a => a -> [DataBox]
contains x = if isAlgType dtyp then concatMap f ctrs else [ ]
  where
  f c = gmapQ DataBox (asTypeOf (fromConstr c) x) 
  ctrs = dataTypeConstrs dtyp
  dtyp = dataTypeOf x