add `--limit` option

lib/kyut/Reordering-inl.hpp

@@ -51,6 +51,7 @@ namespace kyut {
         template <typename RandomAccessIterator, typename Less>
         inline std::size_t embed_by_reordering(
             CircularBitStreamReader& r,
+            std::size_t limit,
             std::size_t chunk_size,
             RandomAccessIterator begin,
             RandomAccessIterator end,
@@ -67,6 +68,10 @@ namespace kyut {
                 const auto chunk_end = chunk_begin + n;
 
                 size_bits += embed_in_chunk(r, chunk_begin, chunk_end, less);
+
+                if (size_bits >= limit) {
+                    break;
+                }
             }
 
             return size_bits;
@@ -170,11 +175,12 @@ namespace kyut {
     template <typename RandomAccessIterator, typename Less>
     inline std::size_t embed_by_reordering(
         CircularBitStreamReader& r,
+        std::size_t limit,
         std::size_t chunk_size,
         RandomAccessIterator begin,
         RandomAccessIterator end,
         Less less) {
-        return detail::embed_by_reordering(r, chunk_size, begin, end, less);
+        return detail::embed_by_reordering(r, limit, chunk_size, begin, end, less);
     }
 
     template <typename RandomAccessIterator, typename Less>

lib/kyut/Reordering.hpp

@@ -12,6 +12,7 @@ namespace kyut {
     template <typename RandomAccessIterator, typename Less>
     std::size_t embed_by_reordering(
         CircularBitStreamReader& r,
+        std::size_t limit,
         std::size_t chunk_size,
         RandomAccessIterator begin,
         RandomAccessIterator end,

lib/kyut/methods/ExportReordering.hpp

@@ -10,9 +10,10 @@ namespace kyut {
 } // namespace kyut
 
 namespace kyut::methods::export_reordering {
-    inline std::size_t embed(CircularBitStreamReader& r, wasm::Module& module, std::size_t chunk_size) {
+    inline std::size_t embed(CircularBitStreamReader& r, wasm::Module& module, std::size_t limit, std::size_t chunk_size) {
         const auto size_bits = embed_by_reordering(
             r,
+            limit,
             chunk_size,
             std::begin(module.exports),
             std::end(module.exports),

lib/kyut/methods/FunctionReordering.hpp

@@ -10,7 +10,7 @@ namespace kyut {
 } // namespace kyut
 
 namespace kyut::methods::function_reordering {
-    inline std::size_t embed(CircularBitStreamReader& r, wasm::Module& module, std::size_t chunk_size) {
+    inline std::size_t embed(CircularBitStreamReader& r, wasm::Module& module, std::size_t limit, std::size_t chunk_size) {
         const auto begin = std::begin(module.functions);
         const auto end = std::end(module.functions);
 
@@ -20,6 +20,7 @@ namespace kyut::methods::function_reordering {
 
         const auto size_bits = embed_by_reordering(
             r,
+            limit,
             chunk_size,
             start,
             end,

lib/kyut/methods/OperandSwapping.cpp

@@ -335,7 +335,7 @@ namespace kyut::methods::operand_swapping {
         };
     } // namespace
 
-    std::size_t embed(CircularBitStreamReader& r, wasm::Module& module) {
+    std::size_t embed(CircularBitStreamReader& r, wasm::Module& module, std::size_t limit) {
         std::vector<wasm::Function*> functions{};
         functions.reserve(module.functions.size());
 
@@ -373,6 +373,10 @@ namespace kyut::methods::operand_swapping {
 
         for (const auto& f : functions) {
             visitor.visitFunction(f);
+
+            if (size_bits >= limit) {
+                break;
+            }
         }
 
         return size_bits;

lib/kyut/methods/OperandSwapping.hpp

@@ -13,7 +13,7 @@ namespace kyut {
 } // namespace kyut
 
 namespace kyut::methods::operand_swapping {
-    std::size_t embed(CircularBitStreamReader& r, wasm::Module& module);
+    std::size_t embed(CircularBitStreamReader& r, wasm::Module& module, std::size_t limit);
 
     std::size_t extract(BitStreamWriter& w, wasm::Module& module);
 } // namespace kyut::methods::operand_swapping

src/snpi.cpp

@@ -20,6 +20,7 @@ int main(int argc, char* argv[]) {
     options.add<std::string>("method", 'm', "Embedding method (function-reorder, export-reorder, operand-swap, null)", true, "", cmdline::oneof<std::string>("function-reorder", "export-reorder", "operand-swap", "null"));
     options.add<std::string>("watermark", 'w', "Watermark to embed", true);
     options.add<std::size_t>("chunk-size", 'c', "Chunk size [2~20]", false, 20, cmdline::range<std::size_t>(2, 20));
+    options.add<std::size_t>("limit", 'l', "Embedding limit", false, std::size_t(-1));
     options.add("debug", 'd', "Preserve debug info");
 
     options.set_program_name(program_name);
@@ -61,6 +62,7 @@ int main(int argc, char* argv[]) {
     const auto method = options.get<std::string>("method");
     const auto watermark = options.get<std::string>("watermark");
     const auto chunk_size = options.get<std::size_t>("chunk-size");
+    const auto limit = options.get<std::size_t>("limit");
     const auto preserve_debug = options.exist("debug");
 
     try {
@@ -71,11 +73,11 @@ int main(int argc, char* argv[]) {
 
         std::size_t size_bits;
         if (method == "function-reorder") {
-            size_bits = kyut::methods::function_reordering::embed(r, module, chunk_size);
+            size_bits = kyut::methods::function_reordering::embed(r, module, limit, chunk_size);
         } else if (method == "export-reorder") {
-            size_bits = kyut::methods::export_reordering::embed(r, module, chunk_size);
+            size_bits = kyut::methods::export_reordering::embed(r, module, limit, chunk_size);
         } else if (method == "operand-swap") {
-            size_bits = kyut::methods::operand_swapping::embed(r, module);
+            size_bits = kyut::methods::operand_swapping::embed(r, module, limit);
         } else if (method == "null") {
             size_bits = 0; /* Don't do anything */
         } else {

test/test_Reordering.cpp

@@ -13,6 +13,7 @@ namespace {
 
         const auto size_bits = kyut::embed_by_reordering(
             r,
+            std::size_t(-1),
             chunk_size,
             std::begin(data),
             std::end(data),
@@ -88,6 +89,7 @@ namespace {
 
         const auto size_bits_embedded = kyut::embed_by_reordering(
             r,
+            std::size_t(-1),
             chunk_size,
             std::begin(data),
             std::end(data),