From Lua 5.2: Add string.rep(s, n, sep).

6 files changed, 40 insertions, 13 deletions

diff --git a/src/lib_string.c b/src/lib_string.c
index b36fcd6d..9aa74d5c 100644
--- a/src/lib_string.c
+++ b/src/lib_string.c
@@ -86,22 +86,48 @@ LJLIB_ASM(string_sub)		LJLIB_REC(string_range 1)
 LJLIB_ASM(string_rep)
 {
   GCstr *s = lj_lib_checkstr(L, 1);
-  int32_t len = (int32_t)s->len;
   int32_t k = lj_lib_checkint(L, 2);
-  int64_t tlen = (int64_t)k * len;
+  GCstr *sep = lj_lib_optstr(L, 3);
+  int32_t len = (int32_t)s->len;
+  global_State *g = G(L);
+  int64_t tlen;
   const char *src;
   char *buf;
-  if (k <= 0) return FFH_RETRY;
-  if (tlen > LJ_MAX_STR)
-    lj_err_caller(L, LJ_ERR_STROV);
-  buf = lj_str_needbuf(L, &G(L)->tmpbuf, (MSize)tlen);
-  if (len <= 1) return FFH_RETRY;  /* ASM code only needed buffer resize. */
+  if (k <= 0) {
+  empty:
+    setstrV(L, L->base-1, &g->strempty);
+    return FFH_RES(1);
+  }
+  if (sep) {
+    tlen = (int64_t)len + sep->len;
+    if (tlen > LJ_MAX_STR)
+      lj_err_caller(L, LJ_ERR_STROV);
+    tlen *= k;
+    if (tlen > LJ_MAX_STR)
+      lj_err_caller(L, LJ_ERR_STROV);
+  } else {
+    tlen = (int64_t)k * len;
+    if (tlen > LJ_MAX_STR)
+      lj_err_caller(L, LJ_ERR_STROV);
+  }
+  if (tlen == 0) goto empty;
+  buf = lj_str_needbuf(L, &g->tmpbuf, (MSize)tlen);
   src = strdata(s);
+  if (sep) {
+    tlen -= sep->len;  /* Ignore trailing separator. */
+    if (k > 1) {  /* Paste one string and one separator. */
+      int32_t i;
+      i = 0; while (i < len) *buf++ = src[i++];
+      src = strdata(sep); len = sep->len;
+      i = 0; while (i < len) *buf++ = src[i++];
+      src = g->tmpbuf.buf; len += s->len; k--;  /* Now copy that k-1 times. */
+    }
+  }
   do {
     int32_t i = 0;
     do { *buf++ = src[i++]; } while (i < len);
   } while (--k > 0);
-  setstrV(L, L->base-1, lj_str_new(L, G(L)->tmpbuf.buf, (size_t)tlen));
+  setstrV(L, L->base-1, lj_str_new(L, g->tmpbuf.buf, (size_t)tlen));
   return FFH_RES(1);
 }
 
diff --git a/src/vm_arm.dasc b/src/vm_arm.dasc
index 331a1b70..4909d827 100644
--- a/src/vm_arm.dasc
+++ b/src/vm_arm.dasc
@@ -1778,7 +1778,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  ldrd CARG12, [BASE]
   |   ldrd CARG34, [BASE, #8]
   |    cmp NARGS8:RC, #16
-  |    blo ->fff_fallback
+  |    bne ->fff_fallback               // Exactly 2 arguments
   |  checktp CARG2, LJ_TSTR
   |   checktpeq CARG4, LJ_TISNUM
   |   bne ->fff_fallback
diff --git a/src/vm_mips.dasc b/src/vm_mips.dasc
index 719d9252..f82cefe8 100644
--- a/src/vm_mips.dasc
+++ b/src/vm_mips.dasc
@@ -1696,7 +1696,7 @@ static void build_subroutines(BuildCtx *ctx)
   |.ffunc string_rep                    // Only handle the 1-char case inline.
   |  ffgccheck
   |  lw TMP0, HI(BASE)
-  |   sltiu AT, NARGS8:RC, 16
+  |   addiu AT, NARGS8:RC, -16          // Exactly 2 arguments.
   |  lw CARG4, 8+HI(BASE)
   |   lw STR:CARG1, LO(BASE)
   |  addiu TMP0, TMP0, -LJ_TSTR
diff --git a/src/vm_ppc.dasc b/src/vm_ppc.dasc
index 6dbfb90d..7eafebe1 100644
--- a/src/vm_ppc.dasc
+++ b/src/vm_ppc.dasc
@@ -2127,7 +2127,7 @@ static void build_subroutines(BuildCtx *ctx)
   |.else
   |    lfd FARG2, 8(BASE)
   |.endif
-  |  blt ->fff_fallback
+  |  bne ->fff_fallback                 // Exactly 2 arguments.
   |  checkstr TMP0; bne ->fff_fallback
   |.if DUALNUM
   |  checknum CARG4; bne ->fff_fallback
diff --git a/src/vm_ppcspe.dasc b/src/vm_ppcspe.dasc
index 2af06494..ab2d7670 100644
--- a/src/vm_ppcspe.dasc
+++ b/src/vm_ppcspe.dasc
@@ -1630,7 +1630,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  cmplwi NARGS8:RC, 16
   |   evldd CARG1, 0(BASE)
   |   evldd CARG2, 8(BASE)
-  |  blt ->fff_fallback
+  |  bne ->fff_fallback                 // Exactly 2 arguments.
   |  checknum CARG2
   |  checkfail ->fff_fallback
   |  checkstr STR:CARG1
diff --git a/src/vm_x86.dasc b/src/vm_x86.dasc
index 75683225..92e98f98 100644
--- a/src/vm_x86.dasc
+++ b/src/vm_x86.dasc
@@ -2437,8 +2437,9 @@ static void build_subroutines(BuildCtx *ctx)
   |  xor RC, RC                         // Zero length. Any ptr in RB is ok.
   |  jmp <4
   |
-  |.ffunc_2 string_rep                  // Only handle the 1-char case inline.
+  |.ffunc string_rep                    // Only handle the 1-char case inline.
   |  ffgccheck
+  |  cmp NARGS:RD, 2+1; jne ->fff_fallback      // Exactly 2 arguments.
   |  cmp dword [BASE+4], LJ_TSTR;  jne ->fff_fallback
   |  cmp dword [BASE+12], LJ_TISNUM
   |  mov STR:RB, [BASE]