2 files changed, 46 insertions, 0 deletions

diff --git a/dht-client.cabal b/dht-client.cabal
index de6eb33f..6bd680cc 100644
--- a/dht-client.cabal
+++ b/dht-client.cabal
@@ -134,6 +134,7 @@ library
                     , monad-control
                     , transformers-base
                     , mtl
+                    , ghc-prim
 
   if impl(ghc < 8)
       Build-depends:  transformers
diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index 8a65dfb4..9f7c5e16 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -7,6 +7,7 @@
 {-# LANGUAGE DeriveTraversable #-}
 {-# LANGUAGE ExplicitNamespaces #-}
 {-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE MagicHash, UnboxedTuples #-}
 module Crypto.Tox
     ( PublicKey
     , publicKey
@@ -33,6 +34,7 @@ module Crypto.Tox
     , decrypt
     , Nonce8(..)
     , Nonce24(..)
+    , incrementNonce24
     , Nonce32(..)
     , getRemainingEncrypted
     , putEncrypted
@@ -82,6 +84,9 @@ import qualified Data.ByteString.Internal
 import Control.Concurrent.STM
 import Crypto.Error.Types                     (CryptoFailable (..), throwCryptoError)
 import Network.Socket (SockAddr)
+import GHC.Exts (Word(..))
+import GHC.Prim
+import Data.Word64Map (fitsInInt)
 
 -- | A 16-byte mac and an arbitrary-length encrypted stream.
 newtype Encrypted a = Encrypted ByteString
@@ -253,6 +258,46 @@ 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
+    where
+        init :: Ptr Word -> IO ()
+        init ptr | fitsInInt (Proxy :: Proxy Word64) = do
+            let frmBE64 = fromIntegral . fromBE64 . fromIntegral
+                tBE64   = fromIntegral . toBE64   . fromIntegral
+            W# w1 <- frmBE64 <$> peek ptr
+            W# w2 <- frmBE64 <$> peekElemOff ptr 1
+            W# w3 <- frmBE64 <$> peekElemOff ptr 2
+            let (# overflw,  sum   #) = plusWord2# w3 (int2Word# 1#)
+                (# overflw', sum'  #) = plusWord2# w2 overflw
+                (# discard,  sum'' #) = plusWord2# w1 overflw'
+            poke        ptr   $ tBE64 (W# sum'')
+            pokeElemOff ptr 1 $ tBE64 (W# sum')
+            pokeElemOff ptr 2 $ tBE64 (W# sum)
+
+        init ptr | fitsInInt (Proxy :: Proxy Word32) = do
+            let frmBE32 = fromIntegral . fromBE32 . fromIntegral
+                tBE32   = fromIntegral . toBE32   . fromIntegral
+            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#)
+                (# overflw__,  sum__    #) = plusWord2# w5 overflw_
+                (# overflw___, sum___   #) = plusWord2# w6 overflw__
+                (# overflw,  sum   #) = plusWord2# w3 overflw___
+                (# overflw', sum'  #) = plusWord2# w2 overflw
+                (# discard,  sum'' #) = plusWord2# w1 overflw'
+            poke        ptr   $ tBE32 (W# sum'')
+            pokeElemOff ptr 1 $ tBE32 (W# sum')
+            pokeElemOff ptr 2 $ tBE32 (W# sum)
+            pokeElemOff ptr 3 $ tBE32 (W# sum___)
+            pokeElemOff ptr 4 $ tBE32 (W# sum__)
+            pokeElemOff ptr 5 $ tBE32 (W# sum_)
+        init _ = error "incrementNonce24: I only support 64 and 32 bits"
+
 quoted :: ShowS -> ShowS
 quoted shows s = '"':shows ('"':s)