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

{-# 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.
--
-- Please refer to "Shape.SYB" for descriptions of these functions.
-- Here, in addition to the extra 'SYB.Stage' argument, most of the
-- @ghom@ combinator functions also take an extra argument of type @r@,
-- the default value to use wherever staging holes are encountered.
--
-----------------------------------------------------------------------------

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

-- (See comment below.)
#if 0
    Stage(..) ,
--  GHC_AST_HOLE ,
#endif

    -- * Staged shape functions

    ghom_Staged ,
    ghomK_Staged ,
    ghomP_Staged ,
    ghomE_Staged ,
    ghomDyn_Staged ,
    ghomBi_Staged ,

    shapeOf_Staged ,
    shapeOf_Staged_ ,
    sizeOf_Staged ,

    symmorphic_Staged ,

    weightedShapeOf_Staged ,
    weightedShapeOf_Staged_ ,

  ) 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 ( HomoM )

  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 Data.Generics.Aliases ( mkQ )
  import Data.Generics.Aliases ( extQ )

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

  import qualified Data.Generics as SYB

  import Data.Dynamic

  import qualified GHC.SYB.Utils as SYB

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

#if 0
-- Note: It works to do this (and then no ghc-syb-utils dependence
-- for this project) ... but if try to use ghc-syb-utils by a
-- project which also uses this, you run afoul!...
  -- | 'Stage' is extracted from <http://hackage.haskell.org/package/ghc-syb-utils ghc-syb-utils> package, including the following comment.
  --
  -- "/GHC AST types tend to have undefined holes, to be filled by later compiler phases. We tag ASTs with their source, so that we can avoid such holes based on who generated the ASTs./"
  data Stage = Parser | Renamer | TypeChecker deriving (Eq,Ord,Show)

  newtype GHC_AST_HOLE = GHC_AST_HOLE Stage deriving ( Typeable )
#else
  newtype GHC_AST_HOLE = GHC_AST_HOLE SYB.Stage deriving ( Typeable )
#endif

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

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

  ghomK_Staged :: forall r d. Data d =>
                     SYB.Stage
                  -> r
                  -> (r -> r -> r)
                  -> GenericQ r
                  -> d
                  -> Homo r
  ghomK_Staged stage z k f x
    | checkItemStage stage x = z'
    | otherwise = foldl k' b (gmapQ (ghomK_Staged 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])

  ghomP_Staged :: forall r s d. Data d =>
                     SYB.Stage
                  -> r              -- default for holes
                  -> GenericQ Bool  -- defines stop nodes
                  -> GenericQ r     -- typically, specialised for stop nodes
                  -> d
                  -> Homo r
  ghomP_Staged stage z p f x
    | checkItemStage stage x = z'
    | p x        = Node (f x) []
    | otherwise  = foldl k b (gmapQ (ghomP_Staged stage z p f) x)
   where
     b = Node (f x) []
     z' = Node z []
     k (Node r chs) nod = Node r (chs++[nod])

  ghomE_Staged :: forall r s d. Data d =>
                     SYB.Stage
                  -> r              -- default for holes
                  -> GenericQ Bool  -- defines stop nodes
                  -> GenericQ r     -- Left : what to do with non-stop nodes
                  -> GenericQ s     -- Right : what to do with stop nodes
                  -> d
                  -> Homo (Either r s)
  ghomE_Staged stage z p f f_stop x
    | checkItemStage stage x = z'
    | p x        = Node (Right (f_stop x)) []
    | otherwise  = foldl k b (gmapQ (ghomE_Staged stage z p f f_stop) x)
   where
     b = Node (Left (f x)) []
     z' = Node (Left z) []
     k (Node r chs) nod = Node r (chs++[nod])

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

  -- | @'ghomBi_Staged' s f x = 'zipRose' ('ghomDyn_Staged' s x) ('ghom_Staged' s f x)@
  ghomBi_Staged :: forall r d. Data d => SYB.Stage -> r -> GenericQ r -> d -> Bi r
  ghomBi_Staged stage z f x = zipRose (ghomDyn_Staged stage x) $ ghom_Staged stage z f x

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

  shapeOf_Staged :: forall d. Data d => SYB.Stage -> d -> Shape
  shapeOf_Staged stage = ghom_Staged stage () (const ())

  -- | Treat 'String' and 'GHC.FastString' as atomic values.
  shapeOf_Staged_ :: forall d. Data d => SYB.Stage -> d -> Shape
  shapeOf_Staged_ stage x = ghomP_Staged stage () pg fg x
   where
    pg :: forall d'. Data d' => d' -> Bool
    pg = mkQ False p_String `extQ` p_FastString
    p_String :: String -> Bool
    p_String _ = True
    p_FastString :: GHC.FastString -> Bool
    p_FastString _ = True
    fg :: forall d''. Data d'' => d'' -> ()
    fg = const ()

  sizeOf_Staged :: forall d. Data d => SYB.Stage -> d -> Int
  sizeOf_Staged stage = sizeOfRose . (shapeOf_Staged stage)

  weightedShapeOf_Staged :: forall d. Data d => SYB.Stage -> d -> Homo Int
  weightedShapeOf_Staged stage = ghomK_Staged stage 1 (+) (const 1)

  -- | Treat 'String' and 'GHC.FastString' as atomic values.
  weightedShapeOf_Staged_ :: forall d. Data d => SYB.Stage -> d -> Homo Int
  weightedShapeOf_Staged_ stage x = weightedRoseSpecial $ ghomP_Staged stage 1 pg fg x
   where
    pg :: forall d'. Data d' => d' -> Bool
    pg = mkQ False p_String `extQ` p_FastString
    p_String :: String -> Bool
    p_String _ = True
    p_FastString :: GHC.FastString -> Bool
    p_FastString _ = True
    fg :: forall d''. Data d'' => d'' -> Int
    fg = mkQ 1 f_String `extQ` f_FastString
    f_String :: String -> Int
    f_String x = length x
    f_FastString :: GHC.FastString -> Int
    f_FastString x = GHC.lengthFS x
    weightedRoseSpecial :: Rose Int -> Rose Int
    weightedRoseSpecial (Node r chs) = foldl k' b (map weightedRoseSpecial chs)
     where
       b = Node r []
       k' (Node rw chs) nod@(Node rw' _) = Node (rw + rw') (chs++[nod])

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

-- Borrowed from HaRe:

-- From @frsoares

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

-- Check the Typeable items
  checkItemStage1 :: (Typeable a) => SYB.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 => SYB.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 Renamer x
#endif

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

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

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