Lex.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns #-}

-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Token lexing and comment preservation for the Swarm language.
module Swarm.Language.Parser.Lex (
  -- * Parsing with source locations
  parseLoc,
  parseLocG,

  -- * Whitespace + comments
  getCommentSituation,
  lineComment,
  blockComment,
  sc,

  -- * Tokens

  -- ** Lexemes
  lexeme,

  -- ** Specific token types
  symbol,
  operator,
  reservedWords,
  reservedCS,
  reserved,
  IdentifierType (..),
  locIdentifier,
  locTmVar,
  locTyName,
  identifier,
  tyVar,
  tyName,
  tmVar,
  textLiteral,
  integer,

  -- ** Combinators
  braces,
  parens,
  brackets,
) where

import Control.Lens (use, (%=), (.=))
import Control.Monad (void)
import Data.Char (isLower, isUpper)
import Data.Containers.ListUtils (nubOrd)
import Data.List.Extra (enumerate)
import Data.Sequence qualified as Seq
import Data.Set (Set)
import Data.Set qualified as S
import Data.Text (Text)
import Data.Text qualified as T
import Swarm.Language.Parser.Core
import Swarm.Language.Syntax
import Swarm.Language.Syntax.Direction
import Swarm.Language.TDVar (TDVar, mkTDVar)
import Swarm.Language.Types (baseTyName)
import Swarm.Util (failT, squote)
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer qualified as L
import Witch (from, into)

------------------------------------------------------------
-- Parsing with source locations

-- | Add 'SrcLoc' to a parser
parseLocG :: Parser a -> Parser (SrcLoc, a)
parseLocG pa = do
  -- Remember the start location.
  start <- getOffset
  a <- pa
  -- Instead of using 'getOffset' to get the end location, which would
  -- include any whitespace consumed at the end of @pa@, get the
  -- @preWSLoc@ which was set by 'sc' at the /beginning/ of the
  -- consumed whitespace.
  end <- use preWSLoc
  pure (SrcLoc start end, a)

-- | Add 'SrcLoc' to a 'Term' parser
parseLoc :: Parser Term -> Parser Syntax
parseLoc pterm = uncurry Syntax <$> parseLocG pterm

------------------------------------------------------------
-- Whitespace

-- Approach for preserving comments taken from https://www.reddit.com/r/haskell/comments/ni4gpm/comment/gz0ipmp/

-- | If we see a comment starting now, is it the first non-whitespace
--   thing on the current line so far, or were there other
--   non-whitespace tokens previously?
getCommentSituation :: Parser CommentSituation
getCommentSituation = do
  fl <- use freshLine
  return $ if fl then StandaloneComment else SuffixComment

-- | Parse a line comment, while appending it out-of-band to the list of
--   comments saved in the custom state.
lineComment :: Text -> Parser ()
lineComment start = do
  cs <- getCommentSituation
  -- Note we must manually get the start and end offset rather than
  -- using parseLocG, since parseLocG explicitly does not include the
  -- source span of trailing whitespace.
  s <- getOffset
  t <- string start *> takeWhileP (Just "character") (/= '\n')
  e <- getOffset
  comments %= (Seq.|> Comment (SrcLoc s e) LineComment cs t)

-- | Parse a block comment, while appending it out-of-band to the list of
--   comments saved in the custom state.
blockComment :: Text -> Text -> Parser ()
blockComment start end = do
  cs <- getCommentSituation
  -- Note we must manually get the start and end offset rather than
  -- using parseLocG, since parseLocG explicitly does not include the
  -- source span of trailing whitespace.
  s <- getOffset
  void $ string start
  t <- manyTill anySingle (string end)
  e <- getOffset
  comments %= (Seq.|> Comment (SrcLoc s e) BlockComment cs (into @Text t))

-- | Skip spaces and comments.
sc :: Parser ()
sc = do
  -- Remember where we were before we started consuming whitespace.
  l <- getOffset
  preWSLoc .= l
  -- Typically we would use L.space here, but we have to inline its
  -- definition and use our own slight variant, since we need to treat
  -- end-of-line specially.
  skipMany . choice . map hidden $
    [ hspace1
    , eol *> (freshLine .= True) -- If we see a newline, reset freshLine to True.
    , lineComment "//"
    , blockComment "/*" "*/"
    ]

------------------------------------------------------------
-- Tokens

-- | In general, we follow the convention that every token parser
--   assumes no leading whitespace and consumes all trailing
--   whitespace.  Concretely, we achieve this by wrapping every token
--   parser using 'lexeme'.
--
--   Also sets freshLine to False every time we see a non-whitespace
--   token.
lexeme :: Parser a -> Parser a
lexeme p = (freshLine .= False) *> L.lexeme sc p

-- | A lexeme consisting of a literal string.
symbol :: Text -> Parser Text
symbol s = (freshLine .= False) *> L.symbol sc s

-- | A lexeme consisting of a specific string, not followed by any other
--   operator character.
operator :: Text -> Parser Text
operator n = (lexeme . try) (string n <* notFollowedBy operatorChar)

-- | Recognize a single character which is one of the characters used
--   by a built-in operator.
operatorChar :: Parser Text
operatorChar = T.singleton <$> oneOf opChars
 where
  isOp = \case { ConstMFunc {} -> False; _ -> True } . constMeta
  opChars = nubOrd . concatMap (from . syntax) . filter isOp $ map constInfo allConst

-- | Names of base types built into the language.
baseTypeNames :: [Text]
baseTypeNames = map baseTyName enumerate

-- | Names of types built into the language.
primitiveTypeNames :: [Text]
primitiveTypeNames = "Cmd" : baseTypeNames

-- | List of keywords built into the language.
keywords :: [Text]
keywords = T.words "let in def tydef end true false forall require stock requirements rec"

-- | List of reserved words that cannot be used as variable names.
reservedWords :: Set Text
reservedWords =
  S.fromList $
    map (syntax . constInfo) (filter isUserFunc allConst)
      ++ map directionSyntax allDirs
      ++ primitiveTypeNames
      ++ keywords

-- | Parse a reserved word, given a string recognizer (which can
--   /e.g./ be case sensitive or not), making sure it is not a prefix
--   of a longer variable name, and allowing the parser to backtrack
--   if it fails.
reservedGen :: (Text -> Parser a) -> Text -> Parser ()
reservedGen str w = (lexeme . try) $ str w *> notFollowedBy (alphaNumChar <|> char '_')

-- | Parse a case-sensitive reserved word.
reservedCS :: Text -> Parser ()
reservedCS = reservedGen string

-- | Parse a case-insensitive reserved word.
reserved :: Text -> Parser ()
reserved = reservedGen string'

-- | What kind of identifier are we parsing?
data IdentifierType = IDTyVar | IDTyName | IDTmVar
  deriving (Eq, Ord, Show)

-- | Parse an identifier together with its source location info.
locIdentifier :: IdentifierType -> Parser LocVar
locIdentifier idTy =
  uncurry LV <$> parseLocG ((lexeme . try) (p >>= check) <?> "variable name")
 where
  p = (:) <$> (letterChar <|> char '_') <*> many (alphaNumChar <|> char '_' <|> char '\'')
  check (into @Text -> t)
    | IDTyVar <- idTy
    , T.toTitle t `S.member` reservedWords =
        failT ["Reserved type name", squote t, "cannot be used as a type variable name; perhaps you meant", squote (T.toTitle t) <> "?"]
    | IDTyName <- idTy
    , t `S.member` reservedWords =
        failT ["Reserved type name", squote t, "cannot be redefined."]
    | t `S.member` reservedWords || T.toLower t `S.member` reservedWords =
        failT ["Reserved word", squote t, "cannot be used as a variable name"]
    | IDTyName <- idTy
    , isLower (T.head t) =
        failT ["Type synonym names must start with an uppercase letter"]
    | IDTyVar <- idTy
    , isUpper (T.head t) =
        failT ["Type variable names must start with a lowercase letter"]
    | otherwise = pure t

-- | Parse a term variable together with its source location info.
locTmVar :: Parser LocVar
locTmVar = locIdentifier IDTmVar

-- | Parse a user-defined type name together with its source location
--   info.
locTyName :: Parser (Located TDVar)
locTyName = (fmap . fmap) mkTDVar (locIdentifier IDTyName)

-- | Parse an identifier, i.e. any non-reserved string containing
--   alphanumeric characters and underscores, not starting with a
--   digit. The Bool indicates whether we are parsing a type variable.
identifier :: IdentifierType -> Parser Var
identifier = fmap lvVar . locIdentifier

-- | Parse a type variable, which must start with an underscore or
--   lowercase letter and cannot be the lowercase version of a type
--   name.
tyVar :: Parser Var
tyVar = identifier IDTyVar

-- | Parse a (user-defined) type constructor name, which must start
--   with an uppercase letter.
--
--   Note that this always produces a variable with version number 0,
--   via 'mkTDVar'.  We cannot properly version variables at parsing
--   time since we don't know what else is in scope.  There is a
--   separate name resolution pass later that assigns correct version
--   numbers to user type names.
tyName :: Parser TDVar
tyName = mkTDVar <$> identifier IDTyName

-- | Parse a term variable, which can start in any case and just
--   cannot be the same (case-insensitively) as a lowercase reserved
--   word.
tmVar :: Parser Var
tmVar = identifier IDTmVar

-- | Parse a text literal (including escape sequences) in double quotes.
textLiteral :: Parser Text
textLiteral = into <$> lexeme (char '"' >> manyTill L.charLiteral (char '"'))

-- | Parse a positive integer literal token, in decimal, binary,
--   octal, or hexadecimal notation.  Note that negation is handled as
--   a separate operator.
integer :: Parser Integer
integer =
  label "integer literal" $
    lexeme $ do
      n <-
        string "0b"
          *> L.binary
          <|> string "0o"
            *> L.octal
          <|> string "0x"
            *> L.hexadecimal
          <|> L.decimal
      notFollowedBy alphaNumChar
      return n

------------------------------------------------------------
-- Combinators

braces :: Parser a -> Parser a
braces = between (symbol "{") (symbol "}")

parens :: Parser a -> Parser a
parens = between (symbol "(") (symbol ")")

brackets :: Parser a -> Parser a
brackets = between (symbol "[") (symbol "]")