From: PHO <pho@cielonegro.org>
Date: Wed, 20 May 2009 03:24:31 +0000 (+0900)
Subject: Clean-up with binary-strict's BitGet and BitPut.
X-Git-Url: http://git.cielonegro.org/gitweb.cgi?p=haskell-dns.git;a=commitdiff_plain;h=44ecef06a0a5d639ceefc761857558c1e1ac08a7

Clean-up with binary-strict's BitGet and BitPut.
---

diff --git a/Network/DNS/Message.hs b/Network/DNS/Message.hs
index 71fec55..2bb5a80 100644
--- a/Network/DNS/Message.hs
+++ b/Network/DNS/Message.hs
@@ -11,10 +11,12 @@ module Network.DNS.Message
     )
     where
 
+import           Control.Monad
 import           Data.Binary
-import           Data.Binary.Get
+import           Data.Binary.BitPut
+import           Data.Binary.Get as G
 import           Data.Binary.Put
-import           Data.Bits
+import           Data.Binary.Strict.BitGet as BG
 import           Data.Word
 
 
@@ -39,10 +41,10 @@ data Header
       , hdResponseCode          :: !ResponseCode
 
       -- These fields are supressed in this data structure:
-      -- * QDCOUNT
-      -- * ANCOUNT
-      -- * NSCOUNT
-      -- * ARCOUNT
+      -- + QDCOUNT
+      -- + ANCOUNT
+      -- + NSCOUNT
+      -- + ARCOUNT
       }
 
 data MessageType
@@ -120,45 +122,40 @@ data RecordType
 
 instance Binary Header where
     put h = do putWord16be $ hdMessageID h
-               let qr    = fromIntegral $ fromEnum $ hdMessageType h
-                   op    = fromIntegral $ fromEnum $ hdOpcode h
-                   aa    = boolToNum $ hdIsAuthoritativeAnswer h
-                   tc    = boolToNum $ hdIsTruncated h
-                   rd    = boolToNum $ hdIsRecursionDesired h
-                   ra    = boolToNum $ hdIsRecursionAvailable h
-                   rc    = fromIntegral $ fromEnum $ hdResponseCode h
-                   flags = ((qr `shiftL` 15) .&. 0x01) .|.
-                           ((op `shiftL` 11) .&. 0x0F) .|.
-                           ((aa `shiftL` 10) .&. 0x01) .|.
-                           ((tc `shiftL`  9) .&. 0x01) .|.
-                           ((rd `shiftL`  8) .&. 0x01) .|.
-                           ((ra `shiftL`  7) .&. 0x01) .|.
-                           ((rc `shiftL`  0) .&. 0x0F)
-               putWord16be flags
+               putLazyByteString flags
         where
-          boolToNum :: Num a => Bool -> a
-          boolToNum True  = 1
-          boolToNum False = 0
-
-    get = do mID   <- getWord16be
-             flags <- getWord16be
-             let qr = toEnum $ fromIntegral ((flags `shiftR` 15) .&. 0x01)
-                 op = toEnum $ fromIntegral ((flags `shiftR` 11) .&. 0x0F)
-                 aa = testBit flags 10
-                 tc = testBit flags 9
-                 rd = testBit flags 8
-                 ra = testBit flags 7
-                 rc = toEnum $ fromIntegral (flags .&. 0x0F)
-                 hd = Header {
-                         hdMessageID             = mID
-                       , hdMessageType           = qr
-                       , hdOpcode                = op
-                       , hdIsAuthoritativeAnswer = aa
-                       , hdIsTruncated           = tc
-                       , hdIsRecursionDesired    = rd
-                       , hdIsRecursionAvailable  = ra
-                       , hdResponseCode          = rc
-                       }
+          flags = runBitPut $
+                  do putNBits 1 $ fromEnum $ hdMessageType h
+                     putNBits 4 $ fromEnum $ hdOpcode h
+                     putBit $ hdIsAuthoritativeAnswer h
+                     putBit $ hdIsTruncated h
+                     putBit $ hdIsRecursionDesired h
+                     putBit $ hdIsRecursionAvailable h
+                     putNBits 3 (0 :: Int)
+                     putNBits 4 $ fromEnum $ hdResponseCode h
+
+    get = do mID   <- G.getWord16be
+             flags <- getByteString 2
+             let Right hd
+                     = runBitGet flags $
+                       do qr <- liftM (toEnum . fromIntegral) $ getAsWord8 1
+                          op <- liftM (toEnum . fromIntegral) $ getAsWord8 4
+                          aa <- getBit
+                          tc <- getBit
+                          rd <- getBit
+                          ra <- getBit
+                          BG.skip 3
+                          rc <- liftM (toEnum . fromIntegral) $ getAsWord8 4
+                          return Header {
+                                       hdMessageID             = mID
+                                     , hdMessageType           = qr
+                                     , hdOpcode                = op
+                                     , hdIsAuthoritativeAnswer = aa
+                                     , hdIsTruncated           = tc
+                                     , hdIsRecursionDesired    = rd
+                                     , hdIsRecursionAvailable  = ra
+                                     , hdResponseCode          = rc
+                                     }
              return hd
 
 instance Enum MessageType where
@@ -239,7 +236,7 @@ instance Enum RecordType where
     toEnum 255 = AnyType
     toEnum _  = undefined
 
-instance Enum RecordClass
+instance Enum RecordClass where
     fromEnum IN       = 1
     fromEnum CS       = 2
     fromEnum CH       = 3
diff --git a/dns.cabal b/dns.cabal
index 93ab364..92cb60c 100644
--- a/dns.cabal
+++ b/dns.cabal
@@ -12,7 +12,7 @@ Build-Type:          Simple
 
 Library
     Build-Depends:
-        base, binary
+        base, binary, binary-strict
 
     Exposed-Modules:
         Network.DNS.Message