module Store.AV
( AVError(..)
, avErrorString
, AV
, initAV
, AVMediaType(..)
, AVProbe(..)
, avProbe
, avProbeLength
, avProbeHas
, avFrame
) where
import Control.Applicative ((<|>))
import Control.Arrow (first)
import Control.Concurrent.MVar (MVar, newMVar, takeMVar, putMVar)
import Control.Exception (Exception, throwIO, bracket, bracket_, finally, onException)
import Control.Monad ((<=<), void, when, forM, forM_)
import qualified Data.ByteString as BS
import Data.Int (Int32, Int64)
import Data.Maybe (isNothing)
import Data.Ratio (Ratio, (%), numerator, denominator)
import Data.Time.Clock (DiffTime)
import Data.Time.Format (parseTimeM, defaultTimeLocale)
import Data.Time.LocalTime (ZonedTime)
import Data.Typeable (Typeable)
import Data.Word (Word16, Word32)
import Foreign.C.Error (throwErrnoIfNull)
import Foreign.C.String (CString, CStringLen, peekCAString, withCAString)
import Foreign.C.Types (CInt(..), CUInt(..), CSize(..))
import Foreign.Marshal.Alloc (allocaBytes)
import Foreign.Marshal.Utils (with, maybePeek)
import Foreign.Ptr (Ptr, FunPtr, nullPtr, plusPtr, castPtr)
import Foreign.StablePtr
import Foreign.Storable
import System.IO.Unsafe (unsafeDupablePerformIO)
import Ops
import Files
import Model.Offset
type AVLockOp = Word32
type AVLockmgr a = Ptr (Ptr a) -> AVLockOp -> IO CInt
type AVPixelFormat = Int32
type AVCodecID = Word32
data AVDictionary
data AVDictionaryEntry
data AVInputFormat
data AVFrame
data AVPacket
data AVCodec
data AVCodecContext
data AVStream
data AVFormatContext
data AVIOContext
data AVOutputFormat
foreign import ccall unsafe "libavutil/error.h av_strerror"
avStrerror :: CInt -> CString -> CSize -> IO CInt
foreign import ccall unsafe "libavutil/mem.h av_free"
avFree :: Ptr a -> IO ()
foreign import ccall unsafe "libavutil/dict.h av_dict_get"
avDictGet :: Ptr AVDictionary -> CString -> Ptr AVDictionaryEntry -> CInt -> IO (Ptr AVDictionaryEntry)
foreign import ccall unsafe "libavutil/dict.h av_dict_set"
avDictSet :: Ptr (Ptr AVDictionary) -> CString -> CString -> CInt -> IO CInt
foreign import ccall unsafe "libavutil/dict.h av_dict_free"
avDictFree :: Ptr (Ptr AVDictionary) -> IO ()
foreign import ccall unsafe "libavutil/frame.h av_frame_alloc"
avFrameAlloc :: IO (Ptr AVFrame)
foreign import ccall unsafe "libavutil/frame.h av_frame_free"
avFrameFree :: Ptr (Ptr AVFrame) -> IO ()
foreign import ccall unsafe "libavutil/frame.h av_frame_unref"
avFrameUnref :: Ptr AVFrame -> IO ()
foreign import ccall unsafe "libavutil/frame.h av_frame_get_best_effort_timestamp"
avFrameGetBestEffortTimestamp :: Ptr AVFrame -> IO Int64
foreign import ccall "libavcodec/avcodec.h av_lockmgr_register"
avLockmgrRegister :: FunPtr (AVLockmgr a) -> IO CInt
foreign import ccall "wrapper"
mkAVLockmgr :: AVLockmgr a -> IO (FunPtr (AVLockmgr a))
foreign import ccall unsafe "libavcodec/avcodec.h avcodec_get_name"
avcodecGetName :: AVCodecID -> IO CString
foreign import ccall unsafe "libavcodec/avcodec.h av_init_packet"
avInitPacket :: Ptr AVPacket -> IO ()
foreign import ccall unsafe "libavcodec/avcodec.h av_free_packet"
avFreePacket :: Ptr AVPacket -> IO ()
foreign import ccall "libavcodec/avcodec.h avcodec_open2"
avcodecOpen2 :: Ptr AVCodecContext -> Ptr AVCodec -> Ptr (Ptr AVDictionary) -> IO CInt
foreign import ccall "libavcodec/avcodec.h avcodec_close"
avcodecClose :: Ptr AVCodecContext -> IO CInt
foreign import ccall "libavcodec/avcodec.h avcodec_decode_video2"
avcodecDecodeVideo2 :: Ptr AVCodecContext -> Ptr AVFrame -> Ptr CInt -> Ptr AVPacket -> IO CInt
foreign import ccall "libavcodec/avcodec.h avcodec_encode_video2"
avcodecEncodeVideo2 :: Ptr AVCodecContext -> Ptr AVPacket -> Ptr AVFrame -> Ptr CInt -> IO CInt
foreign import ccall "libavcodec/avcodec.h avcodec_find_encoder"
avcodecFindEncoder :: AVCodecID -> IO (Ptr AVCodec)
foreign import ccall "libavcodec/avcodec.h avcodec_find_best_pix_fmt_of_list"
avcodecFindBestPixFmtOfList :: Ptr AVPixelFormat -> AVPixelFormat -> CInt -> Ptr CInt -> IO AVPixelFormat
foreign import ccall unsafe "libavformat/avio.h avio_open_dyn_buf"
avioOpenDynBuf :: Ptr (Ptr AVIOContext) -> IO CInt
foreign import ccall unsafe "libavformat/avio.h avio_close_dyn_buf"
avioCloseDynBuf :: Ptr AVIOContext -> Ptr CString -> IO CInt
foreign import ccall "libavformat/avformat.h av_register_all"
avRegisterAll :: IO ()
foreign import ccall "libavformat/avformat.h avformat_open_input"
avformatOpenInput :: Ptr (Ptr AVFormatContext) -> CString -> Ptr AVInputFormat -> Ptr (Ptr AVDictionary) -> IO CInt
foreign import ccall "libavformat/avformat.h avformat_close_input"
avformatCloseInput :: Ptr (Ptr AVFormatContext) -> IO ()
foreign import ccall "libavformat/avformat.h avformat_find_stream_info"
avformatFindStreamInfo :: Ptr AVFormatContext -> Ptr (Ptr AVDictionary) -> IO CInt
foreign import ccall unsafe "libavformat/avformat.h av_find_input_format"
avFindInputFormat :: CString -> IO (Ptr AVInputFormat)
foreign import ccall "libavformat/avformat.h av_find_best_stream"
avFindBestStream :: Ptr AVFormatContext -> Int32 -> CInt -> CInt -> Ptr (Ptr AVCodec) -> CInt -> IO CInt
foreign import ccall "libavformat/avformat.h avformat_seek_file"
avformatSeekFile :: Ptr AVFormatContext -> CInt -> Int64 -> Int64 -> Int64 -> CInt -> IO CInt
foreign import ccall "libavformat/avformat.h av_read_frame"
avReadFrame :: Ptr AVFormatContext -> Ptr AVPacket -> IO CInt
foreign import ccall "libavformat/avformat.h avformat_alloc_output_context2"
avformatAllocOutputContext2 :: Ptr (Ptr AVFormatContext) -> Ptr AVOutputFormat -> CString -> CString -> IO CInt
foreign import ccall "libavformat/avformat.h avformat_free_context"
avformatFreeContext :: Ptr AVFormatContext -> IO ()
foreign import ccall "libavformat/avformat.h avformat_new_stream"
avformatNewStream :: Ptr AVFormatContext -> Ptr AVCodec -> IO (Ptr AVStream)
foreign import ccall "libavformat/avformat.h avformat_write_header"
avformatWriteHeader :: Ptr AVFormatContext -> Ptr (Ptr AVDictionary) -> IO CInt
foreign import ccall "libavformat/avformat.h av_write_frame"
avWriteFrame :: Ptr AVFormatContext -> Ptr AVPacket -> IO CInt
foreign import ccall "libavformat/avformat.h av_write_trailer"
avWriteTrailer :: Ptr AVFormatContext -> IO CInt
newtype AVRational = AVRational (Ratio CInt) deriving (Eq, Ord, Num, Fractional, Real, RealFrac)
instance Storable AVRational where
sizeOf _ = (8)
alignment (AVRational x) = alignment $ numerator x
peek p = do
num <- (\hsc_ptr -> peekByteOff hsc_ptr 0) p
den <- (\hsc_ptr -> peekByteOff hsc_ptr 4) p
return $ AVRational $ num % den
poke p (AVRational x) = do
(\hsc_ptr -> pokeByteOff hsc_ptr 0) p (numerator x)
(\hsc_ptr -> pokeByteOff hsc_ptr 4) p (denominator x)
avShowError :: CInt -> String
avShowError e = unsafeDupablePerformIO $
allocaBytes len $ \p -> do
_ <- avStrerror e p (fromIntegral len)
peekCAString p
where len = 256
data AVError = AVError
{ avErrorCode :: CInt
, avErrorFunction :: String
, avErrorFile :: Maybe RawFilePath
} deriving (Typeable)
instance Exception AVError
avErrorString :: AVError -> String
avErrorString = avShowError . avErrorCode
instance Show AVError where
showsPrec p (AVError e c f) = showParen (p > 10) $
showString "AVError "
. showString c
. maybe id (((' ' :) .) . shows) f
. if e == 0 then id else
showString ": " . showString (avShowError e)
data ErrorFile
= NoFile
| FileName !RawFilePath
| FileContext !(Ptr AVFormatContext)
errorFile :: ErrorFile -> IO (Maybe RawFilePath)
errorFile NoFile = return Nothing
errorFile (FileName f) = return $ Just f
errorFile (FileContext a) =
(n /= nullPtr) `thenReturn` (BS.packCString n)
where n = (\hsc_ptr -> hsc_ptr `plusPtr` 56) a
throwAVError :: CInt -> String -> ErrorFile -> IO a
throwAVError e c f = throwIO . AVError e c =<< errorFile f
throwAVErrorIf :: Integral a => String -> ErrorFile -> IO a -> IO a
throwAVErrorIf c f g = do
r <- g
when (r < 0) $
throwAVError (fromIntegral r) c f
return r
throwAVErrorIf_ :: Integral a => String -> ErrorFile -> IO a -> IO ()
throwAVErrorIf_ c f = void . throwAVErrorIf c f
throwAVErrorIfNull :: String -> ErrorFile -> IO (Ptr a) -> IO (Ptr a)
throwAVErrorIfNull c f g = do
r <- g
when (r == nullPtr) $
throwAVError 0 c f
return r
withAVDictionary :: (Ptr (Ptr AVDictionary) -> IO a) -> IO a
withAVDictionary f = with nullPtr $ \d ->
f d `finally` avDictFree d
getAVDictionary :: Ptr AVDictionary -> String -> IO (Maybe String)
getAVDictionary dict key =
withCAString key $ \ckey ->
maybePeek (peekCAString <=< (\hsc_ptr -> peekByteOff hsc_ptr 8)) =<< avDictGet dict ckey nullPtr 0
setAVDictionary :: Ptr (Ptr AVDictionary) -> String -> String -> IO ()
setAVDictionary dict key val =
withCAString key $ \ckey ->
withCAString val $ \cval ->
throwAVErrorIf_ "av_dict_set" NoFile $ avDictSet dict ckey cval 0
closeAVIODynBuf :: Ptr AVIOContext -> (CStringLen -> IO a) -> IO a
closeAVIODynBuf c g =
with nullPtr $ \p ->
bracket
(do
l <- avioCloseDynBuf c p
b <- peek p
return (b, fromIntegral l))
(avFree . fst)
g
withAVIODynBuf :: ErrorFile -> Ptr (Ptr AVIOContext) -> IO a -> IO (BS.ByteString, a)
withAVIODynBuf ec pb g = do
throwAVErrorIf_ "avio_open_dyn_buf" ec $
avioOpenDynBuf pb
buf <- peek pb
r <- g `onException`
closeAVIODynBuf buf return
b <- closeAVIODynBuf buf BS.packCStringLen
return (b, r)
withAVFrame :: (Ptr AVFrame -> IO a) -> IO a
withAVFrame =
bracket
(throwErrnoIfNull "av_frame_alloc" avFrameAlloc)
(`with` avFrameFree)
withAVInput :: RawFilePath -> Maybe String -> (Ptr AVFormatContext -> IO a) -> IO a
withAVInput f fmt g =
BS.useAsCString f $ \cf ->
with nullPtr $ \a ->
bracket_
(do
avfmt <- maybe (return nullPtr) (`withCAString` avFindInputFormat) fmt
throwAVErrorIf "avformat_open_input" (FileName f) $
avformatOpenInput a cf avfmt nullPtr)
(avformatCloseInput a)
(g =<< peek a)
findAVStreamInfo :: Ptr AVFormatContext -> IO ()
findAVStreamInfo a = throwAVErrorIf_ "avformat_find_stream_info" (FileContext a) $
avformatFindStreamInfo a nullPtr
withAVOutput :: Maybe RawFilePath -> String -> (Ptr AVFormatContext -> IO a) -> IO (Maybe BS.ByteString, a)
withAVOutput f fmt g =
maybe ($ nullPtr) BS.useAsCString f $ \cf ->
withCAString fmt $ \cfmt ->
with nullPtr $ \a ->
bracket_
(throwAVErrorIf_ "avformat_alloc_output_context2" (maybe NoFile FileName f) $
avformatAllocOutputContext2 a nullPtr cfmt cf)
(avformatFreeContext =<< peek a)
(do
c <- peek a
case f of
Just _ -> (,) Nothing <$> g c
Nothing -> first Just <$>
withAVIODynBuf (FileContext c) ((\hsc_ptr -> hsc_ptr `plusPtr` 32) c)
(g c))
withAVCodec :: Ptr AVFormatContext -> Ptr AVStream -> Ptr AVCodec -> Ptr (Ptr AVDictionary) -> IO a -> IO a
withAVCodec a s c o =
bracket_
(do
sc <- (\hsc_ptr -> peekByteOff hsc_ptr 8) s
throwAVErrorIf_ "avcodec_open2" (FileContext a) $
avcodecOpen2 sc c o)
(do
sc <- (\hsc_ptr -> peekByteOff hsc_ptr 8) s
avcodecClose sc)
withAVPacket :: (Ptr AVPacket -> IO a) -> IO a
withAVPacket f = allocaBytes (88) $ \p -> do
bracket_
(avInitPacket p)
(avFreePacket p)
(do
(\hsc_ptr -> pokeByteOff hsc_ptr 24) p nullPtr
(\hsc_ptr -> pokeByteOff hsc_ptr 32) p (0 :: CInt)
f p)
avLockmgr :: AVLockmgr ()
avLockmgr p o
| o == 0 = (<$) 0 $
poke p . castStablePtrToPtr =<< newStablePtr =<< newMVar ()
| o == 1 = (<$) 0 $
takeMVar =<< deRefStablePtr =<< s
| o == 2 = (<$) 0 $
(`putMVar` ()) =<< deRefStablePtr =<< s
| o == 3 = (<$) 0 $
freeStablePtr =<< s
| otherwise = return (1)
where
s = castPtrToStablePtr <$> peek p
s :: IO (StablePtr (MVar ()))
data AV = AV
initAV :: IO AV
initAV = do
avRegisterAll
mgr <- mkAVLockmgr avLockmgr
throwAVErrorIf_ "av_lockmgr_register" NoFile $ avLockmgrRegister mgr
return AV
data AVMediaType
= AVMediaTypeUnknown
| AVMediaTypeVideo
| AVMediaTypeAudio
| AVMediaTypeData
| AVMediaTypeSubtitle
| AVMediaTypeAttachment
deriving (Eq, Show, Enum, Bounded)
instance Storable AVMediaType where
sizeOf _ = sizeOf (undefined :: Int32)
alignment _ = alignment (undefined :: Int32)
peek p = do
v <- peek (castPtr p :: Ptr Int32)
return $ case v of
0 -> AVMediaTypeVideo
1 -> AVMediaTypeAudio
2 -> AVMediaTypeData
3 -> AVMediaTypeSubtitle
4 -> AVMediaTypeAttachment
_ -> AVMediaTypeUnknown
poke p v = poke (castPtr p :: Ptr Int32) $ case v of
AVMediaTypeUnknown -> 1
AVMediaTypeVideo -> 0
AVMediaTypeAudio -> 1
AVMediaTypeData -> 2
AVMediaTypeSubtitle -> 3
AVMediaTypeAttachment -> 4
data AVProbe = AVProbe
{ avProbeFormat :: BS.ByteString
, avProbeDuration :: DiffTime
, avProbeStreams :: [(AVMediaType, BS.ByteString)]
, avProbeDate :: Maybe ZonedTime
}
avProbeHas :: AVMediaType -> AVProbe -> Bool
avProbeHas t = any ((t ==) . fst) . avProbeStreams
avProbeLength :: AVProbe -> Maybe Offset
avProbeLength AVProbe{ avProbeDuration = o } = (o > 0) `thenUse` (diffTimeOffset o)
avTime :: Int64 -> DiffTime
avTime t = realToFrac $ t % 1000000
avProbe :: RawFilePath -> AV -> IO AVProbe
avProbe f AV = withAVInput f Nothing $ \ic -> do
findAVStreamInfo ic
meta <- (\hsc_ptr -> peekByteOff hsc_ptr 1192) ic
AVProbe
<$> (BS.packCString =<< (\hsc_ptr -> peekByteOff hsc_ptr 0) =<< (\hsc_ptr -> peekByteOff hsc_ptr 8) ic)
<*> (avTime <$> (\hsc_ptr -> peekByteOff hsc_ptr 1088) ic)
<*> do
nb :: CUInt <- (\hsc_ptr -> peekByteOff hsc_ptr 44) ic
ss <- (\hsc_ptr -> peekByteOff hsc_ptr 48) ic
forM [0..pred (fromIntegral nb)] $ \i -> do
c <- (\hsc_ptr -> peekByteOff hsc_ptr 8) =<< peekElemOff ss i
t <- (\hsc_ptr -> peekByteOff hsc_ptr 12) c
n <- BS.packCString =<< avcodecGetName =<< (\hsc_ptr -> peekByteOff hsc_ptr 56) c
return (t, n)
<*> ((=<<) (\t -> parseTimeM True defaultTimeLocale "%F %T%Z" t <|> parseTimeM True defaultTimeLocale "%F %T %Z" t) <$>
((`orElseM` getAVDictionary meta "creation_time") =<< getAVDictionary meta "date"))
avSeekStream :: Ptr AVFormatContext -> Ptr AVStream -> Ptr AVFrame -> Maybe DiffTime -> IO ()
avSeekStream ctx s frame offset = do
off <- forM offset $ \o -> do
den :: CInt <- (\hsc_ptr -> peekByteOff hsc_ptr 52) s
num :: CInt <- (\hsc_ptr -> peekByteOff hsc_ptr 48) s
let off = floor $ o * (fromIntegral den) / (fromIntegral num)
throwAVErrorIf_ "avformat_seek_file" (FileContext ctx) $
avformatSeekFile ctx 0 (9223372036854775808) off off 0
return off
si :: CInt <- (\hsc_ptr -> peekByteOff hsc_ptr 0) s
codec <- (\hsc_ptr -> peekByteOff hsc_ptr 8) s
let
seek = withAVPacket $ \pkt -> do
avFrameUnref frame
throwAVErrorIf_ "av_read_frame" (FileContext ctx) $
avReadFrame ctx pkt
psi <- (\hsc_ptr -> peekByteOff hsc_ptr 36) pkt
if psi /= si then seek else with 0 $ \gpp -> do
throwAVErrorIf_ "avcodec_decode_video2" (FileContext ctx) $
avcodecDecodeVideo2 codec frame gpp pkt
gp <- peek gpp
if gp == 0 then seek else do
pts <- avFrameGetBestEffortTimestamp frame
if any (pts <) off then seek else
(\hsc_ptr -> pokeByteOff hsc_ptr 136) frame pts
seek
foreign import ccall unsafe "av.h avFrame_initialize_stream"
avFrameInitializeStream :: Ptr AVStream -> Ptr AVFormatContext -> Ptr AVStream -> Ptr AVFrame -> CInt -> CInt -> IO ()
foreign import ccall "av.h avFrame_rescale"
avFrameRescale :: Ptr AVCodecContext -> Ptr AVFrame -> IO CInt
avFrame :: RawFilePath -> Maybe DiffTime -> Maybe Word16 -> Maybe Word16 -> Maybe RawFilePath -> AV -> IO (Maybe BS.ByteString)
avFrame infile offset width height outfile AV =
withAVInput infile (isimg `thenUse` "image2") $ \inctx ->
with nullPtr $ \icodecp ->
withAVDictionary $ \opts -> do
when isimg $
(\hsc_ptr -> pokeByteOff hsc_ptr 1160) inctx (8 :: AVCodecID)
findAVStreamInfo inctx
si <- throwAVErrorIf "av_find_best_stream" (FileContext inctx) $
avFindBestStream inctx 0 (1) (1) icodecp 0
nb :: CUInt <- (\hsc_ptr -> peekByteOff hsc_ptr 44) inctx
isl <- (\hsc_ptr -> peekByteOff hsc_ptr 48) inctx
forM_ [0..pred (fromIntegral nb)] $ \i ->
when (i /= si) $ do
is <- peekElemOff isl (fromIntegral i)
(\hsc_ptr -> pokeByteOff hsc_ptr 84) is (48 :: Int32)
is <- peekElemOff isl (fromIntegral si)
setAVDictionary opts "threads" "1"
icodec <- peek icodecp
withAVCodec inctx is icodec opts $ withAVFrame $ \frame -> do
avSeekStream inctx is frame offset
ffmt <- (\hsc_ptr -> peekByteOff hsc_ptr 116) frame
fwidth :: CInt <- (\hsc_ptr -> peekByteOff hsc_ptr 104) frame
fheight :: CInt <- (\hsc_ptr -> peekByteOff hsc_ptr 108) frame
fmap fst $ withAVOutput outfile (maybe "image2pipe" (const "image2") outfile) $ \outctx -> do
ocodec <- throwAVErrorIfNull "avcodec_find_encoder(AV_CODEC_ID_MJPEG)" (FileContext outctx) $
avcodecFindEncoder 8
os <- throwAVErrorIfNull "avformat_new_stream" (FileContext outctx) $
avformatNewStream outctx ocodec
oc :: Ptr AVCodecContext <- (\hsc_ptr -> peekByteOff hsc_ptr 8) os
avFrameInitializeStream os inctx is frame (maybe (1) fromIntegral width) (maybe (1) fromIntegral height)
owidth <- (\hsc_ptr -> peekByteOff hsc_ptr 156) oc
oheight <- (\hsc_ptr -> peekByteOff hsc_ptr 160) oc
fmts <- (\hsc_ptr -> peekByteOff hsc_ptr 40) ocodec
fmt <- throwAVErrorIf "avcodec_find_best_pix_fmt_of_list" (FileContext outctx) $
avcodecFindBestPixFmtOfList fmts ffmt 0 nullPtr
(\hsc_ptr -> pokeByteOff hsc_ptr 176) oc fmt
when (fmt /= ffmt || owidth /= fwidth || oheight /= fheight) $
throwAVErrorIf_ "av_frame_get_buffer" (FileContext outctx) $
avFrameRescale oc frame
setAVDictionary opts "threads" "1"
withAVCodec outctx os ocodec opts $
withAVPacket $ \pkt -> with 0 $ \gpp -> do
throwAVErrorIf_ "avformat_write_header" (FileContext outctx) $
avformatWriteHeader outctx nullPtr
throwAVErrorIf_ "avcodec_encode_video2" (FileContext outctx) $
avcodecEncodeVideo2 oc pkt frame gpp
gp <- peek gpp
when (gp == 0) $
throwAVError 0 "avcodec_encode_video2 packet" (FileContext outctx)
throwAVErrorIf_ "av_write_frame" (FileContext outctx) $
avWriteFrame outctx pkt
throwAVErrorIf_ "av_write_trailer" (FileContext outctx) $
avWriteTrailer outctx
where
isimg = isNothing offset