Add BitStream pool

CMakeLists.txt

@@ -19,6 +19,7 @@ add_library(${PROJECT_NAME} MODULE
   src/plugin.cc
   src/native_param.h
   src/config.h
+  src/bitstream_pool.h
   src/internal_packet_channel.h
   src/script.h
   src/script.cc

src/bitstream_pool.h

@@ -0,0 +1,65 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2016-2021 katursis
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+#ifndef PAWNRAKNET_BITSTREAM_POOL_H_
+#define PAWNRAKNET_BITSTREAM_POOL_H_
+
+class BitStreamPool {
+ public:
+  BitStream *New() {
+    for (auto &item : items_) {
+      if (!item.second) {
+        item.second = true;
+
+        return item.first.get();
+      }
+    }
+
+    const auto bs = std::make_shared<BitStream>();
+
+    items_.emplace_back(std::make_pair(bs, true));
+
+    return bs.get();
+  }
+
+  void Delete(BitStream *bs) {
+    for (auto &item : items_) {
+      if (item.first.get() == bs) {
+        item.second = false;
+
+        item.first->Reset();
+
+        return;
+      }
+    }
+  }
+
+ private:
+  using Item =
+      std::pair<std::shared_ptr<BitStream> /* bs */, bool /* is_occupied */>;
+
+  std::vector<Item> items_;
+};
+
+#endif  // PAWNRAKNET_BITSTREAM_POOL_H_

src/internal_packet_channel.h

@@ -22,6 +22,9 @@
  * SOFTWARE.
  */
 
+#ifndef PAWNRAKNET_INTERNAL_PACKET_CHANNEL_H_
+#define PAWNRAKNET_INTERNAL_PACKET_CHANNEL_H_
+
 class InternalPacketChannel {
  public:
   void PushPacket(InternalPacket *packet, const PlayerID &player_id,
@@ -84,3 +87,5 @@ class InternalPacketChannel {
 
   bool result_{};
 };
+
+#endif  // PAWNRAKNET_INTERNAL_PACKET_CHANNEL_H_

src/main.h

@@ -44,6 +44,7 @@
 #include <queue>
 #include <thread>
 #include <atomic>
+#include <vector>
 
 #include "Pawn.RakNet.inc"
 
@@ -58,6 +59,7 @@
 #endif
 
 #include "config.h"
+#include "bitstream_pool.h"
 #include "internal_packet_channel.h"
 #include "rakserver.h"
 #include "script.h"

src/script.cc

@@ -111,14 +111,20 @@ cell Script::PR_EmulateIncomingRPC(BitStream *bs, int player_id,
 }
 
 // native BitStream:BS_New();
-cell Script::BS_New() { return NewBitStream(); }
+cell Script::BS_New() { return reinterpret_cast<cell>(bitstream_pool_.New()); }
 
 // native BitStream:BS_NewCopy(BitStream:bs);
-cell Script::BS_NewCopy(BitStream *bs) { return NewBitStreamCopy(bs); }
+cell Script::BS_NewCopy(BitStream *bs) {
+  const auto bs_copy = bitstream_pool_.New();
+
+  bs_copy->Write(bs);
+
+  return reinterpret_cast<cell>(bs_copy);
+}
 
 // native BS_Delete(&BitStream:bs);
 cell Script::BS_Delete(cell *bs) {
-  DeleteBitStream(*bs);
+  bitstream_pool_.Delete(GetBitStream(*bs));
 
   *bs = 0;
 
@@ -585,23 +591,6 @@ void Script::InitHandlers() {
   }
 }
 
-cell Script::NewBitStream() {
-  const auto bs = std::make_shared<BitStream>();
-
-  bitstreams_.insert(bs);
-
-  return reinterpret_cast<cell>(bs.get());
-}
-
-cell Script::NewBitStreamCopy(BitStream *src) {
-  const auto bs = std::make_shared<BitStream>(
-      src->GetData(), src->GetNumberOfBytesUsed(), true);
-
-  bitstreams_.insert(bs);
-
-  return reinterpret_cast<cell>(bs.get());
-}
-
 BitStream *Script::GetBitStream(cell handle) {
   const auto bs = reinterpret_cast<BitStream *>(handle);
   if (!bs) {
@@ -611,23 +600,6 @@ BitStream *Script::GetBitStream(cell handle) {
   return bs;
 }
 
-std::unordered_set<std::shared_ptr<BitStream>>::iterator Script::FindBitStream(
-    cell handle) {
-  const auto bs =
-      std::find_if(bitstreams_.begin(), bitstreams_.end(), [=](const auto &bs) {
-        return reinterpret_cast<cell>(bs.get()) == handle;
-      });
-  if (bs == bitstreams_.end()) {
-    throw std::runtime_error{"Invalid BitStream handle"};
-  }
-
-  return bs;
-}
-
-void Script::DeleteBitStream(cell handle) {
-  bitstreams_.erase(FindBitStream(handle));
-}
-
 template <typename T, bool compressed>
 void Script::WriteValue(BitStream *bs, cell value) {
   T prepared_value{};

src/script.h

@@ -153,17 +153,8 @@ class Script : public ptl::AbstractScript<Script> {
 
   void InitHandlers();
 
-  cell NewBitStream();
-
-  cell NewBitStreamCopy(BitStream *src);
-
-  std::unordered_set<std::shared_ptr<BitStream>>::iterator FindBitStream(
-      cell handle);
-
   BitStream *GetBitStream(cell handle);
 
-  void DeleteBitStream(cell handle);
-
   template <typename T, bool compressed = false>
   void WriteValue(BitStream *bs, cell value);
 
@@ -187,7 +178,7 @@ class Script : public ptl::AbstractScript<Script> {
   PublicPtr public_on_outcoming_packet_;
   PublicPtr public_on_outcoming_rpc_;
 
-  std::unordered_set<std::shared_ptr<BitStream>> bitstreams_;
+  BitStreamPool bitstream_pool_;
 };
 
 #endif  // PAWNRAKNET_SCRIPT_H_