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.Exit                    ( ExitCode(..)
                                                , exitWith
                                                )
import           System.IO                      ( hPutStr
                                                , hPutStrLn
                                                , stderr
                                                )
import           Options.Applicative
import qualified Language.Fortran.Extras.ProgramFile
                                               as P
import qualified Language.Fortran.Extras.Analysis
                                               as A
import           Language.Fortran.Extras.ModFiles
                                                ( 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 :: ProgramFile a -> [Block a]
allB = ProgramFile a -> [Block a]
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 :: ProgramFile a -> [Statement a]
allS = ProgramFile a -> [Statement a]
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 :: ProgramFile a -> [ProgramUnit a]
allPU = ProgramFile a -> [ProgramUnit a]
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 :: ProgramUnit a -> [Block a]
allPUB = ProgramUnit a -> [Block a]
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 :: ProgramUnit a -> [Statement a]
allPUS = ProgramUnit a -> [Statement a]
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' :: ProgramUnitName
-> ProgramFile (Analysis a) -> Maybe (ProgramUnit (Analysis a))
findPU' ProgramUnitName
n = (ProgramUnit (Analysis a) -> Bool)
-> [ProgramUnit (Analysis a)] -> Maybe (ProgramUnit (Analysis a))
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\ProgramUnit (Analysis a)
pu -> ProgramUnit (Analysis a) -> ProgramUnitName
forall a. ProgramUnit (Analysis a) -> ProgramUnitName
puSrcName ProgramUnit (Analysis a)
pu ProgramUnitName -> ProgramUnitName -> Bool
forall a. Eq a => a -> a -> Bool
== ProgramUnitName
n) ([ProgramUnit (Analysis a)] -> Maybe (ProgramUnit (Analysis a)))
-> (ProgramFile (Analysis a) -> [ProgramUnit (Analysis a)])
-> ProgramFile (Analysis a)
-> Maybe (ProgramUnit (Analysis a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ProgramFile (Analysis a) -> [ProgramUnit (Analysis a)]
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 :: String
-> ProgramFile (Analysis a) -> Maybe (ProgramUnit (Analysis a))
findPU String
n = ProgramUnitName
-> ProgramFile (Analysis a) -> Maybe (ProgramUnit (Analysis a))
forall a.
Data a =>
ProgramUnitName
-> ProgramFile (Analysis a) -> Maybe (ProgramUnit (Analysis a))
findPU' (ProgramUnitName
 -> ProgramFile (Analysis a) -> Maybe (ProgramUnit (Analysis a)))
-> ProgramUnitName
-> ProgramFile (Analysis a)
-> Maybe (ProgramUnit (Analysis a))
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
_ -> ProgramFile A0 -> IO (ProgramFile A0)
forall (m :: * -> *) a. Monad m => a -> m a
return (ProgramFile A0 -> IO (ProgramFile A0))
-> ProgramFile A0 -> IO (ProgramFile A0)
forall a b. (a -> b) -> a -> b
$ FortranVersion -> String -> ByteString -> ProgramFile A0
P.versionedProgramFile FortranVersion
fVersion 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 [String] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [String]
pfIncludes
      then ProgramFile (Analysis A0) -> IO (ProgramFile (Analysis A0))
forall (m :: * -> *) a. Monad m => a -> m a
return (ProgramFile (Analysis A0) -> IO (ProgramFile (Analysis A0)))
-> ProgramFile (Analysis A0) -> IO (ProgramFile (Analysis A0))
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
        ProgramFile (Analysis A0) -> IO (ProgramFile (Analysis A0))
forall (m :: * -> *) a. Monad m => a -> m a
return (ProgramFile (Analysis A0) -> IO (ProgramFile (Analysis A0)))
-> ProgramFile (Analysis A0) -> IO (ProgramFile (Analysis A0))
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 the general entry-point for the `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 (String -> IO A0) -> String -> IO A0
forall a b. (a -> b) -> a -> b
$ String
"Caught exception in file: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
filename
  Handle -> String -> IO A0
hPutStr Handle
stderr (String -> IO A0) -> (String -> String) -> String -> IO A0
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unlines ([String] -> String) -> (String -> [String]) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> String) -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (String
"    " String -> String -> String
forall a. [a] -> [a] -> [a]
++) ([String] -> [String])
-> (String -> [String]) -> String -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
lines (String -> IO A0) -> String -> IO A0
forall a b. (a -> b) -> a -> b
$ SomeException -> String
forall a. Show a => a -> String
show SomeException
e
  ExitCode -> IO A0
forall a. ExitCode -> IO a
exitWith (ExitCode -> IO A0) -> ExitCode -> IO A0
forall a b. (a -> b) -> a -> b
$ Int -> ExitCode
ExitFailure Int
2
errorHandler String
_ (Right A0
_) = A0 -> IO 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 <- IO A0 -> IO (Either SomeException A0)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO A0 -> IO (Either SomeException A0))
-> IO A0 -> IO (Either SomeException A0)
forall a b. (a -> b) -> a -> b
$ FortranSrcRunOptions -> IO (ProgramFile A0)
programFile FortranSrcRunOptions
options IO (ProgramFile A0) -> (ProgramFile A0 -> IO A0) -> IO A0
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 <- IO A0 -> IO (Either SomeException A0)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO A0 -> IO (Either SomeException A0))
-> IO A0 -> IO (Either SomeException A0)
forall a b. (a -> b) -> a -> b
$ FortranSrcRunOptions -> IO (ProgramFile (Analysis A0))
programAnalysis FortranSrcRunOptions
options IO (ProgramFile (Analysis A0))
-> (ProgramFile (Analysis A0) -> IO A0) -> IO A0
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 :: 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 <- String -> String -> Parser a -> IO (RunOptions a)
forall a. String -> String -> Parser a -> IO (RunOptions a)
getRunOptions String
programDescription String
programHeader Parser a
toolOptsParser
    let (FortranSrcRunOptions
fortranSrcOptions, a
toolOptions) =
          (RunOptions a -> FortranSrcRunOptions
forall a. RunOptions a -> FortranSrcRunOptions
fortranSrcOpts RunOptions a
options, RunOptions a -> a
forall a. RunOptions a -> a
toolOpts RunOptions a
options)
    Either SomeException A0
results <- IO A0 -> IO (Either SomeException A0)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO A0 -> IO (Either SomeException A0))
-> IO A0 -> IO (Either SomeException A0)
forall a b. (a -> b) -> a -> b
$ FortranSrcRunOptions -> IO (ProgramFile (Analysis A0))
programAnalysis FortranSrcRunOptions
fortranSrcOptions IO (ProgramFile (Analysis A0))
-> (ProgramFile (Analysis A0) -> IO A0) -> IO A0
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 '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 :: ProgramUnit a -> Maybe (String, ProgramUnit a)
namedProgramUnit pu :: ProgramUnit a
pu@(PUMain a
_ SrcSpan
_ Maybe String
mn [Block a]
_ Maybe [ProgramUnit a]
_) = (String, ProgramUnit a) -> Maybe (String, ProgramUnit a)
forall a. a -> Maybe a
Just (String -> Maybe String -> String
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]
_) = (String, ProgramUnit a) -> Maybe (String, 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]
_) = (String, ProgramUnit a) -> Maybe (String, ProgramUnit a)
forall a. a -> Maybe a
Just (String
n, ProgramUnit a
pu)
namedProgramUnit ProgramUnit a
_ = Maybe (String, 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 :: ProgramFile a -> [(String, ProgramUnit a)]
namedProgramUnits = (ProgramUnit a -> Maybe (String, ProgramUnit a))
-> [ProgramUnit a] -> [(String, ProgramUnit a)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ProgramUnit a -> Maybe (String, ProgramUnit a)
forall a. Data a => ProgramUnit a -> Maybe (String, ProgramUnit a)
namedProgramUnit ([ProgramUnit a] -> [(String, ProgramUnit a)])
-> (ProgramFile a -> [ProgramUnit a])
-> ProgramFile a
-> [(String, ProgramUnit a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ProgramFile a -> [ProgramUnit a]
forall a. Data a => ProgramFile a -> [ProgramUnit a]
allPU