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

{-# LANGUAGE CPP #-}

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

{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# LANGUAGE DeriveDataTypeable #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  SAI.Data.Generics.Shape.SYB.GHC
-- Copyright   :  (c) Andrew Seniuk, 2014
-- License     :  BSD-style (see the LICENSE file)
--  
-- Maintainer  :  rasfar@gmail.com
-- Stability   :  experimental
-- Portability :  non-portable (uses Data.Generics.Basics)
--
-- This package provides SYB shape support: generic mapping to
-- homogeneous types, and related features.  Complements existing
-- Uniplate and TH shape libraries.  See <http://www.fremissant.net/shape-syb>
-- for more information.
--
-- The present module provides support for staged GHC AST types.
-- Once you have a @'Homo' r@, 'Hetero', or @'Bi' r@, the rest of
-- the API in "Shape.SYB" and "Shape.SYB.Filter" is applicable.
--
-----------------------------------------------------------------------------

  module SAI.Data.Generics.Shape.SYB.GHC (

    -- * Staged shape functions

    ghomStaged ,
    ghomStagedK ,
    ghomDynStaged ,
    ghomBiStaged ,

    GHC_AST_HOLE ,

    shapeOfStaged ,
    shapeOfStaged_ ,
    sizeOfStaged ,

    symmorphicStaged ,

    weightedShapeOfStaged ,

    -- * Re-exported from the ghc-syb-utils package

    Stage(..) ,

  ) where

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

  import Data.Data ( gfoldl )
  import Data.Data ( gmapQ )
  import Data.Data ( Data )
  import Data.Data ( Typeable )
  import Data.Generics.Aliases ( GenericQ )

#if USE_DATA_TREE
  import SAI.Data.Generics.Shape.SYB ( Rose, Tree(Node) )
#else
  import SAI.Data.Generics.Shape.SYB ( Rose(..) )
#endif
  import SAI.Data.Generics.Shape.SYB ( Homo, Shape, Hetero, Bi )
  import SAI.Data.Generics.Shape.SYB ( zipRose )

  import SAI.Data.Generics.Shape.SYB.Filter ( filterHomoMM )
--import SAI.Data.Generics.Shape.SYB.Filter ( filterHomoM_' )
  import SAI.Data.Generics.Shape.SYB ( shapeOf )
  import SAI.Data.Generics.Shape.SYB ( unliftHomoM )
  import SAI.Data.Generics.Shape.SYB ( sizeOfRose )

  import qualified GHC     as GHC
  import qualified NameSet as GHC
  import qualified FastString as GHC

  import qualified Data.Generics as SYB
  import qualified GHC.SYB.Utils as SYB
  import GHC.SYB.Utils ( Stage(..) )

  import Data.Dynamic

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

  newtype GHC_AST_HOLE = GHC_AST_HOLE Stage deriving ( Typeable )

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

  ghomStaged :: forall r d. Data d =>
             Stage
          -> r
          -> GenericQ r
          -> d
          -> Homo r
  ghomStaged stage z f x
    | checkItemStage stage x = z'
    | otherwise = foldl k b (gmapQ (ghomStaged stage z f) x)
   where
     b = Node (f x) []
     z' = Node z []
     k (Node r chs) nod@(Node r' _) = Node r (chs++[nod])

  ghomStagedK :: forall r d. Data d =>
             Stage
          -> r
          -> (r -> r -> r)
          -> GenericQ r
          -> d
          -> Homo r
  ghomStagedK stage z k f x
    | checkItemStage stage x = z'
    | otherwise = foldl k' b (gmapQ (ghomStagedK stage z k f) x)
   where
     b = Node (f x) []
     z' = Node z []
     k' (Node r chs) nod@(Node r' _) = Node (r `k` r') (chs++[nod])

  -- | Uses "Data.Dynamic" to support mutiple types homogeneously.
  ghomDynStaged :: forall d. Data d => Stage -> d -> Hetero
  ghomDynStaged stage x
    | checkItemStage stage x = Node (toDyn $ GHC_AST_HOLE stage) []
    | otherwise              = foldl k b (gmapQ (ghomDynStaged stage) x)
   where
     b = Node (toDyn x) []
     k (Node r chs) nod = Node r (chs++[nod])

  -- | @'ghomBiStaged' s f x = 'zipRose' ('ghomDynStaged' s x) ('ghomStaged' s f x)@
  ghomBiStaged :: forall r d. Data d => Stage -> r -> GenericQ r -> d -> Bi r
  ghomBiStaged stage z f x = zipRose (ghomDynStaged stage x) $ ghomStaged stage z f x

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

  shapeOfStaged :: forall d. Data d => Stage -> d -> Shape
  shapeOfStaged stage = ghomStaged stage () (const ())

  shapeOfStaged_ :: forall d. Data d => Stage -> d -> Shape
  shapeOfStaged_ stage x = unliftHomoM () $ filterHomoMM $ ghomStaged stage Nothing fg x
--shapeOfStaged_ stage x = filterHomoM_' () $ ghomStaged stage Nothing fg x
   where
     fg :: forall d'. Data d' => d' -> Maybe ()
     fg = (const (Just ())) `SYB.extQ` f_String `SYB.extQ` f_FastString
---  fg = Just `SYB.extQ` f_String `SYB.extQ` f_FastString
     f_String :: String -> Maybe ()
     f_String x = Nothing
     f_FastString :: GHC.FastString -> Maybe ()
     f_FastString x = Nothing

  sizeOfStaged :: forall d. Data d => Stage -> d -> Int
  sizeOfStaged stage = sizeOfRose . (shapeOfStaged stage)

  weightedShapeOfStaged :: forall d. Data d => Stage -> d -> Homo Int
  weightedShapeOfStaged stage = ghomStagedK stage 1 (+) (const 1)

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

-- Borrowed from HaRe:

-- From @frsoares

--- | Checks whether the current item is undesirable for analysis in the current
-- AST Stage.
  checkItemStage :: (Typeable a, Data a) => Stage -> a -> Bool
  checkItemStage stage x = (checkItemStage1 stage x)
#if __GLASGOW_HASKELL__ > 704
                        || (checkItemStage2 stage x)
#endif

-- Check the Typeable items
  checkItemStage1 :: (Typeable a) => Stage -> a -> Bool
  checkItemStage1 stage x = (const False `SYB.extQ` postTcType `SYB.extQ` fixity `SYB.extQ` nameSet) x
    where nameSet     = const (stage `elem` [SYB.Parser,SYB.TypeChecker]) :: GHC.NameSet       -> Bool
          postTcType  = const (stage < SYB.TypeChecker                  ) :: GHC.PostTcType    -> Bool
          fixity      = const (stage < SYB.Renamer                      ) :: GHC.Fixity        -> Bool

#if __GLASGOW_HASKELL__ > 704
--- | Check the Typeable1 items
  checkItemStage2 :: Data a => Stage -> a -> Bool
  checkItemStage2 stage x = (const False `SYB.ext1Q` hsWithBndrs) x
    where
          hsWithBndrs = const (stage < SYB.Renamer) :: GHC.HsWithBndrs a -> Bool
#endif

#if 0
  checkItemRenamer :: (Data a, Typeable a) => a -> Bool
  checkItemRenamer x = checkItemStage SYB.Renamer x
#endif

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

  -- | Compare two GHC ASTs for shape equality.
  symmorphicStaged :: forall d1 d2. (Data d1,Data d2) =>
                      Stage -> d1 -> d2 -> Bool
  symmorphicStaged stage x y = shapeOfStaged stage x == shapeOfStaged stage y

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