From 76a4e82eb11e515577c251248f1409538611d14a Mon Sep 17 00:00:00 2001
From: Dmitry Ivanov <ethercrow@gmail.com>
Date: Sun, 29 Jan 2017 19:55:54 +0100
Subject: [PATCH] Faster masking, also add a benchmark

---
 benchmarks/mask.hs                    | 64 +++++++++++++++++++++++++++
 src/Network/WebSockets/Hybi13/Mask.hs | 16 +++++--
 websockets.cabal                      | 10 +++++
 3 files changed, 86 insertions(+), 4 deletions(-)
 create mode 100644 benchmarks/mask.hs

diff --git a/benchmarks/mask.hs b/benchmarks/mask.hs
new file mode 100644
index 0000000..baa52f5
--- /dev/null
+++ b/benchmarks/mask.hs
@@ -0,0 +1,64 @@
+{-# language BangPatterns #-}
+{-# language OverloadedStrings #-}
+
+import Criterion
+import Criterion.Main
+
+import Network.WebSockets.Hybi13.Mask
+
+import Data.Bits (shiftR, xor)
+import qualified Data.ByteString as B
+import qualified Data.ByteString.Lazy as BL
+
+setupEnv = do
+    let kilo = BL.replicate 1024 37
+        mega = BL.replicate (1024 * 1024) 37
+    return (kilo, mega)
+
+maskPayload' :: Mask -> BL.ByteString -> BL.ByteString
+maskPayload' Nothing     = id
+maskPayload' (Just mask) = snd . BL.mapAccumL f (cycle $ B.unpack mask)
+  where
+    f []     !c = ([], c)
+    f (m:ms) !c = (ms, m `xor` c)
+
+main = defaultMain [
+    env setupEnv $ \ ~(kilo, mega) -> bgroup "main"
+        [ bgroup "kilobyte payload"
+            [ bgroup "zero_mask"
+                [ bench "current" $ nf (maskPayload (Just "\x00\x00\x00\x00")) kilo
+                , bench "old" $ nf (maskPayload' (Just "\x00\x00\x00\x00")) kilo
+                ]
+            ,  bgroup "full_mask"
+                [ bench "current" $ nf (maskPayload (Just "\xFF\xFF\xFF\xFF")) kilo
+                , bench "old" $ nf (maskPayload' (Just "\xFF\xFF\xFF\xFF")) kilo
+                ]
+            ,  bgroup "one_byte_mask"
+                [ bench "current" $ nf (maskPayload (Just "\xCC\xCC\xCC\xCC")) kilo
+                , bench "old" $ nf (maskPayload' (Just "\xCC\xCC\xCC\xCC")) kilo
+                ]
+            ,  bgroup "other_mask"
+                [ bench "current" $ nf (maskPayload (Just "\xB0\xA2\xB0\xA2")) kilo
+                , bench "old" $ nf (maskPayload' (Just "\xB0\xA2\xB0\xA2")) kilo
+                ]
+            ]
+        , bgroup "megabyte payload"
+            [ bgroup "zero_mask"
+                [ bench "current" $ nf (maskPayload (Just "\x00\x00\x00\x00")) mega
+                , bench "old" $ nf (maskPayload' (Just "\x00\x00\x00\x00")) mega
+                ]
+            ,  bgroup "full_mask"
+                [ bench "current" $ nf (maskPayload (Just "\xFF\xFF\xFF\xFF")) mega
+                , bench "old" $ nf (maskPayload' (Just "\xFF\xFF\xFF\xFF")) mega
+                ]
+            ,  bgroup "one_byte_mask"
+                [ bench "current" $ nf (maskPayload (Just "\xCC\xCC\xCC\xCC")) mega
+                , bench "old" $ nf (maskPayload' (Just "\xCC\xCC\xCC\xCC")) mega
+                ]
+            ,  bgroup "other_mask"
+                [ bench "current" $ nf (maskPayload (Just "\xB0\xA2\xB0\xA2")) mega
+                , bench "old" $ nf (maskPayload' (Just "\xB0\xA2\xB0\xA2")) mega
+                ]
+            ]
+        ]
+    ]
\ No newline at end of file
diff --git a/src/Network/WebSockets/Hybi13/Mask.hs b/src/Network/WebSockets/Hybi13/Mask.hs
index df57fcf..7a6256f 100644
--- a/src/Network/WebSockets/Hybi13/Mask.hs
+++ b/src/Network/WebSockets/Hybi13/Mask.hs
@@ -2,6 +2,7 @@
 -- | Masking of fragmes using a simple XOR algorithm
 {-# LANGUAGE BangPatterns        #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# language OverloadedStrings   #-}
 module Network.WebSockets.Hybi13.Mask
     ( Mask
     , maskPayload
@@ -24,11 +25,18 @@ type Mask = Maybe B.ByteString
 --------------------------------------------------------------------------------
 -- | Apply mask
 maskPayload :: Mask -> BL.ByteString -> BL.ByteString
-maskPayload Nothing     = id
-maskPayload (Just mask) = snd . BL.mapAccumL f (cycle $ B.unpack mask)
+maskPayload Nothing = id
+maskPayload (Just "\x00\x00\x00\x00") = id
+maskPayload (Just mask) =
+  BL.fromChunks . go (cycle (B.unpack mask)) . BL.toChunks
   where
-    f []     !c = ([], c)
-    f (m:ms) !c = (ms, m `xor` c)
+  go _ [] = []
+  go ms (chunk : chunks) =
+      let (ms', chunk') = B.mapAccumL f ms chunk
+      in chunk' : go ms' chunks
+  f (m : ms) c = (ms, m `xor` c)
+  f [] _ = error "impossible, we have infinite stream of mask bytes"
+
 
 --------------------------------------------------------------------------------
 -- | Create a random mask
diff --git a/websockets.cabal b/websockets.cabal
index c703915..fcc1c25 100644
--- a/websockets.cabal
+++ b/websockets.cabal
@@ -156,3 +156,13 @@ Executable websockets-example
 Source-repository head
   Type:     git
   Location: https://github.com/jaspervdj/websockets
+
+Benchmark bench-mask
+    type:       exitcode-stdio-1.0
+    main-is:    Mask.hs
+    hs-source-dirs: benchmarks, src
+    build-depends:
+        base,
+        bytestring,
+        criterion,
+        random
\ No newline at end of file