-- |
-- Description: Utility functions for transforming and manipulating
--              'ghc-exactprint' 'Language.Haskell.GHC.ExactPrint.Anns'
module Smuggler2.Anns
  ( mkLocWithAnns,
    mkLoc,
    mkParenT,
    setAnnsForT,
    swapEntryDPT,
  )
where

import Data.Generics as SYB (Data)
import qualified Data.Map.Strict as Map (alter, fromList, insert, lookup, toList, union)
import Data.Maybe (fromMaybe)
import GHC (AnnKeywordId (AnnCloseP, AnnOpenP))
import GhcPlugins (GenLocated (L), Located)
import Language.Haskell.GHC.ExactPrint
  ( Annotation (annEntryDelta, annPriorComments, annsDP),
    TransformT,
    modifyAnnsT,
    uniqueSrcSpanT,
  )
import Language.Haskell.GHC.ExactPrint.Types (DeltaPos (..), KeywordId (G), annNone, mkAnnKey)

-- Inspired by retrie

-- | Generates a unique location and wraps the given ast chunk with that location
-- Also adds a DP and an annotation at that location
mkLocWithAnns :: (Data e, Monad m) => e -> DeltaPos -> [(KeywordId, DeltaPos)] -> TransformT m (Located e)
mkLocWithAnns :: e
-> DeltaPos -> [(KeywordId, DeltaPos)] -> TransformT m (Located e)
mkLocWithAnns e
e DeltaPos
dp [(KeywordId, DeltaPos)]
anns = do
  Located e
le <- SrcSpan -> e -> Located e
forall l e. l -> e -> GenLocated l e
L (SrcSpan -> e -> Located e)
-> TransformT m SrcSpan -> TransformT m (e -> Located e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TransformT m SrcSpan
forall (m :: * -> *). Monad m => TransformT m SrcSpan
uniqueSrcSpanT TransformT m (e -> Located e)
-> TransformT m e -> TransformT m (Located e)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> e -> TransformT m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure e
e
  Located e
-> DeltaPos -> [(KeywordId, DeltaPos)] -> TransformT m (Located e)
forall e (m :: * -> *).
(Data e, Monad m) =>
Located e
-> DeltaPos -> [(KeywordId, DeltaPos)] -> TransformT m (Located e)
setAnnsForT Located e
le DeltaPos
dp [(KeywordId, DeltaPos)]
anns

-- | `mkLoc` generates a unique location and wraps the given ast chunk with that location
-- Also adds an empty annotation at that location
mkLoc :: (Data e, Monad m) => e -> TransformT m (Located e)
mkLoc :: e -> TransformT m (Located e)
mkLoc e
e = e
-> DeltaPos -> [(KeywordId, DeltaPos)] -> TransformT m (Located e)
forall e (m :: * -> *).
(Data e, Monad m) =>
e
-> DeltaPos -> [(KeywordId, DeltaPos)] -> TransformT m (Located e)
mkLocWithAnns e
e ((Int, Int) -> DeltaPos
DP (Int
0, Int
0)) []

-- | Add an open and close paren annotation to a located thing
mkParenT ::
  (Data x, Monad m) =>
  (Located x -> x) ->
  Located x ->
  TransformT m (Located x)
mkParenT :: (Located x -> x) -> Located x -> TransformT m (Located x)
mkParenT Located x -> x
k Located x
e = do
  Located x
pe <- x -> TransformT m (Located x)
forall e (m :: * -> *).
(Data e, Monad m) =>
e -> TransformT m (Located e)
mkLoc (Located x -> x
k Located x
e)
  Located x
_ <- Located x
-> DeltaPos -> [(KeywordId, DeltaPos)] -> TransformT m (Located x)
forall e (m :: * -> *).
(Data e, Monad m) =>
Located e
-> DeltaPos -> [(KeywordId, DeltaPos)] -> TransformT m (Located e)
setAnnsForT Located x
pe ((Int, Int) -> DeltaPos
DP (Int
0, Int
0)) [(AnnKeywordId -> KeywordId
G AnnKeywordId
GHC.AnnOpenP, (Int, Int) -> DeltaPos
DP (Int
0, Int
1)), (AnnKeywordId -> KeywordId
G AnnKeywordId
GHC.AnnCloseP, (Int, Int) -> DeltaPos
DP (Int
0, Int
1))]
  Located x -> Located x -> TransformT m ()
forall a b (m :: * -> *).
(Data a, Data b, Monad m) =>
Located a -> Located b -> TransformT m ()
swapEntryDPT Located x
e Located x
pe
  Located x -> TransformT m (Located x)
forall (m :: * -> *) a. Monad m => a -> m a
return Located x
pe

-- | Set the `ghc-exactprint` annotations for a 'Located' thing
setAnnsForT ::
  (Data e, Monad m) =>
  Located e ->
  DeltaPos ->
  [(KeywordId, DeltaPos)] ->
  TransformT m (Located e)
setAnnsForT :: Located e
-> DeltaPos -> [(KeywordId, DeltaPos)] -> TransformT m (Located e)
setAnnsForT Located e
e DeltaPos
dp [(KeywordId, DeltaPos)]
anns = (Anns -> Anns) -> TransformT m ()
forall (m :: * -> *). Monad m => (Anns -> Anns) -> TransformT m ()
modifyAnnsT ((Maybe Annotation -> Maybe Annotation) -> AnnKey -> Anns -> Anns
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter Maybe Annotation -> Maybe Annotation
f (Located e -> AnnKey
forall a. Constraints a => a -> AnnKey
mkAnnKey Located e
e)) TransformT m ()
-> TransformT m (Located e) -> TransformT m (Located e)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Located e -> TransformT m (Located e)
forall (m :: * -> *) a. Monad m => a -> m a
return Located e
e
  where
    f :: Maybe Annotation -> Maybe Annotation
f Maybe Annotation
Nothing = Annotation -> Maybe Annotation
forall a. a -> Maybe a
Just Annotation
annNone {annEntryDelta :: DeltaPos
annEntryDelta = DeltaPos
dp, annsDP :: [(KeywordId, DeltaPos)]
annsDP = [(KeywordId, DeltaPos)]
anns}
    f (Just Annotation
a) =
      Annotation -> Maybe Annotation
forall a. a -> Maybe a
Just
        Annotation
a
          { annEntryDelta :: DeltaPos
annEntryDelta = DeltaPos
dp,
            annsDP :: [(KeywordId, DeltaPos)]
annsDP =
              Map KeywordId DeltaPos -> [(KeywordId, DeltaPos)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map KeywordId DeltaPos -> [(KeywordId, DeltaPos)])
-> Map KeywordId DeltaPos -> [(KeywordId, DeltaPos)]
forall a b. (a -> b) -> a -> b
$
                Map KeywordId DeltaPos
-> Map KeywordId DeltaPos -> Map KeywordId DeltaPos
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union ([(KeywordId, DeltaPos)] -> Map KeywordId DeltaPos
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(KeywordId, DeltaPos)]
anns) ([(KeywordId, DeltaPos)] -> Map KeywordId DeltaPos
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (Annotation -> [(KeywordId, DeltaPos)]
annsDP Annotation
a))
          }

-- | Swap two 'Located' things' relative position tage ('DeltaPos')
swapEntryDPT ::
  (Data a, Data b, Monad m) => Located a -> Located b -> TransformT m ()
swapEntryDPT :: Located a -> Located b -> TransformT m ()
swapEntryDPT Located a
a Located b
b = (Anns -> Anns) -> TransformT m ()
forall (m :: * -> *). Monad m => (Anns -> Anns) -> TransformT m ()
modifyAnnsT ((Anns -> Anns) -> TransformT m ())
-> (Anns -> Anns) -> TransformT m ()
forall a b. (a -> b) -> a -> b
$ \Anns
anns ->
  let akey :: AnnKey
akey = Located a -> AnnKey
forall a. Constraints a => a -> AnnKey
mkAnnKey Located a
a
      bkey :: AnnKey
bkey = Located b -> AnnKey
forall a. Constraints a => a -> AnnKey
mkAnnKey Located b
b
      aann :: Annotation
aann = Annotation -> Maybe Annotation -> Annotation
forall a. a -> Maybe a -> a
fromMaybe Annotation
annNone (Maybe Annotation -> Annotation) -> Maybe Annotation -> Annotation
forall a b. (a -> b) -> a -> b
$ AnnKey -> Anns -> Maybe Annotation
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup AnnKey
akey Anns
anns
      bann :: Annotation
bann = Annotation -> Maybe Annotation -> Annotation
forall a. a -> Maybe a -> a
fromMaybe Annotation
annNone (Maybe Annotation -> Annotation) -> Maybe Annotation -> Annotation
forall a b. (a -> b) -> a -> b
$ AnnKey -> Anns -> Maybe Annotation
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup AnnKey
bkey Anns
anns
   in AnnKey -> Annotation -> Anns -> Anns
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert
        AnnKey
akey
        Annotation
aann
          { annEntryDelta :: DeltaPos
annEntryDelta = Annotation -> DeltaPos
annEntryDelta Annotation
bann,
            annPriorComments :: [(Comment, DeltaPos)]
annPriorComments = Annotation -> [(Comment, DeltaPos)]
annPriorComments Annotation
bann
          }
        (Anns -> Anns) -> Anns -> Anns
forall a b. (a -> b) -> a -> b
$ AnnKey -> Annotation -> Anns -> Anns
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert
          AnnKey
bkey
          Annotation
bann
            { annEntryDelta :: DeltaPos
annEntryDelta = Annotation -> DeltaPos
annEntryDelta Annotation
aann,
              annPriorComments :: [(Comment, DeltaPos)]
annPriorComments = Annotation -> [(Comment, DeltaPos)]
annPriorComments Annotation
aann
            }
          Anns
anns