Merge branch 'dht-rewrite' of jerkface.net:repo/bittorrent into dht-rewrite

1 files changed, 11 insertions, 5 deletions

diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index a25f9f4f..9f86f6a4 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -7,7 +7,7 @@
 {-# LANGUAGE DeriveTraversable #-}
 {-# LANGUAGE ExplicitNamespaces #-}
 {-# LANGUAGE TypeOperators #-}
-{-# LANGUAGE MagicHash, UnboxedTuples #-}
+{-# LANGUAGE MagicHash, UnboxedTuples, BangPatterns #-}
 module Crypto.Tox
     ( PublicKey
     , publicKey
@@ -35,6 +35,7 @@ module Crypto.Tox
     , Nonce8(..)
     , Nonce24(..)
     , incrementNonce24
+    , addtoNonce24
     , Nonce32(..)
     , getRemainingEncrypted
     , putEncrypted
@@ -258,17 +259,18 @@ hsalsa20 k n = BA.append a b
 newtype Nonce24 = Nonce24 ByteString
  deriving (Eq, Ord, ByteArrayAccess,Data)
 
-incrementNonce24 :: Nonce24 -> IO Nonce24
-incrementNonce24 (Nonce24 n24) = Nonce24 <$> BA.copy n24 init
+addtoNonce24 :: Nonce24 -> Word -> IO Nonce24
+addtoNonce24 (Nonce24 n24) n = Nonce24 <$> BA.copy n24 init
     where
         init :: Ptr Word -> IO ()
         init ptr | fitsInInt (Proxy :: Proxy Word64) = do
             let frmBE64 = fromIntegral . fromBE64 . fromIntegral
                 tBE64   = fromIntegral . toBE64   . fromIntegral
+                !(W# input) = n
             W# w1 <- frmBE64 <$> peek ptr
             W# w2 <- frmBE64 <$> peekElemOff ptr 1
             W# w3 <- frmBE64 <$> peekElemOff ptr 2
-            let (# overflw,  sum   #) = plusWord2# w3 (int2Word# 1#)
+            let (# overflw,  sum   #) = plusWord2# w3 input
                 (# overflw', sum'  #) = plusWord2# w2 overflw
                 (# discard,  sum'' #) = plusWord2# w1 overflw'
             poke        ptr   $ tBE64 (W# sum'')
@@ -278,13 +280,14 @@ incrementNonce24 (Nonce24 n24) = Nonce24 <$> BA.copy n24 init
         init ptr | fitsInInt (Proxy :: Proxy Word32) = do
             let frmBE32 = fromIntegral . fromBE32 . fromIntegral
                 tBE32   = fromIntegral . toBE32   . fromIntegral
+                !(W# input) = n
             W# w1 <- frmBE32 <$> peek ptr
             W# w2 <- frmBE32 <$> peekElemOff ptr 1
             W# w3 <- frmBE32 <$> peekElemOff ptr 2
             W# w4 <- frmBE32 <$> peekElemOff ptr 3
             W# w5 <- frmBE32 <$> peekElemOff ptr 4
             W# w6 <- frmBE32 <$> peekElemOff ptr 5
-            let (# overflw_,   sum_     #) = plusWord2# w6 (int2Word# 1#)
+            let (# overflw_,   sum_     #) = plusWord2# w6 input
                 (# overflw__,  sum__    #) = plusWord2# w5 overflw_
                 (# overflw___, sum___   #) = plusWord2# w6 overflw__
                 (# overflw,  sum   #) = plusWord2# w3 overflw___
@@ -298,6 +301,9 @@ incrementNonce24 (Nonce24 n24) = Nonce24 <$> BA.copy n24 init
             pokeElemOff ptr 5 $ tBE32 (W# sum_)
         init _ = error "incrementNonce24: I only support 64 and 32 bits"
 
+incrementNonce24 :: Nonce24 -> IO Nonce24
+incrementNonce24 nonce24 = addtoNonce24 nonce24 1
+
 quoted :: ShowS -> ShowS
 quoted shows s = '"':shows ('"':s)