content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE CPP #-}
  module AutoSplit.Split
    ( NonExhaustivePatterns(..)
    , modifyModule
    , splitName
    ) where
  
  import qualified Control.Monad.Trans.Writer.CPS as Writer
  import qualified Data.Generics as Syb
  import qualified Data.List as List
  #if MIN_VERSION_ghc(9,12,0)
  import qualified Data.List.NonEmpty as NE
  #endif
  import           Data.Maybe
  import           Data.Monoid (Any(..))
  import           Data.String (IsString, fromString)
  import qualified Language.Haskell.GHC.ExactPrint as EP
  
  import qualified AutoSplit.GhcFacade as Ghc
  import           AutoSplit.Shared
  
  data NonExhaustivePatterns = NonExhaustivePatterns
    { patternIds :: [Ghc.Id]
    , patternNablas :: [Ghc.Nabla]
    , srcCodeLoc :: Ghc.RealSrcSpan
    }
  
  -- | Applies pattern split transformation and updates the module file.
  modifyModule
    :: Ghc.GlobalRdrEnv
    -> Ghc.ParsedSource
    -> Bool
    -> Ghc.Bag NonExhaustivePatterns
    -> FilePath
    -> IO Bool -- ^ True if module was updated
  modifyModule gblRdrEnv parsedMod usesCpp missingPatWarns filePath = do
    let ast = EP.makeDeltaAst
            $ searchAndMark (Ghc.bagToList missingPatWarns) parsedMod
    case splitPattern gblRdrEnv (Ghc.bagToList missingPatWarns) ast of
      (ps, Any True) ->
        -- If the source contains CPP, newlines are appended
        -- to the end of the file when exact printing. The simple
        -- solution is to remove trailing newlines after exact printing
        -- if the source contains CPP comments.
        let removeTrailingNewlines
              | usesCpp =
                  reverse . ('\n' :) . dropWhile (== '\n') . reverse
              | otherwise = id
            printed = removeTrailingNewlines $ EP.exactPrint ps
         in True <$ writeFile filePath printed
      _ -> pure False
  
  -- | Before applying delta transformation, find the expressions that go with
  -- non exhaustive patterns and mark them with a special comment containing the
  -- index of that pattern. This must be done first because source code locations
  -- are removed by delta transformation.
  -- Problematic if delta moves comments to a different node, hopefully it won't.
  searchAndMark
    :: [NonExhaustivePatterns]
    -> Ghc.ParsedSource
    -> Ghc.ParsedSource
  searchAndMark nePats =
      Syb.everywhere (Syb.mkT goExpr `Syb.extT` goDecl `Syb.extT` goBind)
    where
    goExpr :: Ghc.LHsExpr Ghc.GhcPs -> Ghc.LHsExpr Ghc.GhcPs
    goExpr (Ghc.L ann c@Ghc.HsCase{})
      | Just caseLoc <- Ghc.srcSpanToRealSrcSpan $ Ghc.locA ann
      , Just neIdx <- List.findIndex ((caseLoc ==) . srcCodeLoc) nePats
      = Ghc.L (addIndexComment ann neIdx) c
  #if MIN_VERSION_ghc(9,10,0)
    goExpr (Ghc.L ann l@(Ghc.HsLam _ lamType _))
      | lamType /= Ghc.LamSingle
      , Just caseLoc <- Ghc.srcSpanToRealSrcSpan $ Ghc.locA ann
      , Just neIdx <- List.findIndex ((caseLoc ==) . srcCodeLoc) nePats
      = Ghc.L (addIndexComment ann neIdx) l
  #elif MIN_VERSION_ghc(9,6,0)
    goExpr (Ghc.L ann l@(Ghc.HsLamCase _ _ _))
      | Just caseLoc <- Ghc.srcSpanToRealSrcSpan $ Ghc.locA ann
      , Just neIdx <- List.findIndex ((caseLoc ==) . srcCodeLoc) nePats
      = Ghc.L (addIndexComment ann neIdx) l
  #endif
    goExpr x = x
  
    goDecl :: Ghc.LHsDecl Ghc.GhcPs -> Ghc.LHsDecl Ghc.GhcPs
    goDecl (Ghc.L ann f@(Ghc.ValD _ Ghc.FunBind{}))
      | Just caseLoc <- Ghc.srcSpanToRealSrcSpan $ Ghc.locA ann
      , Just neIdx <- List.findIndex ((caseLoc ==) . srcCodeLoc) nePats
      = Ghc.L (addIndexComment ann neIdx) f
    goDecl x = x
  
    goBind :: Ghc.LHsBind Ghc.GhcPs -> Ghc.LHsBind Ghc.GhcPs
    goBind (Ghc.L ann f@Ghc.FunBind{})
      | Just caseLoc <- Ghc.srcSpanToRealSrcSpan $ Ghc.locA ann
      , Just neIdx <- List.findIndex ((caseLoc ==) . srcCodeLoc) nePats
      = Ghc.L (addIndexComment ann neIdx) f
    goBind x = x
  
  -- | Finds the target pattern and splits it. Returns the modified source and True if successful.
  -- Applied post delta transformation.
  splitPattern
    :: Ghc.GlobalRdrEnv
    -> [NonExhaustivePatterns]
    -> Ghc.ParsedSource
    -> (Ghc.ParsedSource, Any)
  splitPattern gblRdrEnv nePats ps =
      Writer.runWriter $
        Syb.everywhereM
          ( Syb.mkM (Writer.writer . goExpr)
            `Syb.extM` (Writer.writer . goDecl)
            `Syb.extM` (Writer.writer . goBind)
          ) ps
    where
    goExpr :: Ghc.LHsExpr Ghc.GhcPs -> (Ghc.LHsExpr Ghc.GhcPs, Any)
    goExpr (Ghc.L ann (Ghc.HsCase x scrut matchGroup))
      | Just (neIdx, otherComms) <- extractIdxComment (Ghc.getComments ann)
      , Just nePat <- listToMaybe $ drop neIdx nePats
      , Just newMG <- splitMG gblRdrEnv False False 0 nePat matchGroup
      = ( Ghc.L (Ghc.setComments otherComms mempty ann) (Ghc.HsCase x scrut newMG)
        , Any True
        )
  #if MIN_VERSION_ghc(9,10,0)
    goExpr (Ghc.L ann (Ghc.HsLam x lamType matchGroup@(Ghc.MG _ (Ghc.L matchesAnn _))))
      | lamType /= Ghc.LamSingle
      , Just (neIdx, otherComms) <- extractIdxComment (Ghc.comments ann)
      , Just nePat <- nePats List.!? neIdx
      , Just newMG <- splitMG gblRdrEnv (lamType == Ghc.LamCases) False (Ghc.colDelta matchesAnn) nePat matchGroup
      = ( Ghc.L ann {Ghc.comments = otherComms} (Ghc.HsLam x lamType newMG)
        , Any True
        )
  #elif MIN_VERSION_ghc(9,6,0)
    goExpr (Ghc.L ann (Ghc.HsLamCase x lamType matchGroup@(Ghc.MG _ (Ghc.L matchesAnn _))))
      | Just (neIdx, otherComms) <- extractIdxComment (Ghc.getComments ann)
      , Just nePat <- listToMaybe $ drop neIdx nePats
      , Just newMG <- splitMG gblRdrEnv (lamType == Ghc.LamCases) False (Ghc.colDelta matchesAnn) nePat matchGroup
      = ( Ghc.L (Ghc.setComments otherComms mempty ann) (Ghc.HsLamCase x lamType newMG)
        , Any True
        )
  #endif
    goExpr expr = (expr, Any False)
  
    goDecl :: Ghc.LHsDecl Ghc.GhcPs -> (Ghc.LHsDecl Ghc.GhcPs, Any)
    goDecl (Ghc.L ann (Ghc.ValD x (Ghc.FunBind x2 fid matchGroup)))
      | Just (neIdx, otherComms) <- extractIdxComment (Ghc.getComments ann)
      , Just nePat <- listToMaybe $ drop neIdx nePats
      , Just newMG <- splitMG gblRdrEnv True True 0 nePat matchGroup
      = ( Ghc.L (Ghc.setComments otherComms mempty ann) (Ghc.ValD x (Ghc.FunBind x2 fid newMG))
        , Any True
        )
    goDecl decl = (decl, Any False)
  
    goBind :: Ghc.LHsBind Ghc.GhcPs -> (Ghc.LHsBind Ghc.GhcPs, Any)
    goBind (Ghc.L ann (Ghc.FunBind x2 fid matchGroup))
      | Just (neIdx, otherComms) <- extractIdxComment (Ghc.getComments ann)
      , Just nePat <- listToMaybe $ drop neIdx nePats
      , Just newMG <- splitMG gblRdrEnv True True 0 nePat matchGroup
      = ( Ghc.L (Ghc.setComments otherComms mempty ann) (Ghc.FunBind x2 fid newMG)
        , Any True
        )
    goBind bind = (bind, Any False)
  
  splitMG
    :: Ghc.GlobalRdrEnv
    -> Bool -- True => match group can have multiple patterns
    -> Bool -- True => add left padding to each new pattern
    -> Int -- Number of horizontal spaces at the front of inserted pattern match
    -> NonExhaustivePatterns
    -> Ghc.MatchGroup Ghc.GhcPs (Ghc.LHsExpr Ghc.GhcPs)
    -> Maybe (Ghc.MatchGroup Ghc.GhcPs (Ghc.LHsExpr Ghc.GhcPs))
  splitMG gblRdrEnv multiplePats offsetFirstPat colOffset nePat (Ghc.MG x2 (Ghc.L ann2 matches))
    | (beforeSplit, Ghc.L splitAnn targetMatch@(Ghc.Match _ ctx _ rhs) : afterSplit)
        <- break matchHasSplit matches
    , let mUpdatedMatch = Ghc.L splitAnn <$> removeSplitFromOrPat targetMatch
          -- If splitting the first match, trailing matches need to have a delta
          -- putting it on a new line
          correctDeltas [] = []
          correctDeltas (x : xs) | isNothing mUpdatedMatch =
            x :
              (Ghc.L (Ghc.nextLine colOffset) . Ghc.unLoc
               <$> xs)
          correctDeltas xs = Ghc.L (Ghc.nextLine colOffset) . Ghc.unLoc <$> xs
          newMatches = correctDeltas $ do
            nabla <- patternNablas nePat
            let pats =
                  zipWith (mkPat gblRdrEnv nabla $ not multiplePats)
                          (offsetFirstPat : repeat True)
                          (patternIds nePat)
            [ Ghc.L splitAnn $ Ghc.Match Ghc.noExtFieldCpp ctx (Ghc.noLocCpp pats) rhs ]
          removeSplits m =
            case traverse removeSplitFromOrPat m of
              Nothing -> if matchHasSplit m then Nothing else Just m
              Just newM -> Just newM
          afterSplitUpdated = mapMaybe removeSplits afterSplit
          newMatchGroup
            = beforeSplit
           ++ maybeToList mUpdatedMatch
           ++ newMatches
           ++ afterSplitUpdated
    = Just $ Ghc.MG x2 (Ghc.L ann2 newMatchGroup)
    | otherwise = Nothing
  
  matchHasSplit :: Ghc.LMatch Ghc.GhcPs (Ghc.LHsExpr Ghc.GhcPs) -> Bool
  matchHasSplit (Ghc.L _ (Ghc.Match _ _ pats _)) =
      Syb.everything (||) (Syb.mkQ False isSplitCon) pats
  
  isSplitCon :: Ghc.Pat Ghc.GhcPs -> Bool
  isSplitCon (Ghc.ConPat _ (Ghc.L _ rdr) _) =
    Ghc.rdrNameOcc rdr == Ghc.mkDataOcc splitName
  isSplitCon _ = False
  
  -- | Remove SPLIT pattern from OrPat groups. Returns Just if the match was modified.
  removeSplitFromOrPat
    :: Ghc.Match Ghc.GhcPs (Ghc.LHsExpr Ghc.GhcPs)
    -> Maybe (Ghc.Match Ghc.GhcPs (Ghc.LHsExpr Ghc.GhcPs))
  #if MIN_VERSION_ghc(9,12,0)
  removeSplitFromOrPat (Ghc.Match a b pats c) =
    let mNewPats = Syb.everywhereM (Syb.mkM (Writer.writer . go)) pats
        patHasSplit = Syb.everything (||) (Syb.mkQ False isSplitCon)
        dropSemicolon ann = ann { Ghc.anns = Ghc.noAnn }
        go :: Ghc.Pat Ghc.GhcPs -> (Ghc.Pat Ghc.GhcPs, Any)
        go (Ghc.OrPat x oPats) =
          case reverse (NE.filter (not . patHasSplit) oPats) of
            [] -> (Ghc.WildPat Ghc.noExtField, Any True)
            (Ghc.L patAnn lastPat : fPats)
              | length fPats + 1 /= NE.length oPats ->
                  ( Ghc.OrPat x
                      (NE.reverse $ Ghc.L (dropSemicolon patAnn) lastPat NE.:| fPats)
                  , Any True)
              | otherwise -> (Ghc.OrPat x oPats, Any False)
        go other = (other, Any False)
     in case Writer.runWriter mNewPats of
          (newPats, Any True) -> Just (Ghc.Match a b newPats c)
          _ -> Nothing
  #else
  removeSplitFromOrPat _ = Nothing
  #endif
  
  -- | Produce a 'Pat' for a missing pattern
  mkPat
    :: Ghc.GlobalRdrEnv
    -> Ghc.Nabla
    -> Bool -- ^ True => is a singular pattern which doesn't need outer parens
    -> Bool -- ^ True => needs left padding to separate it from another pattern
    -> Ghc.Id
    -> Ghc.LPat Ghc.GhcPs
  mkPat gblRdrEnv nabla isOutermost needsLeftPad x = Ghc.L delta $
    case Ghc.lookupSolution nabla x of
      Nothing -> Ghc.WildPat Ghc.noExtField
      Just (Ghc.PACA (Ghc.PmAltConLike con) _tvs args) -> paren con args $
        case Ghc.conLikeIsInfix con of
          True | [arg1, arg2] <- args ->
            Ghc.ConPat Ghc.noAnn
              ( Ghc.L Ghc.nameAnchorD1
              . nameToRdrName gblRdrEnv
              $ Ghc.conLikeName con
              )
            $ Ghc.InfixCon (mkPat gblRdrEnv nabla False False arg1)
                           (mkPat gblRdrEnv nabla False True arg2)
          _ | Ghc.RealDataCon dc <- con
            , Ghc.isUnboxedTupleDataCon dc
            -> Ghc.TuplePat Ghc.parenHashAnns
                 (addCommaAnns $ zipWith (mkPat gblRdrEnv nabla True) (False : repeat True) args)
                 Ghc.Unboxed
          _ | Ghc.RealDataCon dc <- con
            , Ghc.isTupleDataCon dc
            -> Ghc.TuplePat Ghc.parenAnns
                 (addCommaAnns $ zipWith (mkPat gblRdrEnv nabla True) (False : repeat True) args)
                 Ghc.Boxed
          _ ->
            -- If GHC tries to use SPLIT as a missing pattern, replace it with wildcard
            if Ghc.occName (Ghc.conLikeName con) == Ghc.mkDataOcc splitName
            then Ghc.WildPat Ghc.noExtField
            else Ghc.ConPat Ghc.noAnn
                  ( Ghc.L Ghc.nameAnchorD0
                  . nameToRdrName gblRdrEnv $ Ghc.conLikeName con
                  )
               $ Ghc.PrefixCon [] (mkPat gblRdrEnv nabla False True <$> args)
      Just (Ghc.PACA (Ghc.PmAltLit lit) _tvs _args) ->
        case Ghc.pm_lit_val lit of
          Ghc.PmLitInt integer ->
            Ghc.NPat Ghc.noAnn (Ghc.noLocA $ Ghc.OverLit Ghc.noExtField $ Ghc.HsIntegral $ Ghc.IL (Ghc.SourceText . fromString $ show integer) (integer < 0) integer) Nothing Ghc.noExtField
          Ghc.PmLitRat rational ->
            Ghc.NPat Ghc.noAnn (Ghc.noLocA $ Ghc.OverLit Ghc.noExtField $ Ghc.HsFractional $ Ghc.mkTHFractionalLit rational) Nothing Ghc.noExtField
          Ghc.PmLitChar char -> Ghc.LitPat Ghc.noExtField $ Ghc.HsChar Ghc.NoSourceText char
          Ghc.PmLitString fastString -> Ghc.LitPat Ghc.noExtField $ Ghc.HsString Ghc.NoSourceText fastString
          Ghc.PmLitOverInt _minuses integer ->
            Ghc.NPat Ghc.noAnn (Ghc.noLocA $ Ghc.OverLit Ghc.noExtField $ Ghc.HsIntegral $ Ghc.IL (Ghc.SourceText . fromString $ show integer) (integer < 0) integer) Nothing Ghc.noExtField
          Ghc.PmLitOverRat _minuses fractionalLit ->
            Ghc.NPat Ghc.noAnn (Ghc.noLocA $ Ghc.OverLit Ghc.noExtField $ Ghc.HsFractional fractionalLit) Nothing Ghc.noExtField
          Ghc.PmLitOverString fastString -> Ghc.LitPat Ghc.noExtField $ Ghc.HsString Ghc.NoSourceText fastString
    where
      delta = if needsLeftPad
                 then Ghc.anchorD1
                 else Ghc.anchorD0
      paren _ [] inner = inner
      paren (Ghc.RealDataCon dc) _ inner | Ghc.isTupleDataCon dc = inner -- No parens for tuple pats
      paren _ _ inner =
        if not isOutermost
        then Ghc.mkParPat' (Ghc.L Ghc.anchorD0 inner)
        else inner
  
  splitName :: IsString a => a
  splitName = "SPLIT"
  