Imports.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}

-- | Putting a module's imports in order.
module Tilia.Imports
  ( normalizeImports,
  )
where

import Data.Char (isAlphaNum)
import Data.Choice (Choice, isTrue)
import Data.Function (on, (&))
import Data.List (groupBy, sortOn)
import Data.Text (Text)
import Data.Text qualified as T
import GHC.Data.FastString (unpackFS)
import GHC.Hs
import GHC.Types.Name.Occurrence (occNameString)
import GHC.Types.Name.Reader (RdrName, rdrNameOcc)
import GHC.Types.PkgQual (RawPkgQual (..))
import GHC.Types.SourceText (StringLiteral (..))
import GHC.Types.SrcLoc
import Tilia.Comments (Comment (..), commentTrailing, commentsWithin)
import Tilia.Span (endPoint, isSingleLine, startPoint)
import Tilia.Span.Ghc (bracketsSpan, spanOf, spanOfSrcSpan, tokenSpan)

-- | Sort and fold together a module's imports.
normalizeImports ::
  -- | Whether @ImplicitPrelude@ is on.
  Choice "implicitPrelude" ->
  -- | Source lines the block must not be sorted across.
  [Int] ->
  -- | Source lines the names of a list must not be sorted across.
  [Int] ->
  -- | The module's comments.
  [Comment] ->
  -- | Original imports.
  [LImportDecl GhcPs] ->
  -- | Normalized imports.
  [LImportDecl GhcPs]
normalizeImports implicitPrelude barriers directives written imports =
  concatMap stretch (cutAtBarriers barriers tidied)
  where
    tidied = fmap (fmap (tidyImportList directives written)) imports
    stretch is =
      foldRuns
        (combineImports directives written)
        [((importIdentity implicitPrelude i, hiddenNames i, alone i), i) | i <- is]
    alone i
      | any strands (spanOf i) = startLineOf i
      | otherwise = 0
      where
        strands s =
          any
            (unanchored (itemStarts i))
            (filter loose (commentsWithin s written))
        loose = not . commentTrailing
    unanchored starts c = not (any (> endPoint (commentSpan c)) starts)
    startLineOf i = case srcSpanStart (getLocA i) of
      RealSrcLoc l _ -> srcLocLine l
      _ -> 0

-- | The names an import hides, where it hides any, so that only imports
-- hiding the same names are folded together.
hiddenNames ::
  LImportDecl GhcPs ->
  Maybe [(NameIdentity, Maybe (Bool, [NameIdentity]))]
hiddenNames (L _ decl) = case ideclImportList decl of
  Just (EverythingBut, L _ items) -> Just (fmap (hidden . unLoc) items)
  _ -> Nothing
  where
    hidden item = (nameIdentity item, subordinates item)
    subordinates = \case
      IEThingAll{} -> Just (True, [])
      IEThingWith _ _ NoIEWildcard subs _ -> Just (False, names subs)
      IEThingWith _ _ (IEWildcard _) subs _ -> Just (True, names subs)
      _ -> Nothing
    names = fmap (wrappedNameIdentity . unLoc)

-- | Where every name an import lists begins, each as a line and a column.
--
-- The names inside a thing's own brackets count too, so @T (A, B)@ gives
-- three positions: the start of @T@, of @A@ and of @B@. An import with no
-- list has none at all.
itemStarts :: LImportDecl GhcPs -> [(Int, Int)]
itemStarts (L _ decl) = case ideclImportList decl of
  Nothing -> []
  Just (_, L _ items) -> concatMap starts items
  where
    starts item =
      foldMap ((: []) . startPoint) (spanOf item)
        <> inside (unLoc item)
    inside = \case
      IEThingWith _ _ _ members _ ->
        concatMap (foldMap ((: []) . startPoint) . spanOf) members
      _ -> []

-- | Cut imports, or the names of a list, into the stretches the barriers
-- leave between them.
cutAtBarriers ::
  -- | Lines not to be sorted across, in ascending order.
  [Int] ->
  -- | What to cut, as it was written.
  [GenLocated SrcSpanAnnA a] ->
  -- | One stretch per run between barriers, in the same order.
  [[GenLocated SrcSpanAnnA a]]
cutAtBarriers [] xs = [xs]
cutAtBarriers barriers xs =
  groupBy ((==) `on` fst) [(between i, i) | i <- xs] & fmap (fmap snd)
  where
    between i = length (takeWhile (< lineOf i) barriers)
    lineOf i = case srcSpanStart (getLocA i) of
      RealSrcLoc l _ -> srcLocLine l
      _ -> 0

-- | Sort by the keys, then replace each run of equal keys by one value
-- folded out of it.
foldRuns ::
  (Ord k) =>
  -- | How to fold two values that share a key.
  (a -> a -> a) ->
  -- | The values, each under the key it sorts and folds by.
  [(k, a)] ->
  -- | One value per distinct key, in key order.
  [a]
foldRuns fold' =
  fmap (foldl1 fold' . fmap snd) . groupBy ((==) `on` fst) . sortOn fst

-- | What has to agree before two imports may be folded together.
--
-- The derived ordering is the whole of the policy: the fields are compared
-- in the order they are declared, which is the order imports are printed
-- in. The two leading fields are about reading rather than about identity.
-- A @Prelude@ that only refines what the module already has goes at the
-- end, since looking for it among the @D@s would be looking for the least
-- interesting line in the block. The package goes before the module name so
-- that the imports from one package stay in one run; sorting by module
-- first would interleave them and hide who provides what.
data ImportIdentity = ImportIdentity
  { -- | A @Prelude@ that only refines what the module already has.
    iiRefinesPrelude :: Bool,
    -- | The package it names, where it names one.
    iiPackage :: (Int, Text),
    -- | The module it brings in.
    iiModule :: Text,
    -- | Written @{-\# SOURCE \#-}@?
    iiBoot :: Bool,
    -- | Written @safe@?
    iiSafe :: Bool,
    -- | Written @qualified@?
    iiQualified :: Bool,
    -- | The name it was given with @as@, where it was given one.
    iiAlias :: Maybe Text,
    -- | Whether its list hides rather than names, where it has a list.
    iiHides :: Maybe Bool,
    -- | Whether it is a @quote@ rather than a @splice@ import, where it is
    -- levelled at all.
    iiLevel :: Maybe Bool
  }
  deriving (Eq, Ord, Show)

-- | The identity of one import.
importIdentity ::
  -- | Whether @ImplicitPrelude@ is on.
  Choice "implicitPrelude" ->
  -- | The import.
  LImportDecl GhcPs ->
  ImportIdentity
importIdentity implicitPrelude (L _ decl) =
  ImportIdentity
    { iiRefinesPrelude =
        isTrue implicitPrelude && named (ideclName decl) == T.pack "Prelude",
      iiPackage = package (ideclPkgQual decl),
      iiModule = named (ideclName decl),
      iiBoot = ideclSource decl == IsBoot,
      iiSafe = ideclSafe decl,
      iiQualified = isImportDeclQualified (ideclQualified decl),
      iiAlias = named <$> ideclAs decl,
      iiHides = hides . fst <$> ideclImportList decl,
      iiLevel = lifted (ideclLevelSpec decl)
    }
  where
    named = T.pack . moduleNameString . unLoc
    package = \case
      NoRawPkgQual -> (0, T.empty)
      RawPkgQual (sl_fs -> fs)
        | name == T.pack "this" -> (2, T.empty)
        | otherwise -> (1, name)
        where
          name = T.pack (unpackFS fs)
    hides = \case
      Exactly -> False
      EverythingBut -> True
    lifted = \case
      NotLevelled -> Nothing
      LevelStylePre l -> Just (quoted l)
      LevelStylePost l -> Just (quoted l)
    quoted = \case
      ImportDeclSplice -> False
      ImportDeclQuote -> True

-- | Keep the first import and give it everything the second named.
combineImports ::
  -- | Source lines the names of a list must not be sorted across.
  [Int] ->
  -- | The module's comments, which decide how far the span may grow.
  [Comment] ->
  -- | The import kept, whose spelling the result takes.
  LImportDecl GhcPs ->
  -- | The import folded in, which gives up its list and its span.
  LImportDecl GhcPs ->
  -- | The first, listing what both named, across both their spans.
  LImportDecl GhcPs
combineImports directives written (L ann kept) (L other folded) =
  L
    ann{entry = EpaSpan (combineSrcSpans (locA ann) (locA other))}
    kept{ideclImportList = both (ideclImportList kept) (ideclImportList folded)}
  where
    both (Just (interpretation, L l xs)) (Just (_, L l' ys)) =
      Just
        ( interpretation,
          L
            (bracketsAcross (widened written l l') l')
            (tidyImportItems directives written interpretation (xs <> ys))
        )
    both _ _ = Nothing

-- | Take the second list's brackets where only they were written across
-- lines, so that a list folded out of several breaks if any of them did.
bracketsAcross :: EpAnn (AnnList a) -> EpAnn (AnnList a) -> EpAnn (AnnList a)
bracketsAcross kept folded
  | onOneLine kept && not (onOneLine folded) =
      kept{anns = (anns kept){al_brackets = al_brackets (anns folded)}}
  | otherwise = kept
  where
    onOneLine = all isSingleLine . bracketsSpan . al_brackets . anns

-- | An import with its list sorted and the entries naming one thing folded
-- together. An import with no list is left as it is.
tidyImportList :: [Int] -> [Comment] -> ImportDecl GhcPs -> ImportDecl GhcPs
tidyImportList directives written decl =
  decl
    { ideclImportList = tidied <$> ideclImportList decl
    }
  where
    tidied (interpretation, items) =
      ( interpretation,
        tidyImportItems directives written interpretation <$> items
      )

-- | Sort an import list and fold together the entries naming one thing,
-- each stretch between directives on its own.
--
-- @import M (T (A), T (B))@ names one type twice and comes out as @import M
-- (T (A, B))@.
tidyImportItems ::
  [Int] ->
  [Comment] ->
  ImportListInterpretation ->
  [LIE GhcPs] ->
  [LIE GhcPs]
tidyImportItems directives written interpretation items
  | any (unnameable . unLoc) items = items
  | otherwise = concatMap tidied (cutAtBarriers directives items)
  where
    tidied stretch =
      foldRuns
        (wider written)
        [ ((nameIdentity (unLoc i), apart (unLoc i)), fmap sortSubnames i)
        | i <- stretch
        ]
    -- A bare name in a hiding list also hides any data constructor of that
    -- name, which the name with its own parentheses does not.
    apart item = case (interpretation, item) of
      (EverythingBut, IEThingAll{}) -> True
      (EverythingBut, IEThingWith{}) -> True
      _ -> False
    unnameable = \case
      IEVar{} -> False
      IEThingAbs{} -> False
      IEThingAll{} -> False
      IEThingWith{} -> False
      _ -> True

-- | Cover both of these regions, if anything was written between them.
widened :: [Comment] -> EpAnn ann -> EpAnn ann -> EpAnn ann
widened written a b
  | any holdsComment (spanOfSrcSpan combined) = a{entry = EpaSpan combined}
  | otherwise = a
  where
    combined = combineSrcSpans (locA a) (locA b)
    holdsComment s = not (null (commentsWithin s written))

-- | Of two entries naming one thing, the one that brings in more of it.
wider :: [Comment] -> LIE GhcPs -> LIE GhcPs -> LIE GhcPs
wider written (L ann kept) (L other folded) =
  L (widened written ann other) (combine kept folded)
  where
    combine a b = case (a, b) of
      (IEThingAll x n _, _) -> IEThingAll x n Nothing
      (_, IEThingAll x n _) -> IEThingAll x n Nothing
      (IEThingWith x n wildcard subs _, IEThingWith x' _ wildcard' subs' _) ->
        IEThingWith
          (parensAcross x x')
          n
          (eitherWildcard wildcard wildcard')
          (dedupeSubnames (subs <> subs'))
          Nothing
      (IEThingWith x n wildcard subs _, _) ->
        IEThingWith x n wildcard subs Nothing
      (_, IEThingWith x n wildcard subs _) ->
        IEThingWith x n wildcard subs Nothing
      (IEVar _ n _, _) -> IEVar Nothing n Nothing
      _ -> a

    eitherWildcard a b = case (a, b) of
      (NoIEWildcard, NoIEWildcard) -> NoIEWildcard
      _ -> IEWildcard 0

    parensAcross (warning, (open, dots, comma, close)) (_, (open', _, _, close'))
      | onOneLine open close && not (onOneLine open' close') =
          (warning, (open', dots, comma, close'))
      | otherwise = (warning, (open, dots, comma, close))
    onOneLine open close = all isSingleLine (tokenSpan open <> tokenSpan close)

-- | An entry with the names in its own brackets sorted and deduplicated.
sortSubnames :: IE GhcPs -> IE GhcPs
sortSubnames = \case
  IEThingWith x n wildcard subs doc ->
    IEThingWith x n wildcard (dedupeSubnames subs) doc
  other -> other

-- | Deduplicate names in sub-lists.
dedupeSubnames :: [LIEWrappedName GhcPs] -> [LIEWrappedName GhcPs]
dedupeSubnames subs = foldRuns const [(wrappedNameIdentity (unLoc s), s) | s <- subs]

-- | Where a name sorts, and what makes two of them name one thing.
--
-- The derived ordering is the whole of the policy: the fields are compared
-- in the order they are declared. Names are grouped first by what kind of
-- thing is named, so that the @pattern@s and the @type@s of an import list
-- stay together. Then the ones spelled with letters come before the ones
-- spelled with punctuation, which gathers the operators at the end where
-- they are easy to find; ordering by character code would scatter them,
-- some before the letters and some after.
data NameIdentity = NameIdentity
  { -- | Which kind of thing is named: a plain name, then a @default@, a
    -- @pattern@, a @type@ and a @data@.
    niKind :: Int,
    -- | Is it spelled with punctuation rather than letters?
    niPunctuation :: Bool,
    -- | The name as it was written.
    niName :: String
  }
  deriving (Eq, Ord, Show)

-- | The identity of the name an entry is about.
nameIdentity :: IE GhcPs -> NameIdentity
nameIdentity = \case
  IEVar _ x _ -> wrappedNameIdentity (unLoc x)
  IEThingAbs _ x _ -> wrappedNameIdentity (unLoc x)
  IEThingAll _ x _ -> wrappedNameIdentity (unLoc x)
  IEThingWith _ x _ _ _ -> wrappedNameIdentity (unLoc x)
  _ ->
    NameIdentity
      { niKind = maxBound,
        niPunctuation = True,
        niName = ""
      }

-- | The identity of one name, as an import list writes it.
wrappedNameIdentity :: IEWrappedName GhcPs -> NameIdentity
wrappedNameIdentity = \case
  IEName _ x -> spelled 0 x
  IEDefault _ x -> spelled 1 x
  IEPattern _ x -> spelled 2 x
  IEType _ x -> spelled 3 x
  IEData _ x -> spelled 4 x
  where
    spelled :: Int -> LocatedN RdrName -> NameIdentity
    spelled kind (unLoc -> name) =
      NameIdentity
        { niKind = kind,
          niPunctuation = punctuation text,
          niName = text
        }
      where
        text = occNameString (rdrNameOcc name)
    punctuation = \case
      (c : _) -> not (isAlphaNum c)
      [] -> False