-- | Definition of the core compiler driver building blocks.  The
-- spine of the compiler is the 'FutharkM' monad, although note that
-- individual passes are pure functions, and do not use the 'FutharkM'
-- monad (see "Futhark.Pass").
--
-- Running the compiler involves producing an initial IR program (see
-- "Futhark.Compiler.Program"), running a 'Pipeline' to produce a
-- final program (still in IR), then running an 'Action', which is
-- usually a code generator.
module Futhark.Pipeline
  ( Pipeline,
    PipelineConfig (..),
    Action (..),
    FutharkM,
    runFutharkM,
    Verbosity (..),
    module Futhark.Error,
    onePass,
    passes,
    condPipeline,
    runPipeline,
  )
where

import Control.Category
import Control.Exception (SomeException, catch, throwIO)
import Control.Monad
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Control.Parallel
import Data.Text qualified as T
import Data.Text.IO qualified as T
import Data.Time.Clock
import Futhark.Analysis.Alias qualified as Alias
import Futhark.Compiler.Config (Verbosity (..))
import Futhark.Error
import Futhark.IR (PrettyRep, Prog)
import Futhark.IR.TypeCheck
import Futhark.MonadFreshNames
import Futhark.Pass
import Futhark.Util.Log
import Futhark.Util.Pretty (prettyText)
import System.IO
import Text.Printf
import Prelude hiding (id, (.))

newtype FutharkEnv = FutharkEnv {futharkVerbose :: Verbosity}

data FutharkState = FutharkState
  { futharkPrevLog :: UTCTime,
    futharkNameSource :: VNameSource
  }

-- | The main Futhark compiler driver monad - basically some state
-- tracking on top if 'IO'.
newtype FutharkM a = FutharkM (ExceptT CompilerError (StateT FutharkState (ReaderT FutharkEnv IO)) a)
  deriving
    ( Applicative,
      Functor,
      Monad,
      MonadError CompilerError,
      MonadState FutharkState,
      MonadReader FutharkEnv,
      MonadIO
    )

instance MonadFreshNames FutharkM where
  getNameSource = gets futharkNameSource
  putNameSource src = modify $ \s -> s {futharkNameSource = src}

instance MonadLogger FutharkM where
  addLog = mapM_ perLine . T.lines . toText
    where
      perLine msg = do
        verb <- asks $ (>= Verbose) . futharkVerbose
        prev <- gets futharkPrevLog
        now <- liftIO getCurrentTime
        let delta :: Double
            delta = fromRational $ toRational (now `diffUTCTime` prev)
            prefix = printf "[  +%.6f] " delta
        modify $ \s -> s {futharkPrevLog = now}
        when verb $ liftIO $ T.hPutStrLn stderr $ T.pack prefix <> msg

runFutharkM' ::
  FutharkM a ->
  FutharkState ->
  FutharkEnv ->
  IO (Either CompilerError a, FutharkState)
runFutharkM' (FutharkM m) s =
  runReaderT (runStateT (runExceptT m) s)

-- | Run a 'FutharkM' action.
runFutharkM :: FutharkM a -> Verbosity -> IO (Either CompilerError a)
runFutharkM m verbose = do
  s <- FutharkState <$> getCurrentTime <*> pure blankNameSource
  fst <$> runFutharkM' m s (FutharkEnv verbose)

catchIO :: FutharkM a -> (SomeException -> FutharkM a) -> FutharkM a
catchIO m f = FutharkM $ do
  s <- get
  env <- ask
  (x, s') <-
    liftIO $
      runFutharkM' m s env `catch` \e ->
        runFutharkM' (f e) s env
  put s'
  case x of
    Left e -> throwError e
    Right x' -> pure x'

-- | A compilation always ends with some kind of action.
data Action rep = Action
  { actionName :: String,
    actionDescription :: String,
    actionProcedure :: Prog rep -> FutharkM ()
  }

-- | Configuration object for running a compiler pipeline.
data PipelineConfig = PipelineConfig
  { pipelineVerbose :: Bool,
    pipelineValidate :: Bool
  }

-- | A compiler pipeline is conceptually a function from programs to
-- programs, where the actual representation may change.  Pipelines
-- can be composed using their 'Category' instance.
newtype Pipeline fromrep torep = Pipeline
  { unPipeline ::
      forall a.
      PipelineConfig ->
      Prog fromrep ->
      FutharkM ((Prog torep -> FutharkM a) -> FutharkM a)
  }

instance Category Pipeline where
  id = Pipeline $ \_ prog -> pure $ \c -> c prog
  p2 . p1 = Pipeline perform
    where
      perform cfg prog = do
        rc <- unPipeline p1 cfg prog
        rc $ unPipeline p2 cfg

-- | Run the pipeline on the given program.
runPipeline ::
  Pipeline fromrep torep ->
  PipelineConfig ->
  Prog fromrep ->
  FutharkM (Prog torep)
runPipeline p cfg prog = do
  rc <- unPipeline p cfg prog
  rc pure

-- | Construct a pipeline from a single compiler pass.
onePass ::
  Checkable torep =>
  Pass fromrep torep ->
  Pipeline fromrep torep
onePass pass = Pipeline perform
  where
    perform cfg prog = do
      when (pipelineVerbose cfg) . logMsg $
        "Running pass " <> T.pack (passName pass)
      prog' <- runPass pass prog
      -- Spark validation in a separate task and speculatively execute
      -- next pass.  If the next pass throws an exception, we better
      -- be ready to catch it and check if it might be because the
      -- program was actually ill-typed.
      let check =
            if pipelineValidate cfg
              then validate prog'
              else Right ()
      par check $ pure $ \c ->
        (errorOnError check pure =<< c prog')
          `catchIO` errorOnError check (liftIO . throwIO)
    validate prog =
      let prog' = Alias.aliasAnalysis prog
       in case checkProg prog' of
            Left err -> Left (prog', err)
            Right () -> Right ()
    errorOnError (Left (prog, err)) _ _ =
      validationError pass prog $ show err
    errorOnError _ c x = c x

-- | Conditionally run pipeline if predicate is true.
condPipeline ::
  (Prog rep -> Bool) -> Pipeline rep rep -> Pipeline rep rep
condPipeline cond (Pipeline f) =
  Pipeline $ \cfg prog ->
    if cond prog
      then f cfg prog
      else pure $ \c -> c prog

-- | Create a pipeline from a list of passes.
passes ::
  Checkable rep =>
  [Pass rep rep] ->
  Pipeline rep rep
passes = foldl (>>>) id . map onePass

validationError ::
  PrettyRep rep =>
  Pass fromrep torep ->
  Prog rep ->
  String ->
  FutharkM a
validationError pass prog err =
  throwError $ InternalError msg (prettyText prog) CompilerBug
  where
    msg = "Type error after pass '" <> T.pack (passName pass) <> "':\n" <> T.pack err

runPass ::
  Pass fromrep torep ->
  Prog fromrep ->
  FutharkM (Prog torep)
runPass pass prog = do
  (prog', logged) <- runPassM (passFunction pass prog)
  verb <- asks $ (>= VeryVerbose) . futharkVerbose
  when verb $ addLog logged
  pure prog'