packages feed

aihc-parser-1.0.0.2: src/Aihc/Parser/Lex/Trivia.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}

module Aihc.Parser.Lex.Trivia
  ( consumeBlockCommentTokenOrError,
    consumeLineCommentToken,
    isHaskellWhitespace,
    isLineComment,
    tryConsumeControlPragma,
    tryConsumeLineDirective,
  )
where

import Aihc.Parser.Lex.Pragmas (parseControlPragma)
import Aihc.Parser.Lex.Types
import Data.Char (isDigit, isSpace)
import Data.Maybe (fromMaybe)
import Data.Text (Text, pattern (:<))
import Data.Text qualified as T
import Text.Read (readMaybe)

isHaskellWhitespace :: Char -> Bool
isHaskellWhitespace c =
  c == ' ' || c == '\t' || c == '\r' || c == '\f' || c == '\v' || isSpace c

tryConsumeLineDirective :: LexerState -> Maybe (Maybe LexToken, LexerState)
tryConsumeLineDirective st
  | not (lexerAtLineStart st) = Nothing
  | otherwise =
      let inp = lexerInput st
          (spaces, rest) = T.span (\c -> c == ' ' || c == '\t') inp
          -- CPP #line directives require '#' at column 1 (no leading
          -- whitespace).  GHC enforces the same rule: an indented '#' is
          -- never a line directive.  Shebangs, however, are allowed with
          -- optional leading whitespace.
          --
          -- Note: skipTrivia may have already consumed leading whitespace
          -- before calling us, so we check lexerCol rather than whether
          -- 'spaces' is empty.  The column after consuming 'spaces' is
          -- where '#' would sit.
          hashCol = lexerCol st + T.length spaces
          atColumn1 = hashCol == 1
       in case rest of
            '#' :< more ->
              let lineText = "#" <> takeLineRemainder more
                  consumed = spaces <> lineText
                  isFirstLine = lexerLine st == 1
               in case classifyHashLineTrivia atColumn1 isFirstLine lineText of
                    Just (HashLineDirective update) ->
                      Just (Nothing, applyDirectiveAdvance consumed update st)
                    Just HashLineShebang ->
                      let st' = advanceChars consumed st
                       in Just (Nothing, st')
                    Just HashLineMalformed ->
                      let st' = advanceChars consumed st
                       in Just (Just (mkToken st st' consumed (TkError "malformed line directive")), st')
                    Nothing -> Nothing
            _ -> Nothing

tryConsumeControlPragma :: LexerState -> Maybe (Maybe LexToken, LexerState)
tryConsumeControlPragma st =
  let inp = lexerInput st
   in case parseControlPragma inp of
        Just (consumedT, Right update0) ->
          let (consumedT', update) =
                case directiveLine update0 of
                  Just lineNo ->
                    case T.drop (T.length consumedT) (lexerInput st) of
                      '\n' :< _ ->
                        (consumedT <> "\n", update0 {directiveLine = Just lineNo, directiveCol = Just 1})
                      _ -> (consumedT, update0)
                  Nothing -> (consumedT, update0)
           in Just (Nothing, applyDirectiveAdvance consumedT' update st)
        Just (consumedT, Left msg) ->
          let st' = advanceChars consumedT st
           in Just (Just (mkToken st st' consumedT (TkError msg)), st')
        Nothing -> Nothing

consumeLineCommentToken :: LexerState -> (LexToken, LexerState)
consumeLineCommentToken st =
  let inp = lexerInput st
      rest = T.drop 2 inp
      consumed = "--" <> T.takeWhile (/= '\n') rest
      st' = advanceChars consumed st
   in (mkToken st st' consumed TkLineComment, st')

consumeBlockCommentToken :: LexerState -> Maybe (LexToken, LexerState)
consumeBlockCommentToken st =
  case scanNestedBlockComment 1 (T.drop 2 (lexerInput st)) of
    Just consumedTail ->
      let consumed = "{-" <> consumedTail
          st' = advanceChars consumed st
          st'' =
            if T.any (== '\n') consumed
              then st'
              else st' {lexerAtLineStart = lexerAtLineStart st}
       in Just (mkToken st st'' consumed TkBlockComment, st'')
    Nothing -> Nothing

consumeBlockCommentTokenOrError :: LexerState -> Either (LexToken, LexerState) (LexToken, LexerState)
consumeBlockCommentTokenOrError st =
  case consumeBlockCommentToken st of
    Just result -> Right result
    Nothing ->
      let consumed = lexerInput st
          st' = advanceChars consumed st
          tok = mkToken st st' consumed (TkError "unterminated block comment")
       in Left (tok, st')

scanNestedBlockComment :: Int -> Text -> Maybe Text
scanNestedBlockComment depth0 input = go depth0 0 input
  where
    -- Skip characters that can't start a nesting change in bulk, then
    -- inspect the stopping character.  Allocation is O(nesting changes)
    -- instead of O(length).
    go depth !n remaining
      | depth <= 0 = Just (T.take n input)
      | otherwise =
          let (prefix, rest0) = T.span (\c -> c /= '{' && c /= '-') remaining
              n' = n + T.length prefix
           in case T.uncons rest0 of
                Nothing -> Nothing
                Just (c, rest1) ->
                  case T.uncons rest1 of
                    Nothing -> Nothing -- truncated escape sequence, unterminated
                    Just (c2, rest2)
                      | c == '{' && c2 == '-' -> go (depth + 1) (n' + 2) rest2
                      | c == '-' && c2 == '}' ->
                          if depth == 1
                            then Just (T.take (n' + 2) input)
                            else go (depth - 1) (n' + 2) rest2
                      | otherwise ->
                          -- c was '{' or '-' but not a nesting pair; advance by 1
                          -- and re-examine rest1 (which still starts with c2)
                          go depth (n' + 1) rest1

applyDirectiveAdvance :: Text -> DirectiveUpdate -> LexerState -> LexerState
applyDirectiveAdvance consumed update st =
  let hasTrailingNewline = T.isSuffixOf "\n" consumed
      st' = advanceChars consumed st
   in st'
        { lexerLogicalSourceName = fromMaybe (lexerLogicalSourceName st') (directiveSourceName update),
          lexerLine = maybe (lexerLine st') (max 1) (directiveLine update),
          lexerCol = maybe (lexerCol st') (max 1) (directiveCol update),
          lexerAtLineStart = hasTrailingNewline || (Just 1 == directiveCol update)
        }

-- | Classify a line beginning with @#@.
--
-- @atColumn1@ is 'True' when @#@ sits at column 1 of the physical source
-- line (no leading whitespace).  @isFirstLine@ is 'True' on line 1.
-- Shebangs are accepted at column 1 on any line (for mid-file shebangs
-- as in some polyglot scripts) or anywhere on line 1 (GHC also accepts
-- an optional leading space before the initial shebang).
-- An indented @#!@ past line 1 is left for the operator lexer.
-- CPP @#line@ directives are only recognised at column 1, matching GHC.
classifyHashLineTrivia :: Bool -> Bool -> Text -> Maybe HashLineTrivia
classifyHashLineTrivia atColumn1 isFirstLine raw
  | (atColumn1 || isFirstLine) && isHashBangLine raw = Just HashLineShebang
  | not atColumn1 = Nothing
  | looksLikeHashLineDirective raw =
      case parseHashLineDirective raw of
        Just update -> Just (HashLineDirective update)
        Nothing -> Just HashLineMalformed
  | otherwise = Nothing

parseHashLineDirective :: Text -> Maybe DirectiveUpdate
parseHashLineDirective raw =
  let trimmed = T.dropWhile isSpace (T.drop 1 (T.dropWhile isSpace raw))
      trimmed' =
        case T.stripPrefix "line" trimmed of
          Just afterLine -> T.dropWhile isSpace afterLine
          Nothing -> trimmed
      (digits, rest) = T.span isDigit trimmed'
   in if T.null digits
        then Nothing
        else do
          lineNo <- readBoundedInt digits
          Just
            DirectiveUpdate
              { directiveLine = Just lineNo,
                directiveCol = Just 1,
                directiveSourceName = parseDirectiveSourceName rest
              }

isHashBangLine :: Text -> Bool
isHashBangLine raw =
  "#!" `T.isPrefixOf` T.dropWhile isSpace raw

looksLikeHashLineDirective :: Text -> Bool
looksLikeHashLineDirective raw =
  let afterHash = T.dropWhile isSpace (T.drop 1 (T.dropWhile isSpace raw))
   in case afterHash of
        c :< _ | isDigit c -> True
        _ -> "line" `T.isPrefixOf` afterHash

parseDirectiveSourceName :: Text -> Maybe FilePath
parseDirectiveSourceName rest =
  let rest' = T.dropWhile isSpace rest
   in case rest' of
        '"' :< more ->
          let (name, trailing) = T.break (== '"') more
           in case trailing of
                '"' :< _ -> Just (T.unpack name)
                _ -> Nothing
        _ -> Nothing

takeLineRemainder :: Text -> Text
takeLineRemainder chars =
  let (prefix, rest) = T.break (== '\n') chars
   in case rest of
        '\n' :< _ -> prefix <> "\n"
        _ -> prefix

isLineComment :: Text -> Bool
isLineComment rest =
  case rest of
    c :< _
      | c == '-' -> isLineComment (T.dropWhile (== '-') rest)
      | isSymbolicOpChar c -> False
      | otherwise -> True
    _ -> True

readBoundedInt :: Text -> Maybe Int
readBoundedInt txt = do
  n <- readMaybe (T.unpack txt) :: Maybe Integer
  if n >= fromIntegral (minBound :: Int) && n <= fromIntegral (maxBound :: Int)
    then Just (fromInteger n)
    else Nothing