From e9663b3371ca0d138210d11683fed1ad184ea991 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Sun, 7 Jan 2024 22:47:29 +0000
Subject: [PATCH] refactor

---
 src/Simplex/Messaging/Client.hs   | 29 ++++++++++++-----------
 src/Simplex/Messaging/Protocol.hs | 38 ++++++++++++++++---------------
 2 files changed, 34 insertions(+), 33 deletions(-)

diff --git a/src/Simplex/Messaging/Client.hs b/src/Simplex/Messaging/Client.hs
index 996e0e3c3..471a9dbaa 100644
--- a/src/Simplex/Messaging/Client.hs
+++ b/src/Simplex/Messaging/Client.hs
@@ -676,30 +676,29 @@ data ClientBatch err msg
   | CBLargeTransmission (Request err msg)
 
 batchClientTransmissions :: forall err msg. Bool -> Int -> NonEmpty (PCTransmission err msg) -> [ClientBatch err msg]
-batchClientTransmissions batch blkSize ts
-  | batch =
-      let (bs, b, _, n, rs) = foldr addToBatch ([], mempty, 0, 0, []) ts
-       in if n == 0 then bs else CBTransmissions b n rs : bs
-  | otherwise = map mkBatch1 $ L.toList ts
+batchClientTransmissions batch bSize
+  | batch = addBatch . foldr addTransmission ([], mempty, 0, 0, [])
+  | otherwise = map mkBatch1 . L.toList
   where
     mkBatch1 :: PCTransmission err msg -> ClientBatch err msg
     mkBatch1 (t, r)
       -- 2 bytes are reserved for pad size
-      | LB.length s <= fromIntegral (blkSize - 2) = CBTransmission (lazyByteString s) r
+      | LB.length s <= fromIntegral (bSize - 2) = CBTransmission (lazyByteString s) r
       | otherwise = CBLargeTransmission r
       where
         s = tEncode t
-    addToBatch :: PCTransmission err msg -> ([ClientBatch err msg], Builder, Int, Int, [Request err msg]) -> ([ClientBatch err msg], Builder, Int, Int, [Request err msg])
-    addToBatch (t, r) (bs, b, len, n, rs)
-      | len' <= blkSize - 3 && n < 255 = (bs, s <> b, len', 1 + n, r : rs)
-      | sLen <= blkSize - 3 = (bs', s, sLen, 1, [r])
-      | otherwise = (CBLargeTransmission r : (if n == 0 then bs else bs'), mempty, 0, 0, [])
+    addTransmission :: PCTransmission err msg -> ([ClientBatch err msg], Builder, Int, Int, [Request err msg]) -> ([ClientBatch err msg], Builder, Int, Int, [Request err msg])
+    addTransmission (t, r) acc@(bs, b, len, n, rs)
+      | len' <= bSize - 3 && n < 255 = (bs, s <> b, len', 1 + n, r : rs)
+      | sLen <= bSize - 3 = (addBatch acc, s, sLen, 1, [r])
+      | otherwise = (CBLargeTransmission r : addBatch acc, mempty, 0, 0, [])
       where
-        s = encodeLarge s'
-        sLen = 2 + (fromIntegral $ LB.length s')
-        s' = tEncode t
+        s = encodeLarge t'
+        sLen = 2 + fromIntegral (LB.length t') -- 2-bytes length is added by encodeLarge
+        t' = tEncode t
         len' = sLen + len
-        bs' = CBTransmissions b n rs : bs
+    addBatch :: ([ClientBatch err msg], Builder, Int, Int, [Request err msg]) -> [ClientBatch err msg]
+    addBatch (bs, b, _, n, rs) = if n == 0 then bs else CBTransmissions b n rs : bs
 
 -- | Send Protocol command
 sendProtocolCommand :: forall err msg. ProtocolEncoding err (ProtoCommand msg) => ProtocolClient err msg -> Maybe C.APrivateSignKey -> EntityId -> ProtoCommand msg -> ExceptT (ProtocolClientError err) IO msg
diff --git a/src/Simplex/Messaging/Protocol.hs b/src/Simplex/Messaging/Protocol.hs
index 8afd44e56..6a16ea557 100644
--- a/src/Simplex/Messaging/Protocol.hs
+++ b/src/Simplex/Messaging/Protocol.hs
@@ -1308,27 +1308,29 @@ data TransportBatch = TBTransmissions Int Builder | TBTransmission Builder | TBL
 
 -- | encodes and batches transmissions into blocks,
 batchTransmissions :: Bool -> Int -> NonEmpty SentRawTransmission -> [TransportBatch]
-batchTransmissions batch bSize ts
-  | batch =
-      let (bs, b, _, n) = foldr addToBatch ([], mempty, 0, 0) ts
-        in if n == 0 then bs else TBTransmissions n b : bs
-  | otherwise = map (mkBatch1 . tEncode) (L.toList ts)
+batchTransmissions batch bSize
+  | batch = addBatch . foldr addTransmission ([], mempty, 0, 0)
+  | otherwise = map mkBatch1 . L.toList
   where
-    mkBatch1 :: LB.ByteString -> TransportBatch
-    mkBatch1 s
-      | LB.length s > fromIntegral (bSize - 2) = TBLargeTransmission
-      | otherwise = TBTransmission $ lazyByteString s
-    addToBatch :: SentRawTransmission -> ([TransportBatch], Builder, Int, Int) -> ([TransportBatch], Builder, Int, Int)
-    addToBatch t (bs, b, len, n)
-      | len' <= bSize - 3 && n < 255 = (bs, s <> b, len', 1 + n)
-      | sLen <= bSize - 3 = (bs', s, sLen, 1)
-      | otherwise = (TBLargeTransmission : (if n == 0 then bs else bs'), mempty, 0, 0)
+    mkBatch1 :: SentRawTransmission -> TransportBatch
+    mkBatch1 t
+      -- 2 bytes are reserved for pad size
+      | LB.length s <= fromIntegral (bSize - 2) = TBTransmission (lazyByteString s)
+      | otherwise = TBLargeTransmission
       where
-        s = encodeLarge s'
-        sLen = 2 + fromIntegral (LB.length s')
-        s' = tEncode t
+        s = tEncode t
+    addTransmission :: SentRawTransmission -> ([TransportBatch], Builder, Int, Int) -> ([TransportBatch], Builder, Int, Int)
+    addTransmission t acc@(bs, b, len, n)
+      | len' <= bSize - 3 && n < 255 = (bs, s <> b, len', 1 + n)
+      | sLen <= bSize - 3 = (addBatch acc, s, sLen, 1)
+      | otherwise = (TBLargeTransmission : addBatch acc, mempty, 0, 0)
+      where
+        s = encodeLarge t'
+        sLen = 2 + fromIntegral (LB.length t') -- 2-bytes length is added by encodeLarge
+        t' = tEncode t
         len' = sLen + len
-        bs' = TBTransmissions n b : bs
+    addBatch :: ([TransportBatch], Builder, Int, Int) -> [TransportBatch]
+    addBatch (bs, b, _, n) = if n == 0 then bs else TBTransmissions n b : bs
 
 tEncode :: SentRawTransmission -> LB.ByteString
 tEncode (sig, t) = LB.chunk (smpEncode $ C.signatureBytes sig) (LB.fromStrict t)