Spec.hs

-- | Definition and parsing of a test specification.
module Futhark.Test.Spec
  ( testSpecFromProgram,
    testSpecFromProgramOrDie,
    testSpecsFromPaths,
    testSpecsFromPathsOrDie,
    testSpecFromFile,
    testSpecFromFileOrDie,
    ProgramTest (..),
    StructureTest (..),
    StructurePipeline (..),
    WarningTest (..),
    TestAction (..),
    ExpectedError (..),
    InputOutputs (..),
    PropertyCase (..),
    TestRun (..),
    ExpectedResult (..),
    Success (..),
    Values (..),
    GenValue (..),
    genValueType,
    parseEmbeddedExp,
  )
where

import Control.Applicative
import Control.Exception (catch)
import Control.Monad
import Data.Char
import Data.Either (partitionEithers)
import Data.Functor
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Text qualified as T
import Data.Text.IO qualified as T
import Data.Void
import Futhark.Analysis.Metrics.Type
import Futhark.Data.Parser
import Futhark.Data.Parser qualified as V
import Futhark.Test.Values qualified as V
import Futhark.Util (directoryContents, nubOrd, showText)
import Futhark.Util.Loc qualified as Loc
import Futhark.Util.Pretty (prettyTextOneLine)
import Language.Futhark.Parser (SyntaxError (..), parseExpAt)
import Language.Futhark.Pretty ()
import Language.Futhark.Prop (UncheckedExp)
import System.Exit
import System.FilePath
import System.IO
import System.IO.Error
import Text.Megaparsec hiding (many, some)
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer (charLiteral)
import Text.Regex.TDFA
import Prelude

-- | Description of a test to be carried out on a Futhark program.
-- The Futhark program is stored separately.
data ProgramTest = ProgramTest
  { testDescription :: T.Text,
    testTags :: [T.Text],
    testAction :: TestAction
  }
  deriving (Show)

-- | How to test a program.
data TestAction
  = CompileTimeFailure ExpectedError
  | RunCases [InputOutputs] [StructureTest] [WarningTest] [PropertyCase]
  deriving (Show)

-- | Input and output pairs for some entry point(s).
data InputOutputs = InputOutputs
  { iosEntryPoint :: T.Text,
    iosTestRuns :: [TestRun]
  }
  deriving (Show)

newtype PropertyCase = PropertyCase
  {propertyEntryPoint :: T.Text}
  deriving (Show)

-- | The error expected for a negative test.
data ExpectedError
  = AnyError
  | ThisError T.Text Regex

instance Show ExpectedError where
  show AnyError = "AnyError"
  show (ThisError r _) = "ThisError " ++ show r

-- | How a program can be transformed.
data StructurePipeline
  = GpuPipeline
  | MCPipeline
  | SOACSPipeline
  | SeqMemPipeline
  | GpuMemPipeline
  | MCMemPipeline
  | NoPipeline
  deriving (Show)

-- | A structure test specifies a compilation pipeline, as well as
-- metrics for the program coming out the other end.
data StructureTest = StructureTest StructurePipeline AstMetrics
  deriving (Show)

-- | A warning test requires that a warning matching the regular
-- expression is produced.  The program must also compile succesfully.
data WarningTest = ExpectedWarning T.Text Regex

instance Show WarningTest where
  show (ExpectedWarning r _) = "ExpectedWarning " ++ T.unpack r

-- | A condition for execution, input, and expected result.
data TestRun = TestRun
  { runTags :: [T.Text],
    runInput :: Values,
    runExpectedResult :: ExpectedResult Success,
    runIndex :: Int,
    runDescription :: T.Text
  }
  deriving (Show)

-- | Several values - either literally, or by reference to a file, or
-- to be generated on demand.  All paths are relative to test program.
data Values
  = Values [V.Value]
  | InFile FilePath
  | GenValues [GenValue]
  | -- | A Futhark expression, evaluated by the interpreter in the
    -- context of the program, with calls to entry points run on the
    -- compiled program.
    ScriptValues UncheckedExp
  | -- | A file containing such an expression.
    ScriptFile FilePath
  deriving (Show)

-- | How to generate a single random value.
data GenValue
  = -- | Generate a value of the given rank and primitive
    -- type.  Scalars are considered 0-ary arrays.
    GenValue V.ValueType
  | -- | A fixed non-randomised primitive value.
    GenPrim V.Value
  deriving (Show)

-- | A prettyprinted representation of type of value produced by a
-- 'GenValue'.
genValueType :: GenValue -> T.Text
genValueType (GenValue (V.ValueType ds t)) =
  foldMap (\d -> "[" <> showText d <> "]") ds <> V.primTypeText t
genValueType (GenPrim v) =
  V.valueText v

-- | How a test case is expected to terminate.
data ExpectedResult values
  = -- | Execution suceeds, with or without
    -- expected result values.
    Succeeds (Maybe values)
  | -- | Execution fails with this error.
    RunTimeFailure ExpectedError
  deriving (Show)

-- | The result expected from a succesful execution.
data Success
  = -- | These values are expected.
    SuccessValues Values
  | -- | Compute expected values from executing a known-good
    -- reference implementation.
    SuccessGenerateValues
  deriving (Show)

type Parser = Parsec Void T.Text

lexeme :: Parser () -> Parser a -> Parser a
lexeme sep p = p <* sep

-- Like 'lexeme', but does not consume trailing linebreaks.
lexeme' :: Parser a -> Parser a
lexeme' p = p <* hspace

-- Like 'lexstr', but does not consume trailing linebreaks.
lexstr' :: T.Text -> Parser ()
lexstr' = void . try . lexeme' . string

inBraces :: Parser () -> Parser a -> Parser a
inBraces sep = between (lexeme sep "{") (lexeme sep "}")

parseNatural :: Parser () -> Parser Int
parseNatural sep =
  lexeme sep $ L.foldl' addDigit 0 . map num <$> some digitChar
  where
    addDigit acc x = acc * 10 + x
    num c = ord c - ord '0'

restOfLine :: Parser T.Text
restOfLine = do
  l <- restOfLine_
  if T.null l then void eol else void eol <|> eof
  pure l

restOfLine_ :: Parser T.Text
restOfLine_ = takeWhileP Nothing (/= '\n')

parseDescription :: Parser () -> Parser T.Text
parseDescription sep =
  T.unlines <$> pDescLine `manyTill` pDescriptionSeparator
  where
    pDescLine = restOfLine <* sep
    pDescriptionSeparator = void $ "==" *> sep

lTagName :: Parser () -> Parser T.Text
lTagName sep =
  lexeme sep $
    takeWhile1P (Just "tag-constituent character") tagConstituent

parseTags :: Parser () -> Parser [T.Text]
parseTags sep =
  choice
    [ lexeme' "tags" *> inBraces sep (many (lTagName sep)),
      pure []
    ]

tagConstituent :: Char -> Bool
tagConstituent c = isAlphaNum c || c == '_' || c == '-'

parseAction :: Parser () -> Parser TestAction
parseAction sep =
  choice
    [ CompileTimeFailure <$> (lexstr' "error:" *> parseExpectedError sep),
      try parseAsProperty,
      parseAsRunCase
    ]
  where
    parseAsProperty = do
      props <- parseProperty sep
      pure $ RunCases [] [] [] props

    parseAsRunCase =
      RunCases
        <$> parseInputOutputs sep
        <*> many (parseExpectedStructure sep)
        <*> many (parseWarning sep)
        <*> pure []

parseProperty :: Parser () -> Parser [PropertyCase]
parseProperty sep = do
  entryss <- some (parseEntryPointsProp sep)
  pure $ map PropertyCase (concat entryss)

parseEntryPointsProp :: Parser () -> Parser [T.Text]
parseEntryPointsProp sep =
  lexeme' "property:" *> some entry <* sep
  where
    constituent c = not (isSpace c)
    entry = lexeme' $ takeWhile1P Nothing constituent

parseInputOutputs :: Parser () -> Parser [InputOutputs]
parseInputOutputs sep = do
  entrys <- parseEntryPoints sep
  cases <- parseRunCases sep
  pure $
    if null cases
      then []
      else map (`InputOutputs` cases) entrys

parseEntryPoints :: Parser () -> Parser [T.Text]
parseEntryPoints sep =
  (lexeme' "entry:" *> many entry <* sep) <|> pure ["main"]
  where
    constituent c = not (isSpace c) && c /= '}'
    entry = lexeme' $ takeWhile1P Nothing constituent

parseRunTags :: Parser () -> Parser [T.Text]
parseRunTags sep = many . try . lexeme' $ do
  s <- lTagName sep
  guard $ s `notElem` ["input", "structure", "warning"]
  pure s

parseStringLiteral :: Parser () -> Parser T.Text
parseStringLiteral sep =
  lexeme sep . fmap T.pack $ char '"' >> manyTill charLiteral (char '"')

parseRunCases :: Parser () -> Parser [TestRun]
parseRunCases sep = parseRunCases' (0 :: Int)
  where
    parseRunCases' i =
      (:) <$> parseRunCase i <*> parseRunCases' (i + 1)
        <|> pure []
    parseRunCase i = do
      name <- optional $ parseStringLiteral sep
      tags <- parseRunTags sep
      void $ lexeme sep "input"
      input <-
        if "random" `elem` tags
          then parseRandomValues sep
          else
            if "script" `elem` tags
              then parseScriptValues sep
              else parseValues sep
      expr <- parseExpectedResult sep
      pure $ TestRun tags input expr i $ fromMaybe (desc i input) name

    -- If the file is gzipped, we strip the 'gz' extension from
    -- the dataset name.  This makes it more convenient to rename
    -- from 'foo.in' to 'foo.in.gz', as the reported dataset name
    -- does not change (which would make comparisons to historical
    -- data harder).
    desc _ (InFile path)
      | takeExtension path == ".gz" = T.pack $ dropExtension path
      | otherwise = T.pack path
    desc i (Values vs) =
      -- Turn linebreaks into space.
      "#" <> showText i <> " (\"" <> T.unwords (T.lines vs') <> "\")"
      where
        vs' = case T.unwords $ map V.valueText vs of
          s
            | T.length s > 50 -> T.take 50 s <> "..."
            | otherwise -> s
    desc _ (GenValues gens) =
      T.unwords $ map genValueType gens
    desc _ (ScriptValues e) =
      prettyTextOneLine e
    desc _ (ScriptFile path) =
      T.pack path

parseExpectedResult :: Parser () -> Parser (ExpectedResult Success)
parseExpectedResult sep =
  choice
    [ lexeme sep "auto" *> lexeme sep "output" $> Succeeds (Just SuccessGenerateValues),
      Succeeds . Just . SuccessValues <$> (lexeme sep "output" *> parseValues sep),
      RunTimeFailure <$> (lexeme sep "error:" *> parseExpectedError sep),
      pure (Succeeds Nothing)
    ]

parseExpectedError :: Parser () -> Parser ExpectedError
parseExpectedError sep = lexeme sep $ do
  s <- T.strip <$> restOfLine_ <* sep
  if T.null s
    then pure AnyError
    else -- blankCompOpt creates a regular expression that treats
    -- newlines like ordinary characters, which is what we want.
      ThisError s <$> makeRegexOptsM blankCompOpt defaultExecOpt (T.unpack s)

parseScriptValues :: Parser () -> Parser Values
parseScriptValues sep =
  choice
    [ ScriptValues <$> inBraces sep parseScriptExp,
      ScriptFile . T.unpack <$> (lexeme sep "@" *> lexeme sep nextWord)
    ]
  where
    nextWord = takeWhileP Nothing $ not . isSpace

-- | Parse a Futhark expression embedded in a comment block, consisting of the
-- text consumed by the given parser. We slice out that text and hand it to the
-- Futhark parser. This is somewhat clumsy because the Futhark parser is not
-- written with parser combinators.
parseEmbeddedExp :: Parser () -> Parser UncheckedExp
parseEmbeddedExp extent = do
  pos <- sourcePos
  s <- getInput
  bef <- getOffset
  extent
  aft <- getOffset
  -- To get the right source positions, we replace comment prefixes with spaces.
  case parseExpAt pos $ blankCommentPrefix $ T.take (aft - bef) s of
    Left (SyntaxError loc msg) -> do
      case loc of
        Loc.Loc start _ -> setOffset $ Loc.posCoff start
        Loc.NoLoc -> pure ()
      fail $ T.unpack $ T.strip msg
    Right e -> pure e

-- | The current position, in the form used by the Futhark parser.
sourcePos :: Parser Loc.Pos
sourcePos = do
  p <- getSourcePos
  Loc.Pos (sourceName p) (unPos (sourceLine p)) (unPos (sourceColumn p)) <$> getOffset

-- | A script expression extends to the matching closing brace. Braces inside
-- string literals are not counted.
parseScriptExp :: Parser UncheckedExp
parseScriptExp = parseEmbeddedExp balanced
  where
    balanced =
      skipMany $
        choice
          [ void $ takeWhile1P Nothing (`notElem` ("{}\"" :: String)),
            void $ char '"' *> manyTill charLiteral (char '"'),
            void $ char '{' *> balanced *> char '}'
          ]

-- | Replace the comment marker on every line but the first with spaces.
blankCommentPrefix :: T.Text -> T.Text
blankCommentPrefix s =
  case T.lines s of
    [] -> s
    l : ls -> T.intercalate "\n" $ l : map onLine ls
  where
    onLine l = maybe l ("  " <>) $ T.stripPrefix "--" l

parseRandomValues :: Parser () -> Parser Values
parseRandomValues sep = GenValues <$> inBraces sep (many (parseGenValue sep))

parseGenValue :: Parser () -> Parser GenValue
parseGenValue sep =
  choice
    [ GenValue <$> lexeme sep parseType,
      GenPrim <$> lexeme sep V.parsePrimValue
    ]

parseValues :: Parser () -> Parser Values
parseValues sep =
  choice
    [ Values <$> inBraces sep (many $ parseValue sep),
      InFile . T.unpack <$> (lexeme sep "@" *> lexeme sep nextWord)
    ]
  where
    nextWord = takeWhileP Nothing $ not . isSpace

parseWarning :: Parser () -> Parser WarningTest
parseWarning sep = lexeme sep "warning:" >> parseExpectedWarning
  where
    parseExpectedWarning = lexeme sep $ do
      s <- T.strip <$> restOfLine_
      ExpectedWarning s <$> makeRegexOptsM blankCompOpt defaultExecOpt (T.unpack s)

parseExpectedStructure :: Parser () -> Parser StructureTest
parseExpectedStructure sep =
  lexeme sep "structure" *> (StructureTest <$> optimisePipeline sep <*> parseMetrics sep)

optimisePipeline :: Parser () -> Parser StructurePipeline
optimisePipeline sep =
  choice
    [ lexeme sep "gpu-mem" $> GpuMemPipeline,
      lexeme sep "gpu" $> GpuPipeline,
      lexeme sep "mc-mem" $> MCMemPipeline,
      lexeme sep "mc" $> MCPipeline,
      lexeme sep "seq-mem" $> SeqMemPipeline,
      lexeme sep "internalised" $> NoPipeline,
      pure SOACSPipeline
    ]

parseMetrics :: Parser () -> Parser AstMetrics
parseMetrics sep =
  inBraces sep . fmap (AstMetrics . M.fromList) . many $
    (,) <$> lexeme sep (takeWhile1P Nothing constituent) <*> parseNatural sep
  where
    constituent c = isAlpha c || c == '/'

testSpec :: Parser () -> Parser ProgramTest
testSpec sep =
  ProgramTest <$> parseDescription sep <*> parseTags sep <*> parseAction sep

couldNotRead :: IOError -> IO (Either String a)
couldNotRead = pure . Left . show

pProgramTest :: Parser ProgramTest
pProgramTest = do
  void $ many pNonTestLine
  maybe_spec <-
    optional ("--" *> sep *> testSpec sep) <* pEndOfTestBlock <* many pNonTestLine
  case maybe_spec of
    Just spec
      | RunCases old_cases structures warnings old_properties <- testAction spec -> do
          let pAnyBlock = (Left <$> pInputOutputs) <|> (Right <$> pPropertyCases)
          restBlocks <- many (pAnyBlock <* many pNonTestLine)
          let (restCases, restProperties) = partitionEithers restBlocks
          pure
            spec
              { testAction =
                  RunCases
                    (old_cases ++ concat restCases)
                    structures
                    warnings
                    (old_properties ++ concat restProperties)
              }
      | otherwise ->
          many pNonTestLine
            *> notFollowedBy "-- =="
            *> pure spec
            <?> "no more test blocks, since first test block specifies type error."
    Nothing ->
      eof $> noTest
  where
    sep = void $ hspace *> optional (try $ eol *> "--" *> sep)

    noTest =
      ProgramTest mempty mempty (RunCases mempty mempty mempty mempty)

    pEndOfTestBlock =
      (void eol <|> eof) *> notFollowedBy "--"
    pNonTestLine =
      void $ notFollowedBy "-- ==" *> restOfLine
    pInputOutputs =
      try $
        "--" *> sep *> parseDescription sep *> parseInputOutputs sep <* pEndOfTestBlock
    pPropertyCases =
      "--" *> sep *> parseDescription sep *> parseProperty sep <* pEndOfTestBlock

validate :: FilePath -> ProgramTest -> Either String ProgramTest
validate path pt = do
  case testAction pt of
    CompileTimeFailure {} -> pure pt
    RunCases ios _ _ _ -> do
      mapM_ (noDups . map runDescription . iosTestRuns) ios
      Right pt
  where
    noDups xs =
      let xs' = nubOrd xs
       in -- Works because \\ only removes first instance.
          case xs L.\\ xs' of
            [] -> Right ()
            x : _ -> Left $ path <> ": multiple datasets with name " <> show (T.unpack x)

-- | Read the test specification from the given Futhark program.
testSpecFromProgram :: FilePath -> IO (Either String ProgramTest)
testSpecFromProgram path =
  ( either (Left . errorBundlePretty) (validate path) . parse pProgramTest path
      <$> T.readFile path
  )
    `catch` couldNotRead

-- | Like 'testSpecFromProgram', but exits the process on error.
testSpecFromProgramOrDie :: FilePath -> IO ProgramTest
testSpecFromProgramOrDie prog = do
  spec_or_err <- testSpecFromProgram prog
  case spec_or_err of
    Left err -> do
      hPutStrLn stderr err
      exitFailure
    Right spec -> pure spec

testPrograms :: FilePath -> IO [FilePath]
testPrograms dir = filter isFut <$> directoryContents dir
  where
    isFut = (== ".fut") . takeExtension

-- | Read test specifications from the given path, which can be a file
-- or directory containing @.fut@ files and further directories.
testSpecsFromPath :: FilePath -> IO (Either String [(FilePath, ProgramTest)])
testSpecsFromPath path = do
  programs_or_err <- (Right <$> testPrograms path) `catch` couldNotRead
  case programs_or_err of
    Left err -> pure $ Left err
    Right programs -> do
      specs_or_errs <- mapM testSpecFromProgram programs
      pure $ zip programs <$> sequence specs_or_errs

-- | Read test specifications from the given paths, which can be a
-- files or directories containing @.fut@ files and further
-- directories.
testSpecsFromPaths ::
  [FilePath] ->
  IO (Either String [(FilePath, ProgramTest)])
testSpecsFromPaths = fmap (fmap concat . sequence) . mapM testSpecsFromPath

-- | Like 'testSpecsFromPaths', but kills the process on errors.
testSpecsFromPathsOrDie ::
  [FilePath] ->
  IO [(FilePath, ProgramTest)]
testSpecsFromPathsOrDie dirs = do
  specs_or_err <- testSpecsFromPaths dirs
  case specs_or_err of
    Left err -> do
      hPutStrLn stderr err
      exitFailure
    Right specs -> pure specs

-- | Read a test specification from a file.  Expects only a single
-- block, and no comment prefixes.
testSpecFromFile :: FilePath -> IO (Either String ProgramTest)
testSpecFromFile path =
  ( either (Left . errorBundlePretty) Right . parse (testSpec space) path
      <$> T.readFile path
  )
    `catch` couldNotRead

-- | Like 'testSpecFromFile', but kills the process on errors.
testSpecFromFileOrDie :: FilePath -> IO ProgramTest
testSpecFromFileOrDie dirs = do
  spec_or_err <- testSpecFromFile dirs
  case spec_or_err of
    Left err -> do
      hPutStrLn stderr err
      exitFailure
    Right spec -> pure spec