diff options
Diffstat (limited to 'src/Data/ByteString/Base32/Internal.hs')
-rw-r--r-- | src/Data/ByteString/Base32/Internal.hs | 230 |
1 files changed, 230 insertions, 0 deletions
diff --git a/src/Data/ByteString/Base32/Internal.hs b/src/Data/ByteString/Base32/Internal.hs new file mode 100644 index 0000000..c34e80c --- /dev/null +++ b/src/Data/ByteString/Base32/Internal.hs | |||
@@ -0,0 +1,230 @@ | |||
1 | -- | | ||
2 | -- Copyright : (c) Sam Truzjan 2013 | ||
3 | -- License : BSD3 | ||
4 | -- Maintainer : pxqr.sta@gmail.com | ||
5 | -- Stability : stable | ||
6 | -- Portability : portable | ||
7 | -- | ||
8 | -- (Word5 <-> Word8) and (Word8 -> Word5) bytestring packers using | ||
9 | -- lookup table. | ||
10 | -- | ||
11 | {-# LANGUAGE BangPatterns #-} | ||
12 | module Data.ByteString.Base32.Internal | ||
13 | ( Word5 | ||
14 | , Word8 | ||
15 | |||
16 | , EncTable | ||
17 | , unpack5 | ||
18 | |||
19 | , DecTable | ||
20 | , pack5 | ||
21 | , invIx | ||
22 | ) where | ||
23 | |||
24 | import Data.Bits.Extras | ||
25 | import Data.ByteString as BS | ||
26 | import Data.ByteString.Internal as BS | ||
27 | import Data.Word | ||
28 | import Foreign hiding (unsafePerformIO) | ||
29 | import System.IO.Unsafe (unsafePerformIO) | ||
30 | import System.Endian | ||
31 | |||
32 | |||
33 | {----------------------------------------------------------------------- | ||
34 | -- Utils | ||
35 | -----------------------------------------------------------------------} | ||
36 | |||
37 | type Word5 = Word8 | ||
38 | |||
39 | -- System.Endian.toBE32 is slower because toBE32 implemented using | ||
40 | -- cbits shuffle functions while toBE32' implemented used gcc | ||
41 | -- intrinsics | ||
42 | -- | ||
43 | toBE64' :: Word64 -> Word64 | ||
44 | toBE64' = if getSystemEndianness == BigEndian then id else byteSwap | ||
45 | {-# INLINE toBE64' #-} | ||
46 | |||
47 | toBE32' :: Word32 -> Word32 | ||
48 | toBE32' = if getSystemEndianness == BigEndian then id else byteSwap | ||
49 | {-# INLINE toBE32' #-} | ||
50 | |||
51 | fromBE32' :: Word32 -> Word32 | ||
52 | fromBE32' = toBE32' | ||
53 | {-# INLINE fromBE32' #-} | ||
54 | |||
55 | -- n = 2 ^ d | ||
56 | padCeilN :: Int -> Int -> Int | ||
57 | padCeilN !n !x | ||
58 | | remd == 0 = x | ||
59 | | otherwise = (x - remd) + n | ||
60 | where mask = n - 1 | ||
61 | remd = x .&. mask | ||
62 | |||
63 | {----------------------------------------------------------------------- | ||
64 | -- Encoding | ||
65 | -----------------------------------------------------------------------} | ||
66 | |||
67 | unpack5Ptr :: Ptr Word8 -> ByteString -> ByteString | ||
68 | unpack5Ptr !tbl bs @ (PS fptr off sz) = | ||
69 | unsafePerformIO $ do | ||
70 | let unpackedSize = dstSize $ BS.length bs | ||
71 | BS.create unpackedSize $ \ dst -> do | ||
72 | withForeignPtr fptr $ \ ptr -> do | ||
73 | dst_end <- bigStep dst (advancePtr ptr off) sz | ||
74 | _ <- fillPadding dst_end (unpackedSize - (dst_end `minusPtr` dst)) | ||
75 | return () | ||
76 | where | ||
77 | dstSize x = padCeilN 8 (d + if m == 0 then 0 else 1) | ||
78 | where (d, m) = (x * 8) `quotRem` 5 | ||
79 | |||
80 | fillPadding dst s = memset dst (c2w '=') (fromIntegral s) | ||
81 | |||
82 | bigStep !dst !src !s | ||
83 | | s >= 5 = do | ||
84 | unpack5_40 dst src | ||
85 | bigStep (dst `advancePtr` 8) (src `advancePtr` 5) (s - 5) | ||
86 | | otherwise = smallStep dst src s 0 0 | ||
87 | |||
88 | unpack5_40 !dst !src = do | ||
89 | w32he <- peek (castPtr src) :: IO Word32 | ||
90 | let w32 = toBE32' w32he | ||
91 | fill8_32 0 (w32 `unsafeShiftR` 27) | ||
92 | fill8_32 1 (w32 `unsafeShiftR` 22) | ||
93 | fill8_32 2 (w32 `unsafeShiftR` 17) | ||
94 | fill8_32 3 (w32 `unsafeShiftR` 12) | ||
95 | fill8_32 4 (w32 `unsafeShiftR` 7) | ||
96 | fill8_32 5 (w32 `unsafeShiftR` 2) | ||
97 | |||
98 | w8 <- peekElemOff src 4 | ||
99 | fill8_32 6 ( (w32 `unsafeShiftL` 3) | ||
100 | .|. fromIntegral (w8 `unsafeShiftR` 5)) | ||
101 | fill8_32 7 (fromIntegral w8) | ||
102 | where | ||
103 | fill8_32 :: Int -> Word32 -> IO () | ||
104 | fill8_32 !i !w32 = do | ||
105 | w8 <- peekByteOff tbl (fromIntegral w32 .&. 0x1f) | ||
106 | poke (dst `advancePtr` i) w8 | ||
107 | |||
108 | smallStep !dst !src !s !unused !un_cnt | ||
109 | | un_cnt >= 5 = do | ||
110 | let ix = unused `unsafeShiftR` 3 | ||
111 | peekByteOff tbl (fromIntegral ix) >>= poke dst | ||
112 | smallStep (advancePtr dst 1) | ||
113 | src s | ||
114 | (unused `unsafeShiftL` 5) | ||
115 | (un_cnt - 5) | ||
116 | |||
117 | | s == 0 = do | ||
118 | if un_cnt == 0 | ||
119 | then return dst | ||
120 | else do | ||
121 | let ix = unused `unsafeShiftR` 3 | ||
122 | peekByteOff tbl (fromIntegral ix) >>= poke dst | ||
123 | return (dst `advancePtr` 1) | ||
124 | |||
125 | | otherwise = do | ||
126 | w8 <- peek src | ||
127 | let usd_cnt = 5 - un_cnt | ||
128 | let bits = w8 .&. complement (bit (8 - usd_cnt) - 1) | ||
129 | let ix = (unused .|. bits `shiftR` un_cnt) `unsafeShiftR` 3 | ||
130 | peekByteOff tbl (fromIntegral ix) >>= poke dst | ||
131 | smallStep (advancePtr dst 1) | ||
132 | (advancePtr src 1) (pred s) | ||
133 | (w8 `shiftL` usd_cnt) (8 - usd_cnt) | ||
134 | |||
135 | type EncTable = ByteString | ||
136 | |||
137 | unpack5 :: EncTable -> ByteString -> ByteString | ||
138 | unpack5 (PS fptr off len) bs | ||
139 | | len /= 32 | ||
140 | = error $ "base32: unpack5: invalid lookup table size " ++ show len | ||
141 | | otherwise = | ||
142 | unsafePerformIO $ do | ||
143 | withForeignPtr fptr $ \ptr -> do | ||
144 | return $ unpack5Ptr (ptr `advancePtr` off) bs | ||
145 | |||
146 | {----------------------------------------------------------------------- | ||
147 | -- Decoding | ||
148 | -----------------------------------------------------------------------} | ||
149 | |||
150 | invIx :: Word5 | ||
151 | invIx = 255 | ||
152 | |||
153 | pack5Ptr :: Ptr Word5 -> ByteString -> ByteString | ||
154 | pack5Ptr !tbl bs @ (PS fptr off sz) = | ||
155 | unsafePerformIO $ do | ||
156 | let packedSize = dstSize $ BS.length bs | ||
157 | BS.createAndTrim packedSize $ \ dst -> do | ||
158 | withForeignPtr fptr $ \ ptr -> do | ||
159 | dst_end <- bigStep dst (advancePtr ptr off) sz | ||
160 | return (dst_end `minusPtr` dst) | ||
161 | where | ||
162 | lookupTable :: Word8 -> Word5 | ||
163 | lookupTable ix | ||
164 | | x == invIx = error $ "base32: decode: invalid character" ++ show ix | ||
165 | | otherwise = x | ||
166 | where x = inlinePerformIO (peekByteOff tbl (fromIntegral ix)) | ||
167 | {-# INLINE lookupTable #-} | ||
168 | |||
169 | dstSize x = d + if m == 0 then 0 else 1 | ||
170 | where (d, m) = (x * 5) `quotRem` 8 | ||
171 | |||
172 | bigStep !dst !src !s | ||
173 | | s > 8 = do | ||
174 | pack5_40 dst src | ||
175 | bigStep (dst `advancePtr` 5) (src `advancePtr` 8) (s - 8) | ||
176 | | otherwise = smallStep dst src s (0 :: Word64) 0 | ||
177 | |||
178 | pack5_40 !dst !src = do | ||
179 | w64he <- peek (castPtr src) :: IO Word64 | ||
180 | let w64 = toBE64' w64he | ||
181 | let w40 = putAsW5 (w64 `unsafeShiftR` 00) $ | ||
182 | putAsW5 (w64 `unsafeShiftR` 08) $ | ||
183 | putAsW5 (w64 `unsafeShiftR` 16) $ | ||
184 | putAsW5 (w64 `unsafeShiftR` 24) $ | ||
185 | putAsW5 (w64 `unsafeShiftR` 32) $ | ||
186 | putAsW5 (w64 `unsafeShiftR` 40) $ | ||
187 | putAsW5 (w64 `unsafeShiftR` 48) $ | ||
188 | putAsW5 (w64 `unsafeShiftR` 56) 0 | ||
189 | pokeW40 w40 | ||
190 | where | ||
191 | putAsW5 :: Word64 -> Word64 -> Word64 | ||
192 | {-# INLINE putAsW5 #-} | ||
193 | putAsW5 !w8 !acc = (acc `unsafeShiftL` 5) | ||
194 | .|. fromIntegral (lookupTable (fromIntegral w8)) | ||
195 | |||
196 | pokeW40 :: Word64 -> IO () | ||
197 | {-# INLINE pokeW40 #-} | ||
198 | pokeW40 !w40 = do | ||
199 | poke dst (fromIntegral (w40 `unsafeShiftR` 32) :: Word8) | ||
200 | poke (castPtr (dst `advancePtr` 1)) | ||
201 | (fromBE32' (fromIntegral w40 :: Word32)) | ||
202 | |||
203 | smallStep !dst !src !s !unused !un_cnt | ||
204 | | un_cnt >= 8 = do | ||
205 | poke dst $ fromIntegral (unused `unsafeShiftR` (un_cnt - 8)) | ||
206 | smallStep (dst `advancePtr` 1) src s unused (un_cnt - 8) | ||
207 | |||
208 | | s == 0 = return dst | ||
209 | | otherwise = do | ||
210 | w8 <- peek src | ||
211 | if w2c w8 == '=' | ||
212 | then if (bit un_cnt - 1) .&. unused == 0 | ||
213 | then smallStep dst src 0 0 0 | ||
214 | else smallStep dst src 0 (unused `shiftL` (8 - un_cnt)) 8 | ||
215 | else smallStep dst | ||
216 | (src `advancePtr` 1) (pred s) | ||
217 | ((unused `unsafeShiftL` 5) | ||
218 | .|. fromIntegral (lookupTable (fromIntegral w8))) | ||
219 | (un_cnt + 5) | ||
220 | |||
221 | type DecTable = ByteString | ||
222 | |||
223 | pack5 :: DecTable -> ByteString -> ByteString | ||
224 | pack5 (PS fptr off len) bs | ||
225 | | len /= 256 | ||
226 | = error $ "base32: pack5: invalid lookup table size " ++ show len | ||
227 | | otherwise = | ||
228 | unsafePerformIO $ do | ||
229 | withForeignPtr fptr $ \ptr -> | ||
230 | return $ pack5Ptr (ptr `advancePtr` off) bs | ||