1
// Copyright (c) 2023 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
#pragma once
5
 
6
#include <kj/string.h>
7
#include <kj/vector.h>
8
 
9
#include <cstring>
10
 
11
namespace workerd {
12
 
13
// String buffer optimized for appending a lot of strings together.
14
// Allocates StackSize chunk on the stack and uses that until full.
15
// Keeps allocating new chunks of at least HeapChunkSize as needed.
16
// Doesn't perform any heap allocations if string stays within
17
// StackSize bytes (without \0)
18
template <size_t StackSize>
19
class StringBuffer {
20
 
21
 public:
22
  KJ_DISALLOW_COPY_AND_MOVE(StringBuffer);
23
 
24
  explicit StringBuffer(size_t heapChunkSize)
25
      : heapChunkSize(heapChunkSize),
26
        tail(&arr[0]),
27
        cap(StackSize) {}
28
 
29
  void append() {}
30
 
31
  template <typename First, typename... Rest>
32
  void append(First&& first, Rest&&... rest) {
33
    appendImpl(kj::fwd<First>(first));
34
    append(kj::fwd<Rest>(rest)...);
35
  }
36
 
37
  kj::String toString() {
38
    auto result = kj::heapString(len);
39
    copyTo(result.begin());
40
    return result;
41
  }
42
 
43
 private:
44
  // minimum heap chunk size
45
  const size_t heapChunkSize;
46
 
47
  // chunk on the stack
48
  char arr[StackSize];
49
 
50
  // on the heap chunks
51
  kj::Vector<kj::Array<char>> chunks;
52
 
53
  // points after the last used bytes in current chunk
54
  char* tail;
55
 
56
  // number of bytes available in current chunk
57
  size_t cap;
58
 
59
  // total length of the data appended so far
60
  size_t len = 0;
61
 
62
  void appendImpl(const char* ptr, size_t size) {
63
    size_t toCopy = kj::min(size, cap);
64
    memcpy(tail, ptr, toCopy);
65
    tail += toCopy;
66
    cap -= toCopy;
67
 
68
    if (toCopy != size) {
69
      // prepare new chunk
70
      size_t remaining = size - toCopy;
71
      size_t chunkSize = kj::max(remaining, heapChunkSize);  // don't chunk large strings
72
      auto chunk = kj::heapArray<char>(chunkSize);
73
 
74
      // copy the rest of the string to the new chunk
75
      memcpy(chunk.begin(), ptr + toCopy, remaining);
76
      tail = chunk.begin() + remaining;
77
      cap = chunk.size() - remaining;
78
 
79
      chunks.add(kj::mv(chunk));
80
    }
81
 
82
    len += size;
83
  }
84
 
85
  void appendImpl(const kj::StringPtr& str) {
86
    appendImpl(str.begin(), str.size());
87
  }
88
 
89
  template <size_t size>
90
  void appendImpl(const char (&arr)[size]) {
91
    appendImpl(arr, size - 1 /* assume 0-terminated strings */);
92
  }
93
 
94
  inline void appendImpl(const kj::ArrayPtr<const char>& arr) {
95
    appendImpl(arr.begin(), arr.size());
96
  }
97
 
98
  inline void appendImpl(const kj::String& str) {
99
    appendImpl(str.asPtr());
100
  }
101
 
102
  void copyTo(char* dest) {
103
    // copy stack portion first
104
    size_t onStack = kj::min(len, StackSize);
105
    memcpy(dest, arr, onStack);
106
    dest += onStack;
107
 
108
    // copy from heap chunks
109
    if (onStack < len) {
110
      size_t remaining = len - onStack;
111
      for (auto& chunk: chunks) {
112
        size_t inChunk = kj::min(remaining, chunk.size());  // last chunk won't be full
113
        memcpy(dest, chunk.begin(), inChunk);
114
        dest += inChunk;
115
        remaining -= inChunk;
116
      }
117
 
118
      KJ_IREQUIRE(remaining == 0);
119
    }
120
  }
121
};
122
 
123
}  // namespace workerd