diff options
author | Sam T <sta.cs.vsu@gmail.com> | 2013-03-31 15:44:10 +0400 |
---|---|---|
committer | Sam T <sta.cs.vsu@gmail.com> | 2013-03-31 15:44:10 +0400 |
commit | 3c82f6740cfab315892d6cf0186ec0b8188d8d57 (patch) | |
tree | 5ee48746be8a470905a64ef12bea915de3ac98e1 /src/Data/BEncode.hs | |
parent | d9d8be772d574236a5c1e33625e804973be3b7fb (diff) |
add src
Diffstat (limited to 'src/Data/BEncode.hs')
-rw-r--r-- | src/Data/BEncode.hs | 181 |
1 files changed, 181 insertions, 0 deletions
diff --git a/src/Data/BEncode.hs b/src/Data/BEncode.hs new file mode 100644 index 0000000..fceb88b --- /dev/null +++ b/src/Data/BEncode.hs | |||
@@ -0,0 +1,181 @@ | |||
1 | -- | This module is intented to be imported qualified. | ||
2 | module Data.BEncode | ||
3 | ( -- ^ Datatype | ||
4 | BEncode(..) | ||
5 | |||
6 | -- ^ Construction | ||
7 | , string, integer, list, dict | ||
8 | , int, charstring, dictAssoc | ||
9 | |||
10 | -- ^ Destructuring | ||
11 | -- ^ Serialization | ||
12 | , encode, decode | ||
13 | |||
14 | -- ^ Extra | ||
15 | , builder, parser, printPretty | ||
16 | |||
17 | -- ^ Predicates | ||
18 | , isInteger, isString, isList, isDict | ||
19 | ) where | ||
20 | |||
21 | |||
22 | import Control.Applicative | ||
23 | import Data.Int | ||
24 | import Data.Foldable | ||
25 | import Data.Monoid ((<>)) | ||
26 | import Data.Map (Map) | ||
27 | import qualified Data.Map as M | ||
28 | import Data.Attoparsec.ByteString.Char8 (Parser) | ||
29 | import qualified Data.Attoparsec.ByteString.Char8 as P | ||
30 | import Data.ByteString (ByteString) | ||
31 | import qualified Data.ByteString as B | ||
32 | import qualified Data.ByteString.Lazy as Lazy | ||
33 | import Data.ByteString.Internal as B (c2w, w2c) | ||
34 | import qualified Data.ByteString.Builder as B | ||
35 | import qualified Data.ByteString.Builder.Prim as BP () | ||
36 | import Text.PrettyPrint.ANSI.Leijen (Pretty, Doc, pretty, (<+>), (</>)) | ||
37 | import qualified Text.PrettyPrint.ANSI.Leijen as PP | ||
38 | |||
39 | type Dict = Map ByteString BEncode | ||
40 | |||
41 | -- | 'BEncode' is straightforward AST for b-encoded values. | ||
42 | -- Please note that since dictionaries are sorted, in most cases we can | ||
43 | -- compare BEncoded values without serialization and vice versa. | ||
44 | -- Lists is not required to be sorted through. | ||
45 | -- Also note that 'BEncode' have JSON-like instance for 'Pretty'. | ||
46 | -- | ||
47 | data BEncode = BInteger Int64 | ||
48 | | BString ByteString | ||
49 | | BList [BEncode] | ||
50 | | BDict Dict | ||
51 | deriving (Show, Read, Eq, Ord) | ||
52 | |||
53 | integer :: Integer -> BEncode | ||
54 | integer = BInteger . fromIntegral | ||
55 | {-# INLINE integer #-} | ||
56 | |||
57 | string :: ByteString -> BEncode | ||
58 | string = BString | ||
59 | {-# INLINE string #-} | ||
60 | |||
61 | list :: [BEncode] -> BEncode | ||
62 | list = BList | ||
63 | {-# INLINE list #-} | ||
64 | |||
65 | dict :: Dict -> BEncode | ||
66 | dict = BDict | ||
67 | {-# INLINE dict #-} | ||
68 | |||
69 | |||
70 | int :: Int -> BEncode | ||
71 | int = integer . fromIntegral | ||
72 | {-# INLINE int #-} | ||
73 | |||
74 | charstring :: String -> BEncode | ||
75 | charstring = string . B.pack . map (toEnum . fromEnum) | ||
76 | {-# INLINE charstring #-} | ||
77 | |||
78 | dictAssoc :: [(ByteString, BEncode)] -> BEncode | ||
79 | dictAssoc = dict . M.fromList | ||
80 | {-# INLINE dictAssoc #-} | ||
81 | |||
82 | |||
83 | isInteger :: BEncode -> Bool | ||
84 | isInteger (BInteger _) = True | ||
85 | isInteger _ = False | ||
86 | {-# INLINE isInteger #-} | ||
87 | |||
88 | isString :: BEncode -> Bool | ||
89 | isString (BString _) = True | ||
90 | isString _ = False | ||
91 | {-# INLINE isString #-} | ||
92 | |||
93 | isList :: BEncode -> Bool | ||
94 | isList (BList _) = True | ||
95 | isList _ = False | ||
96 | {-# INLINE isList #-} | ||
97 | |||
98 | isDict :: BEncode -> Bool | ||
99 | isDict (BList _) = True | ||
100 | isDict _ = False | ||
101 | {-# INLINE isDict #-} | ||
102 | |||
103 | encode :: BEncode -> Lazy.ByteString | ||
104 | encode = B.toLazyByteString . builder | ||
105 | |||
106 | decode :: ByteString -> Either String BEncode | ||
107 | decode = P.parseOnly parser | ||
108 | |||
109 | |||
110 | builder :: BEncode -> B.Builder | ||
111 | builder = go | ||
112 | where | ||
113 | go (BInteger i) = B.word8 (c2w 'i') <> | ||
114 | B.intDec (fromIntegral i) <> -- TODO FIXME | ||
115 | B.word8 (c2w 'e') | ||
116 | go (BString s) = buildString s | ||
117 | go (BList l) = B.word8 (c2w 'l') <> | ||
118 | foldMap go l <> | ||
119 | B.word8 (c2w 'e') | ||
120 | go (BDict d) = B.word8 (c2w 'd') <> | ||
121 | foldMap mkKV (M.toAscList d) <> | ||
122 | B.word8 (c2w 'e') | ||
123 | where | ||
124 | mkKV (k, v) = buildString k <> go v | ||
125 | |||
126 | buildString s = B.intDec (B.length s) <> | ||
127 | B.word8 (c2w ':') <> | ||
128 | B.byteString s | ||
129 | {-# INLINE buildString #-} | ||
130 | |||
131 | |||
132 | parser :: Parser BEncode | ||
133 | parser = valueP | ||
134 | where | ||
135 | valueP = do | ||
136 | mc <- P.peekChar | ||
137 | case mc of | ||
138 | Nothing -> fail "end of input" | ||
139 | Just c -> | ||
140 | case c of | ||
141 | -- if we have digit it always should be string length | ||
142 | di | di <= '9' -> BString <$> stringP | ||
143 | 'i' -> P.anyChar *> ((BInteger <$> integerP) <* P.anyChar) | ||
144 | 'l' -> P.anyChar *> ((BList <$> many valueP) <* P.anyChar) | ||
145 | 'd' -> do | ||
146 | P.anyChar | ||
147 | (BDict . M.fromDistinctAscList <$> many ((,) <$> stringP <*> valueP)) | ||
148 | <* P.anyChar | ||
149 | _ -> fail "" | ||
150 | |||
151 | stringP :: Parser ByteString | ||
152 | stringP = do | ||
153 | n <- P.decimal :: Parser Int | ||
154 | P.char ':' | ||
155 | P.take n | ||
156 | {-# INLINE stringP #-} | ||
157 | |||
158 | integerP :: Parser Int64 | ||
159 | integerP = negate <$> (P.char8 '-' *> P.decimal) | ||
160 | <|> P.decimal | ||
161 | {-# INLINE integerP #-} | ||
162 | |||
163 | |||
164 | printPretty :: BEncode -> IO () | ||
165 | printPretty = print . pretty | ||
166 | |||
167 | ppBS :: ByteString -> Doc | ||
168 | ppBS = PP.string . map w2c . B.unpack | ||
169 | |||
170 | instance Pretty BEncode where | ||
171 | pretty (BInteger i) = PP.integer (fromIntegral i) | ||
172 | pretty (BString s) = ppBS s | ||
173 | pretty (BList l) = PP.lbracket <+> | ||
174 | PP.hsep (PP.punctuate PP.comma (map PP.pretty l)) <+> | ||
175 | PP.rbracket | ||
176 | pretty (BDict d) = | ||
177 | PP.align $ PP.lbrace <+> | ||
178 | PP.vsep (PP.punctuate PP.comma (map ppKV (M.toAscList d))) </> | ||
179 | PP.rbrace | ||
180 | where | ||
181 | ppKV (k, v) = ppBS k <+> PP.colon <+> PP.pretty v | ||