{-# LANGUAGE CPP                   #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
------------------------------------------------------------
-- |
-- Module      :  FileEmbedLzma
-- Copyright   :  (c) 2015-2018 Futurice, 2018 Oleg Grenrus
-- License     :  BSD-3-Clause
-- Maintainer  :  Oleg Grenrus <oleg.grenrus@iki.fi>
----------------------------------------------------------------------------
module FileEmbedLzma (
    -- * Embed files
    embedByteString,
    embedLazyByteString,
    embedText,
    embedLazyText,
    -- * Embed directories
    embedDir,
    embedRecursiveDir,
    -- * Internal
    -- ** Directory listing
    listDirectoryFiles,
    listRecursiveDirectoryFiles,
    listDirectoryFilesF,
    -- ** Template Haskell
    lazyBytestringE,
    ) where

import Control.Arrow                    (first)
import Control.Monad                    (forM)
import Control.Monad.Trans.State.Strict (runState, state)
import Data.Foldable                    (for_)
import Data.Functor.Compose             (Compose (..))
import Data.Int                         (Int64)
import Data.List                        (sort)
import Language.Haskell.TH
import Language.Haskell.TH.Syntax       (qAddDependentFile)
import System.Directory
       (doesDirectoryExist, getDirectoryContents)
import System.FilePath                  (makeRelative, (</>))
import System.IO.Unsafe                 (unsafePerformIO)

import qualified Codec.Compression.Lzma  as LZMA
import qualified Data.ByteString         as BS
import qualified Data.ByteString.Lazy    as LBS
import qualified Data.ByteString.Unsafe  as BS.Unsafe
import qualified Data.Text               as T
import qualified Data.Text.Lazy          as LT
import qualified Data.Text.Lazy.Encoding as LTE

#if MIN_VERSION_template_haskell(2,16,0)
import qualified Data.ByteString.Internal as BS.Internal

import Language.Haskell.TH.Syntax (Bytes (..))
#endif

-- $setup
-- >>> :set -XTemplateHaskell -dppr-cols=9999
-- >>> import qualified Data.ByteString.Lazy as LBS
-- >>> import qualified Data.ByteString as BS
-- >>> import qualified Data.Text.Lazy as LT
-- >>> import qualified Data.Text as T

listRecursiveDirectoryFiles :: FilePath -> IO [(FilePath, LBS.ByteString)]
listRecursiveDirectoryFiles :: FilePath -> IO [(FilePath, ByteString)]
listRecursiveDirectoryFiles = (FilePath -> IO [(FilePath, ByteString)])
-> FilePath -> IO [(FilePath, ByteString)]
listDirectoryFilesF FilePath -> IO [(FilePath, ByteString)]
listRecursiveDirectoryFiles

listDirectoryFiles  :: FilePath -> IO [(FilePath, LBS.ByteString)]
listDirectoryFiles :: FilePath -> IO [(FilePath, ByteString)]
listDirectoryFiles = (FilePath -> IO [(FilePath, ByteString)])
-> FilePath -> IO [(FilePath, ByteString)]
listDirectoryFilesF (\FilePath
_ -> [(FilePath, ByteString)] -> IO [(FilePath, ByteString)]
forall (m :: * -> *) a. Monad m => a -> m a
return [])

listDirectoryFilesF
    :: (FilePath -> IO [(FilePath, LBS.ByteString)]) -- ^ what to do with a sub-directory
    -> FilePath -> IO [(FilePath, LBS.ByteString)]
listDirectoryFilesF :: (FilePath -> IO [(FilePath, ByteString)])
-> FilePath -> IO [(FilePath, ByteString)]
listDirectoryFilesF FilePath -> IO [(FilePath, ByteString)]
go FilePath
topdir = do
    [FilePath]
names <- FilePath -> IO [FilePath]
getDirectoryContents FilePath
topdir
    let properNames :: [FilePath]
properNames = (FilePath -> Bool) -> [FilePath] -> [FilePath]
forall a. (a -> Bool) -> [a] -> [a]
filter (FilePath -> [FilePath] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [FilePath
".", FilePath
".."]) [FilePath]
names
    [[(FilePath, ByteString)]]
paths <- [FilePath]
-> (FilePath -> IO [(FilePath, ByteString)])
-> IO [[(FilePath, ByteString)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [FilePath]
properNames ((FilePath -> IO [(FilePath, ByteString)])
 -> IO [[(FilePath, ByteString)]])
-> (FilePath -> IO [(FilePath, ByteString)])
-> IO [[(FilePath, ByteString)]]
forall a b. (a -> b) -> a -> b
$ \FilePath
name -> do
        let path :: FilePath
path = FilePath
topdir FilePath -> FilePath -> FilePath
</> FilePath
name
        Bool
isDirectory <- FilePath -> IO Bool
doesDirectoryExist FilePath
path
        if Bool
isDirectory
        then FilePath -> IO [(FilePath, ByteString)]
go FilePath
path
        else do
            ByteString
contents <- FilePath -> IO ByteString
LBS.readFile FilePath
path
            [(FilePath, ByteString)] -> IO [(FilePath, ByteString)]
forall (m :: * -> *) a. Monad m => a -> m a
return [(FilePath
path, ByteString
contents)]
    [(FilePath, ByteString)] -> IO [(FilePath, ByteString)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([[(FilePath, ByteString)]] -> [(FilePath, ByteString)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(FilePath, ByteString)]]
paths)

makeAllRelative :: FilePath -> [(FilePath, a)] -> [(FilePath, a)]
makeAllRelative :: FilePath -> [(FilePath, a)] -> [(FilePath, a)]
makeAllRelative FilePath
topdir = ((FilePath, a) -> (FilePath, a))
-> [(FilePath, a)] -> [(FilePath, a)]
forall a b. (a -> b) -> [a] -> [b]
map ((FilePath -> FilePath) -> (FilePath, a) -> (FilePath, a)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first ((FilePath
"/" FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++) (FilePath -> FilePath)
-> (FilePath -> FilePath) -> FilePath -> FilePath
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FilePath -> FilePath -> FilePath
makeRelative FilePath
topdir))

-- | Makes lazy 'LBS.ByteString' expression.
-- Embedded value is compressed with LZMA.
lazyBytestringE :: LBS.ByteString -> Q Exp
lazyBytestringE :: ByteString -> Q Exp
lazyBytestringE ByteString
lbs =
    [| LZMA.decompress . LBS.fromStrict . unsafePerformIO |] Q Exp -> Q Exp -> Q Exp
`appE`
    ([| BS.Unsafe.unsafePackAddressLen |] Q Exp -> Q Exp -> Q Exp
`appE` Q Exp
l Q Exp -> Q Exp -> Q Exp
`appE` Q Exp
s)
  where
    bs :: ByteString
bs = ByteString -> ByteString
LBS.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ CompressParams -> ByteString -> ByteString
LZMA.compressWith CompressParams
params ByteString
lbs
#if MIN_VERSION_template_haskell(2,16,0)
    s :: Q Exp
s = Lit -> Q Exp
litE (Lit -> Q Exp) -> Lit -> Q Exp
forall a b. (a -> b) -> a -> b
$ Bytes -> Lit
bytesPrimL (Bytes -> Lit) -> Bytes -> Lit
forall a b. (a -> b) -> a -> b
$ ByteString -> Bytes
bsToBytes ByteString
bs
#else
    s = litE $ stringPrimL $ BS.unpack bs
#endif
    l :: Q Exp
l = Lit -> Q Exp
litE (Lit -> Q Exp) -> Lit -> Q Exp
forall a b. (a -> b) -> a -> b
$ Integer -> Lit
integerL (Integer -> Lit) -> Integer -> Lit
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
bs

    params :: CompressParams
params = CompressParams
LZMA.defaultCompressParams
        {- doesn't seem to affect much
        { LZMA.compressLevel = LZMA.CompressionLevel9
        , LZMA.compressLevelExtreme = True
        }
        -}

#if MIN_VERSION_template_haskell(2,16,0)
bsToBytes :: BS.ByteString -> Bytes
bsToBytes :: ByteString -> Bytes
bsToBytes (BS.Internal.PS ForeignPtr Word8
fptr Int
off Int
len) = ForeignPtr Word8 -> Word -> Word -> Bytes
Bytes ForeignPtr Word8
fptr (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
off) (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
#endif

makeEmbeddedEntry :: Name -> (FilePath, (Int64, Int64)) -> Q Exp
makeEmbeddedEntry :: Name -> (FilePath, (Int64, Int64)) -> Q Exp
makeEmbeddedEntry Name
name (FilePath
path, (Int64
off, Int64
len)) = do
    let y :: Q Exp
y = [| LBS.toStrict . LBS.take len . LBS.drop off |] Q Exp -> Q Exp -> Q Exp
`appE` Name -> Q Exp
varE Name
name
    [| (,) path |] Q Exp -> Q Exp -> Q Exp
`appE` Q Exp
y

concatEntries :: Traversable t => t LBS.ByteString -> (LBS.ByteString, t (Int64, Int64))
concatEntries :: t ByteString -> (ByteString, t (Int64, Int64))
concatEntries t ByteString
xs = (ByteString -> ByteString
bslEndo ByteString
LBS.empty, t (Int64, Int64)
ys)
  where
    (t (Int64, Int64)
ys, (Int64
_, ByteString -> ByteString
bslEndo)) = State (Int64, ByteString -> ByteString) (t (Int64, Int64))
-> (Int64, ByteString -> ByteString)
-> (t (Int64, Int64), (Int64, ByteString -> ByteString))
forall s a. State s a -> s -> (a, s)
runState ((ByteString
 -> StateT
      (Int64, ByteString -> ByteString) Identity (Int64, Int64))
-> t ByteString
-> State (Int64, ByteString -> ByteString) (t (Int64, Int64))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (((Int64, ByteString -> ByteString)
 -> ((Int64, Int64), (Int64, ByteString -> ByteString)))
-> StateT (Int64, ByteString -> ByteString) Identity (Int64, Int64)
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state (((Int64, ByteString -> ByteString)
  -> ((Int64, Int64), (Int64, ByteString -> ByteString)))
 -> StateT
      (Int64, ByteString -> ByteString) Identity (Int64, Int64))
-> (ByteString
    -> (Int64, ByteString -> ByteString)
    -> ((Int64, Int64), (Int64, ByteString -> ByteString)))
-> ByteString
-> StateT (Int64, ByteString -> ByteString) Identity (Int64, Int64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString
-> (Int64, ByteString -> ByteString)
-> ((Int64, Int64), (Int64, ByteString -> ByteString))
single) t ByteString
xs) (Int64
0, ByteString -> ByteString
forall a. a -> a
id)

    single
        :: LBS.ByteString                             -- file bytestring
        -> (Int64, LBS.ByteString -> LBS.ByteString)  -- current offset, buffer so far
        -> ((Int64, Int64), (Int64, LBS.ByteString -> LBS.ByteString))
    single :: ByteString
-> (Int64, ByteString -> ByteString)
-> ((Int64, Int64), (Int64, ByteString -> ByteString))
single ByteString
bsl (Int64
off, ByteString -> ByteString
endo) = ((Int64
off, Int64
l), (Int64
off Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
l, ByteString -> ByteString
endo (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString -> ByteString
LBS.append ByteString
bsl))
      where
        l :: Int64
l = Int64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int64) -> Int64 -> Int64
forall a b. (a -> b) -> a -> b
$ ByteString -> Int64
LBS.length ByteString
bsl

-------------------------------------------------------------------------------
-- Directories
-------------------------------------------------------------------------------

-- | Embed a @[('FilePath', 'Data.ByteString.ByteString')]@ list, traversing given directory.
embedDir :: FilePath -> Q Exp
embedDir :: FilePath -> Q Exp
embedDir FilePath
topdir = do
    [(FilePath, ByteString)]
pairs' <- IO [(FilePath, ByteString)] -> Q [(FilePath, ByteString)]
forall a. IO a -> Q a
runIO (IO [(FilePath, ByteString)] -> Q [(FilePath, ByteString)])
-> IO [(FilePath, ByteString)] -> Q [(FilePath, ByteString)]
forall a b. (a -> b) -> a -> b
$ FilePath -> IO [(FilePath, ByteString)]
listDirectoryFiles FilePath
topdir
    [(FilePath, ByteString)]
-> ((FilePath, ByteString) -> Q ()) -> Q ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [(FilePath, ByteString)]
pairs' (((FilePath, ByteString) -> Q ()) -> Q ())
-> ((FilePath, ByteString) -> Q ()) -> Q ()
forall a b. (a -> b) -> a -> b
$ FilePath -> Q ()
forall (m :: * -> *). Quasi m => FilePath -> m ()
qAddDependentFile (FilePath -> Q ())
-> ((FilePath, ByteString) -> FilePath)
-> (FilePath, ByteString)
-> Q ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FilePath, ByteString) -> FilePath
forall a b. (a, b) -> a
fst
    let pairs :: [(FilePath, ByteString)]
pairs = FilePath -> [(FilePath, ByteString)] -> [(FilePath, ByteString)]
forall a. FilePath -> [(FilePath, a)] -> [(FilePath, a)]
makeAllRelative FilePath
topdir [(FilePath, ByteString)]
pairs'
    [(FilePath, ByteString)] -> Q Exp
embedPairs [(FilePath, ByteString)]
pairs

embedPairs :: [(FilePath, LBS.ByteString)] -> Q Exp
embedPairs :: [(FilePath, ByteString)] -> Q Exp
embedPairs [(FilePath, ByteString)]
pairs = do
    -- we do a hop to only embed single big bytestring.
    -- it's beneficial as lzma have more stuff to compress
    let (ByteString
bsl, Compose [(FilePath, (Int64, Int64))]
offsets) = Compose [] ((,) FilePath) ByteString
-> (ByteString, Compose [] ((,) FilePath) (Int64, Int64))
forall (t :: * -> *).
Traversable t =>
t ByteString -> (ByteString, t (Int64, Int64))
concatEntries ([(FilePath, ByteString)] -> Compose [] ((,) FilePath) ByteString
forall k k1 (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose [(FilePath, ByteString)]
pairs)
    Name
bslName <- FilePath -> Q Name
newName FilePath
"embedBsl"
    Exp
bslExpr <- ByteString -> Q Exp
lazyBytestringE ByteString
bsl
    let e :: Q Exp
e = [DecQ] -> Q Exp -> Q Exp
letE [ Dec -> DecQ
forall (m :: * -> *) a. Monad m => a -> m a
return (Dec -> DecQ) -> Dec -> DecQ
forall a b. (a -> b) -> a -> b
$ Pat -> Body -> [Dec] -> Dec
ValD (Name -> Pat
VarP Name
bslName) (Exp -> Body
NormalB Exp
bslExpr) [] ] (Q Exp -> Q Exp) -> Q Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$
                [Q Exp] -> Q Exp
listE ([Q Exp] -> Q Exp) -> [Q Exp] -> Q Exp
forall a b. (a -> b) -> a -> b
$ ((FilePath, (Int64, Int64)) -> Q Exp)
-> [(FilePath, (Int64, Int64))] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> (FilePath, (Int64, Int64)) -> Q Exp
makeEmbeddedEntry Name
bslName) [(FilePath, (Int64, Int64))]
offsets
    Q Exp -> TypeQ -> Q Exp
sigE Q Exp
e [t| [(FilePath, BS.ByteString)] |]

-- | Embed a @[('FilePath', 'Data.ByteString.ByteString')]@ list, recursively traversing given directory path.
--
-- For example, with @wai-static-app@ this can be used as:
--
-- @
-- staticApp $ embeddedSettings $('embedRecursiveDir' "static")
-- -- is an embedded (no data-files!) equivalent of
-- staticApp $ defaultFileServerSettings "static"
-- @
--
--
-- >>> $(embedRecursiveDir "example")
-- [("/Example.hs","..."),("/example.txt","Hello from the inside.\n")]
--
-- >>> :t $(embedRecursiveDir "example")
-- $(embedRecursiveDir "example") :: [(FilePath, BS.ByteString)]
--
embedRecursiveDir :: FilePath -> Q Exp
embedRecursiveDir :: FilePath -> Q Exp
embedRecursiveDir FilePath
topdir = do
    [(FilePath, ByteString)]
pairs' <- IO [(FilePath, ByteString)] -> Q [(FilePath, ByteString)]
forall a. IO a -> Q a
runIO (IO [(FilePath, ByteString)] -> Q [(FilePath, ByteString)])
-> IO [(FilePath, ByteString)] -> Q [(FilePath, ByteString)]
forall a b. (a -> b) -> a -> b
$ FilePath -> IO [(FilePath, ByteString)]
listRecursiveDirectoryFiles FilePath
topdir
    [(FilePath, ByteString)]
-> ((FilePath, ByteString) -> Q ()) -> Q ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [(FilePath, ByteString)]
pairs' (((FilePath, ByteString) -> Q ()) -> Q ())
-> ((FilePath, ByteString) -> Q ()) -> Q ()
forall a b. (a -> b) -> a -> b
$ FilePath -> Q ()
forall (m :: * -> *). Quasi m => FilePath -> m ()
qAddDependentFile (FilePath -> Q ())
-> ((FilePath, ByteString) -> FilePath)
-> (FilePath, ByteString)
-> Q ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FilePath, ByteString) -> FilePath
forall a b. (a, b) -> a
fst
    let pairs :: [(FilePath, ByteString)]
pairs = [(FilePath, ByteString)] -> [(FilePath, ByteString)]
forall a. Ord a => [a] -> [a]
sort (FilePath -> [(FilePath, ByteString)] -> [(FilePath, ByteString)]
forall a. FilePath -> [(FilePath, a)] -> [(FilePath, a)]
makeAllRelative FilePath
topdir [(FilePath, ByteString)]
pairs')
    [(FilePath, ByteString)] -> Q Exp
embedPairs [(FilePath, ByteString)]
pairs

-------------------------------------------------------------------------------
-- Strings
-------------------------------------------------------------------------------

-- | Embed a lazy 'Data.ByteString.Lazy.ByteString' from a file.
--
-- >>> :t $(embedLazyByteString "file-embed-lzma.cabal")
-- $(embedLazyByteString "file-embed-lzma.cabal") :: LBS.ByteString
--
embedLazyByteString :: FilePath -> Q Exp
embedLazyByteString :: FilePath -> Q Exp
embedLazyByteString FilePath
fp = do
    FilePath -> Q ()
forall (m :: * -> *). Quasi m => FilePath -> m ()
qAddDependentFile FilePath
fp
    ByteString
bsl <- IO ByteString -> Q ByteString
forall a. IO a -> Q a
runIO (IO ByteString -> Q ByteString) -> IO ByteString -> Q ByteString
forall a b. (a -> b) -> a -> b
$ FilePath -> IO ByteString
LBS.readFile FilePath
fp
    ByteString -> Q Exp
lazyBytestringE ByteString
bsl

-- | Embed a strict 'Data.ByteString.ByteString' from a file.
--
-- >>> :t $(embedByteString "file-embed-lzma.cabal")
-- $(embedByteString "file-embed-lzma.cabal") :: BS.ByteString
--
embedByteString :: FilePath -> Q Exp
embedByteString :: FilePath -> Q Exp
embedByteString FilePath
fp = [| LBS.toStrict |] Q Exp -> Q Exp -> Q Exp
`appE` FilePath -> Q Exp
embedLazyByteString FilePath
fp

-- | Embed a lazy 'Data.Text.Lazy.Text' from a UTF8-encoded file.
--
-- >>> :t $(embedLazyText "file-embed-lzma.cabal")
-- $(embedLazyText "file-embed-lzma.cabal") :: LT.Text
--
embedLazyText :: FilePath -> Q Exp
embedLazyText :: FilePath -> Q Exp
embedLazyText FilePath
fp = do
    FilePath -> Q ()
forall (m :: * -> *). Quasi m => FilePath -> m ()
qAddDependentFile FilePath
fp
    ByteString
bsl <- IO ByteString -> Q ByteString
forall a. IO a -> Q a
runIO (IO ByteString -> Q ByteString) -> IO ByteString -> Q ByteString
forall a b. (a -> b) -> a -> b
$ FilePath -> IO ByteString
LBS.readFile FilePath
fp
    case ByteString -> Either UnicodeException Text
LTE.decodeUtf8' ByteString
bsl of
        Left UnicodeException
e  -> FilePath -> Q ()
reportError (UnicodeException -> FilePath
forall a. Show a => a -> FilePath
show UnicodeException
e)
        Right Text
_ -> () -> Q ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    [| LTE.decodeUtf8 |] Q Exp -> Q Exp -> Q Exp
`appE` ByteString -> Q Exp
lazyBytestringE ByteString
bsl

-- | Embed a strict 'Data.Text.Text' from a UTF8-encoded file.
--
-- >>> :t $(embedText "file-embed-lzma.cabal")
-- $(embedText "file-embed-lzma.cabal") :: T.Text
--
embedText :: FilePath -> Q Exp
embedText :: FilePath -> Q Exp
embedText FilePath
fp = [| LT.toStrict |] Q Exp -> Q Exp -> Q Exp
`appE` FilePath -> Q Exp
embedLazyText FilePath
fp