content: {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  -- | A module for parsing of pragmas from comments.
  module Ormolu.Parser.Pragma
    ( Pragma (..),
      parsePragma,
    )
  where
  
  import Data.Char (isSpace)
  import Data.Text (Text)
  import Data.Text qualified as T
  import Data.Text.Encoding qualified as T
  import GHC.Data.FastString (bytesFS, mkFastString)
  import GHC.Driver.Config.Parser (initParserOpts)
  import GHC.DynFlags (baseDynFlags)
  import GHC.Parser.Lexer qualified as L
  import GHC.Types.SrcLoc
  import Ormolu.Utils (textToStringBuffer)
  
  -- | Ormolu's representation of pragmas.
  data Pragma
    = -- | Language pragma
      PragmaLanguage [Text]
    | -- | GHC options pragma
      PragmaOptionsGHC Text
    | -- | Haddock options pragma
      PragmaOptionsHaddock Text
    deriving (Show, Eq)
  
  -- | Extract a pragma from a comment if possible, or return 'Nothing'
  -- otherwise.
  parsePragma ::
    -- | Comment to try to parse
    Text ->
    Maybe Pragma
  parsePragma input = do
    contents <- T.stripSuffix "#-}" =<< T.stripPrefix "{-#" input
    let (pragmaName, cs) = (T.break isSpace . T.dropWhile isSpace) contents
    case T.toLower pragmaName of
      "language" -> PragmaLanguage <$> parseExtensions cs
      "options_ghc" -> Just $ PragmaOptionsGHC (T.strip cs)
      "options_haddock" -> Just $ PragmaOptionsHaddock (T.strip cs)
      _ -> Nothing
  
  -- | Assuming the input consists of a series of tokens from a language
  -- pragma, return the set of enabled extensions.
  parseExtensions :: Text -> Maybe [Text]
  parseExtensions str = tokenize str >>= go
    where
      go = \case
        [L.ITconid ext] -> return [fsToText ext]
        (L.ITconid ext : L.ITcomma : xs) -> (fsToText ext :) <$> go xs
        _ -> Nothing
      fsToText = T.decodeUtf8 . bytesFS
  
  -- | Tokenize a given input using GHC's lexer.
  tokenize :: Text -> Maybe [L.Token]
  tokenize input =
    case L.unP pLexer parseState of
      L.PFailed {} -> Nothing
      L.POk _ x -> Just x
    where
      location = mkRealSrcLoc (mkFastString "") 1 1
      buffer = textToStringBuffer input
      parseState = L.initParserState parserOpts buffer location
      parserOpts = initParserOpts baseDynFlags
  
  -- | Haskell lexer.
  pLexer :: L.P [L.Token]
  pLexer = go
    where
      go = do
        r <- L.lexer False return
        case unLoc r of
          L.ITeof -> return []
          x -> (x :) <$> go
  