/
Base85.hs
179 lines (164 loc) · 6.93 KB
/
Base85.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
{-# LANGUAGE ForeignFunctionInterface #-}
-- |
-- Module: Codec.Binary.Base85
-- Copyright: (c) 2012 Magnus Therning
-- License: BSD3
--
-- Implemented as described at <http://en.wikipedia.org/wiki/Ascii85>.
module Codec.Binary.Base85
( b85_encode_part
, b85_encode_final
, b85_decode_part
, b85_decode_final
, encode
, decode
) where
import qualified Data.ByteString as BS
import Foreign
import Foreign.C.Types
import System.IO.Unsafe as U
import Data.ByteString.Unsafe
castEnum :: (Enum a, Enum b) => a -> b
castEnum = toEnum . fromEnum
foreign import ccall "static b85.h b85_enc_part"
c_b85_enc_part :: Ptr Word8 -> CSize -> Ptr Word8 -> Ptr CSize -> Ptr (Ptr Word8) -> Ptr CSize -> IO ()
foreign import ccall "static b85.h b85_enc_final"
c_b85_enc_final :: Ptr Word8 -> CSize -> Ptr Word8 -> Ptr CSize -> IO CInt
foreign import ccall "static b85.h b85_dec_part"
c_b85_dec_part :: Ptr Word8 -> CSize -> Ptr Word8 -> Ptr CSize -> Ptr (Ptr Word8) -> Ptr CSize -> IO CInt
foreign import ccall "static b85.h b85_dec_final"
c_b85_dec_final :: Ptr Word8 -> CSize -> Ptr Word8 -> Ptr CSize -> IO CInt
-- | Encoding function.
--
-- Encodes as large a part as possible of the indata.
--
-- >>> b85_encode_part $ Data.ByteString.Char8.pack "foobar"
-- ("AoDTs","ar")
--
-- It supports special handling of both all-zero groups and all-space groups.
--
-- >>> b85_encode_part $ Data.ByteString.Char8.pack " "
-- ("y", "")
-- >>> b85_encode_part $ Data.ByteString.Char8.pack "\0\0\0\0"
-- ("z", "")
b85_encode_part :: BS.ByteString -> (BS.ByteString, BS.ByteString)
b85_encode_part bs = U.unsafePerformIO $ unsafeUseAsCStringLen bs $ \ (inBuf, inLen) -> do
let maxOutLen = inLen `div` 4 * 5
outBuf <- mallocBytes maxOutLen
alloca $ \ pOutLen ->
alloca $ \ pRemBuf ->
alloca $ \ pRemLen -> do
poke pOutLen (castEnum maxOutLen)
c_b85_enc_part (castPtr inBuf) (castEnum inLen) outBuf pOutLen pRemBuf pRemLen
outLen <- peek pOutLen
newOutBuf <- reallocBytes outBuf (castEnum outLen)
remBuf <- peek pRemBuf
remLen <- peek pRemLen
remBs <- BS.packCStringLen (castPtr remBuf, castEnum remLen)
outBs <- unsafePackCStringFinalizer newOutBuf (castEnum outLen) (free newOutBuf)
return (outBs, remBs)
-- | Encoding function for the final block.
--
-- >>> b85_encode_final $ Data.ByteString.Char8.pack "ar"
-- Just "@<)"
b85_encode_final :: BS.ByteString -> Maybe BS.ByteString
b85_encode_final bs = U.unsafePerformIO $ unsafeUseAsCStringLen bs $ \ (inBuf, inLen) -> do
outBuf <- mallocBytes 5
alloca $ \ pOutLen -> do
r <- c_b85_enc_final (castPtr inBuf) (castEnum inLen) outBuf pOutLen
if r == 0
then do
outLen <- peek pOutLen
newOutBuf <- reallocBytes outBuf (castEnum outLen)
outBs <- unsafePackCStringFinalizer newOutBuf (castEnum outLen) (free newOutBuf)
return $ Just outBs
else free outBuf >> return Nothing
-- | Decoding function.
--
-- Decode as large a portion of the input as possible.
--
-- >>> b85_decode_part $ Data.ByteString.Char8.pack "AoDTs"
-- Right ("foob","")
-- >>> b85_decode_part $ Data.ByteString.Char8.pack "AoDTs@<)"
-- Right ("foob","@<)")
-- >>> b85_decode_part $ Data.ByteString.Char8.pack "@<)"
-- Right ("","@<)")
--
-- At least 512 bytes of data is allocated for the output, but because of the
-- special handling of all-zero and all-space groups it is possible that the
-- space won't be enough. (To be sure to always fit the output one would have
-- to allocate 5 times the length of the input. It seemed a good trade-off to
-- sometimes have to call the function more than once instead.)
--
-- >>> either snd snd $ b85_decode_part $ Data.ByteString.Char8.pack $ Prelude.take 129 $ repeat 'y'
-- "y"
b85_decode_part :: BS.ByteString -> Either (BS.ByteString, BS.ByteString) (BS.ByteString, BS.ByteString)
b85_decode_part bs = U.unsafePerformIO $ unsafeUseAsCStringLen bs $ \ (inBuf, inLen) -> do
let maxOutLen = max 512 $ inLen `div` 5 * 4
outBuf <- mallocBytes maxOutLen
alloca $ \ pOutLen ->
alloca $ \ pRemBuf ->
alloca $ \ pRemLen -> do
poke pOutLen (castEnum maxOutLen)
r <- c_b85_dec_part (castPtr inBuf) (castEnum inLen) outBuf pOutLen pRemBuf pRemLen
outLen <- peek pOutLen
newOutBuf <- reallocBytes outBuf (castEnum outLen)
remBuf <- peek pRemBuf
remLen <- peek pRemLen
remBs <- BS.packCStringLen (castPtr remBuf, castEnum remLen)
outBs <- unsafePackCStringFinalizer newOutBuf (castEnum outLen) (free newOutBuf)
if r == 0
then return $ Right (outBs, remBs)
else return $ Left (outBs, remBs)
-- | Decoding function for the final block.
--
-- >>> b85_decode_final $ Data.ByteString.Char8.pack "@<)"
-- Just "ar"
-- >>> b85_decode_final $ Data.ByteString.Char8.pack ""
-- Just ""
-- >>> b85_decode_final $ Data.ByteString.Char8.pack "AoDTs"
-- Nothing
b85_decode_final :: BS.ByteString -> Maybe BS.ByteString
b85_decode_final bs = U.unsafePerformIO $ unsafeUseAsCStringLen bs $ \ (inBuf, inLen) -> do
outBuf <- mallocBytes 4
alloca $ \ pOutLen -> do
r <- c_b85_dec_final (castPtr inBuf) (castEnum inLen) outBuf pOutLen
if r == 0
then do
outLen <- peek pOutLen
newOutBuf <- reallocBytes outBuf (castEnum outLen)
outBs <- unsafePackCStringFinalizer newOutBuf (castEnum outLen) (free newOutBuf)
return $ Just outBs
else free outBuf >> return Nothing
-- | Convenience function that combines 'b85_encode_part' and
-- 'b85_encode_final' to encode a complete string.
--
-- >>> encode $ Data.ByteString.Char8.pack "foob"
-- "AoDTs"
-- >>> encode $ Data.ByteString.Char8.pack "foobar"
-- "AoDTs@<)"
encode :: BS.ByteString -> BS.ByteString
encode bs = first `BS.append` final
where
(first, rest) = b85_encode_part bs
Just final = b85_encode_final rest
-- | Convenience function that combines 'b85_decode_part' and
-- 'b85_decode_final' to decode a complete string.
--
-- >>> decode $ Data.ByteString.Char8.pack "AoDTs"
-- "foob"
-- >>> encode $ Data.ByteString.Char8.pack "AoDTs@<)"
-- "foobar"
decode :: BS.ByteString -> Either (BS.ByteString, BS.ByteString) BS.ByteString
decode bs = either Left handleFinal (iterateDecode [] bs)
where
iterateDecode bss re = case b85_decode_part re of
Right (d, r) ->
if BS.null d
then Right (BS.concat (reverse bss), r)
else iterateDecode (d : bss) r
Left (d, r) -> Left (BS.concat $ reverse $ d : bss, r)
handleFinal a@(first, rest) = maybe
(Left a)
(\ final -> Right (first `BS.append` final))
(b85_decode_final rest)