content: -- Copyright 2017 Google Inc.
  --
  -- Licensed under the Apache License, Version 2.0 (the "License");
  -- you may not use this file except in compliance with the License.
  -- You may obtain a copy of the License at
  --
  --     http://www.apache.org/licenses/LICENSE-2.0
  --
  -- Unless required by applicable law or agreed to in writing, software
  -- distributed under the License is distributed on an "AS IS" BASIS,
  -- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  -- See the License for the specific language governing permissions and
  -- limitations under the License.
  
  {-# LANGUAGE GADTs, ScopedTypeVariables, TypeOperators #-}
  
  -- | Landmine-proof GHC AST traversal.
  module Language.Haskell.Indexer.Backend.GhcLens
      ( careful
      , carefulUniverse
      -- * Uniplate-like API (see details).
      , universe, children
      , universeBi
      ) where
  
  import Control.Applicative
  import Control.Exception
  import Data.Data
  import Data.Hashable
  import qualified Data.HashSet as HS
  import Data.IORef
  import System.IO.Unsafe
  import System.Mem.StableName
  
  import Panic
  
  data SeenNode where
    SeenNode :: StableName a -> SeenNode
  
  instance Eq SeenNode where
    SeenNode sn == SeenNode sn' = eqStableName sn sn'
  
  instance Hashable SeenNode where
    hashWithSalt salt (SeenNode sn) = hashWithSalt salt sn
  
  type SeenNodes = IORef (HS.HashSet SeenNode)
  
  {-|
  A function that traverses a structure of type @s@ and finds values of
  type @a@, applying a function @f@ to lift them into a `Monoid` @m@.
  
  Does not recurse further into values of type @a@. If it detects a node
  it has already traversed, it is skipped using the given Monoid's 'mempty'
  in order to prevent traversing self-referential structures forever.
  
  Nodes that would immediately throw a GHC 'Panic' are also skipped, but other
  GHC exceptions are re-thrown.
  -}
  careful :: (Monoid m, Typeable a, Data s) => (a -> m) -> s -> m
  careful f x = unsafeDupablePerformIO $ do
    seenNodes <- newIORef HS.empty
    carefulChildren seenNodes f x
  
  -- | Returns all subelements of the same type (including self).
  carefulUniverse :: Data a => a -> [a]
  carefulUniverse a = a:careful carefulUniverse a
  
  carefulChildren :: forall m a s. (Monoid m, Typeable a, Data s)
    => SeenNodes -> (a -> m) -> s -> IO m
  carefulChildren nodesRef f
    = gmapQr (liftA2 mappend) (pure mempty) (careful' nodesRef f)
  
  careful' :: forall m a s. (Monoid m, Typeable a, Data s)
    => SeenNodes -> (a -> m) -> s -> IO m
  careful' nodesRef f node
    = handleGhcException handlePanic $ do
      _ <- evaluate node -- check for landmines before we do anything
      nodeName <- SeenNode <$> makeStableName node
      seenNodes <- readIORef nodesRef
      if HS.member nodeName seenNodes
        then return mempty
        else do
          writeIORef nodesRef $! HS.insert nodeName seenNodes
          case eqT of
            Just (Refl :: s :~: a) -> return (f node)
            Nothing -> carefulChildren nodesRef f node
    where
      handlePanic (Panic _) = return mempty
      handlePanic e = throwGhcExceptionIO e
  
  -- | All values matching the type, deeply from the structure (including root).
  universe :: Data a => a -> [a]
  universe = carefulUniverse
  
  -- | All values of type 'a' below the root of type 's'.
  --
  -- Warning! Edge case behavior differs from uniplate - if 'a'~'s', uniplate
  -- would return only the root, while this returns everything except the root.
  -- Such usecase is not intended for universeBi anyway - use universe or
  -- children.
  universeBi :: forall s a. (Data a, Data s) => s -> [a]
  universeBi = concatMap carefulUniverse . (careful return :: s -> [a])
  
  -- | Nearest descendants, excluding root.
  children :: Data a => a -> [a]
  children = careful return
  
  