{-# LANGUAGE RecordWildCards #-}

module Telescope.Asdf.Encoding.File where

import Control.Monad (replicateM_)
import Data.Binary.Get
import Data.Binary.Put
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as BC
import Data.ByteString.Lazy qualified as BL
import Data.String (IsString)
import Data.Word
import Effectful
import Effectful.Error.Static
import Effectful.Fail
import Effectful.NonDet
import Effectful.State.Static.Local
import Telescope.Asdf.Error
import Telescope.Asdf.Node (Tree)


-- | Split an encoded 'ByteString' into a 'Tree', '[Encoded Block]' and 'Encoded Index'
splitAsdfFile :: (Error AsdfError :> es) => ByteString -> Eff es AsdfFile
splitAsdfFile :: forall (es :: [Effect]).
(Error AsdfError :> es) =>
ByteString -> Eff es AsdfFile
splitAsdfFile ByteString
dat = do
  Either String AsdfFile
res <- Eff (Fail : es) AsdfFile -> Eff es (Either String AsdfFile)
forall (es :: [Effect]) a.
HasCallStack =>
Eff (Fail : es) a -> Eff es (Either String a)
runFail (Eff (Fail : es) AsdfFile -> Eff es (Either String AsdfFile))
-> Eff (Fail : es) AsdfFile -> Eff es (Either String AsdfFile)
forall a b. (a -> b) -> a -> b
$ ByteString
-> Eff (State ByteString : Fail : es) AsdfFile
-> Eff (Fail : es) AsdfFile
forall s (es :: [Effect]) a.
HasCallStack =>
s -> Eff (State s : es) a -> Eff es a
evalState ByteString
dat Eff (State ByteString : Fail : es) AsdfFile
forall (es :: [Effect]).
(State ByteString :> es, Fail :> es) =>
Eff es AsdfFile
parseAsdfFile
  case Either String AsdfFile
res of
    Left String
e -> AsdfError -> Eff es AsdfFile
forall e (es :: [Effect]) a.
(HasCallStack, Error e :> es, Show e) =>
e -> Eff es a
throwError (AsdfError -> Eff es AsdfFile) -> AsdfError -> Eff es AsdfFile
forall a b. (a -> b) -> a -> b
$ String -> AsdfError
ParseError String
e
    Right AsdfFile
a -> AsdfFile -> Eff es AsdfFile
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure AsdfFile
a


-- | Lightweight first-pass parsing using 'State ByteString' instead of Parsec or similar
parseAsdfFile :: (State ByteString :> es, Fail :> es) => Eff es AsdfFile
parseAsdfFile :: forall (es :: [Effect]).
(State ByteString :> es, Fail :> es) =>
Eff es AsdfFile
parseAsdfFile = do
  -- if it's empty, then give an error
  Encoded Tree
tree <- String -> Eff (NonDet : es) (Encoded Tree) -> Eff es (Encoded Tree)
forall {es :: [Effect]} {b}.
(State ByteString :> es, Fail :> es) =>
String -> Eff (NonDet : es) b -> Eff es b
onEmpty String
"tree" Eff (NonDet : es) (Encoded Tree)
forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
Eff es (Encoded Tree)
parseTree
  [Encoded Block]
blocks <- String
-> Eff (NonDet : es) [Encoded Block] -> Eff es [Encoded Block]
forall {es :: [Effect]} {b}.
(State ByteString :> es, Fail :> es) =>
String -> Eff (NonDet : es) b -> Eff es b
onEmpty String
"blocks" Eff (NonDet : es) [Encoded Block]
forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
Eff es [Encoded Block]
parseBlocks
  Encoded Index
index <- Eff es (Encoded Index)
forall (es :: [Effect]).
(State ByteString :> es) =>
Eff es (Encoded Index)
parseIndex
  AsdfFile -> Eff es AsdfFile
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AsdfFile -> Eff es AsdfFile) -> AsdfFile -> Eff es AsdfFile
forall a b. (a -> b) -> a -> b
$ AsdfFile{Encoded Tree
tree :: Encoded Tree
tree :: Encoded Tree
tree, [Encoded Block]
blocks :: [Encoded Block]
blocks :: [Encoded Block]
blocks, Encoded Index
index :: Encoded Index
index :: Encoded Index
index}
 where
  onEmpty :: String -> Eff (NonDet : es) b -> Eff es b
onEmpty String
ex Eff (NonDet : es) b
eff = do
    Either CallStack b
res <- OnEmptyPolicy -> Eff (NonDet : es) b -> Eff es (Either CallStack b)
forall (es :: [Effect]) a.
HasCallStack =>
OnEmptyPolicy -> Eff (NonDet : es) a -> Eff es (Either CallStack a)
runNonDet OnEmptyPolicy
OnEmptyKeep Eff (NonDet : es) b
eff
    case Either CallStack b
res of
      Left CallStack
_ -> do
        ByteString
inp <- forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
Eff es s
get @ByteString
        String -> Eff es b
forall a. String -> Eff es a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Eff es b) -> String -> Eff es b
forall a b. (a -> b) -> a -> b
$ String
"Expected " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
ex String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ ByteString -> String
forall a. Show a => a -> String
show (Int -> ByteString -> ByteString
BS.take Int
100 ByteString
inp)
      Right b
a -> b -> Eff es b
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
a


parseTree :: (State ByteString :> es, NonDet :> es) => Eff es (Encoded Tree)
parseTree :: forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
Eff es (Encoded Tree)
parseTree = do
  ByteString
t <- ByteString -> Eff es ByteString
forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
ByteString -> Eff es ByteString
parseUntil ByteString
blockMagicToken Eff es ByteString -> Eff es ByteString -> Eff es ByteString
forall a. Eff es a -> Eff es a -> Eff es a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ByteString -> Eff es ByteString
forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
ByteString -> Eff es ByteString
parseUntil ByteString
blockIndexHeader Eff es ByteString -> Eff es ByteString -> Eff es ByteString
forall a. Eff es a -> Eff es a -> Eff es a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Eff es ByteString
forall (es :: [Effect]).
(NonDet :> es, State ByteString :> es) =>
Eff es ByteString
remainingBytes
  Encoded Tree -> Eff es (Encoded Tree)
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Encoded Tree -> Eff es (Encoded Tree))
-> Encoded Tree -> Eff es (Encoded Tree)
forall a b. (a -> b) -> a -> b
$ ByteString -> Encoded Tree
forall {k} (a :: k). ByteString -> Encoded a
Encoded ByteString
t


parseIndex :: (State ByteString :> es) => Eff es (Encoded Index)
parseIndex :: forall (es :: [Effect]).
(State ByteString :> es) =>
Eff es (Encoded Index)
parseIndex =
  ByteString -> Encoded Index
forall {k} (a :: k). ByteString -> Encoded a
Encoded (ByteString -> Encoded Index)
-> Eff es ByteString -> Eff es (Encoded Index)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Eff es ByteString
forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
Eff es s
get


parseBlocks :: (State ByteString :> es, NonDet :> es) => Eff es [Encoded Block]
parseBlocks :: forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
Eff es [Encoded Block]
parseBlocks = Eff es (Encoded Block) -> Eff es [Encoded Block]
forall a. Eff es a -> Eff es [a]
forall (f :: * -> *) a. Alternative f => f a -> f [a]
many Eff es (Encoded Block)
forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
Eff es (Encoded Block)
parseBlock


parseBlock :: (State ByteString :> es, NonDet :> es) => Eff es (Encoded Block)
parseBlock :: forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
Eff es (Encoded Block)
parseBlock = do
  ByteString -> Eff es ()
forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
ByteString -> Eff es ()
exactly ByteString
blockMagicToken
  ByteString
b <- ByteString -> Eff es ByteString
forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
ByteString -> Eff es ByteString
parseUntil ByteString
blockMagicToken Eff es ByteString -> Eff es ByteString -> Eff es ByteString
forall a. Eff es a -> Eff es a -> Eff es a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ByteString -> Eff es ByteString
forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
ByteString -> Eff es ByteString
parseUntil ByteString
blockIndexHeader Eff es ByteString -> Eff es ByteString -> Eff es ByteString
forall a. Eff es a -> Eff es a -> Eff es a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Eff es ByteString
forall (es :: [Effect]).
(NonDet :> es, State ByteString :> es) =>
Eff es ByteString
nonEmpty
  case ByteString
b of
    ByteString
"" -> Eff es (Encoded Block)
forall a. Eff es a
forall (f :: * -> *) a. Alternative f => f a
empty
    ByteString
_ -> Encoded Block -> Eff es (Encoded Block)
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Encoded Block -> Eff es (Encoded Block))
-> Encoded Block -> Eff es (Encoded Block)
forall a b. (a -> b) -> a -> b
$ ByteString -> Encoded Block
forall {k} (a :: k). ByteString -> Encoded a
Encoded (ByteString -> Encoded Block) -> ByteString -> Encoded Block
forall a b. (a -> b) -> a -> b
$ ByteString
blockMagicToken ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
b


exactly :: (State ByteString :> es, NonDet :> es) => ByteString -> Eff es ()
exactly :: forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
ByteString -> Eff es ()
exactly ByteString
val = do
  ByteString
inp <- forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
Eff es s
get @ByteString
  if ByteString
val ByteString -> ByteString -> Bool
`BS.isPrefixOf` ByteString
inp
    then do
      ByteString -> Eff es ()
forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
s -> Eff es ()
put (ByteString -> Eff es ()) -> ByteString -> Eff es ()
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
BS.drop (ByteString -> Int
BS.length ByteString
val) ByteString
inp
      () -> Eff es ()
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    else Eff es ()
forall a. Eff es a
forall (f :: * -> *) a. Alternative f => f a
empty


parseUntil :: (State ByteString :> es, NonDet :> es) => ByteString -> Eff es ByteString
parseUntil :: forall (es :: [Effect]).
(State ByteString :> es, NonDet :> es) =>
ByteString -> Eff es ByteString
parseUntil ByteString
val = do
  ByteString
inp <- forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
Eff es s
get @ByteString
  let (ByteString
before, ByteString
rest) = ByteString -> ByteString -> (ByteString, ByteString)
BS.breakSubstring ByteString
val ByteString
inp
  case ByteString
rest of
    ByteString
"" -> Eff es ByteString
forall a. Eff es a
forall (f :: * -> *) a. Alternative f => f a
empty
    ByteString
_ -> do
      ByteString -> Eff es ()
forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
s -> Eff es ()
put ByteString
rest
      ByteString -> Eff es ByteString
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
before


nonEmpty :: (NonDet :> es, State ByteString :> es) => Eff es ByteString
nonEmpty :: forall (es :: [Effect]).
(NonDet :> es, State ByteString :> es) =>
Eff es ByteString
nonEmpty = do
  ByteString
b <- Eff es ByteString
forall (es :: [Effect]).
(NonDet :> es, State ByteString :> es) =>
Eff es ByteString
remainingBytes
  case ByteString
b of
    ByteString
"" -> Eff es ByteString
forall a. Eff es a
forall (f :: * -> *) a. Alternative f => f a
empty
    ByteString
_ -> ByteString -> Eff es ByteString
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
b


remainingBytes :: (NonDet :> es, State ByteString :> es) => Eff es ByteString
remainingBytes :: forall (es :: [Effect]).
(NonDet :> es, State ByteString :> es) =>
Eff es ByteString
remainingBytes = do
  ByteString
inp <- Eff es ByteString
forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
Eff es s
get
  forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
s -> Eff es ()
put @ByteString ByteString
""
  ByteString -> Eff es ByteString
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
inp


breakIndex :: (State ByteString :> es) => Eff es (Encoded Index)
breakIndex :: forall (es :: [Effect]).
(State ByteString :> es) =>
Eff es (Encoded Index)
breakIndex = ByteString -> Encoded Index
forall {k} (a :: k). ByteString -> Encoded a
Encoded (ByteString -> Encoded Index)
-> Eff es ByteString -> Eff es (Encoded Index)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Eff es ByteString
forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
Eff es s
get


concatAsdfFile :: AsdfFile -> ByteString
concatAsdfFile :: AsdfFile -> ByteString
concatAsdfFile AsdfFile
a =
  [ByteString] -> ByteString
forall a. Monoid a => [a] -> a
mconcat [AsdfFile
a.tree.bytes, [Encoded Block] -> ByteString
blocks AsdfFile
a.blocks, Encoded Index -> ByteString
forall {r} {a}. HasField "bytes" r a => r -> a
index AsdfFile
a.index]
 where
  blocks :: [Encoded Block] -> ByteString
  blocks :: [Encoded Block] -> ByteString
blocks [Encoded Block]
ebks = [ByteString] -> ByteString
forall a. Monoid a => [a] -> a
mconcat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ (Encoded Block -> ByteString) -> [Encoded Block] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (.bytes) [Encoded Block]
ebks
  -- case mconcat $ fmap (.bytes) ebks of
  --   "" -> ""
  --   s -> s <> "\n"
  index :: r -> a
index r
ix = r
ix.bytes


encodeTree :: ByteString -> Encoded Tree
encodeTree :: ByteString -> Encoded Tree
encodeTree ByteString
tr =
  ByteString -> Encoded Tree
forall {k} (a :: k). ByteString -> Encoded a
Encoded (ByteString -> Encoded Tree) -> ByteString -> Encoded Tree
forall a b. (a -> b) -> a -> b
$ ByteString -> [ByteString] -> ByteString
BS.intercalate ByteString
"\n" ([ByteString]
headers [ByteString] -> [ByteString] -> [ByteString]
forall a. Semigroup a => a -> a -> a
<> ByteString -> [ByteString]
forall {a}. (Semigroup a, IsString a) => a -> [a]
formatDoc ByteString
tr) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"\n" -- has a trailing newline
 where
  formatDoc :: a -> [a]
formatDoc a
doc = [a
"--- " a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
doc, a
"..."]
  headers :: [ByteString]
headers = [ByteString
"#ASDF 1.0.0", ByteString
"#ASDF_STANDARD 1.5.0", ByteString
"%YAML 1.1", ByteString
tagDirective]
  tagDirective :: ByteString
tagDirective = ByteString
"%TAG ! tag:stsci.edu:asdf/"


encodeBlocks :: [BlockData] -> [Encoded Block]
encodeBlocks :: [BlockData] -> [Encoded Block]
encodeBlocks = (BlockData -> Encoded Block) -> [BlockData] -> [Encoded Block]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap BlockData -> Encoded Block
encodeBlock


encodeBlock :: BlockData -> Encoded Block
encodeBlock :: BlockData -> Encoded Block
encodeBlock BlockData
b =
  ByteString -> Encoded Block
forall {k} (a :: k). ByteString -> Encoded a
Encoded (ByteString -> Encoded Block) -> ByteString -> Encoded Block
forall a b. (a -> b) -> a -> b
$ LazyByteString -> ByteString
BL.toStrict (LazyByteString -> ByteString) -> LazyByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Put -> LazyByteString
runPut (Put -> LazyByteString) -> Put -> LazyByteString
forall a b. (a -> b) -> a -> b
$ BlockData -> Put
putBlock BlockData
b


encodeIndex :: BlockIndex -> Encoded Index
encodeIndex :: BlockIndex -> Encoded Index
encodeIndex (BlockIndex [Int]
ns) =
  ByteString -> Encoded Index
forall {k} (a :: k). ByteString -> Encoded a
Encoded (ByteString -> Encoded Index) -> ByteString -> Encoded Index
forall a b. (a -> b) -> a -> b
$ ByteString -> [ByteString] -> ByteString
BS.intercalate ByteString
"\n" ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString
"#ASDF BLOCK INDEX", ByteString
"%YAML 1.1", ByteString
"---"] [ByteString] -> [ByteString] -> [ByteString]
forall a. Semigroup a => a -> a -> a
<> (Int -> ByteString) -> [Int] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> ByteString
forall {a}. Show a => a -> ByteString
indexEntry [Int]
ns [ByteString] -> [ByteString] -> [ByteString]
forall a. Semigroup a => a -> a -> a
<> [ByteString
"..."]
 where
  indexEntry :: a -> ByteString
indexEntry a
n = ByteString
"- " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
BC.pack (a -> String
forall a. Show a => a -> String
show a
n)


data Index
data Block


-- | values that have been encoded to the file format: padding, etc
newtype Encoded a = Encoded {forall {k} (a :: k). Encoded a -> ByteString
bytes :: ByteString}
  deriving (Int -> Encoded a -> String -> String
[Encoded a] -> String -> String
Encoded a -> String
(Int -> Encoded a -> String -> String)
-> (Encoded a -> String)
-> ([Encoded a] -> String -> String)
-> Show (Encoded a)
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
forall k (a :: k). Int -> Encoded a -> String -> String
forall k (a :: k). [Encoded a] -> String -> String
forall k (a :: k). Encoded a -> String
$cshowsPrec :: forall k (a :: k). Int -> Encoded a -> String -> String
showsPrec :: Int -> Encoded a -> String -> String
$cshow :: forall k (a :: k). Encoded a -> String
show :: Encoded a -> String
$cshowList :: forall k (a :: k). [Encoded a] -> String -> String
showList :: [Encoded a] -> String -> String
Show, Encoded a -> Encoded a -> Bool
(Encoded a -> Encoded a -> Bool)
-> (Encoded a -> Encoded a -> Bool) -> Eq (Encoded a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (a :: k). Encoded a -> Encoded a -> Bool
$c== :: forall k (a :: k). Encoded a -> Encoded a -> Bool
== :: Encoded a -> Encoded a -> Bool
$c/= :: forall k (a :: k). Encoded a -> Encoded a -> Bool
/= :: Encoded a -> Encoded a -> Bool
Eq)
  deriving newtype (String -> Encoded a
(String -> Encoded a) -> IsString (Encoded a)
forall a. (String -> a) -> IsString a
forall k (a :: k). String -> Encoded a
$cfromString :: forall k (a :: k). String -> Encoded a
fromString :: String -> Encoded a
IsString)


-- | Decompressed block data
newtype BlockData = BlockData {BlockData -> ByteString
bytes :: ByteString}
  deriving (BlockData -> BlockData -> Bool
(BlockData -> BlockData -> Bool)
-> (BlockData -> BlockData -> Bool) -> Eq BlockData
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BlockData -> BlockData -> Bool
== :: BlockData -> BlockData -> Bool
$c/= :: BlockData -> BlockData -> Bool
/= :: BlockData -> BlockData -> Bool
Eq)


newtype BlockSource = BlockSource Int
  deriving (BlockSource -> BlockSource -> Bool
(BlockSource -> BlockSource -> Bool)
-> (BlockSource -> BlockSource -> Bool) -> Eq BlockSource
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BlockSource -> BlockSource -> Bool
== :: BlockSource -> BlockSource -> Bool
$c/= :: BlockSource -> BlockSource -> Bool
/= :: BlockSource -> BlockSource -> Bool
Eq)


instance Show BlockData where
  show :: BlockData -> String
show (BlockData ByteString
bs) = String
"BlockData " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show (ByteString -> Int
BS.length ByteString
bs)


data BlockHeader = BlockHeader
  { BlockHeader -> Word16
headerSize :: Word16
  , BlockHeader -> Word32
flags :: Word32
  , BlockHeader -> Compression
compression :: Compression
  , BlockHeader -> Word64
allocatedSize :: Word64
  , BlockHeader -> Word64
usedSize :: Word64
  , BlockHeader -> Word64
dataSize :: Word64
  , BlockHeader -> Checksum
checksum :: Checksum
  }
  deriving (Int -> BlockHeader -> String -> String
[BlockHeader] -> String -> String
BlockHeader -> String
(Int -> BlockHeader -> String -> String)
-> (BlockHeader -> String)
-> ([BlockHeader] -> String -> String)
-> Show BlockHeader
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> BlockHeader -> String -> String
showsPrec :: Int -> BlockHeader -> String -> String
$cshow :: BlockHeader -> String
show :: BlockHeader -> String
$cshowList :: [BlockHeader] -> String -> String
showList :: [BlockHeader] -> String -> String
Show, BlockHeader -> BlockHeader -> Bool
(BlockHeader -> BlockHeader -> Bool)
-> (BlockHeader -> BlockHeader -> Bool) -> Eq BlockHeader
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BlockHeader -> BlockHeader -> Bool
== :: BlockHeader -> BlockHeader -> Bool
$c/= :: BlockHeader -> BlockHeader -> Bool
/= :: BlockHeader -> BlockHeader -> Bool
Eq)


newtype BlockIndex = BlockIndex [Int]
  deriving (BlockIndex -> BlockIndex -> Bool
(BlockIndex -> BlockIndex -> Bool)
-> (BlockIndex -> BlockIndex -> Bool) -> Eq BlockIndex
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BlockIndex -> BlockIndex -> Bool
== :: BlockIndex -> BlockIndex -> Bool
$c/= :: BlockIndex -> BlockIndex -> Bool
/= :: BlockIndex -> BlockIndex -> Bool
Eq)


data Compression
  = NoCompression -- "\0\0\0\0"
  | ZLib -- "zlib"
  | BZip2 -- "bzp2"
  deriving (Int -> Compression -> String -> String
[Compression] -> String -> String
Compression -> String
(Int -> Compression -> String -> String)
-> (Compression -> String)
-> ([Compression] -> String -> String)
-> Show Compression
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> Compression -> String -> String
showsPrec :: Int -> Compression -> String -> String
$cshow :: Compression -> String
show :: Compression -> String
$cshowList :: [Compression] -> String -> String
showList :: [Compression] -> String -> String
Show, Compression -> Compression -> Bool
(Compression -> Compression -> Bool)
-> (Compression -> Compression -> Bool) -> Eq Compression
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Compression -> Compression -> Bool
== :: Compression -> Compression -> Bool
$c/= :: Compression -> Compression -> Bool
/= :: Compression -> Compression -> Bool
Eq)


newtype Checksum = Checksum ByteString
  deriving (Int -> Checksum -> String -> String
[Checksum] -> String -> String
Checksum -> String
(Int -> Checksum -> String -> String)
-> (Checksum -> String)
-> ([Checksum] -> String -> String)
-> Show Checksum
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> Checksum -> String -> String
showsPrec :: Int -> Checksum -> String -> String
$cshow :: Checksum -> String
show :: Checksum -> String
$cshowList :: [Checksum] -> String -> String
showList :: [Checksum] -> String -> String
Show, Checksum -> Checksum -> Bool
(Checksum -> Checksum -> Bool)
-> (Checksum -> Checksum -> Bool) -> Eq Checksum
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Checksum -> Checksum -> Bool
== :: Checksum -> Checksum -> Bool
$c/= :: Checksum -> Checksum -> Bool
/= :: Checksum -> Checksum -> Bool
Eq)


noChecksum :: Checksum
noChecksum :: Checksum
noChecksum = ByteString -> Checksum
Checksum (ByteString -> Checksum) -> ByteString -> Checksum
forall a b. (a -> b) -> a -> b
$ Int -> Char -> ByteString
BC.replicate Int
16 Char
'0'


data AsdfFile = AsdfFile
  { AsdfFile -> Encoded Tree
tree :: Encoded Tree
  , AsdfFile -> [Encoded Block]
blocks :: [Encoded Block]
  , AsdfFile -> Encoded Index
index :: Encoded Index
  }
  deriving (Int -> AsdfFile -> String -> String
[AsdfFile] -> String -> String
AsdfFile -> String
(Int -> AsdfFile -> String -> String)
-> (AsdfFile -> String)
-> ([AsdfFile] -> String -> String)
-> Show AsdfFile
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> AsdfFile -> String -> String
showsPrec :: Int -> AsdfFile -> String -> String
$cshow :: AsdfFile -> String
show :: AsdfFile -> String
$cshowList :: [AsdfFile] -> String -> String
showList :: [AsdfFile] -> String -> String
Show, AsdfFile -> AsdfFile -> Bool
(AsdfFile -> AsdfFile -> Bool)
-> (AsdfFile -> AsdfFile -> Bool) -> Eq AsdfFile
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: AsdfFile -> AsdfFile -> Bool
== :: AsdfFile -> AsdfFile -> Bool
$c/= :: AsdfFile -> AsdfFile -> Bool
/= :: AsdfFile -> AsdfFile -> Bool
Eq)


getBlock :: Get BlockData
getBlock :: Get BlockData
getBlock = do
  BlockHeader
h <- Get BlockHeader
getBlockHeader
  BlockHeader -> Get BlockData
getBlockData BlockHeader
h


-- | Skip along blocks and create a list of all of them
getBlocks :: Get [BlockData]
getBlocks :: Get [BlockData]
getBlocks = do
  Bool
isBlock <- Get Bool
checkMagicToken
  if Bool -> Bool
not Bool
isBlock
    then [BlockData] -> Get [BlockData]
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    else do
      BlockData
b <- Get BlockData
getBlock
      [BlockData]
bs <- Get [BlockData]
getBlocks
      [BlockData] -> Get [BlockData]
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BlockData
b BlockData -> [BlockData] -> [BlockData]
forall a. a -> [a] -> [a]
: [BlockData]
bs)


getBlockHeader :: Get BlockHeader
getBlockHeader :: Get BlockHeader
getBlockHeader = do
  Get ()
expectMagicToken
  Word16
headerSize <- String -> Get Word16 -> Get Word16
forall a. String -> Get a -> Get a
label String
"header_size" Get Word16
getWord16be
  Int64
start <- Get Int64
bytesRead

  -- the remainder is inside the headerSize
  Word32
flags <- String -> Get Word32 -> Get Word32
forall a. String -> Get a -> Get a
label String
"flags" Get Word32
getWord32be -- 4
  Compression
compression <- String -> Get Compression -> Get Compression
forall a. String -> Get a -> Get a
label String
"compression" Get Compression
getCompression -- 4
  Word64
allocatedSize <- String -> Get Word64 -> Get Word64
forall a. String -> Get a -> Get a
label String
"allocated_size" Get Word64
getWord64be -- 8
  Word64
usedSize <- String -> Get Word64 -> Get Word64
forall a. String -> Get a -> Get a
label String
"used_size" Get Word64
getWord64be -- 8
  Word64
dataSize <- String -> Get Word64 -> Get Word64
forall a. String -> Get a -> Get a
label String
"data_size" Get Word64
getWord64be -- 8
  Checksum
checksum <- String -> Get Checksum -> Get Checksum
forall a. String -> Get a -> Get a
label String
"checksum" Get Checksum
getChecksum -- 2
  Int64
end <- Get Int64
bytesRead

  -- skip until the end of headerSize
  let usedHead :: Int
usedHead = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> Int64 -> Int
forall a b. (a -> b) -> a -> b
$ Int64
end Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
start
  Int -> Get ()
skip (Int -> Get ()) -> Int -> Get ()
forall a b. (a -> b) -> a -> b
$ Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
headerSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
usedHead

  BlockHeader -> Get BlockHeader
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BlockHeader -> Get BlockHeader) -> BlockHeader -> Get BlockHeader
forall a b. (a -> b) -> a -> b
$ BlockHeader{Word16
Word32
Word64
Checksum
Compression
headerSize :: Word16
flags :: Word32
compression :: Compression
allocatedSize :: Word64
usedSize :: Word64
dataSize :: Word64
checksum :: Checksum
headerSize :: Word16
flags :: Word32
compression :: Compression
allocatedSize :: Word64
usedSize :: Word64
dataSize :: Word64
checksum :: Checksum
..}
 where
  getCompression :: Get Compression
getCompression = do
    ByteString
val <- Int -> Get ByteString
getByteString Int
4
    case ByteString
val of
      ByteString
"\0\0\0\0" -> Compression -> Get Compression
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Compression
NoCompression
      -- "zlib" -> pure ZLib
      -- "bzp2" -> pure BZip2
      ByteString
_ -> String -> Get Compression
forall a. String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get Compression) -> String -> Get Compression
forall a b. (a -> b) -> a -> b
$ String
"BlockHeader compression invalid, found " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ByteString -> String
forall a. Show a => a -> String
show ByteString
val

  getChecksum :: Get Checksum
getChecksum = ByteString -> Checksum
Checksum (ByteString -> Checksum) -> Get ByteString -> Get Checksum
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get ByteString
getByteString Int
16

  expectMagicToken :: Get ()
expectMagicToken = do
    Either ByteString ()
m <- Get (Either ByteString ())
getMagicToken
    case Either ByteString ()
m of
      Right ()
a -> () -> Get ()
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
a
      Left ByteString
str -> String -> Get ()
forall a. String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get ()) -> String -> Get ()
forall a b. (a -> b) -> a -> b
$ String
"BlockHeader magic token invalid: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ ByteString -> String
forall a. Show a => a -> String
show ByteString
str


putBlockHeader :: BlockHeader -> Put
putBlockHeader :: BlockHeader -> Put
putBlockHeader BlockHeader
h = do
  ByteString -> Put
putByteString ByteString
blockMagicToken
  Word16 -> Put
putWord16be BlockHeader
h.headerSize
  Int
size <- PutM Int
putHeaderContent
  let emptyBytes :: Int
emptyBytes = Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral BlockHeader
h.headerSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
size :: Int
  Int -> Put -> Put
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
emptyBytes (Put -> Put) -> Put -> Put
forall a b. (a -> b) -> a -> b
$ Word8 -> Put
putWord8 Word8
0x0
 where
  putCompression :: p -> Put
putCompression p
_ = ByteString -> Put
putByteString ByteString
"\0\0\0\0"
  putChecksum :: Checksum -> Put
putChecksum (Checksum ByteString
cs) = ByteString -> Put
putByteString ByteString
cs
  putHeaderContent :: PutM Int
putHeaderContent = do
    let bs :: LazyByteString
bs = Put -> LazyByteString
runPut (Put -> LazyByteString) -> Put -> LazyByteString
forall a b. (a -> b) -> a -> b
$ do
          Word32 -> Put
putWord32be Word32
0 -- flags
          Compression -> Put
forall {p}. p -> Put
putCompression Compression
NoCompression
          Word64 -> Put
putWord64be BlockHeader
h.allocatedSize
          Word64 -> Put
putWord64be BlockHeader
h.usedSize
          Word64 -> Put
putWord64be BlockHeader
h.dataSize
          Checksum -> Put
putChecksum BlockHeader
h.checksum
    LazyByteString -> Put
putLazyByteString LazyByteString
bs
    Int -> PutM Int
forall a. a -> PutM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> PutM Int) -> Int -> PutM Int
forall a b. (a -> b) -> a -> b
$ Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> Int64 -> Int
forall a b. (a -> b) -> a -> b
$ LazyByteString -> Int64
BL.length LazyByteString
bs


blockIndex :: Encoded Tree -> [Encoded Block] -> BlockIndex
blockIndex :: Encoded Tree -> [Encoded Block] -> BlockIndex
blockIndex (Encoded ByteString
bytes) [Encoded Block]
ebs =
  let ns :: [Int]
ns = (Int -> Encoded Block -> Int) -> Int -> [Encoded Block] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Encoded Block -> Int
go (ByteString -> Int
BS.length ByteString
bytes) [Encoded Block]
ebs
   in [Int] -> BlockIndex
BlockIndex ([Int] -> BlockIndex) -> [Int] -> BlockIndex
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ns Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Int]
ns
 where
  go :: Int -> Encoded Block -> Int
  go :: Int -> Encoded Block -> Int
go Int
n Encoded Block
eb = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
BS.length Encoded Block
eb.bytes


putBlock :: BlockData -> Put
putBlock :: BlockData -> Put
putBlock bd :: BlockData
bd@(BlockData ByteString
bs) = do
  BlockHeader -> Put
putBlockHeader (BlockHeader -> Put) -> BlockHeader -> Put
forall a b. (a -> b) -> a -> b
$ BlockData -> BlockHeader
blockHeader BlockData
bd
  ByteString -> Put
putByteString ByteString
bs


blockHeader :: BlockData -> BlockHeader
blockHeader :: BlockData -> BlockHeader
blockHeader (BlockData ByteString
bs) =
  let bytes :: Word64
bytes = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> Int -> Word64
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
bs
   in BlockHeader
        { headerSize :: Word16
headerSize = Word16
48 -- minimum allowed size. Our encoding uses fewer bytes than this
        , compression :: Compression
compression = Compression
NoCompression
        , allocatedSize :: Word64
allocatedSize = Word64
bytes
        , usedSize :: Word64
usedSize = Word64
bytes
        , dataSize :: Word64
dataSize = Word64
bytes
        , checksum :: Checksum
checksum = Checksum
noChecksum
        , flags :: Word32
flags = Word32
0
        }


getBlockData :: BlockHeader -> Get BlockData
getBlockData :: BlockHeader -> Get BlockData
getBlockData BlockHeader
h = do
  -- LATER: handle compression
  ByteString
bytes <- Int -> Get ByteString
getByteString (Int -> Get ByteString) -> Int -> Get ByteString
forall a b. (a -> b) -> a -> b
$ Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral BlockHeader
h.usedSize
  ByteString
_empty <- Int -> Get ByteString
getByteString (Int -> Get ByteString) -> Int -> Get ByteString
forall a b. (a -> b) -> a -> b
$ Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int) -> Word64 -> Int
forall a b. (a -> b) -> a -> b
$ BlockHeader
h.allocatedSize Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- BlockHeader
h.usedSize
  BlockData -> Get BlockData
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BlockData -> Get BlockData) -> BlockData -> Get BlockData
forall a b. (a -> b) -> a -> b
$ ByteString -> BlockData
BlockData ByteString
bytes


blockMagicToken :: ByteString
blockMagicToken :: ByteString
blockMagicToken = [Word8] -> ByteString
BS.pack [Word8
0xd3, Word8
0x42, Word8
0x4c, Word8
0x4b]


blockIndexHeader :: ByteString
blockIndexHeader :: ByteString
blockIndexHeader = ByteString
"#ASDF BLOCK INDEX"


-- | Verify that the magic token is next without consuming any input
checkMagicToken :: Get Bool
checkMagicToken :: Get Bool
checkMagicToken = do
  Bool
emp <- Get Bool
isEmpty
  if Bool
emp
    then Bool -> Get Bool
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
    else do
      Either ByteString ()
eb <- Get (Either ByteString ()) -> Get (Either ByteString ())
forall a. Get a -> Get a
lookAhead Get (Either ByteString ())
getMagicToken
      Bool -> Get Bool
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> Get Bool) -> Bool -> Get Bool
forall a b. (a -> b) -> a -> b
$ (ByteString -> Bool)
-> (() -> Bool) -> Either ByteString () -> Bool
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Bool -> ByteString -> Bool
forall a b. a -> b -> a
const Bool
False) (Bool -> () -> Bool
forall a b. a -> b -> a
const Bool
True) Either ByteString ()
eb


-- | consume the magic token if available
getMagicToken :: Get (Either ByteString ())
getMagicToken :: Get (Either ByteString ())
getMagicToken = do
  -- this still fails if it is empty
  ByteString
str <- Int -> Get ByteString
getByteString Int
4
  Either ByteString () -> Get (Either ByteString ())
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either ByteString () -> Get (Either ByteString ()))
-> Either ByteString () -> Get (Either ByteString ())
forall a b. (a -> b) -> a -> b
$
    if ByteString
str ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
blockMagicToken
      then () -> Either ByteString ()
forall a b. b -> Either a b
Right ()
      else ByteString -> Either ByteString ()
forall a b. a -> Either a b
Left ByteString
str