{-# LANGUAGE TupleSections #-}

module Language.Fortran.Extras where

import           Control.Exception              ( try
                                                , SomeException
                                                )
import           Data.Data                      ( Data )
import           Data.List                      ( find
                                                )
import           Data.Maybe                     ( fromMaybe
                                                , mapMaybe
                                                )
import           Data.Generics.Uniplate.Data    ( universeBi )
import           Language.Fortran.AST           ( A0
                                                , Block
                                                , ProgramFile
                                                , Statement
                                                , ProgramUnit(..)
                                                , ProgramUnitName(..)
                                                )
import           Language.Fortran.Analysis      ( Analysis
                                                , puSrcName
                                                )
import           Language.Fortran.Version       ( FortranVersion(..) )
import           System.FilePath                ( takeExtension )
import           System.Exit                    ( ExitCode(..)
                                                , exitWith
                                                )
import           System.IO                      ( hPutStr
                                                , hPutStrLn
                                                , stderr
                                                )
import           Options.Applicative
import qualified Language.Fortran.Parser as Parser
import qualified Language.Fortran.Extras.ProgramFile
                                               as P
import qualified Language.Fortran.Extras.Analysis
                                               as A
import           Language.Fortran.Util.ModFile  ( decodeModFiles' )
import           Language.Fortran.Extras.RunOptions
                                                ( unwrapFortranSrcOptions
                                                , getFortranSrcRunOptions
                                                , getRunOptions
                                                , FortranSrcRunOptions(..)
                                                , RunOptions(..)
                                                )

-- | Get a list of all 'Block's in a 'ProgramFile'
allB :: Data a => ProgramFile a -> [Block a]
allB :: forall a. Data a => ProgramFile a -> [Block a]
allB = forall from to. Biplate from to => from -> [to]
universeBi

-- | Get a list of all 'Statement's in a 'ProgramFile'
allS :: Data a => ProgramFile a -> [Statement a]
allS :: forall a. Data a => ProgramFile a -> [Statement a]
allS = forall from to. Biplate from to => from -> [to]
universeBi

-- | Get a list of all 'ProgramUnit's in a 'ProgramFile'
allPU :: Data a => ProgramFile a -> [ProgramUnit a]
allPU :: forall a. Data a => ProgramFile a -> [ProgramUnit a]
allPU = forall from to. Biplate from to => from -> [to]
universeBi

-- | Get a list of all 'Block's in a 'ProgramUnit'
allPUB :: Data a => ProgramUnit a -> [Block a]
allPUB :: forall a. Data a => ProgramUnit a -> [Block a]
allPUB = forall from to. Biplate from to => from -> [to]
universeBi

-- | Get a list of all 'Statement's in a 'ProgramUnit'
allPUS :: Data a => ProgramUnit a -> [Statement a]
allPUS :: forall a. Data a => ProgramUnit a -> [Statement a]
allPUS = forall from to. Biplate from to => from -> [to]
universeBi

-- | Given a 'ProgramFile' find a 'ProgramUnit' with a particular 'ProgramUnitName'
findPU'
  :: Data a
  => ProgramUnitName
  -> ProgramFile (Analysis a)
  -> Maybe (ProgramUnit (Analysis a))
findPU' :: forall a.
Data a =>
ProgramUnitName
-> ProgramFile (Analysis a) -> Maybe (ProgramUnit (Analysis a))
findPU' ProgramUnitName
n = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\ProgramUnit (Analysis a)
pu -> forall a. ProgramUnit (Analysis a) -> ProgramUnitName
puSrcName ProgramUnit (Analysis a)
pu forall a. Eq a => a -> a -> Bool
== ProgramUnitName
n) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Data a => ProgramFile a -> [ProgramUnit a]
allPU

-- | Given a 'ProgramFile' find a 'ProgramUnit' with a particular name
findPU
  :: Data a
  => String
  -> ProgramFile (Analysis a)
  -> Maybe (ProgramUnit (Analysis a))
findPU :: forall a.
Data a =>
String
-> ProgramFile (Analysis a) -> Maybe (ProgramUnit (Analysis a))
findPU String
n = forall a.
Data a =>
ProgramUnitName
-> ProgramFile (Analysis a) -> Maybe (ProgramUnit (Analysis a))
findPU' forall a b. (a -> b) -> a -> b
$ String -> ProgramUnitName
Named String
n

-- | Get a 'ProgramFile' from version and path specified in 'FortranSrcRunOptions'
programFile :: FortranSrcRunOptions -> IO (ProgramFile A0)
programFile :: FortranSrcRunOptions -> IO (ProgramFile A0)
programFile FortranSrcRunOptions
options = do
  (String
pfPath, ByteString
pfContents, [String]
pfIncludes, FortranVersion
fVersion) <- FortranSrcRunOptions
-> IO (String, ByteString, [String], FortranVersion)
unwrapFortranSrcOptions FortranSrcRunOptions
options
  case FortranVersion
fVersion of
    FortranVersion
Fortran77Legacy ->
      FortranVersion
-> [String] -> String -> ByteString -> IO (ProgramFile A0)
P.versionedExpandedProgramFile FortranVersion
fVersion [String]
pfIncludes String
pfPath ByteString
pfContents
    FortranVersion
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ FortranVersion -> String -> ByteString -> ProgramFile A0
P.versionedProgramFile FortranVersion
fVersion String
pfPath ByteString
pfContents

incFile :: FortranSrcRunOptions -> IO [Block A0]
incFile :: FortranSrcRunOptions -> IO [Block A0]
incFile FortranSrcRunOptions
options = do
  (String
pfPath, ByteString
pfContents, [String]
_pfIncludes, FortranVersion
_fVersion) <- FortranSrcRunOptions
-> IO (String, ByteString, [String], FortranVersion)
unwrapFortranSrcOptions FortranSrcRunOptions
options
  forall e (m :: * -> *) a.
(Exception e, MonadIO m) =>
Either e a -> m a
Parser.throwIOLeft forall a b. (a -> b) -> a -> b
$ Parser [Block A0]
Parser.f77lIncludesNoTransform String
pfPath ByteString
pfContents

-- | Get a 'ProgramFile' with 'Analysis' from version and path specified
-- in 'FortranSrcRunOptions'
programAnalysis :: FortranSrcRunOptions -> IO (ProgramFile (Analysis A0))
programAnalysis :: FortranSrcRunOptions -> IO (ProgramFile (Analysis A0))
programAnalysis FortranSrcRunOptions
options = do
  (String
pfPath, ByteString
pfContents, [String]
pfIncludes, FortranVersion
fVersion) <- FortranSrcRunOptions
-> IO (String, ByteString, [String], FortranVersion)
unwrapFortranSrcOptions FortranSrcRunOptions
options
  case FortranVersion
fVersion of
    FortranVersion
Fortran77Legacy ->
      FortranVersion
-> [String]
-> String
-> ByteString
-> IO (ProgramFile (Analysis A0))
A.versionedExpandedProgramAnalysis FortranVersion
fVersion [String]
pfIncludes String
pfPath ByteString
pfContents
    FortranVersion
_ -> if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [String]
pfIncludes
      then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ FortranVersion -> String -> ByteString -> ProgramFile (Analysis A0)
A.versionedProgramAnalysis FortranVersion
fVersion String
pfPath ByteString
pfContents
      else do
        ModFiles
pfMods <- [String] -> IO ModFiles
decodeModFiles' [String]
pfIncludes
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ FortranVersion
-> ModFiles -> String -> ByteString -> ProgramFile (Analysis A0)
A.versionedProgramAnalysisWithMods FortranVersion
fVersion
                                                    ModFiles
pfMods
                                                    String
pfPath
                                                    ByteString
pfContents

-- | Parse arguments and return 'ProgramFile'
--
-- This function has the purpose of being a general entry-point for `fortran-src-tools` tools.
-- It handles parsing the common arguments and then returns a 'ProgramFile' that the tool can then use to
-- do further processing.
-- 
-- This function takes in two arguments, namely:
--
-- * A description of the program that shows up when the program is invoked incorrectly
-- @
-- $ some-fortran-tool
-- Missing: (-v|--fortranVersion VERSION) PATH
-- 
-- Usage: vars (-v|--fortranVersion VERSION) [-I|--include DIRECTORY] PATH
--   THIS IS WHERE THE DESCRIPTION GOES
-- @
--
-- * A header that is shown when the user passes the `--help` argument to the tool (note that this does
--  not show up when the program is invoked incorrectly)
-- @
-- $ some-fortran-tool --help
-- THIS IS WHERE THE HEADER GOES
-- 
-- Usage: vars (-v|--fortranVersion VERSION) [-I|--include DIRECTORY] PATH
--   THIS IS WHERE THE DESCRIPTION GOES
-- 
-- Available options:
--    -h,--help                Show this help text
--    -v,--fortranVersion VERSION
--                             Fortran version to use, format:
--                             Fortran[66/77/BigIron/77Legacy/77Extended/90]
--    -I,--include DIRECTORY   Directory to include files from
-- @
getProgramFile :: String -> String -> IO (ProgramFile A0)
getProgramFile :: String -> String -> IO (ProgramFile A0)
getProgramFile String
programDescription String
programHeader = do
  FortranSrcRunOptions
options <- String -> String -> IO FortranSrcRunOptions
getFortranSrcRunOptions String
programDescription String
programHeader
  FortranSrcRunOptions -> IO (ProgramFile A0)
programFile FortranSrcRunOptions
options

-- | Parse arguments and return a 'ProgramFile' with 'Analysis'
--
-- This function takes the same arguments as 'getProgramFile', however it will return an 'Analysis' object
-- within the 'ProgramFile' monad. If any `-I DIR` arguments were specified when invoking the tool, this
-- function will ensure that any module files located in this directory are loaded and incorporated into the
-- analysis.
getProgramAnalysis :: String -> String -> IO (ProgramFile (Analysis A0))
getProgramAnalysis :: String -> String -> IO (ProgramFile (Analysis A0))
getProgramAnalysis String
programDescription String
programHeader = do
  FortranSrcRunOptions
options <- String -> String -> IO FortranSrcRunOptions
getFortranSrcRunOptions String
programDescription String
programHeader
  FortranSrcRunOptions -> IO (ProgramFile (Analysis A0))
programAnalysis FortranSrcRunOptions
options

-- | Helper to print out exceptions with the name of the file being processed
errorHandler :: String -> Either SomeException () -> IO ()
errorHandler :: String -> Either SomeException A0 -> IO A0
errorHandler String
filename (Left SomeException
e) = do
  Handle -> String -> IO A0
hPutStrLn Handle
stderr forall a b. (a -> b) -> a -> b
$ String
"Caught exception in file: " forall a. [a] -> [a] -> [a]
++ String
filename
  Handle -> String -> IO A0
hPutStr Handle
stderr forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unlines forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (String
"    " forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
lines forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show SomeException
e
  forall a. ExitCode -> IO a
exitWith forall a b. (a -> b) -> a -> b
$ Int -> ExitCode
ExitFailure Int
2
errorHandler String
_ (Right A0
_) = forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Given a program description, a program header, and a handler that
-- takes a 'ProgramFile', this function generates the 'ProgramFile'
-- and passes it to the handler, while catching any exceptions that
-- occur within either the parsing of the 'ProgramFile' itself
-- or while the handler is processing
withProgramFile :: String -> String -> (ProgramFile A0 -> IO ()) -> IO ()
withProgramFile :: String -> String -> (ProgramFile A0 -> IO A0) -> IO A0
withProgramFile String
programDescription String
programHeader ProgramFile A0 -> IO A0
handler = do
  FortranSrcRunOptions
options <- String -> String -> IO FortranSrcRunOptions
getFortranSrcRunOptions String
programDescription String
programHeader
  Either SomeException A0
results <- forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ FortranSrcRunOptions -> IO (ProgramFile A0)
programFile FortranSrcRunOptions
options forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ProgramFile A0 -> IO A0
handler
  String -> Either SomeException A0 -> IO A0
errorHandler (FortranSrcRunOptions -> String
path FortranSrcRunOptions
options) Either SomeException A0
results

-- | Given a program description, a program header, and a handler that
-- takes a 'ProgramFile', this function generates the 'ProgramFile'
-- annotated with 'Analysis' and passes it to the handler,
-- while catching any exceptions that occur within either the parsing
-- of the 'ProgramFile' itself or while the handler is processing
withProgramAnalysis
  :: String -> String -> (ProgramFile (Analysis A0) -> IO ()) -> IO ()
withProgramAnalysis :: String -> String -> (ProgramFile (Analysis A0) -> IO A0) -> IO A0
withProgramAnalysis String
programDescription String
programHeader ProgramFile (Analysis A0) -> IO A0
handler = do
  FortranSrcRunOptions
options <- String -> String -> IO FortranSrcRunOptions
getFortranSrcRunOptions String
programDescription String
programHeader
  Either SomeException A0
results <- forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ FortranSrcRunOptions -> IO (ProgramFile (Analysis A0))
programAnalysis FortranSrcRunOptions
options forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ProgramFile (Analysis A0) -> IO A0
handler
  String -> Either SomeException A0 -> IO A0
errorHandler (FortranSrcRunOptions -> String
path FortranSrcRunOptions
options) Either SomeException A0
results

-- | Given a program description, a program header, a parser for
-- tool CLI options, and a handler that takes tool CLI options object,
-- and a 'ProgramFile', this function generates tool CLI options object,
-- and a 'ProgramFile' annotated with 'Analysis', and passes them to the
-- handler, while catching any exceptions that occurs during handler
-- processing
withToolOptionsAndProgramAnalysis
  :: String
  -> String
  -> Parser a
  -> (a -> ProgramFile (Analysis A0) -> IO ())
  -> IO ()
withToolOptionsAndProgramAnalysis :: forall a.
String
-> String
-> Parser a
-> (a -> ProgramFile (Analysis A0) -> IO A0)
-> IO A0
withToolOptionsAndProgramAnalysis String
programDescription String
programHeader Parser a
toolOptsParser a -> ProgramFile (Analysis A0) -> IO A0
handler
  = do
    RunOptions a
options <- forall a. String -> String -> Parser a -> IO (RunOptions a)
getRunOptions String
programDescription String
programHeader Parser a
toolOptsParser
    let (FortranSrcRunOptions
fortranSrcOptions, a
toolOptions) =
          (forall a. RunOptions a -> FortranSrcRunOptions
fortranSrcOpts RunOptions a
options, forall a. RunOptions a -> a
toolOpts RunOptions a
options)
    Either SomeException A0
results <- forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ FortranSrcRunOptions -> IO (ProgramFile (Analysis A0))
programAnalysis FortranSrcRunOptions
fortranSrcOptions forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> ProgramFile (Analysis A0) -> IO A0
handler a
toolOptions
    String -> Either SomeException A0 -> IO A0
errorHandler (FortranSrcRunOptions -> String
path FortranSrcRunOptions
fortranSrcOptions) Either SomeException A0
results

-- | Given a program description, a program header, cli options parser, and a
-- handler which takes the cli options and either an include files '[Block A0]'
-- or 'Programfile A0', run the handler on the appropriately parsed source,
-- parsing anything that has the ".inc" extension as an include
withToolOptionsAndProgramOrBlock
  :: String
  -> String
  -> Parser a
  -> (a -> Either (FilePath, [Block A0]) (ProgramFile A0) -> IO ())
  -> IO ()
withToolOptionsAndProgramOrBlock :: forall a.
String
-> String
-> Parser a
-> (a -> Either (String, [Block A0]) (ProgramFile A0) -> IO A0)
-> IO A0
withToolOptionsAndProgramOrBlock String
programDescription String
programHeader Parser a
optsParser a -> Either (String, [Block A0]) (ProgramFile A0) -> IO A0
handler = do
  RunOptions FortranSrcRunOptions
srcOptions a
toolOptions <-
    forall a. String -> String -> Parser a -> IO (RunOptions a)
getRunOptions String
programDescription String
programHeader Parser a
optsParser
  Either (String, [Block A0]) (ProgramFile A0)
ast <- if String -> String
takeExtension (FortranSrcRunOptions -> String
path FortranSrcRunOptions
srcOptions) forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String
".inc", String
".ins"]
    then forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FortranSrcRunOptions -> String
path FortranSrcRunOptions
srcOptions, ) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FortranSrcRunOptions -> IO [Block A0]
incFile FortranSrcRunOptions
srcOptions
    else forall a b. b -> Either a b
Right forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FortranSrcRunOptions -> IO (ProgramFile A0)
programFile FortranSrcRunOptions
srcOptions
  a -> Either (String, [Block A0]) (ProgramFile A0) -> IO A0
handler a
toolOptions Either (String, [Block A0]) (ProgramFile A0)
ast

-- | Given a 'ProgramUnit' return a pair of the name of the unit as well as the unit itself
-- only if the 'ProgramUnit' is a 'PUMain', 'PUSubroutine', or a 'PUFunction'
namedProgramUnit :: Data a => ProgramUnit a -> Maybe (String, ProgramUnit a)
namedProgramUnit :: forall a. Data a => ProgramUnit a -> Maybe (String, ProgramUnit a)
namedProgramUnit pu :: ProgramUnit a
pu@(PUMain a
_ SrcSpan
_ Maybe String
mn [Block a]
_ Maybe [ProgramUnit a]
_) = forall a. a -> Maybe a
Just (forall a. a -> Maybe a -> a
fromMaybe String
"MAIN" Maybe String
mn, ProgramUnit a
pu)
namedProgramUnit pu :: ProgramUnit a
pu@(PUSubroutine a
_ SrcSpan
_ PrefixSuffix a
_ String
n Maybe (AList Expression a)
_ [Block a]
_ Maybe [ProgramUnit a]
_) = forall a. a -> Maybe a
Just (String
n, ProgramUnit a
pu)
namedProgramUnit pu :: ProgramUnit a
pu@(PUFunction a
_ SrcSpan
_ Maybe (TypeSpec a)
_ PrefixSuffix a
_ String
n Maybe (AList Expression a)
_ Maybe (Expression a)
_ [Block a]
_ Maybe [ProgramUnit a]
_) = forall a. a -> Maybe a
Just (String
n, ProgramUnit a
pu)
namedProgramUnit ProgramUnit a
_ = forall a. Maybe a
Nothing

-- | Given a 'ProgramFile' return all of the named 'ProgramUnits' within the file, i.e.
-- 'PUMain's, 'PUSubroutine's, or 'PUFunction's
namedProgramUnits :: Data a => ProgramFile a -> [(String, ProgramUnit a)]
namedProgramUnits :: forall a. Data a => ProgramFile a -> [(String, ProgramUnit a)]
namedProgramUnits = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe forall a. Data a => ProgramUnit a -> Maybe (String, ProgramUnit a)
namedProgramUnit forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Data a => ProgramFile a -> [ProgramUnit a]
allPU