Bug 1551473 - Provide utility methods to read brotli stream in BinASTTokenReaderContext. r=Yoric

Differential Revision: https://phabricator.services.mozilla.com/D31212

--HG--
extra : moz-landing-system : lando

js/src/frontend/BinASTTokenReaderContext.cpp

@@ -8,7 +8,7 @@
 
 #include "mozilla/Result.h"  // MOZ_TRY*
 
-#include <string.h>  // memcmp
+#include <string.h>  // memcmp, memmove
 
 #include "frontend/BinAST-macros.h"  // BINJS_TRY*, BINJS_MOZ_TRY*
 #include "vm/JSScript.h"             // ScriptSource
@@ -43,6 +43,72 @@ BinASTTokenReaderContext::~BinASTTokenReaderContext() {
   if (metadata_ && metadataOwned_ == MetadataOwnership::Owned) {
     UniqueBinASTSourceMetadataPtr ptr(metadata_);
   }
+  if (decoder_) {
+    BrotliDecoderDestroyInstance(decoder_);
+  }
+}
+
+template <>
+JS::Result<Ok>
+BinASTTokenReaderContext::readBuf<BinASTTokenReaderContext::Compression::No>(
+    uint8_t* bytes, uint32_t len) {
+  return Base::readBuf(bytes, len);
+}
+
+template <>
+JS::Result<Ok>
+BinASTTokenReaderContext::readBuf<BinASTTokenReaderContext::Compression::Yes>(
+    uint8_t* bytes, uint32_t len) {
+  while (availableDecodedLength() < len) {
+    if (availableDecodedLength()) {
+      memmove(bytes, decodedBufferBegin(), availableDecodedLength());
+      bytes += availableDecodedLength();
+      len -= availableDecodedLength();
+    }
+
+    if (isEOF()) {
+      return raiseError("Unexpected end of file");
+    }
+
+    // We have exhausted the in-memory buffer. Start from the beginning.
+    decodedBegin_ = 0;
+
+    size_t inSize = stop_ - current_;
+    size_t outSize = DECODED_BUFFER_SIZE;
+    uint8_t* out = decodedBuffer_;
+
+    BrotliDecoderResult result;
+    result = BrotliDecoderDecompressStream(decoder_, &inSize, &current_,
+                                           &outSize, &out,
+                                           /* total_out = */ nullptr);
+    if (result == BROTLI_DECODER_RESULT_ERROR) {
+      return raiseError("Failed to decompress brotli stream");
+    }
+
+    decodedEnd_ = out - decodedBuffer_;
+  }
+
+  memmove(bytes, decodedBufferBegin(), len);
+  decodedBegin_ += len;
+  return Ok();
+}
+
+bool BinASTTokenReaderContext::isEOF() const {
+  return BrotliDecoderIsFinished(decoder_);
+}
+
+template <>
+JS::Result<uint8_t> BinASTTokenReaderContext::readByte<
+    BinASTTokenReaderContext::Compression::No>() {
+  return Base::readByte();
+}
+
+template <>
+JS::Result<uint8_t> BinASTTokenReaderContext::readByte<
+    BinASTTokenReaderContext::Compression::Yes>() {
+  uint8_t buf;
+  MOZ_TRY(readBuf<Compression::Yes>(&buf, 1));
+  return buf;
 }
 
 BinASTSourceMetadata* BinASTTokenReaderContext::takeMetadata() {
@@ -65,13 +131,20 @@ JS::Result<Ok> BinASTTokenReaderContext::readHeader() {
 
   // Read global headers.
   MOZ_TRY(readConst(CX_MAGIC_HEADER));
-  BINJS_MOZ_TRY_DECL(version, readVarU32());
+  BINJS_MOZ_TRY_DECL(version, readVarU32<Compression::No>());
 
   if (version != MAGIC_FORMAT_VERSION) {
     return raiseError("Format version not implemented");
   }
 
-  // TODO: handle `LinkToSharedDictionary` and remaining things here.
+  decoder_ = BrotliDecoderCreateInstance(/* alloc_func = */ nullptr,
+                                         /* free_func = */ nullptr,
+                                         /* opaque = */ nullptr);
+  if (!decoder_) {
+    return raiseError("Failed to create brotli decoder");
+  }
+
+  // TODO: handle strings and models here.
 
   return raiseError("Not Yet Implemented");
 }
@@ -169,13 +242,14 @@ JS::Result<Ok> BinASTTokenReaderContext::AutoList::done() {
 //
 // Encoded as variable length number.
 
-MOZ_MUST_USE JS::Result<uint32_t> BinASTTokenReaderContext::readVarU32() {
+template <BinASTTokenReaderContext::Compression compression>
+JS::Result<uint32_t> BinASTTokenReaderContext::readVarU32() {
   uint32_t result = 0;
   uint32_t shift = 0;
   while (true) {
     MOZ_ASSERT(shift < 32);
     uint32_t byte;
-    MOZ_TRY_VAR(byte, readByte());
+    MOZ_TRY_VAR(byte, readByte<compression>());
 
     const uint32_t newResult = result | (byte & 0x7f) << shift;
     if (newResult < result) {
@@ -195,6 +269,10 @@ MOZ_MUST_USE JS::Result<uint32_t> BinASTTokenReaderContext::readVarU32() {
   }
 }
 
+JS::Result<uint32_t> BinASTTokenReaderContext::readUnsignedLong(const Context&) {
+  return readVarU32<Compression::Yes>();
+}
+
 BinASTTokenReaderContext::AutoTaggedTuple::AutoTaggedTuple(
     BinASTTokenReaderContext& reader)
     : AutoBase(reader) {}

js/src/frontend/BinASTTokenReaderContext.h

@@ -12,6 +12,8 @@
 
 #include "mozilla/Maybe.h"  // mozilla::Maybe
 
+#include <brotli/decode.h>  // BrotliDecoderState
+
 #include <stddef.h>  // size_t
 #include <stdint.h>  // uint8_t, uint32_t
 
@@ -48,6 +50,8 @@ class ErrorReporter;
  * - the reader does not support lookahead or pushback.
  */
 class MOZ_STACK_CLASS BinASTTokenReaderContext : public BinASTTokenReaderBase {
+  using Base = BinASTTokenReaderBase;
+
  public:
   class AutoList;
   class AutoTaggedTuple;
@@ -82,6 +86,51 @@ class MOZ_STACK_CLASS BinASTTokenReaderContext : public BinASTTokenReaderBase {
 
   ~BinASTTokenReaderContext();
 
+ private:
+  // {readByte, readBuf, readVarU32} are implemented both for uncompressed
+  // stream and brotli-compressed stream.
+  //
+  // Uncompressed variant is for reading the magic header, and compressed
+  // variant is for reading the remaining part.
+  //
+  // Once compressed variant is called, the underlying uncompressed stream is
+  // buffered and uncompressed variant cannot be called.
+  enum class Compression { No, Yes };
+
+  // Buffer that holds already brotli-decoded but not yet used data.
+  // decodedBuffer[decodedBegin, decodedEnd) holds the data.
+  static const size_t DECODED_BUFFER_SIZE = 128;
+  uint8_t decodedBuffer_[DECODED_BUFFER_SIZE];
+  size_t decodedBegin_ = 0;
+  size_t decodedEnd_ = 0;
+
+  // The number of already decoded bytes.
+  size_t availableDecodedLength() const { return decodedEnd_ - decodedBegin_; }
+
+  // The beginning of decoded buffer.
+  const uint8_t* decodedBufferBegin() const {
+    return decodedBuffer_ + decodedBegin_;
+  }
+
+  // Returns true if the brotli stream finished.
+  bool isEOF() const;
+
+  /**
+   * Read a single byte.
+   */
+  template <Compression compression>
+  MOZ_MUST_USE JS::Result<uint8_t> readByte();
+
+  /**
+   * Read several bytes.
+   *
+   * If there is not enough data, or if the tokenizer has previously been
+   * poisoned, return an error.
+   */
+  template <Compression compression>
+  MOZ_MUST_USE JS::Result<Ok> readBuf(uint8_t* bytes, uint32_t len);
+
+ public:
   /**
    * Read the header of the file.
    */
@@ -196,14 +245,13 @@ class MOZ_STACK_CLASS BinASTTokenReaderContext : public BinASTTokenReaderBase {
   /**
    * Read a single unsigned long.
    */
-  MOZ_MUST_USE JS::Result<uint32_t> readUnsignedLong(const Context&) {
-    return readVarU32();
-  }
+  MOZ_MUST_USE JS::Result<uint32_t> readUnsignedLong(const Context&);
 
  private:
   /**
    * Read a single uint32_t.
    */
+  template <Compression compression>
   MOZ_MUST_USE JS::Result<uint32_t> readVarU32();
 
  private:
@@ -219,6 +267,8 @@ class MOZ_STACK_CLASS BinASTTokenReaderContext : public BinASTTokenReaderBase {
 
   const uint8_t* posBeforeTree_;
 
+  BrotliDecoderState* decoder_ = nullptr;
+
  public:
   BinASTTokenReaderContext(const BinASTTokenReaderContext&) = delete;
   BinASTTokenReaderContext(BinASTTokenReaderContext&&) = delete;