Fix: GetOnDemand forgot copying " for parseStringInplace when key is escaped string (#89)

Author: xiegx94

include/sonic/internal/arch/common/base.h

@@ -17,8 +17,6 @@
 
 #pragma once
 
-#include "../common/base.h"
-
 namespace sonic_json {
 namespace internal {
 namespace common {
@@ -268,25 +266,19 @@ sonic_force_inline bool handle_unicode_codepoint(const uint8_t **src_ptr,
   return offset > 0;
 }
 
-template <size_t BLOK_SIZE>
-sonic_force_inline uint64_t GetEscapedBranchless(uint64_t &prev_escaped,
-                                                 uint64_t backslash) {
-  static_assert(BLOK_SIZE == 32 || BLOK_SIZE == 64,
-                "escaped branchless block only support 32 or 64 bytes");
-  backslash &= ~prev_escaped;
-  uint64_t follows_escape = backslash << 1 | prev_escaped;
-  const uint64_t even_bits = 0x5555555555555555ULL;
-  uint64_t odd_sequence_starts = backslash & ~even_bits & ~follows_escape;
-  uint64_t sequences_starting_on_even_bits;
-  if (BLOK_SIZE == 64) {
-    prev_escaped = AddOverflow64(odd_sequence_starts, backslash,
-                                 &sequences_starting_on_even_bits);
-  } else if (BLOK_SIZE == 32) {
-    prev_escaped = AddOverflow32(odd_sequence_starts, backslash,
-                                 &sequences_starting_on_even_bits);
-  }
-  uint64_t invert_mask = sequences_starting_on_even_bits << 1;
-  return (even_bits ^ invert_mask) & follows_escape;
+template <size_t BLOCK_SIZE>
+sonic_force_inline uint64_t GetEscaped(uint64_t &prev_escaped,
+                                       uint64_t backslash) {
+  uint64_t with_prev_backslash = backslash & ~prev_escaped;
+  static const uint64_t odd_bits = 0xAAAAAAAAAAAAAAAAULL;
+  // see: https://github.com/simdjson/simdjson/pull/2042
+  uint64_t escaped =
+      (((with_prev_backslash << 1) | odd_bits) - with_prev_backslash) ^
+      odd_bits;
+  // & backslash to clear escaped char
+  uint64_t escaped_with_prev = (escaped ^ (backslash | prev_escaped));
+  prev_escaped = ((escaped & backslash) >> (BLOCK_SIZE - 1)) & 0x1ULL;
+  return escaped_with_prev;
 }
 
 }  // namespace common

include/sonic/internal/arch/common/unicode_common.h

@@ -21,7 +21,7 @@ sonic_force_inline uint64_t GetStringBits(const uint8_t *data,
   uint64_t escaped = 0;
   uint64_t bs_bits = v.eq('\\');
   if (bs_bits) {
-    escaped = common::GetEscapedBranchless<64>(prev_escaped, bs_bits);
+    escaped = common::GetEscaped<64>(prev_escaped, bs_bits);
   } else {
     escaped = prev_escaped;
     prev_escaped = 0;
@@ -72,12 +72,14 @@ sonic_force_inline int SkipString(const uint8_t *data, size_t &pos,
   bool found = false;
   while (pos + VEC_LEN <= len) {
     const VecUint8Type v(data + pos);
+
+    // bs_bits and quote_bits only has VEC_LEN bits used
     bs_bits = (v == '\\').to_bitmask();
     quote_bits = (v == '"').to_bitmask();
 
     // maybe has escaped quotes
     if (((quote_bits - 1) & bs_bits) || prev_escaped) {
-      escaped = common::GetEscapedBranchless<32>(prev_escaped, bs_bits);
+      escaped = common::GetEscaped<VEC_LEN>(prev_escaped, bs_bits);
       // NOTE: maybe mark the normal string as escaped, example "abc":"\\",
       // abc will marked as escaped.
       found = true;

include/sonic/internal/arch/common/x86_common/skip.inc.h

@@ -184,7 +184,9 @@ class SkipScanner {
       // parse escaped key
       kbuf.resize(sn + 32);
       uint8_t *nsrc = &kbuf[0];
-      std::memcpy(nsrc, sp, sn);
+
+      // parseStringInplace need `"` as the end
+      std::memcpy(nsrc, sp, sn + 1);
       sn = parseStringInplace(nsrc, err);
       if (err) {
         pos = (sp - data) + (nsrc - &kbuf[0]);

include/sonic/internal/arch/simd_skip.h

@@ -79,6 +79,15 @@ TEST(GetOnDemand, SuccessEscapeCharacters) {
                   R"("\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\")");
   TestGetOnDemand(R"({"a":"\"\\\\\\\\\\\\\\\\\\\\\\\\\\\\\""})", {"a"},
                   R"("\"\\\\\\\\\\\\\\\\\\\\\\\\\\\\\"")");
+  TestGetOnDemand(R"({"\"a\"":"\n\tHello,\nworld!\n"})", {"\"a\""},
+                  R"("\n\tHello,\nworld!\n")");
+  TestGetOnDemand(R"({"123456789012345\"123":"sse_string",
+                      "1234567890123456789012345678901\"123":"avx2_string",
+                      "obj\n\t\\":{"name\\\\\\\\":"string\\\\"},
+                      "array\"\t\n\b\r":["\n\tHello,\nworld!\n",
+                        "{\" / \b \f \n \r \t } [景] 测试中文 😀")],
+                      "\"a\"":"\n\tHello,\nworld!\n"})",
+                  {"\"a\""}, R"("\n\tHello,\nworld!\n")");
 }
 
 TEST(GetOnDemand, Failed) {
@@ -93,4 +102,4 @@ TEST(GetOnDemand, Failed) {
 
   TestGetOnDemandFailed(R"("\")", {}, ParseResult(kParseErrorInvalidChar, 3));
 }
-}  // namespace
+}  // namespace