Compile string.rep().

4 files changed, 55 insertions, 16 deletions

diff --git a/src/lib_string.c b/src/lib_string.c
index 598cd93e..e460e834 100644
--- a/src/lib_string.c
+++ b/src/lib_string.c
@@ -85,7 +85,7 @@ LJLIB_ASM(string_sub)		LJLIB_REC(string_range 1)
   return FFH_RETRY;
 }
 
-LJLIB_CF(string_rep)
+LJLIB_CF(string_rep)            LJLIB_REC(.)
 {
   GCstr *s = lj_lib_checkstr(L, 1);
   int32_t rep = lj_lib_checkint(L, 2);
diff --git a/src/lj_ffrecord.c b/src/lj_ffrecord.c
index 1560d3f5..d66aaa53 100644
--- a/src/lj_ffrecord.c
+++ b/src/lj_ffrecord.c
@@ -765,6 +765,34 @@ static void LJ_FASTCALL recff_string_char(jit_State *J, RecordFFData *rd)
   UNUSED(rd);
 }
 
+static void LJ_FASTCALL recff_string_rep(jit_State *J, RecordFFData *rd)
+{
+  TRef str = lj_ir_tostr(J, J->base[0]);
+  TRef rep = lj_opt_narrow_toint(J, J->base[1]);
+  TRef hdr, tr, str2 = 0;
+  if (J->base[2]) {
+    TRef sep = lj_ir_tostr(J, J->base[2]);
+    int32_t vrep = argv2int(J, &rd->argv[1]);
+    emitir(IRTGI(vrep > 1 ? IR_GT : IR_LE), rep, lj_ir_kint(J, 1));
+    if (vrep > 1) {
+      TRef hdr2 = emitir(IRT(IR_BUFHDR, IRT_P32),
+                         lj_ir_kptr(J, &J2G(J)->tmpbuf), IRBUFHDR_RESET);
+      TRef tr2 = emitir(IRT(IR_BUFPUT, IRT_P32), hdr2, sep);
+      tr2 = emitir(IRT(IR_BUFPUT, IRT_P32), tr2, str);
+      str2 = emitir(IRT(IR_BUFSTR, IRT_STR), tr2, hdr2);
+    }
+  }
+  tr = hdr = emitir(IRT(IR_BUFHDR, IRT_P32),
+                    lj_ir_kptr(J, &J2G(J)->tmpbuf), IRBUFHDR_RESET);
+  if (str2) {
+    tr = emitir(IRT(IR_BUFPUT, IRT_P32), tr, str);
+    str = str2;
+    rep = emitir(IRTI(IR_ADD), rep, lj_ir_kint(J, -1));
+  }
+  tr = lj_ir_call(J, IRCALL_lj_buf_putstr_rep, tr, str, rep);
+  J->base[0] = emitir(IRT(IR_BUFSTR, IRT_STR), tr, hdr);
+}
+
 static void LJ_FASTCALL recff_string_op(jit_State *J, RecordFFData *rd)
 {
   TRef str = lj_ir_tostr(J, J->base[0]);
diff --git a/src/lj_ircall.h b/src/lj_ircall.h
index 3e190c80..f652befd 100644
--- a/src/lj_ircall.h
+++ b/src/lj_ircall.h
@@ -114,6 +114,7 @@ typedef struct CCallInfo {
   _(ANY,        lj_buf_putstr_reverse,  2,  FL, P32, 0) \
   _(ANY,        lj_buf_putstr_lower,    2,  FL, P32, 0) \
   _(ANY,        lj_buf_putstr_upper,    2,  FL, P32, 0) \
+  _(ANY,        lj_buf_putstr_rep,      3,   L, P32, 0) \
   _(ANY,        lj_buf_tostr,           1,  FL, STR, 0) \
   _(ANY,        lj_tab_new1,            2,  FS, TAB, CCI_L) \
   _(ANY,        lj_tab_dup,             2,  FS, TAB, CCI_L) \
diff --git a/src/lj_opt_fold.c b/src/lj_opt_fold.c
index fc625b84..94d5702f 100644
--- a/src/lj_opt_fold.c
+++ b/src/lj_opt_fold.c
@@ -596,18 +596,13 @@ LJFOLDF(bufstr_kfold_cse)
   if (LJ_LIKELY(J->flags & JIT_F_OPT_CSE)) {
     IRRef ref = J->chain[IR_BUFSTR];
     while (ref) {
-      IRRef last = fins->op1;
       IRIns *irs = IR(ref), *ira = fleft, *irb = IR(irs->op1);
       while (ira->o == irb->o && ira->op2 == irb->op2) {
-        if (ira->o == IR_BUFHDR && !(ira->op2 & IRBUFHDR_APPEND)) {
+        lua_assert(ira->o == IR_BUFHDR || ira->o == IR_BUFPUT ||
+                   ira->o == IR_CALLL || ira->o == IR_CARG);
+        if (ira->o == IR_BUFHDR && !(ira->op2 & IRBUFHDR_APPEND))
           return ref;  /* CSE succeeded. */
-        } else if (ira->o == IR_CALLL) {
-          ira = IR(ira->op1); irb = IR(irb->op1);
-          lua_assert(ira->o == IR_CARG && irb->o == IR_CARG);
-          if (ira->op2 != irb->op2) break;
-        }
-        last = ira->op1;
-        ira = IR(last);
+        ira = IR(ira->op1);
         irb = IR(irb->op1);
       }
       ref = irs->prev;
@@ -623,17 +618,32 @@ LJFOLDF(bufput_kfold_op)
 {
   if (irref_isk(fleft->op2)) {
     const CCallInfo *ci = &lj_ir_callinfo[fins->op2];
-    SBuf *sb = &J2G(J)->tmpbuf;
-    setsbufL(sb, J->L);
-    lj_buf_reset(sb);
+    SBuf *sb = lj_buf_tmp_(J->L);
     sb = ((SBuf * (LJ_FASTCALL *)(SBuf *, GCstr *))ci->func)(sb,
                                                        ir_kstr(IR(fleft->op2)));
-    fins->op2 = lj_ir_kstr(J, lj_buf_tostr(sb));
-    fins->op1 = fleft->op1;
     fins->o = IR_BUFPUT;
+    fins->op1 = fleft->op1;
+    fins->op2 = lj_ir_kstr(J, lj_buf_tostr(sb));
     return RETRYFOLD;
   }
-  return EMITFOLD;  /* This is a store and always emitted. */
+  return EMITFOLD;  /* Always emit, CSE later. */
+}
+
+LJFOLD(CALLL CARG IRCALL_lj_buf_putstr_rep)
+LJFOLDF(bufput_kfold_rep)
+{
+  if (irref_isk(fleft->op2)) {
+    IRIns *irc = IR(fleft->op1);
+    if (irref_isk(irc->op2)) {
+      SBuf *sb = lj_buf_tmp_(J->L);
+      sb = lj_buf_putstr_rep(sb, ir_kstr(IR(irc->op2)), IR(fleft->op2)->i);
+      fins->o = IR_BUFPUT;
+      fins->op1 = irc->op1;
+      fins->op2 = lj_ir_kstr(J, lj_buf_tostr(sb));
+      return RETRYFOLD;
+    }
+  }
+  return EMITFOLD;  /* Always emit, CSE later. */
 }
 
 /* -- Constant folding of pointer arithmetic ------------------------------ */