-- 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 #-}
module Retrie.ExactPrint.Annotated
  ( -- * Annotated
    Annotated
  , astA
  , annsA
  , seedA
  -- ** Synonyms
  , AnnotatedHsDecl
  , AnnotatedHsExpr
  , AnnotatedHsType
  , AnnotatedImport
  , AnnotatedImports
  , AnnotatedModule
  , AnnotatedPat
  , AnnotatedStmt
  -- ** Operations
  , pruneA
  , graftA
  , transformA
  , trimA
  , printA
    -- * 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.Annotate (Annotate)
import Language.Haskell.GHC.ExactPrint.Types (emptyAnns)

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]
type AnnotatedModule = Annotated (Located HsModule)
type AnnotatedPat = Annotated (Located (Pat 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
  { Annotated ast -> ast
astA :: ast
  -- ^ Examine the actual AST.
  , Annotated ast -> Anns
annsA  :: Anns
  -- ^ Annotations generated/consumed by ghc-exactprint
  , Annotated ast -> Int
seedA  :: Int
  -- ^ Name supply used by ghc-exactprint to generate unique locations.
  }

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

instance Foldable Annotated where
  foldMap :: (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 :: (a -> f b) -> Annotated a -> f (Annotated b)
traverse a -> f b
f Annotated{a
Int
Anns
seedA :: Int
annsA :: Anns
astA :: a
seedA :: forall ast. Annotated ast -> Int
annsA :: forall ast. Annotated ast -> Anns
astA :: forall ast. Annotated ast -> ast
..} =
    (\b
ast -> Annotated :: forall ast. ast -> Anns -> Int -> Annotated ast
Annotated{astA :: b
astA = b
ast, Int
Anns
seedA :: Int
annsA :: Anns
seedA :: Int
annsA :: Anns
..}) (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 -> Anns -> Int -> Annotated ast
forall ast. ast -> Anns -> Int -> Annotated ast
Annotated ast
forall a. Default a => a
D.def Anns
emptyAnns Int
0

instance (Data ast, Monoid ast) => Semigroup (Annotated ast) where
  <> :: Annotated ast -> Annotated ast -> Annotated ast
(<>) = Annotated ast -> Annotated ast -> Annotated ast
forall a. Monoid a => a -> a -> a
mappend

instance (Data ast, Monoid ast) => Monoid (Annotated ast) where
  mempty :: Annotated ast
mempty = ast -> Anns -> Int -> Annotated ast
forall ast. ast -> Anns -> Int -> Annotated ast
Annotated ast
forall a. Monoid a => a
mempty Anns
emptyAnns Int
0
  mappend :: Annotated ast -> Annotated ast -> Annotated ast
mappend Annotated ast
a1 (Annotated ast
ast2 Anns
anns 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
<$> Anns -> ast -> TransformT Identity ast
forall a (m :: * -> *).
(Data a, Monad m) =>
Anns -> a -> TransformT m a
graftT Anns
anns ast
ast2

-- | 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 -> Anns -> Int -> Annotated ast
unsafeMkA :: ast -> Anns -> Int -> Annotated ast
unsafeMkA = ast -> Anns -> Int -> Annotated ast
forall ast. ast -> Anns -> Int -> Annotated ast
Annotated

-- | Transform an 'Annotated' thing.
transformA
  :: Monad m => Annotated ast1 -> (ast1 -> TransformT m ast2) -> m (Annotated ast2)
transformA :: Annotated ast1 -> (ast1 -> TransformT m ast2) -> m (Annotated ast2)
transformA (Annotated ast1
ast Anns
anns Int
seed) ast1 -> TransformT m ast2
f = do
  (ast2
ast',(Anns
anns',Int
seed'),[String]
_) <- Int -> Anns -> TransformT m ast2 -> m (ast2, (Anns, Int), [String])
forall (m :: * -> *) a.
Int -> Anns -> TransformT m a -> m (a, (Anns, Int), [String])
runTransformFromT Int
seed Anns
anns (ast1 -> TransformT m ast2
f ast1
ast)
  Annotated ast2 -> m (Annotated ast2)
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 -> Anns -> Int -> Annotated ast2
forall ast. ast -> Anns -> Int -> Annotated ast
Annotated ast2
ast' Anns
anns' 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 :: Annotated ast -> TransformT m ast
graftA (Annotated ast
x Anns
anns Int
_) = Anns -> ast -> TransformT m ast
forall a (m :: * -> *).
(Data a, Monad m) =>
Anns -> a -> TransformT m a
graftT Anns
anns 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 :: ast -> TransformT m (Annotated ast)
pruneA ast
ast = ast -> Anns -> Int -> Annotated ast
forall ast. ast -> Anns -> Int -> Annotated ast
Annotated ast
ast (Anns -> Int -> Annotated ast)
-> TransformT m Anns -> TransformT m (Int -> Annotated ast)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TransformT m Anns
forall (m :: * -> *). Monad m => TransformT m Anns
getAnnsT TransformT m (Int -> Annotated ast)
-> TransformT m Int -> TransformT m (Annotated ast)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((Anns, Int) -> Int) -> TransformT m Int
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Anns, Int) -> Int
forall a b. (a, b) -> b
snd

-- | 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 :: 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

-- | Exactprint an 'Annotated' thing.
printA :: Annotate ast => Annotated (Located ast) -> String
printA :: Annotated (Located ast) -> String
printA (Annotated Located ast
ast Anns
anns Int
_) = Located ast -> Anns -> String
forall ast. Annotate ast => Located ast -> Anns -> String
exactPrint Located ast
ast Anns
anns