Support caller-provided memory for CompressionContext

Callers that must avoid heap allocation entirely can now construct a
CompressionContext over a caller-provided workspace of WorkspaceSize()
bytes.

Internally, WorkingMemory gains a non-allocating constructor laying out
its scratch space in a caller-provided buffer, and RequiredSize()
factoring out the size computation its allocating constructor already
performed.

snappy-internal.h

diff --git a/snappy-internal.h b/snappy-internal.h
index f68ab77..582e886 100644
--- a/snappy-internal.h
+++ b/snappy-internal.h
@@ -137,8 +137,16 @@
 class WorkingMemory {
  public:
   explicit WorkingMemory(size_t input_size);
+
+  // Non-allocating: lays out the scratch space in the caller-provided
+  // buffer, which must be at least RequiredSize(input_size) bytes, aligned
+  // at least as strictly as uint16_t, and must outlive "*this".
+  WorkingMemory(size_t input_size, char* buffer);
   ~WorkingMemory();
 
+  // The buffer size required by the non-allocating constructor above.
+  static size_t RequiredSize(size_t input_size);
+
   // Allocates and clears a hash table using memory in "*this",
   // stores the number of buckets in "*table_size" and returns a pointer to
   // the base of the hash table.
@@ -149,6 +157,7 @@
  private:
   char* mem_;        // the allocated memory, never nullptr
   size_t size_;      // the size of the allocated memory, never 0
+  bool owns_mem_;    // whether the destructor should free mem_
   uint16_t* table_;  // the pointer to the hashtable
   char* input_;      // the pointer to the input scratch buffer
   char* output_;     // the pointer to the output scratch buffer

snappy.cc

diff --git a/snappy.cc b/snappy.cc
index efca3af..69cc209 100644
--- a/snappy.cc
+++ b/snappy.cc
@@ -76,6 +76,7 @@
 #include <cstring>
 #include <limits>
 #include <memory>
+#include <new>
 #include <string>
 #include <utility>
 #include <vector>
@@ -757,19 +758,36 @@
 }  // namespace
 
 namespace internal {
-WorkingMemory::WorkingMemory(size_t input_size) {
+size_t WorkingMemory::RequiredSize(size_t input_size) {
   const size_t max_fragment_size = std::min(input_size, kBlockSize);
   const size_t table_size = CalculateTableSize(max_fragment_size);
-  size_ = table_size * sizeof(*table_) + max_fragment_size +
-          MaxCompressedLength(max_fragment_size);
-  mem_ = std::allocator<char>().allocate(size_);
+  return table_size * sizeof(uint16_t) + max_fragment_size +
+         MaxCompressedLength(max_fragment_size);
+}
+
+WorkingMemory::WorkingMemory(size_t input_size)
+    : WorkingMemory(input_size,
+                    std::allocator<char>().allocate(RequiredSize(input_size))) {
+  owns_mem_ = true;
+}
+
+WorkingMemory::WorkingMemory(size_t input_size, char* buffer) {
+  assert(buffer != nullptr);
+  assert(reinterpret_cast<uintptr_t>(buffer) % alignof(uint16_t) == 0);
+  const size_t max_fragment_size = std::min(input_size, kBlockSize);
+  const size_t table_size = CalculateTableSize(max_fragment_size);
+  mem_ = buffer;
+  size_ = RequiredSize(input_size);
+  owns_mem_ = false;
   table_ = reinterpret_cast<uint16_t*>(mem_);
   input_ = mem_ + table_size * sizeof(*table_);
   output_ = input_ + max_fragment_size;
 }
 
 WorkingMemory::~WorkingMemory() {
-  std::allocator<char>().deallocate(mem_, size_);
+  if (owns_mem_) {
+    std::allocator<char>().deallocate(mem_, size_);
+  }
 }
 
 uint16_t* WorkingMemory::GetHashTable(size_t fragment_size,
@@ -1953,20 +1971,51 @@
 }
 
 CompressionContext::CompressionContext()
-    : working_memory_(new internal::WorkingMemory(kBlockSize)) {}
+    : working_memory_(new internal::WorkingMemory(kBlockSize)),
+      owns_working_memory_(true) {}
 
-CompressionContext::~CompressionContext() { delete working_memory_; }
+size_t CompressionContext::WorkspaceSize() {
+  return sizeof(internal::WorkingMemory) +
+         internal::WorkingMemory::RequiredSize(kBlockSize);
+}
+
+CompressionContext::CompressionContext(void* workspace, size_t workspace_size)
+    : owns_working_memory_(false) {
+  assert(workspace != nullptr);
+  assert(workspace_size >= WorkspaceSize());
+  assert(reinterpret_cast<uintptr_t>(workspace) %
+             alignof(internal::WorkingMemory) ==
+         0);
+  (void)workspace_size;
+  char* base = static_cast<char*>(workspace);
+  working_memory_ = new (base) internal::WorkingMemory(
+      kBlockSize, base + sizeof(internal::WorkingMemory));
+}
+
+void CompressionContext::Reset() {
+  if (working_memory_ == nullptr) return;
+  if (owns_working_memory_) {
+    delete working_memory_;
+  } else {
+    working_memory_->~WorkingMemory();
+  }
+  working_memory_ = nullptr;
+}
+
+CompressionContext::~CompressionContext() { Reset(); }
 
 CompressionContext::CompressionContext(CompressionContext&& other) noexcept
-    : working_memory_(other.working_memory_) {
+    : working_memory_(other.working_memory_),
+      owns_working_memory_(other.owns_working_memory_) {
   other.working_memory_ = nullptr;
 }
 
 CompressionContext& CompressionContext::operator=(
     CompressionContext&& other) noexcept {
   if (this != &other) {
-    delete working_memory_;
+    Reset();
     working_memory_ = other.working_memory_;
+    owns_working_memory_ = other.owns_working_memory_;
     other.working_memory_ = nullptr;
   }
   return *this;

snappy.h

diff --git a/snappy.h b/snappy.h
index 423142b..5413153 100644
--- a/snappy.h
+++ b/snappy.h
@@ -88,7 +88,18 @@
   // may only be destroyed or assigned to.
   class CompressionContext {
    public:
+    // Allocates the working memory on the heap.
     CompressionContext();
+
+    // Constructs a context whose working memory is placed in the
+    // caller-provided "workspace" instead of being heap-allocated; the
+    // library performs no allocation at all.
+    //
+    // REQUIRES: "workspace" points to at least "workspace_size" bytes with
+    // "workspace_size >= WorkspaceSize()", is suitably aligned for any
+    // object type (as if returned by malloc), and outlives "*this".
+    CompressionContext(void* workspace, size_t workspace_size);
+
     ~CompressionContext();
 
     CompressionContext(CompressionContext&& other) noexcept;
@@ -97,11 +108,20 @@
     CompressionContext(const CompressionContext&) = delete;
     CompressionContext& operator=(const CompressionContext&) = delete;
 
+    // The workspace size required by the non-allocating constructor above.
+    static size_t WorkspaceSize();
+
    private:
     friend size_t Compress(Source* reader, Sink* writer,
                            CompressionOptions options, CompressionContext* ctx);
 
+    // Destroys the working memory as appropriate for how it was created
+    // (delete if heap-allocated, in-place destruction if placement-constructed
+    // in a caller-provided workspace).
+    void Reset();
+
     internal::WorkingMemory* working_memory_;
+    bool owns_working_memory_;
   };
 
   // ------------------------------------------------------------------------

snappy_unittest.cc

diff --git a/snappy_unittest.cc b/snappy_unittest.cc
index 2ee6e22..4f44cbf 100644
--- a/snappy_unittest.cc
+++ b/snappy_unittest.cc
@@ -595,6 +595,58 @@
   }
 }
 
+TEST(Snappy, CompressionContextStaticWorkspace) {
+  // The library performs no allocation for a context constructed over a
+  // caller-provided workspace.
+  std::vector<char> workspace(CompressionContext::WorkspaceSize());
+  CompressionContext static_ctx(workspace.data(), workspace.size());
+  CompressionContext heap_ctx;
+
+  const size_t sizes[] = {0, 1, kBlockSize - 1, kBlockSize + 1,
+                          2 * kBlockSize + 17};
+  for (size_t len : sizes) {
+    std::string input;
+    input.reserve(len);
+    while (input.size() < len) {
+      input.push_back(static_cast<char>('a' + input.size() % 7));
+    }
+
+    std::string with_static(MaxCompressedLength(len), '\0');
+    size_t with_static_len = 0;
+    RawCompress(input.data(), input.size(), &with_static[0], &with_static_len,
+                CompressionOptions{}, &static_ctx);
+    with_static.resize(with_static_len);
+
+    std::string with_heap(MaxCompressedLength(len), '\0');
+    size_t with_heap_len = 0;
+    RawCompress(input.data(), input.size(), &with_heap[0], &with_heap_len,
+                CompressionOptions{}, &heap_ctx);
+    with_heap.resize(with_heap_len);
+
+    EXPECT_EQ(with_static, with_heap) << "len=" << len;
+
+    std::string uncompressed;
+    EXPECT_TRUE(Uncompress(with_static, &uncompressed));
+    EXPECT_EQ(input, uncompressed);
+  }
+
+  // Both context flavors keep working after being moved.
+  CompressionContext moved_static(std::move(static_ctx));
+  CompressionContext moved_heap = std::move(heap_ctx);
+  const std::string input = "the quick brown fox jumps over the lazy dog";
+  std::string a(MaxCompressedLength(input.size()), '\0');
+  std::string b(MaxCompressedLength(input.size()), '\0');
+  size_t a_len = 0;
+  size_t b_len = 0;
+  RawCompress(input.data(), input.size(), &a[0], &a_len, CompressionOptions{},
+              &moved_static);
+  RawCompress(input.data(), input.size(), &b[0], &b_len, CompressionOptions{},
+              &moved_heap);
+  a.resize(a_len);
+  b.resize(b_len);
+  EXPECT_EQ(a, b);
+}
+
 TEST(Snappy, FourByteOffset) {
   // The new compressor cannot generate four-byte offsets since
   // it chops up the input into 32KB pieces.  So we hand-emit the