-- Copyright (c) Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE CPP #-}
module Retrie.ExactPrint.Annotated
  ( -- * Annotated
    Annotated
  , astA
  , seedA
  -- ** Synonyms
  , AnnotatedHsDecl
  , AnnotatedHsExpr
  , AnnotatedHsType
  , AnnotatedImport
  , AnnotatedImports
  , AnnotatedModule
  , AnnotatedPat
  , AnnotatedStmt
  -- ** Operations
  , pruneA
  , graftA
  , transformA
  , trimA
  , setEntryDPA
  , printA
  , printA'
  , showAstA
    -- * Internal
  , unsafeMkA
  ) where

import Control.Monad.State.Lazy hiding (fix)
import Data.Default as D

import Data.Functor.Identity

import Language.Haskell.GHC.ExactPrint hiding
  ( -- cloneT
    -- setEntryDP
  -- , setEntryDPT
  -- , transferEntryDPT
    transferEntryDP
  )
-- import Language.Haskell.GHC.ExactPrint.ExactPrint (ExactPrint(..))
import Language.Haskell.GHC.ExactPrint.Utils

import Retrie.GHC
import Retrie.SYB

-- Annotated -----------------------------------------------------------------

type AnnotatedHsDecl = Annotated (LHsDecl GhcPs)
type AnnotatedHsExpr = Annotated (LHsExpr GhcPs)
type AnnotatedHsType = Annotated (LHsType GhcPs)
type AnnotatedImport = Annotated (LImportDecl GhcPs)
type AnnotatedImports = Annotated [LImportDecl GhcPs]
#if __GLASGOW_HASKELL__ >= 906
type AnnotatedModule = Annotated (Located (HsModule GhcPs))
#else
type AnnotatedModule = Annotated (Located HsModule)
#endif
type AnnotatedPat = Annotated (LPat GhcPs)
type AnnotatedStmt = Annotated (LStmt GhcPs (LHsExpr GhcPs))

-- | 'Annotated' packages an AST fragment with the annotations necessary to
-- 'exactPrint' or 'transform' that AST.
data Annotated ast = Annotated
  { forall ast. Annotated ast -> ast
astA :: ast
  -- ^ Examine the actual AST.
  , forall ast. Annotated ast -> Int
seedA  :: Int
  -- ^ Name supply used by ghc-exactprint to generate unique locations.
  }
deriving instance (Data ast) => Data (Annotated ast)

instance Functor Annotated where
  fmap :: forall a b. (a -> b) -> Annotated a -> Annotated b
fmap a -> b
f Annotated{a
Int
astA :: forall ast. Annotated ast -> ast
seedA :: forall ast. Annotated ast -> Int
astA :: a
seedA :: Int
..} = Annotated{astA :: b
astA = a -> b
f a
astA, Int
seedA :: Int
seedA :: Int
..}

instance Foldable Annotated where
  foldMap :: forall m a. Monoid m => (a -> m) -> Annotated a -> m
foldMap a -> m
f = a -> m
f (a -> m) -> (Annotated a -> a) -> Annotated a -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Annotated a -> a
forall ast. Annotated ast -> ast
astA

instance Traversable Annotated where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Annotated a -> f (Annotated b)
traverse a -> f b
f Annotated{a
Int
astA :: forall ast. Annotated ast -> ast
seedA :: forall ast. Annotated ast -> Int
astA :: a
seedA :: Int
..} =
    (\b
ast -> Annotated{astA :: b
astA = b
ast, Int
seedA :: Int
seedA :: Int
..}) (b -> Annotated b) -> f b -> f (Annotated b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
astA

instance Default ast => Default (Annotated ast) where
  def :: Annotated ast
def = ast -> Int -> Annotated ast
forall ast. ast -> Int -> Annotated ast
Annotated ast
forall a. Default a => a
D.def Int
0

instance (Data ast, Monoid ast) => Semigroup (Annotated ast) where
  Annotated ast
a1 <> :: Annotated ast -> Annotated ast -> Annotated ast
<> (Annotated ast
ast2 Int
_) =
    Identity (Annotated ast) -> Annotated ast
forall a. Identity a -> a
runIdentity (Identity (Annotated ast) -> Annotated ast)
-> Identity (Annotated ast) -> Annotated ast
forall a b. (a -> b) -> a -> b
$ Annotated ast
-> (ast -> TransformT Identity ast) -> Identity (Annotated ast)
forall (m :: * -> *) ast1 ast2.
Monad m =>
Annotated ast1 -> (ast1 -> TransformT m ast2) -> m (Annotated ast2)
transformA Annotated ast
a1 ((ast -> TransformT Identity ast) -> Identity (Annotated ast))
-> (ast -> TransformT Identity ast) -> Identity (Annotated ast)
forall a b. (a -> b) -> a -> b
$ \ ast
ast1 ->
      ast -> ast -> ast
forall a. Monoid a => a -> a -> a
mappend ast
ast1 (ast -> ast) -> TransformT Identity ast -> TransformT Identity ast
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ast -> TransformT Identity ast
forall a. a -> TransformT Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return ast
ast2

instance (Data ast, Monoid ast) => Monoid (Annotated ast) where
  mempty :: Annotated ast
mempty = ast -> Int -> Annotated ast
forall ast. ast -> Int -> Annotated ast
Annotated ast
forall a. Monoid a => a
mempty Int
0

-- | Construct an 'Annotated'.
-- This should really only be used in the parsing functions, hence the scary name.
-- Don't use this unless you know what you are doing.
unsafeMkA :: ast -> Int -> Annotated ast
unsafeMkA :: forall ast. ast -> Int -> Annotated ast
unsafeMkA = ast -> Int -> Annotated ast
forall ast. ast -> Int -> Annotated ast
Annotated

-- | Transform an 'Annotated' thing.
transformA
  :: Monad m => Annotated ast1 -> (ast1 -> TransformT m ast2) -> m (Annotated ast2)
transformA :: forall (m :: * -> *) ast1 ast2.
Monad m =>
Annotated ast1 -> (ast1 -> TransformT m ast2) -> m (Annotated ast2)
transformA (Annotated ast1
ast Int
seed) ast1 -> TransformT m ast2
f = do
  (ast2
ast',Int
seed',[String]
_) <- Int -> TransformT m ast2 -> m (ast2, Int, [String])
forall (m :: * -> *) a.
Int -> TransformT m a -> m (a, Int, [String])
runTransformFromT Int
seed (ast1 -> TransformT m ast2
f ast1
ast)
  Annotated ast2 -> m (Annotated ast2)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Annotated ast2 -> m (Annotated ast2))
-> Annotated ast2 -> m (Annotated ast2)
forall a b. (a -> b) -> a -> b
$ ast2 -> Int -> Annotated ast2
forall ast. ast -> Int -> Annotated ast
Annotated ast2
ast' Int
seed'

-- | Graft an 'Annotated' thing into the current transformation.
-- The resulting AST will have proper annotations within the 'TransformT'
-- computation. For example:
--
-- > mkDeclList :: IO (Annotated [LHsDecl GhcPs])
-- > mkDeclList = do
-- >   ad1 <- parseDecl "myId :: a -> a"
-- >   ad2 <- parseDecl "myId x = x"
-- >   transformA ad1 $ \ d1 -> do
-- >     d2 <- graftA ad2
-- >     return [d1, d2]
--
graftA :: (Data ast, Monad m) => Annotated ast -> TransformT m ast
graftA :: forall ast (m :: * -> *).
(Data ast, Monad m) =>
Annotated ast -> TransformT m ast
graftA (Annotated ast
x Int
_) = ast -> TransformT m ast
forall a. a -> TransformT m a
forall (m :: * -> *) a. Monad m => a -> m a
return ast
x

-- | Encapsulate something in the current transformation into an 'Annotated'
-- thing. This is the inverse of 'graftT'. For example:
--
-- > splitHead :: Monad m => Annotated [a] -> m (Annotated a, Annotated [a])
-- > splitHead l = fmap astA $ transformA l $ \(x:xs) -> do
-- >   y <- pruneA x
-- >   ys <- pruneA xs
-- >   return (y, ys)
--
pruneA :: (Data ast, Monad m) => ast -> TransformT m (Annotated ast)
pruneA :: forall ast (m :: * -> *).
(Data ast, Monad m) =>
ast -> TransformT m (Annotated ast)
pruneA ast
ast = ast -> Int -> Annotated ast
forall ast. ast -> Int -> Annotated ast
Annotated ast
ast (Int -> Annotated ast)
-> TransformT m Int -> TransformT m (Annotated ast)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> Int) -> TransformT m Int
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Int -> Int
forall a. a -> a
id

-- | Trim the annotation data to only include annotations for 'ast'.
-- (Usually, the annotation data is a superset of what is necessary.)
-- Also freshens all source locations, so filename information
-- in annotation keys is discarded.
--
-- Note: not commonly needed, but useful if you want to inspect the annotation
-- data directly and don't want to wade through a mountain of output.
trimA :: Data ast => Annotated ast -> Annotated ast
trimA :: forall ast. Data ast => Annotated ast -> Annotated ast
trimA = Identity (Annotated ast) -> Annotated ast
forall a. Identity a -> a
runIdentity (Identity (Annotated ast) -> Annotated ast)
-> (Annotated ast -> Identity (Annotated ast))
-> Annotated ast
-> Annotated ast
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Annotated ()
-> (() -> TransformT Identity ast) -> Identity (Annotated ast)
forall (m :: * -> *) ast1 ast2.
Monad m =>
Annotated ast1 -> (ast1 -> TransformT m ast2) -> m (Annotated ast2)
transformA Annotated ()
nil ((() -> TransformT Identity ast) -> Identity (Annotated ast))
-> (Annotated ast -> () -> TransformT Identity ast)
-> Annotated ast
-> Identity (Annotated ast)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TransformT Identity ast -> () -> TransformT Identity ast
forall a b. a -> b -> a
const (TransformT Identity ast -> () -> TransformT Identity ast)
-> (Annotated ast -> TransformT Identity ast)
-> Annotated ast
-> ()
-> TransformT Identity ast
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Annotated ast -> TransformT Identity ast
forall ast (m :: * -> *).
(Data ast, Monad m) =>
Annotated ast -> TransformT m ast
graftA
  where
    nil :: Annotated ()
    nil :: Annotated ()
nil = Annotated ()
forall a. Monoid a => a
mempty

setEntryDPA :: (Default an)
            => Annotated (LocatedAn an ast) -> DeltaPos -> Annotated (LocatedAn an ast)
setEntryDPA :: forall an ast.
Default an =>
Annotated (LocatedAn an ast)
-> DeltaPos -> Annotated (LocatedAn an ast)
setEntryDPA (Annotated LocatedAn an ast
ast Int
s) DeltaPos
dp = LocatedAn an ast -> Int -> Annotated (LocatedAn an ast)
forall ast. ast -> Int -> Annotated ast
Annotated (LocatedAn an ast -> DeltaPos -> LocatedAn an ast
forall t a. Default t => LocatedAn t a -> DeltaPos -> LocatedAn t a
setEntryDP LocatedAn an ast
ast DeltaPos
dp) Int
s

-- | Exactprint an 'Annotated' thing.
printA :: (Data ast, ExactPrint ast) => Annotated ast -> String
printA :: forall ast. (Data ast, ExactPrint ast) => Annotated ast -> String
printA (Annotated ast
ast Int
_) = ast -> String
forall ast. ExactPrint ast => ast -> String
exactPrint ast
ast
    String -> String -> String
forall c. c -> String -> c
`debug` (String
"printA:" String -> String -> String
forall a. [a] -> [a] -> [a]
++ ast -> String
forall a. Data a => a -> String
showAst ast
ast)

printA' :: (Data ast, ExactPrint ast) => Annotated ast -> String
printA' :: forall ast. (Data ast, ExactPrint ast) => Annotated ast -> String
printA' (Annotated ast
ast Int
_) = String
"[" String -> String -> String
forall a. [a] -> [a] -> [a]
++ ast -> String
forall ast. ExactPrint ast => ast -> String
exactPrint ast
ast String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"]\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ ast -> String
forall a. Data a => a -> String
showAst ast
ast

-- | showAst an 'Annotated' thing.
showAstA :: (Data ast, ExactPrint ast) => Annotated ast -> String
showAstA :: forall ast. (Data ast, ExactPrint ast) => Annotated ast -> String
showAstA (Annotated ast
ast Int
_) = ast -> String
forall a. Data a => a -> String
showAst ast
ast