Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
355 lines (307 sloc) 12 KB
{-# LANGUAGE LambdaCase #-}
module Read (trySpecificationsOnFile) where
import Control.Monad (when, guard, unless, forM, forM_)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Maybe (MaybeT (MaybeT), runMaybeT)
import Control.Monad.Trans.State.Strict (StateT (StateT), runStateT, gets, modify')
import qualified Data.Binary.Get as B
import Data.Bits ((.&.), shiftR)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LazyBS
import qualified Data.Char as Char
import qualified Data.Map.Strict as Map
import Data.String (fromString)
import Data.Word (Word8)
import Prelude hiding (Read)
import qualified System.Endian as E
import qualified System.IO as IO
import Text.Printf (printf)
import AST
import Util
trySpecificationsOnFile :: Specifications -> FilePath -> Bool -> Bool -> IO (Maybe Values)
trySpecificationsOnFile specs filePath printOutput debug =
let toplevelSpecs = filter toplevel $ Map.elems specs
in
IO.withBinaryFile filePath IO.ReadMode $ \handle -> do
result <- runMaybeT $ runStateT (trySpecifications toplevelSpecs)
$ defaultEnv specs handle debug
case result of
Just (_, env) -> do
when printOutput $ BS.putStr $ output env
return $ Just $ values env
Nothing -> return Nothing
trySpecifications :: [Specification] -> Read ()
trySpecifications = go 0
where
go i = \case
[] -> failRead
s:ss -> (withPrefix i $ readSpecification s) `catch` (go (i+1) ss)
readSpecification :: Specification -> Read ()
readSpecification specification = do
printBlock (name specification) []
mapEnv $ \env -> env { specFilePath = filePath specification }
withIncreasedIndent $ readStatements $ statements specification
readStatements :: [Statement] -> Read ()
readStatements = mapM_ readStatement
readStatement :: Statement -> Read ()
readStatement = \case
StmtExpectConstant t v -> readConstant t v
StmtExpectValue v -> readValue v
StmtExpectEnum e n t f a -> readEnum e n t f a
StmtExpectData n l -> readData n l
StmtExpectAscii n l -> readAscii n l
StmtSequence n l s -> readSequence n l s
StmtIf c t f -> readIf c t f
StmtByteOrder b -> readByteOrder b
StmtLet n e -> readLet n e
StmtTry s -> readTry s
readConstant :: Type -> Value -> Read ()
readConstant t v = readValueOfType t >>= guard . ((==) v)
readValue :: ExpectValue -> Read ()
readValue = \case
EVSingle name type_ format assignment -> do
value <- readValueOfType type_
printValues name type_ format assignment [value]
addValue name value
EVSequence name type_ length format assignment -> do
length <- readLength length
values <- forM [0..length - 1] $ \i -> do
value <- readValueOfType type_
withPrefix i $ addValue name value
return value
printValues name type_ format assignment values
EVPacked assignment type_ format -> do
value <- readValueOfType type_
go value $ reverse assignment
where
go value = \case
[] -> return ()
(numBits, name):assignment ->
let effective = value .&. ((2 ^ numBits) - 1)
value' = value `shiftR` numBits
in do
go value' assignment
printValues name (requiredType numBits) format Nothing [effective]
addValue name effective
readEnum :: [Int] -> BS.ByteString -> Type -> Format -> Maybe Assignment -> Read ()
readEnum enum name type_ format assignment = do
value <- readValueOfType type_
guard $ value `elem` enum
printValues name type_ format assignment [value]
addValue name value
readData :: BS.ByteString -> Length -> Read ()
readData name length = do
length <- readLength length
chunk <- readChunk length
forM_ [0..length - 1] $ \i ->
let value = fromIntegral $ BS.index chunk i
in
withPrefix i $ addValue name value
printData chunk
where
printData data_ =
let hexLines = makeHexLines data_
charLines = makeCharLines data_
in
printBlock name $ mix hexLines charLines
makeHexLines = makeLines $ toLength 23 . map8
where
map8 = BS.intercalate (fromString " ")
. map (fromString . printf "%02x" . fromEnum)
. BS.unpack
makeCharLines = makeLines $ toLength 8 . map8
where
map8 = BS.map $ \c -> if printable c then c else c2w '.'
printable w = Char.isAscii (w2c w) && Char.isPrint (w2c w)
toLength n string = BS.append string $ space $ n - BS.length string
makeLines :: (BS.ByteString -> BS.ByteString) -> BS.ByteString -> [BS.ByteString]
makeLines map8 data_ =
let (data16, rest) = BS.splitAt 16 data_
in
if BS.null data16 then []
else
let (dataL8, dataH8) = BS.splitAt 8 data16
in
BS.concat [map8 dataL8, fromString " ", map8 dataH8] : (makeLines map8 rest)
mix = zipWith $ \x y -> BS.concat [x, fromString " ", y]
readAscii :: BS.ByteString -> Length -> Read ()
readAscii name length = do
length <- readLength length
chunk <- readChunk length
unless (isAscii chunk) failRead
forM_ [0..length - 1] $ \i ->
let value = fromIntegral $ BS.index chunk i
in
withPrefix i $ addValue name value
printData chunk
where
isAscii = all (Char.isAscii . w2c) . BS.unpack
printData = printBlock name . map (BS.filter isPrintable) . BS.splitWith isNewline
isNewline = (==) '\n' . w2c
isPrintable = Char.isPrint . w2c
readSequence :: BS.ByteString -> SequenceLength -> [Statement] -> Read ()
readSequence name length statements = case length of
SeqLengthEOF -> goEOF 0
SeqLengthFixed l -> readLength l >>= goFixed 0
SeqLengthPostCondition c -> goPostCondition 0 c
where
goEOF i = do
isEOF <- fromEnv handle >>= runIO . IO.hIsEOF
unless isEOF $ do
printBlock name []
withPrefix i $ withIncreasedIndent $ readStatements statements
goEOF $ i + 1
goFixed i length =
if i < length
then do printBlock name []
withPrefix i $ withIncreasedIndent $ readStatements statements
goFixed (i + 1) length
else return ()
goPostCondition i condition = do
printBlock name []
withPrefix i $ withIncreasedIndent $ readStatements statements
c <- withPrefix i $ readExpression condition
if c > 0 then goPostCondition (i + 1) condition
else return ()
readIf :: Expression -> [Statement] -> Maybe [Statement] -> Read ()
readIf condition true mFalse = do
c <- readExpression condition
if c > 0 then withPrefix 0 $ readStatements true
else case mFalse of
Nothing -> return ()
Just false -> withPrefix 1 $ readStatements false
readByteOrder :: ByteOrder -> Read ()
readByteOrder = \case
ByteOrderBigEndian -> set E.BigEndian
ByteOrderLittleEndian -> set E.LittleEndian
ByteOrderSystem c -> do
value <- readExpression c
if value > 0 then set E.getSystemEndianness
else case E.getSystemEndianness of
E.BigEndian -> set E.LittleEndian
E.LittleEndian -> set E.BigEndian
where
set e = mapEnv $ \env -> env { byteOrder = e }
readLet :: BS.ByteString -> Expression -> Read ()
readLet name expr = do
value <- readExpression expr
addValue name value
readTry :: [BS.ByteString] -> Read ()
readTry names = do
specs <- forM names $ \n ->
fromEnv specs >>= return . Map.lookup n >>= \case
Just s -> return s
Nothing -> failSpec $ "Could not find specification " ++ (show n)
trySpecifications specs
readLength :: Length -> Read Int
readLength = \case
LengthConstant c -> return c
LengthVariable v -> valueFromEnv v
readExpression :: Expression -> Read Int
readExpression e = do
prefix <- fromEnv prefix
filePos <- fromEnv handle >>= runIO . IO.hTell >>= return . fromIntegral
value <- fromEnv values >>= return . evaluate e prefix filePos
case value of
Nothing -> failSpec "Could not evaluate expression"
Just v -> return v
readValueOfType :: Type -> Read Value
readValueOfType = \case
TypeUInt8 -> readNum 1 B.getWord8 B.getWord8
TypeUInt16 -> readNum 2 B.getWord16le B.getWord16be
TypeUInt32 -> readNum 4 B.getWord32le B.getWord32be
TypeInt8 -> readNum 1 B.getInt8 B.getInt8
TypeInt16 -> readNum 2 B.getInt16le B.getInt16be
TypeInt32 -> readNum 4 B.getInt32le B.getInt32be
readNum :: Integral a => Int -> B.Get a -> B.Get a -> Read Value
readNum numBytes littleEndian bigEndian = do
chunk <- readChunk numBytes >>= return . LazyBS.fromStrict
byteOrder <- fromEnv byteOrder
return $ fromIntegral $ case byteOrder of
E.LittleEndian -> B.runGet littleEndian chunk
E.BigEndian -> B.runGet bigEndian chunk
readChunk :: Int -> Read BS.ByteString
readChunk n = do
handle <- fromEnv handle
chunk <- runIO $ BS.hGet handle n
when (fromIntegral (BS.length chunk) < n) failRead
return chunk
addValue :: BS.ByteString -> Value -> Read ()
addValue name value = do
prefix <- fromEnv prefix
mapEnv $ \env -> env { values = Map.insert (name, prefix) value $ values env }
printValues :: BS.ByteString -> Type -> Format -> Maybe Assignment -> [Value] -> Read ()
printValues label type_ format assignment values = do
indent <- fromEnv indent
toOutput [space indent, label, fromString ": "]
let values' = map (\v -> formatValue v type_ format assignment) values
toOutput [BS.intercalate (space 1) values', newline]
printBlock :: BS.ByteString -> [BS.ByteString] -> Read ()
printBlock label lines = do
indent <- fromEnv indent
toOutput [space indent, label, fromString ": "]
case lines of
[] -> toOutput [newline]
l:ls -> do toOutput [l, newline]
toOutput $ concat $ do
let indent' = indent + (BS.length label) + 2
l <- ls
return [space indent', l, newline]
toOutput :: [BS.ByteString] -> Read ()
toOutput strings = do
debug <- fromEnv debug
when debug $ forM_ strings $ runIO . BS.putStr
mapEnv $ \env -> env { output = BS.concat $ (output env) : strings }
catch :: Read a -> Read a -> Read a
catch try handler = StateT $ \env -> MaybeT $ do
pos <- IO.hTell $ handle env
result <- runMaybeT $ runStateT try env
case result of
Just r -> return result
Nothing -> do
IO.hSeek (handle env) IO.AbsoluteSeek pos
runMaybeT $ runStateT handler env
withIncreasedIndent :: Read a -> Read a
withIncreasedIndent read = do
indent <- gets indent
mapEnv $ \env -> env { indent = indent + 4 }
result <- read
mapEnv $ \env -> env { indent = indent }
return result
withPrefix :: Int -> Read a -> Read a
withPrefix p run = do
mapEnv $ \env -> env { prefix = p : (prefix env) }
result <- run
mapEnv $ \env -> env { prefix = tail $ prefix env }
return result
mapEnv :: (Env -> Env) -> Read ()
mapEnv = modify'
valueFromEnv :: BS.ByteString -> Read Value
valueFromEnv name = do
prefix <- fromEnv prefix
fromEnv values >>= return . getValue name prefix >>= \case
Just value -> return value
Nothing -> failSpec $ "Could not find variable " ++ (show name)
failRead :: Read ()
failRead = lift $ fail ""
failSpec :: String -> Read a
failSpec msg = do
specFilePath <- fromEnv specFilePath
runIO $ error $ concat [specFilePath, ": ", msg]
runIO :: IO a -> Read a
runIO = lift . lift
fromEnv :: (Env -> a) -> Read a
fromEnv = gets
defaultEnv :: Specifications -> IO.Handle -> Bool -> Env
defaultEnv specs handle = Env specs handle 0 Map.empty BS.empty "" E.getSystemEndianness []
type Read a = StateT Env (MaybeT IO) a
data Env = Env { specs :: Specifications
, handle :: IO.Handle
, indent :: Int
, values :: Values
, output :: BS.ByteString
, specFilePath :: FilePath
, byteOrder :: E.Endianness
, prefix :: [Int]
, debug :: Bool
}