#375, Bug 1112537 - Optimize String#split('foo').join('bar') pattern. (includes M1235403) r=nbp

js/src/jit-test/tests/ion/bug977966.js

@@ -72,6 +72,24 @@ function split_join_4(i) {
     return i;
 }
 
+function split_join_5(i) {
+    var s = "abca";
+    assertEq(s.split("a").join("") + i, "bc" + i);
+}
+
+function split_join_two_byte_char(i) {
+    var s1 = "ab";
+    assertEq(s1.split("").join("\u03c0"), "a\u03c0b");
+    var s2 = i + "\u03c0" + i;
+    assertEq(s2.split("\u03c0").join("-"), i + "-" + i);
+}
+
+function split_join_underflow(i)
+{
+    var s = "";
+    assertEq(s.split("").join("x" + i), "");
+}
+
 // Check that we do not consider the string argument of join as a replacement
 // pattern, as the string replace primitive is supposed to do.
 function split_join_pattern(i) {
@@ -104,6 +122,9 @@ for (var i = 0; i < 100; ++i) {
     split_join_2(i);
     split_join_3(i);
     split_join_4(i);
+    split_join_5(i);
     split_join_pattern(i);
     split_join_multiple(i);
+    split_join_two_byte_char(i);
+    split_join_underflow(i);
 }

js/src/jit/CodeGenerator.cpp

@@ -1935,6 +1935,7 @@ CodeGenerator::visitRegExpReplace(LRegExpReplace* lir)
 }
 
 typedef JSString* (*StringReplaceFn)(JSContext*, HandleString, HandleString, HandleString);
+static const VMFunction StringFlatReplaceInfo = FunctionInfo<StringReplaceFn>(js::str_flat_replace_string);
 static const VMFunction StringReplaceInfo = FunctionInfo<StringReplaceFn>(StringReplace);
 
 void
@@ -1955,7 +1956,10 @@ CodeGenerator::visitStringReplace(LStringReplace* lir)
     else
         pushArg(ToRegister(lir->string()));
 
-    callVM(StringReplaceInfo, lir);
+    if (lir->mir()->isFlatReplacement())
+        callVM(StringFlatReplaceInfo, lir);
+    else
+        callVM(StringReplaceInfo, lir);
 }
 
 void

js/src/jit/MIR.cpp

@@ -4843,20 +4843,20 @@ MTableSwitch::foldsTo(TempAllocator& alloc)
 MDefinition*
 MArrayJoin::foldsTo(TempAllocator& alloc)
 {
-    // :TODO: Enable this optimization after fixing Bug 977966 test cases.
-    return this;
-
     MDefinition* arr = array();
 
     if (!arr->isStringSplit())
         return this;
 
-    this->setRecoveredOnBailout();
+    setRecoveredOnBailout();
     if (arr->hasLiveDefUses()) {
-        this->setNotRecoveredOnBailout();
+        setNotRecoveredOnBailout();
         return this;
     }
 
+    // The MStringSplit won't generate any code.
+    arr->setRecoveredOnBailout();
+
     // We're replacing foo.split(bar).join(baz) by
     // foo.replace(bar, baz).  MStringSplit could be recovered by
     // a bailout.  As we are removing its last use, and its result
@@ -4867,7 +4867,9 @@ MArrayJoin::foldsTo(TempAllocator& alloc)
     MDefinition* replacement = sep();
 
     setNotRecoveredOnBailout();
-    return MStringReplace::New(alloc, string, pattern, replacement);
+    MStringReplace *substr = MStringReplace::New(alloc, string, pattern, replacement);
+    substr->setFlatReplacement();
+    return substr;
 }
 
 MConvertUnboxedObjectToNative*

js/src/jit/MIR.h

@@ -7636,8 +7636,10 @@ class MStringReplace
 {
   private:
 
+    bool isFlatReplacement_;
+
     MStringReplace(MDefinition* string, MDefinition* pattern, MDefinition* replacement)
-      : MStrReplace< StringPolicy<1> >(string, pattern, replacement)
+      : MStrReplace< StringPolicy<1> >(string, pattern, replacement), isFlatReplacement_(false)
     {
     }
 
@@ -7648,7 +7650,20 @@ class MStringReplace
         return new(alloc) MStringReplace(string, pattern, replacement);
     }
 
+    void setFlatReplacement() {
+        MOZ_ASSERT(!isFlatReplacement_);
+        isFlatReplacement_ = true;
+    }
+
+    bool isFlatReplacement() const {
+        return isFlatReplacement_;
+    }
+
     bool congruentTo(const MDefinition* ins) const override {
+        if (!ins->isStringReplace())
+            return false;
+        if (isFlatReplacement_ != ins->toStringReplace()->isFlatReplacement())
+            return false;
         return congruentIfOperandsEqual(ins);
     }
 
@@ -7658,6 +7673,11 @@ class MStringReplace
 
     bool writeRecoverData(CompactBufferWriter& writer) const override;
     bool canRecoverOnBailout() const override {
+        if (isFlatReplacement_)
+        {
+            MOZ_ASSERT(!pattern()->isRegExp());
+            return true;
+        }
         if (pattern()->isRegExp())
             return !pattern()->toRegExp()->source()->global();
         return false;

js/src/jit/Recover.cpp

@@ -1500,11 +1500,14 @@ MStringReplace::writeRecoverData(CompactBufferWriter& writer) const
 {
     MOZ_ASSERT(canRecoverOnBailout());
     writer.writeUnsigned(uint32_t(RInstruction::Recover_StringReplace));
+    writer.writeByte(isFlatReplacement_);
     return true;
 }
 
 RStringReplace::RStringReplace(CompactBufferReader& reader)
-{ }
+{
+    isFlatReplacement_ = reader.readByte();
+}
 
 bool RStringReplace::recover(JSContext* cx, SnapshotIterator& iter) const
 {
@@ -1512,7 +1515,9 @@ bool RStringReplace::recover(JSContext* cx, SnapshotIterator& iter) const
     RootedString pattern(cx, iter.read().toString());
     RootedString replace(cx, iter.read().toString());
 
-    JSString* result = js::str_replace_string_raw(cx, string, pattern, replace);
+    JSString* result = isFlatReplacement_ ? js::str_flat_replace_string(cx, string, pattern, replace) :
+                                            js::str_replace_string_raw(cx, string, pattern, replace);
+
     if (!result)
         return false;
 

js/src/jit/Recover.h

@@ -597,6 +597,9 @@ class RRegExpReplace final : public RInstruction
 
 class RStringReplace final : public RInstruction
 {
+  private:
+    bool isFlatReplacement_;
+
   public:
     RINSTRUCTION_HEADER_(StringReplace)
 

js/src/jsstr.cpp

@@ -3406,11 +3406,124 @@ StrReplaceString(JSContext* cx, ReplaceData& rdata, const FlatMatch& fm)
     return BuildFlatReplacement(cx, rdata.str, rdata.repstr, fm);
 }
 
+template <typename StrChar, typename RepChar>
+static bool
+StrFlatReplaceGlobal(JSContext *cx, JSLinearString *str, JSLinearString *pat, JSLinearString *rep,
+                     StringBuffer &sb)
+{
+    MOZ_ASSERT(str->length() > 0);
+
+    AutoCheckCannotGC nogc;
+    const StrChar *strChars = str->chars<StrChar>(nogc);
+    const RepChar *repChars = rep->chars<RepChar>(nogc);
+
+    // The pattern is empty, so we interleave the replacement string in-between
+    // each character.
+    if (!pat->length()) {
+        CheckedInt<uint32_t> strLength(str->length());
+        CheckedInt<uint32_t> repLength(rep->length());
+        CheckedInt<uint32_t> length = repLength * (strLength - 1) + strLength;
+        if (!length.isValid()) {
+            ReportAllocationOverflow(cx);
+            return false;
+        }
+
+        if (!sb.reserve(length.value()))
+            return false;
+
+        for (unsigned i = 0; i < str->length() - 1; ++i, ++strChars) {
+            sb.infallibleAppend(*strChars);
+            sb.infallibleAppend(repChars, rep->length());
+        }
+        sb.infallibleAppend(*strChars);
+        return true;
+    }
+
+    // If it's true, we are sure that the result's length is, at least, the same
+    // length as |str->length()|.
+    if (rep->length() >= pat->length()) {
+        if (!sb.reserve(str->length()))
+            return false;
+    }
+
+    uint32_t start = 0;
+    for (;;) {
+        int match = StringMatch(str, pat, start);
+        if (match < 0)
+            break;
+        if (!sb.append(strChars + start, match - start))
+            return false;
+        if (!sb.append(repChars, rep->length()))
+            return false;
+        start = match + pat->length();
+    }
+
+    if (!sb.append(strChars + start, str->length() - start))
+        return false;
+
+    return true;
+}
+
+// This is identical to "str.split(pattern).join(replacement)" except that we
+// do some deforestation optimization in Ion.
+JSString *
+js::str_flat_replace_string(JSContext *cx, HandleString string, HandleString pattern,
+                            HandleString replacement)
+{
+    MOZ_ASSERT(string);
+    MOZ_ASSERT(pattern);
+    MOZ_ASSERT(replacement);
+
+    if (!string->length())
+        return string;
+
+    RootedLinearString linearRepl(cx, replacement->ensureLinear(cx));
+    if (!linearRepl)
+        return nullptr;
+
+    RootedLinearString linearPat(cx, pattern->ensureLinear(cx));
+    if (!linearPat)
+        return nullptr;
+
+    RootedLinearString linearStr(cx, string->ensureLinear(cx));
+    if (!linearStr)
+        return nullptr;
+
+    StringBuffer sb(cx);
+    if (linearStr->hasTwoByteChars()) {
+        if (!sb.ensureTwoByteChars())
+            return nullptr;
+        if (linearRepl->hasTwoByteChars()) {
+            if (!StrFlatReplaceGlobal<char16_t, char16_t>(cx, linearStr, linearPat, linearRepl, sb))
+                return nullptr;
+        } else {
+            if (!StrFlatReplaceGlobal<char16_t, Latin1Char>(cx, linearStr, linearPat, linearRepl, sb))
+                return nullptr;
+        }
+    } else {
+        if (linearRepl->hasTwoByteChars()) {
+            if (!sb.ensureTwoByteChars())
+                return nullptr;
+            if (!StrFlatReplaceGlobal<Latin1Char, char16_t>(cx, linearStr, linearPat, linearRepl, sb))
+                return nullptr;
+        } else {
+            if (!StrFlatReplaceGlobal<Latin1Char, Latin1Char>(cx, linearStr, linearPat, linearRepl, sb))
+                return nullptr;
+        }
+    }
+
+    JSString *str = sb.finishString();
+    if (!str)
+        return nullptr;
+
+    return str;
+}
+
 static const uint32_t ReplaceOptArg = 2;
 
 JSString*
 js::str_replace_string_raw(JSContext* cx, HandleString string, HandleString pattern,
-                          HandleString replacement)
+                           HandleString replacement)
 {
     ReplaceData rdata(cx);
 

js/src/jsstr.h

@@ -437,6 +437,10 @@ str_split(JSContext* cx, unsigned argc, Value* vp);
 JSObject*
 str_split_string(JSContext* cx, HandleObjectGroup group, HandleString str, HandleString sep);
 
+JSString *
+str_flat_replace_string(JSContext *cx, HandleString string, HandleString pattern,
+                        HandleString replacement);
+
 JSString*
 str_replace_string_raw(JSContext* cx, HandleString string, HandleString pattern,
                        HandleString replacement);