From e32b2a744f70b21c2f0f9c36113ca81a04783cbd Mon Sep 17 00:00:00 2001
From: Daniel Lemire <daniel@lemire.me>
Date: Fri, 22 Nov 2024 19:35:12 -0500
Subject: [PATCH] tweaking

---
 include/simdjson/arm64/implementation.h       |  1 +
 include/simdjson/arm64/simd.h                 |  9 +-
 include/simdjson/fallback/implementation.h    |  1 +
 .../ondemand/json_string_builder-inl.h        |  1 -
 include/simdjson/haswell/implementation.h     |  1 +
 include/simdjson/icelake/implementation.h     |  1 +
 include/simdjson/implementation.h             | 14 +--
 include/simdjson/lasx/implementation.h        |  1 +
 include/simdjson/lsx/implementation.h         |  1 +
 include/simdjson/ppc64/implementation.h       |  1 +
 include/simdjson/westmere/implementation.h    |  1 +
 src/arm64.cpp                                 |  4 +
 src/fallback.cpp                              |  4 +
 src/generic/stage2/stringparsing.h            | 98 +++++++++++++++++++
 src/haswell.cpp                               |  4 +
 src/icelake.cpp                               |  4 +
 src/implementation.cpp                        |  9 ++
 src/lasx.cpp                                  |  4 +
 src/lsx.cpp                                   |  4 +
 src/ppc64.cpp                                 |  4 +
 src/westmere.cpp                              |  4 +
 21 files changed, 161 insertions(+), 10 deletions(-)

diff --git a/include/simdjson/arm64/implementation.h b/include/simdjson/arm64/implementation.h
index c9b7dd753..00c5bf566 100644
--- a/include/simdjson/arm64/implementation.h
+++ b/include/simdjson/arm64/implementation.h
@@ -23,6 +23,7 @@ public:
   ) const noexcept final;
   simdjson_warn_unused error_code minify(const uint8_t *buf, size_t len, uint8_t *dst, size_t &dst_len) const noexcept final;
   simdjson_warn_unused bool validate_utf8(const char *buf, size_t len) const noexcept final;
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final;
 };
 
 } // namespace arm64
diff --git a/include/simdjson/arm64/simd.h b/include/simdjson/arm64/simd.h
index 3b0fa844f..34c5117ad 100644
--- a/include/simdjson/arm64/simd.h
+++ b/include/simdjson/arm64/simd.h
@@ -210,7 +210,7 @@ namespace {
 
     // Bit-specific operations
     simdjson_inline simd8<bool> any_bits_set(simd8<uint8_t> bits) const { return vtstq_u8(*this, bits); }
-    simdjson_inline bool any_bits_set_anywhere() const { return this->max_val() != 0; }
+    simdjson_inline bool any_bits_set_anywhere() const { return vmaxvq_u32(vreinterpretq_u32_u8(*this)) != 0; }
     simdjson_inline bool any_bits_set_anywhere(simd8<uint8_t> bits) const { return (*this & bits).any_bits_set_anywhere(); }
     template<int N>
     simdjson_inline simd8<uint8_t> shr() const { return vshrq_n_u8(*this, N); }
@@ -223,7 +223,12 @@ namespace {
       return lookup_table.apply_lookup_16_to(*this);
     }
 
-
+    // Returns 4-bit out of each byte, alternating between the high 4 bits and low
+    // bits result it is 64 bit.
+    simdjson_inline uint64_t to_bitmask64() const {
+      return vget_lane_u64(
+          vreinterpret_u64_u8(vshrn_n_u16(vreinterpretq_u16_u8(*this), 4)), 0);
+    }
     // Copies to 'output" all bytes corresponding to a 0 in the mask (interpreted as a bitset).
     // Passing a 0 value for mask would be equivalent to writing out every byte to output.
     // Only the first 16 - count_ones(mask) bytes of the result are significant but 16 bytes
diff --git a/include/simdjson/fallback/implementation.h b/include/simdjson/fallback/implementation.h
index 523f06d2e..e607253a0 100644
--- a/include/simdjson/fallback/implementation.h
+++ b/include/simdjson/fallback/implementation.h
@@ -26,6 +26,7 @@ public:
   ) const noexcept final;
   simdjson_warn_unused error_code minify(const uint8_t *buf, size_t len, uint8_t *dst, size_t &dst_len) const noexcept final;
   simdjson_warn_unused bool validate_utf8(const char *buf, size_t len) const noexcept final;
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final;
 };
 
 } // namespace fallback
diff --git a/include/simdjson/generic/ondemand/json_string_builder-inl.h b/include/simdjson/generic/ondemand/json_string_builder-inl.h
index 97981bd1b..7fd96ed92 100644
--- a/include/simdjson/generic/ondemand/json_string_builder-inl.h
+++ b/include/simdjson/generic/ondemand/json_string_builder-inl.h
@@ -218,7 +218,6 @@ simdjson_inline void string_builder::escape_and_append(std::string_view input)
   }
 }
 
-
 simdjson_inline void string_builder::escape_and_append_with_quotes(std::string_view input)  noexcept {
   // escaping might turn a control character into \x00xx so 6 characters.
   if(capacity_check(2 + 6 * input.size())) {
diff --git a/include/simdjson/haswell/implementation.h b/include/simdjson/haswell/implementation.h
index 6861e4298..984d9f7ed 100644
--- a/include/simdjson/haswell/implementation.h
+++ b/include/simdjson/haswell/implementation.h
@@ -28,6 +28,7 @@ public:
   ) const noexcept final;
   simdjson_warn_unused error_code minify(const uint8_t *buf, size_t len, uint8_t *dst, size_t &dst_len) const noexcept final;
   simdjson_warn_unused bool validate_utf8(const char *buf, size_t len) const noexcept final;
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final;
 };
 
 } // namespace haswell
diff --git a/include/simdjson/icelake/implementation.h b/include/simdjson/icelake/implementation.h
index 940c5f992..d197b4309 100644
--- a/include/simdjson/icelake/implementation.h
+++ b/include/simdjson/icelake/implementation.h
@@ -28,6 +28,7 @@ public:
   ) const noexcept final;
   simdjson_warn_unused error_code minify(const uint8_t *buf, size_t len, uint8_t *dst, size_t &dst_len) const noexcept final;
   simdjson_warn_unused bool validate_utf8(const char *buf, size_t len) const noexcept final;
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final;
 };
 
 } // namespace icelake
diff --git a/include/simdjson/implementation.h b/include/simdjson/implementation.h
index eff7562d7..79f34491b 100644
--- a/include/simdjson/implementation.h
+++ b/include/simdjson/implementation.h
@@ -137,13 +137,13 @@ public:
   simdjson_warn_unused virtual bool validate_utf8(const char *buf, size_t len) const noexcept = 0;
 
   /**
-  * Write the string to the output buffer while escaping double-quote, backlash and ascii control characters.
-  *
-  * @param input the string_view to escape
-  * @param out output buffer (for escaped string): to be safe, it should have 6 * input.size() allocated bytes.
-  * @return number of bytes written
-  */
-  simdjson_warn_unused virtual size_t write_string_escaped(const std::string_view input, char *out) noexcept;
+   * Write the string to the output buffer while escaping double-quote, backlash and ascii control characters.
+   *
+   * @param input the string_view to escape
+   * @param out output buffer (for escaped string): to be safe, it should have 6 * input.size() allocated bytes.
+   * @return number of bytes written
+   */
+  simdjson_warn_unused virtual size_t write_string_escaped(const std::string_view input, char *out) const noexcept;
 protected:
   /** @private Construct an implementation with the given name and description. For subclasses. */
   simdjson_inline implementation(
diff --git a/include/simdjson/lasx/implementation.h b/include/simdjson/lasx/implementation.h
index 8aafbb8b8..d64ec3e49 100644
--- a/include/simdjson/lasx/implementation.h
+++ b/include/simdjson/lasx/implementation.h
@@ -23,6 +23,7 @@ public:
   ) const noexcept final;
   simdjson_warn_unused error_code minify(const uint8_t *buf, size_t len, uint8_t *dst, size_t &dst_len) const noexcept final;
   simdjson_warn_unused bool validate_utf8(const char *buf, size_t len) const noexcept final;
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final;
 };
 
 } // namespace lasx
diff --git a/include/simdjson/lsx/implementation.h b/include/simdjson/lsx/implementation.h
index 14468777d..a17fb5532 100644
--- a/include/simdjson/lsx/implementation.h
+++ b/include/simdjson/lsx/implementation.h
@@ -23,6 +23,7 @@ public:
   ) const noexcept final;
   simdjson_warn_unused error_code minify(const uint8_t *buf, size_t len, uint8_t *dst, size_t &dst_len) const noexcept final;
   simdjson_warn_unused bool validate_utf8(const char *buf, size_t len) const noexcept final;
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final;
 };
 
 } // namespace lsx
diff --git a/include/simdjson/ppc64/implementation.h b/include/simdjson/ppc64/implementation.h
index 33436f534..b3445af26 100644
--- a/include/simdjson/ppc64/implementation.h
+++ b/include/simdjson/ppc64/implementation.h
@@ -32,6 +32,7 @@ public:
                                          size_t &dst_len) const noexcept final;
   simdjson_warn_unused bool validate_utf8(const char *buf,
                                           size_t len) const noexcept final;
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final;
 };
 
 } // namespace ppc64
diff --git a/include/simdjson/westmere/implementation.h b/include/simdjson/westmere/implementation.h
index 37392be2a..c7e655b5e 100644
--- a/include/simdjson/westmere/implementation.h
+++ b/include/simdjson/westmere/implementation.h
@@ -24,6 +24,7 @@ public:
   ) const noexcept final;
   simdjson_warn_unused error_code minify(const uint8_t *buf, size_t len, uint8_t *dst, size_t &dst_len) const noexcept final;
   simdjson_warn_unused bool validate_utf8(const char *buf, size_t len) const noexcept final;
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final;
 };
 
 } // namespace westmere
diff --git a/src/arm64.cpp b/src/arm64.cpp
index 436757d53..f641d987c 100644
--- a/src/arm64.cpp
+++ b/src/arm64.cpp
@@ -164,6 +164,10 @@ simdjson_warn_unused error_code dom_parser_implementation::parse(const uint8_t *
   return stage2(_doc);
 }
 
+simdjson_warn_unused size_t implementation::write_string_escaped(const std::string_view input, char *out) const noexcept {
+  return arm64::stringparsing::write_string_escaped(input, out);
+}
+
 } // namespace arm64
 } // namespace simdjson
 
diff --git a/src/fallback.cpp b/src/fallback.cpp
index f8e87be06..139361d5f 100644
--- a/src/fallback.cpp
+++ b/src/fallback.cpp
@@ -402,6 +402,10 @@ simdjson_warn_unused error_code dom_parser_implementation::parse(const uint8_t *
   return stage2(_doc);
 }
 
+simdjson_warn_unused size_t implementation::write_string_escaped(const std::string_view input, char *out) const noexcept {
+  return fallback::stringparsing::write_string_escaped(input, out);
+}
+
 } // namespace fallback
 } // namespace simdjson
 
diff --git a/src/generic/stage2/stringparsing.h b/src/generic/stage2/stringparsing.h
index 07a4daf50..b944c964d 100644
--- a/src/generic/stage2/stringparsing.h
+++ b/src/generic/stage2/stringparsing.h
@@ -236,6 +236,104 @@ simdjson_warn_unused simdjson_inline uint8_t *parse_wobbly_string(const uint8_t
   }
 }
 
+simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) noexcept {
+  // We are making the following assumption: most strings will either be very short or they will not
+  // need escaping.
+  size_t i = 0;
+  size_t pos = 0;
+  if(input.size() >= sizeof(simd8<uint8_t>)) {
+    auto vec_processing = [input,out]() -> size_t {
+      size_t i = 0;
+      size_t pos = 0;
+      for(;input.size() - i >= sizeof(simd8<uint8_t>); i += sizeof(simd8<uint8_t>)) {
+        simd8<uint8_t> vinput(reinterpret_cast<const uint8_t *>(input.data()) + i);
+        // instead of doing it register by register, we could regroup, but consider
+        // that we expect most strings to be short.
+        if(((vinput <= 31) | (vinput == '\\') | (vinput == '"')).any()) {
+          return i; // We have a character that needs escaping
+          // We could be more carefully and identify the character that needs escaping.
+        }
+        vinput.store(reinterpret_cast<uint8_t *>(out) + pos);
+        pos += sizeof(simd8<uint8_t>);
+      }
+      if(i == input.size()) { return input.size(); }
+      simd8<uint8_t> vinput(reinterpret_cast<const uint8_t *>(input.data()) + input.size() - sizeof(simd8<uint8_t>));
+      if(((vinput <= 31) | (vinput == '\\') | (vinput == '"')).any()) {
+        return i; // We have a character that needs escaping
+        // We could be more carefully and identify the character that needs escaping.
+      }
+      vinput.store(reinterpret_cast<uint8_t *>(out) + input.size() - sizeof(simd8<uint8_t>));
+      return input.size();
+    };
+    i = vec_processing();
+    pos = i;
+    if(i == input.size()) { return pos; }
+    // Here we only continue if there was a character that needed escaping.
+  }
+  static std::string_view control_chars[] = {
+    "\\x0000", "\\x0001", "\\x0002", "\\x0003", "\\x0004", "\\x0005", "\\x0006",
+    "\\x0007", "\\x0008", "\\t",     "\\n",     "\\x000b", "\\f",     "\\r",
+    "\\x000e", "\\x000f", "\\x0010", "\\x0011", "\\x0012", "\\x0013", "\\x0014",
+    "\\x0015", "\\x0016", "\\x0017", "\\x0018", "\\x0019", "\\x001a", "\\x001b",
+    "\\x001c", "\\x001d", "\\x001e", "\\x001f"};
+  static std::array<uint8_t, 256> json_quotable_character =
+    []() constexpr {
+      std::array<uint8_t, 256> result{};
+      for (int i = 0; i < 32; i++) {
+        result[i] = 1;
+      }
+      for (int i : {'"', '\\'}) {
+        result[i] = 1;
+      }
+      return result;
+    }();
+  // The rest could possibly be vectorized, but consider that we expect most strings
+  // to be short or not to require escaping.
+  for (; i < input.size(); i++) {
+    uint8_t c = static_cast<uint8_t>(input[i]);
+    if(json_quotable_character[c]) {
+      switch (c) {
+      case '"':
+        out[pos++] = '\\';
+        out[pos++] = '"';
+        break;
+      case '\\':
+        out[pos++] = '\\';
+        out[pos++] = '\\';
+        break;
+      case '\b':
+        out[pos++] = '\\';
+        out[pos++] = 'b';
+        break;
+      case '\f':
+        out[pos++] = '\\';
+        out[pos++] = 'f';
+        break;
+      case '\n':
+        out[pos++] = '\\';
+        out[pos++] = 'n';
+        break;
+      case '\r':
+        out[pos++] = '\\';
+        out[pos++] = 'r';
+        break;
+      case '\t':
+        out[pos++] = '\\';
+        out[pos++] = 't';
+        break;
+      default:
+        control_chars[c].copy(out + pos, 6);
+        pos += 6;
+      }
+    } else {
+      out[pos++] = c;
+    }
+  }
+  return pos;
+}
+
+
+
 } // namespace stringparsing
 } // unnamed namespace
 } // namespace SIMDJSON_IMPLEMENTATION
diff --git a/src/haswell.cpp b/src/haswell.cpp
index f721cac8b..ce15d16a3 100644
--- a/src/haswell.cpp
+++ b/src/haswell.cpp
@@ -161,6 +161,10 @@ simdjson_warn_unused error_code dom_parser_implementation::parse(const uint8_t *
   return stage2(_doc);
 }
 
+simdjson_warn_unused size_t implementation::write_string_escaped(const std::string_view input, char *out) const noexcept {
+  return haswell::stringparsing::write_string_escaped(input, out);
+}
+
 } // namespace SIMDJSON_IMPLEMENTATION
 } // namespace simdjson
 
diff --git a/src/icelake.cpp b/src/icelake.cpp
index 8ec08c69c..5726e8111 100644
--- a/src/icelake.cpp
+++ b/src/icelake.cpp
@@ -207,6 +207,10 @@ simdjson_warn_unused error_code dom_parser_implementation::parse(const uint8_t *
   return stage2(_doc);
 }
 
+simdjson_warn_unused size_t implementation::write_string_escaped(const std::string_view input, char *out) const noexcept {
+  return icelake::stringparsing::write_string_escaped(input, out);
+}
+
 } // namespace icelake
 } // namespace simdjson
 
diff --git a/src/implementation.cpp b/src/implementation.cpp
index 3dba4b605..50feda912 100644
--- a/src/implementation.cpp
+++ b/src/implementation.cpp
@@ -186,6 +186,9 @@ public:
   simdjson_warn_unused bool validate_utf8(const char * buf, size_t len) const noexcept final override {
     return set_best()->validate_utf8(buf, len);
   }
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final {
+    return set_best()->write_string_escaped(input, out);
+  }
   simdjson_inline detect_best_supported_implementation_on_first_use() noexcept : implementation("best_supported_detector", "Detects the best supported implementation and sets it", 0) {}
 private:
   const implementation *set_best() const noexcept;
@@ -236,6 +239,9 @@ public:
   simdjson_warn_unused error_code minify(const uint8_t *, size_t, uint8_t *, size_t &) const noexcept final override {
     return UNSUPPORTED_ARCHITECTURE;
   }
+  simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) const noexcept final {
+    return 0;
+  }
   simdjson_warn_unused bool validate_utf8(const char *, size_t) const noexcept final override {
     return false; // Just refuse to validate. Given that we have a fallback implementation
     // it seems unlikely that unsupported_implementation will ever be used. If it is used,
@@ -319,6 +325,9 @@ simdjson_warn_unused error_code minify(const char *buf, size_t len, char *dst, s
 simdjson_warn_unused bool validate_utf8(const char *buf, size_t len) noexcept {
   return get_active_implementation()->validate_utf8(buf, len);
 }
+simdjson_warn_unused size_t write_string_escaped(const std::string_view input, char *out) noexcept {
+  return get_active_implementation()->write_string_escaped(input, out);
+}
 const implementation * builtin_implementation() {
   static const implementation * builtin_impl = get_available_implementations()[SIMDJSON_STRINGIFY(SIMDJSON_BUILTIN_IMPLEMENTATION)];
   assert(builtin_impl);
diff --git a/src/lasx.cpp b/src/lasx.cpp
index 3be3f080e..f90e2fb0e 100644
--- a/src/lasx.cpp
+++ b/src/lasx.cpp
@@ -124,6 +124,10 @@ simdjson_warn_unused error_code dom_parser_implementation::parse(const uint8_t *
   return stage2(_doc);
 }
 
+simdjson_warn_unused size_t implementation::write_string_escaped(const std::string_view input, char *out) const noexcept {
+  return lasx::stringparsing::write_string_escaped(input, out);
+}
+
 } // namespace lasx
 } // namespace simdjson
 
diff --git a/src/lsx.cpp b/src/lsx.cpp
index 99af664a5..2867987c0 100644
--- a/src/lsx.cpp
+++ b/src/lsx.cpp
@@ -128,6 +128,10 @@ simdjson_warn_unused error_code dom_parser_implementation::parse(const uint8_t *
   return stage2(_doc);
 }
 
+simdjson_warn_unused size_t implementation::write_string_escaped(const std::string_view input, char *out) const noexcept {
+  return lsx::stringparsing::write_string_escaped(input, out);
+}
+
 } // namespace lsx
 } // namespace simdjson
 
diff --git a/src/ppc64.cpp b/src/ppc64.cpp
index 63606380c..c2e93ec07 100644
--- a/src/ppc64.cpp
+++ b/src/ppc64.cpp
@@ -134,6 +134,10 @@ simdjson_warn_unused error_code dom_parser_implementation::parse(const uint8_t *
   return stage2(_doc);
 }
 
+simdjson_warn_unused size_t implementation::write_string_escaped(const std::string_view input, char *out) const noexcept {
+  return ppc64::stringparsing::write_string_escaped(input, out);
+}
+
 } // namespace ppc64
 } // namespace simdjson
 
diff --git a/src/westmere.cpp b/src/westmere.cpp
index 538db42f8..f595ed787 100644
--- a/src/westmere.cpp
+++ b/src/westmere.cpp
@@ -166,6 +166,10 @@ simdjson_warn_unused error_code dom_parser_implementation::parse(const uint8_t *
   return stage2(_doc);
 }
 
+simdjson_warn_unused size_t implementation::write_string_escaped(const std::string_view input, char *out) const noexcept {
+  return westmere::stringparsing::write_string_escaped(input, out);
+}
+
 } // namespace westmere
 } // namespace simdjson