Preprocessor.hs

{-# LANGUAGE OverloadedStrings #-}

module Skeletest.Internal.Preprocessor (
  processFile,
) where

import Control.Monad (guard)
import Data.Char (isDigit, isLower, isUpper)
import Data.List (sort)
import Data.Maybe (mapMaybe)
import Data.Text (Text)
import Data.Text qualified as Text
import Skeletest.Internal.Constants (mainFileSpecsListIdentifier)
import Skeletest.Internal.Error (SkeletestError (..))
import System.Directory (doesDirectoryExist, listDirectory)
import System.FilePath (makeRelative, splitExtensions, takeDirectory, (</>))
import UnliftIO.Exception (throwIO)

-- | Preprocess the given Haskell file. See Main.hs
processFile :: FilePath -> Text -> IO Text
processFile path file = do
  file' <- if isMain file then updateMainFile path file else pure file
  pure
    . addLine pluginPragma
    . addLine linePragma
    $ file'
 where
  addLine line f = line <> "\n" <> f
  quoted s = "\"" <> s <> "\""

  pluginPragma = "{-# OPTIONS_GHC -fplugin=Skeletest.Internal.Plugin #-}"
  linePragma =
    -- this is needed to tell GHC to use original path in error messages
    "{-# LINE 1 " <> quoted (Text.pack path) <> " #-}"

isMain :: Text -> Bool
isMain file =
  case mapMaybe getModuleName $ Text.lines file of
    -- there was a module line
    [name] -> name == "Main"
    -- there were no module lines, it's the main module
    [] -> True
    -- something else? just silently ignore it
    _ -> False
 where
  getModuleName s =
    case Text.words s of
      "module" : name : _ -> Just name
      _ -> Nothing

updateMainFile :: FilePath -> Text -> IO Text
updateMainFile path file = do
  modules <- findTestModules path
  either throwIO pure $
    pure file
      >>= insertImports modules
      >>= pure . addSpecsList modules

-- | Find all test modules using the given path to the Main module.
--
-- >>> findTestModules "test/Main.hs"
-- ["My.Module.Test1", "My.Module.Test2", ...]
findTestModules :: FilePath -> IO [(FilePath, Text)]
findTestModules path = mapMaybe toTestModule <$> listDirectoryRecursive testDir
 where
  testDir = takeDirectory path

  toTestModule fp = do
    guard (fp /= path)
    (fpNoExt, ".hs") <- pure $ splitExtensions fp
    guard ("Spec" `Text.isSuffixOf` Text.pack fpNoExt)
    name <- moduleNameFromPath $ Text.pack $ makeRelative testDir fpNoExt
    pure (fp, name)

  moduleNameFromPath = fmap (Text.intercalate ".") . mapM validateModuleName . Text.splitOn "/"

  -- https://www.haskell.org/onlinereport/syntax-iso.html
  -- large { small | large | digit | ' }
  validateModuleName name = do
    (first, rest) <- Text.uncons name
    guard $ isUpper first
    guard $ Text.all (\c -> isUpper c || isLower c || isDigit c || c == '\'') rest
    pure name

addSpecsList :: [(FilePath, Text)] -> Text -> Text
addSpecsList testModules file =
  Text.unlines
    [ file
    , mainFileSpecsListIdentifier <> " :: [(FilePath, Spec)]"
    , mainFileSpecsListIdentifier <> " = " <> renderSpecList specsList
    ]
 where
  specsList =
    [ (quote $ Text.pack fp, modName <> ".spec")
    | (fp, modName) <- testModules
    ]
  quote s = "\"" <> s <> "\""
  renderSpecList xs = "[" <> (Text.intercalate ", " . map renderSpecInfo) xs <> "]"
  renderSpecInfo (fp, spec) = "(" <> fp <> ", " <> spec <> ")"

-- | Add imports after the Skeletest.Main import, which should always be present in the Main module.
insertImports :: [(FilePath, Text)] -> Text -> Either SkeletestError Text
insertImports testModules file =
  let (pre, post) = break isSkeletestImport $ Text.lines file
   in if null post
        then Left $ CompilationError Nothing "Could not find Skeletest.Main import in Main module"
        else pure . Text.unlines $ pre <> importTests <> post
 where
  isSkeletestImport line =
    case Text.words line of
      "import" : "Skeletest.Main" : _ -> True
      _ -> False

  importTests =
    [ "import qualified " <> name
    | (_, name) <- testModules
    ]

{----- Helpers -----}

listDirectoryRecursive :: FilePath -> IO [FilePath]
listDirectoryRecursive fp = fmap (sort . concat) . mapM (go . (fp </>)) =<< listDirectory fp
 where
  go child = do
    isDir <- doesDirectoryExist child
    if isDir
      then listDirectoryRecursive child
      else pure [child]