Emit optimal-size LEBs in section/subsection/function body sizes (#1128)

* emit optimal-size LEBs in section/subsection/function body sizes, instead of preallocating 5 bytes

2 files changed, 42 insertions, 12 deletions

diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index daa1dc4fa..55a97743f 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -35,6 +35,11 @@
 
 namespace wasm {
 
+enum {
+  // the maximum amount of bytes we emit per LEB
+  MaxLEB32Bytes = 5
+};
+
 template<typename T, typename MiniT>
 struct LEB {
   static_assert(sizeof(MiniT) == 1, "MiniT must be a byte");
@@ -63,7 +68,9 @@ struct LEB {
     } while (more);
   }
 
-  void writeAt(std::vector<uint8_t>* out, size_t at, size_t minimum = 0) {
+  // @minimum: a minimum number of bytes to write, padding as necessary
+  // returns the number of bytes written
+  size_t writeAt(std::vector<uint8_t>* out, size_t at, size_t minimum = 0) {
     T temp = value;
     size_t offset = 0;
     bool more;
@@ -77,6 +84,7 @@ struct LEB {
       (*out)[at + offset] = byte;
       offset++;
     } while (more);
+    return offset;
   }
 
   void read(std::function<MiniT()> get) {
@@ -258,9 +266,18 @@ public:
     (*this)[i+2] = x & 0xff; x >>= 8;
     (*this)[i+3] = x & 0xff;
   }
-  void writeAt(size_t i, U32LEB x) {
+
+  // writes out an LEB to an arbitrary location. this writes the LEB as a full
+  // 5 bytes, the fixed amount that can easily be set aside ahead of time
+  void writeAtFullFixedSize(size_t i, U32LEB x) {
     if (debug) std::cerr << "backpatchU32LEB: " << x.value << " (at " << i << ")" << std::endl;
-    x.writeAt(this, i, 5); // fill all 5 bytes, we have to do this when backpatching
+    x.writeAt(this, i, MaxLEB32Bytes); // fill all 5 bytes, we have to do this when backpatching
+  }
+  // writes out an LEB of normal size
+  // returns how many bytes were written
+  size_t writeAt(size_t i, U32LEB x) {
+    if (debug) std::cerr << "writeAtU32LEB: " << x.value << " (at " << i << ")" << std::endl;
+    return x.writeAt(this, i);
   }
 
   template <typename T>
@@ -647,7 +664,8 @@ public:
   void writeHeader();
   int32_t writeU32LEBPlaceholder();
   void writeResizableLimits(Address initial, Address maximum, bool hasMaximum, bool shared);
-  int32_t startSection(BinaryConsts::Section code);
+  template<typename T>
+  int32_t startSection(T code);
   void finishSection(int32_t start);
   int32_t startSubsection(BinaryConsts::UserSections::Subsection code);
   void finishSubsection(int32_t start);
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 63dddbdd6..eb41ae8a7 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <algorithm>
 #include <fstream>
 
 #include "support/bits.h"
@@ -84,24 +85,29 @@ void WasmBinaryWriter::writeResizableLimits(Address initial, Address maximum,
   }
 }
 
-int32_t WasmBinaryWriter::startSection(BinaryConsts::Section code) {
+template<typename T>
+int32_t WasmBinaryWriter::startSection(T code) {
   o << U32LEB(code);
   return writeU32LEBPlaceholder(); // section size to be filled in later
 }
 
 void WasmBinaryWriter::finishSection(int32_t start) {
-  int32_t size = o.size() - start - 5; // section size does not include the 5 bytes of the size field itself
-  o.writeAt(start, U32LEB(size));
+  int32_t size = o.size() - start - MaxLEB32Bytes; // section size does not include the reserved bytes of the size field itself
+  auto sizeFieldSize = o.writeAt(start, U32LEB(size));
+  if (sizeFieldSize != MaxLEB32Bytes) {
+    // we can save some room, nice
+    assert(sizeFieldSize < MaxLEB32Bytes);
+    std::move(&o[start + MaxLEB32Bytes], &o[start + MaxLEB32Bytes + size], &o[start + sizeFieldSize]);
+    o.resize(o.size() - (MaxLEB32Bytes - sizeFieldSize));
+  }
 }
 
 int32_t WasmBinaryWriter::startSubsection(BinaryConsts::UserSections::Subsection code) {
-  o << U32LEB(code);
-  return writeU32LEBPlaceholder(); // section size to be filled in later
+  return startSection(code);
 }
 
 void WasmBinaryWriter::finishSubsection(int32_t start) {
-  int32_t size = o.size() - start - 5; // section size does not include the 5 bytes of the size field itself
-  o.writeAt(start, U32LEB(size));
+  finishSection(start);
 }
 
 void WasmBinaryWriter::writeStart() {
@@ -270,7 +276,13 @@ void WasmBinaryWriter::writeFunctions() {
     size_t size = o.size() - start;
     assert(size <= std::numeric_limits<uint32_t>::max());
     if (debug) std::cerr << "body size: " << size << ", writing at " << sizePos << ", next starts at " << o.size() << std::endl;
-    o.writeAt(sizePos, U32LEB(size));
+    auto sizeFieldSize = o.writeAt(sizePos, U32LEB(size));
+    if (sizeFieldSize != MaxLEB32Bytes) {
+      // we can save some room, nice
+      assert(sizeFieldSize < MaxLEB32Bytes);
+      std::move(&o[start], &o[start + size], &o[sizePos + sizeFieldSize]);
+      o.resize(o.size() - (MaxLEB32Bytes - sizeFieldSize));
+    }
   }
   currFunction = nullptr;
   finishSection(start);