{- HLINT ignore "Redundant if" -}

{- |
Copyright: (c) 2020 Kowainik
SPDX-License-Identifier: MPL-2.0
Maintainer: Kowainik <xrom.xkov@gmail.com>

Functions to work with @hie@ specific parts.
-}

module Stan.Hie
    ( readHieFiles
    , countLinesOfCode
    , eqAst
    , slice
    ) where

import Colourista (errorMessage, infoMessage, warningMessage)
import System.Directory (doesDirectoryExist, doesFileExist)
import System.Directory.Recursive (getDirRecursive)
import System.FilePath (takeExtension)

import Stan.Core.List (checkWith)
import Stan.Ghc.Compat (RealSrcSpan, srcSpanEndCol, srcSpanStartCol, srcSpanStartLine)
import Stan.Hie.Compat (HieAST (..), HieFile (..), HieFileResult (hie_file_result), NameCache,
                        NodeInfo (..), initNameCache, mkSplitUniqSupply, readHieFile)
import Stan.Hie.Debug ()
import Stan.Pattern.Ast (literalAnns)

import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BS8
import qualified Data.Set as Set


{- | Returns contents of all @.hie@ files recursively in the given
@hie@ directory.
-}
readHieFiles :: FilePath -> IO [HieFile]
readHieFiles :: FilePath -> IO [HieFile]
readHieFiles hieDir :: FilePath
hieDir = do
    IO Bool -> IO () -> IO ()
forall (m :: * -> *). Monad m => m Bool -> m () -> m ()
unlessM (FilePath -> IO Bool
doesDirectoryExist FilePath
hieDir) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        Text -> IO ()
errorMessage (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ "Directory with HIE files doesn't exist: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> FilePath -> Text
forall a. ToText a => a -> Text
toText FilePath
hieDir
        Text -> IO ()
infoMessage "Use the '--hiedir' CLI option to specify path to the directory with HIE files"
        IO ()
forall (m :: * -> *) a. MonadIO m => m a
exitFailure

    NameCache
nameCache <- IO NameCache
createNameCache
    [FilePath]
hieContent <- FilePath -> IO [FilePath]
getDirRecursive FilePath
hieDir
    let isHieFile :: FilePath -> IO Bool
isHieFile f :: FilePath
f = Bool -> Bool -> Bool
(&&) (FilePath -> FilePath
takeExtension FilePath
f FilePath -> FilePath -> Bool
forall a. Eq a => a -> a -> Bool
== ".hie") (Bool -> Bool) -> IO Bool -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> IO Bool
doesFileExist FilePath
f
    [FilePath]
hiePaths <- (FilePath -> IO Bool) -> [FilePath] -> IO [FilePath]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM FilePath -> IO Bool
isHieFile [FilePath]
hieContent

    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([FilePath] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [FilePath]
hiePaths) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> IO ()
warningMessage (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$
        "The directory with HIE files doesn't contain any HIE files: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> FilePath -> Text
forall a. ToText a => a -> Text
toText FilePath
hieDir

    [FilePath] -> (FilePath -> IO HieFile) -> IO [HieFile]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [FilePath]
hiePaths ((FilePath -> IO HieFile) -> IO [HieFile])
-> (FilePath -> IO HieFile) -> IO [HieFile]
forall a b. (a -> b) -> a -> b
$ \hiePath :: FilePath
hiePath -> do
        (hieFileResult :: HieFileResult
hieFileResult, _newCache :: NameCache
_newCache) <- NameCache -> FilePath -> IO (HieFileResult, NameCache)
readHieFile NameCache
nameCache FilePath
hiePath
        HieFile -> IO HieFile
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HieFile -> IO HieFile) -> HieFile -> IO HieFile
forall a b. (a -> b) -> a -> b
$ HieFileResult -> HieFile
hie_file_result HieFileResult
hieFileResult

createNameCache :: IO NameCache
createNameCache :: IO NameCache
createNameCache = do
    UniqSupply
uniqSupply <- Char -> IO UniqSupply
mkSplitUniqSupply 'z'
    NameCache -> IO NameCache
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NameCache -> IO NameCache) -> NameCache -> IO NameCache
forall a b. (a -> b) -> a -> b
$ UniqSupply -> [Name] -> NameCache
initNameCache UniqSupply
uniqSupply []

-- | Get the number of lines of code in the file by analising 'HieFile'.
countLinesOfCode :: HieFile -> Int
countLinesOfCode :: HieFile -> Int
countLinesOfCode HieFile{..} = [ByteString] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([ByteString] -> Int) -> [ByteString] -> Int
forall a b. (a -> b) -> a -> b
$ ByteString -> [ByteString]
BS8.lines ByteString
hie_hs_src

{- | Take sub-bytestring according to a given span.

When the given source is empty returns 'Nothing'.

TODO: currently works only with single-line spans
-}
slice :: RealSrcSpan -> ByteString -> Maybe ByteString
slice :: RealSrcSpan -> ByteString -> Maybe ByteString
slice span :: RealSrcSpan
span =
    (ByteString -> ByteString) -> Maybe ByteString -> Maybe ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
        ( Int -> ByteString -> ByteString
BS.take (RealSrcSpan -> Int
srcSpanEndCol RealSrcSpan
span Int -> Int -> Int
forall a. Num a => a -> a -> a
- RealSrcSpan -> Int
srcSpanStartCol RealSrcSpan
span)
        (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
BS.drop (RealSrcSpan -> Int
srcSpanStartCol RealSrcSpan
span Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
        )
    (Maybe ByteString -> Maybe ByteString)
-> (ByteString -> Maybe ByteString)
-> ByteString
-> Maybe ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([ByteString] -> Int -> Maybe ByteString)
-> Int -> [ByteString] -> Maybe ByteString
forall a b c. (a -> b -> c) -> b -> a -> c
flip [ByteString] -> Int -> Maybe ByteString
forall a. [a] -> Int -> Maybe a
(!!?) (RealSrcSpan -> Int
srcSpanStartLine RealSrcSpan
span Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
    ([ByteString] -> Maybe ByteString)
-> (ByteString -> [ByteString]) -> ByteString -> Maybe ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [ByteString]
BS8.lines

{- | Compare two AST nodes on equality. This is a more relaxed version
of the 'Eq' instance for 'HieAST' because it doesn't compare source
locations. This function is useful if you want to check whether two
AST nodes represent the same AST.

This function needs to take the original 'HieFile' because constants
are not stored in 'HieAST' and to compare constants we need to compare
parts of source code.
-}
eqAst :: forall a . Eq a => HieFile -> HieAST a -> HieAST a -> Bool
eqAst :: HieFile -> HieAST a -> HieAST a -> Bool
eqAst HieFile{..} = HieAST a -> HieAST a -> Bool
eqNodes
  where
    eqNodes :: HieAST a -> HieAST a -> Bool
    eqNodes :: HieAST a -> HieAST a -> Bool
eqNodes (Node info1 :: NodeInfo a
info1 span1 :: RealSrcSpan
span1 children1 :: [HieAST a]
children1) (Node info2 :: NodeInfo a
info2 span2 :: RealSrcSpan
span2 children2 :: [HieAST a]
children2) =
        NodeInfo a -> NodeInfo a -> Bool
eqInfo NodeInfo a
info1 NodeInfo a
info2 Bool -> Bool -> Bool
&& (HieAST a -> HieAST a -> Bool) -> [HieAST a] -> [HieAST a] -> Bool
forall a b. (a -> b -> Bool) -> [a] -> [b] -> Bool
checkWith HieAST a -> HieAST a -> Bool
eqNodes [HieAST a]
children1 [HieAST a]
children2
      where
        eqInfo :: NodeInfo a -> NodeInfo a -> Bool
        eqInfo :: NodeInfo a -> NodeInfo a -> Bool
eqInfo (NodeInfo anns1 :: Set (FastString, FastString)
anns1 types1 :: [a]
types1 ids1 :: NodeIdentifiers a
ids1) (NodeInfo anns2 :: Set (FastString, FastString)
anns2 types2 :: [a]
types2 ids2 :: NodeIdentifiers a
ids2) =
            Set (FastString, FastString)
anns1 Set (FastString, FastString)
-> Set (FastString, FastString) -> Bool
forall a. Eq a => a -> a -> Bool
== Set (FastString, FastString)
anns2 Bool -> Bool -> Bool
&& [a]
types1 [a] -> [a] -> Bool
forall a. Eq a => a -> a -> Bool
== [a]
types2 Bool -> Bool -> Bool
&& NodeIdentifiers a
ids1 NodeIdentifiers a -> NodeIdentifiers a -> Bool
forall a. Eq a => a -> a -> Bool
== NodeIdentifiers a
ids2 Bool -> Bool -> Bool
&&
            if (FastString, FastString) -> Set (FastString, FastString) -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (FastString, FastString)
literalAnns Set (FastString, FastString)
anns1
            then RealSrcSpan -> ByteString -> Maybe ByteString
slice RealSrcSpan
span1 ByteString
hie_hs_src Maybe ByteString -> Maybe ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== RealSrcSpan -> ByteString -> Maybe ByteString
slice RealSrcSpan
span2 ByteString
hie_hs_src
            else Bool
True