Compile string.format().

4 files changed, 133 insertions, 5 deletions

diff --git a/src/lib_string.c b/src/lib_string.c
index cc1affbd..204f6975 100644
--- a/src/lib_string.c
+++ b/src/lib_string.c
@@ -658,10 +658,10 @@ static GCstr *string_fmt_tostring(lua_State *L, int arg, int retry)
   return lj_strfmt_obj(L, o);
 }
 
-LJLIB_CF(string_format)
+LJLIB_CF(string_format)         LJLIB_REC(.)
 {
   int arg, top = (int)(L->top - L->base);
-  GCstr *sfmt;
+  GCstr *fmt;
   SBuf *sb;
   FormatState fs;
   SFormat sf;
@@ -669,8 +669,8 @@ LJLIB_CF(string_format)
 again:
   arg = 1;
   sb = lj_buf_tmp_(L);
-  sfmt = lj_lib_checkstr(L, arg);
-  lj_strfmt_init(&fs, strdata(sfmt), sfmt->len);
+  fmt = lj_lib_checkstr(L, arg);
+  lj_strfmt_init(&fs, strdata(fmt), fmt->len);
   while ((sf = lj_strfmt_parse(&fs)) != STRFMT_EOF) {
     if (sf == STRFMT_LIT) {
       lj_buf_putmem(sb, fs.str, fs.len);
@@ -705,7 +705,7 @@ again:
         if (str == NULL)
           retry = 1;
         else if ((sf & STRFMT_T_QUOTED))
-          lj_strfmt_putquoted(sb, str);
+          lj_strfmt_putquoted(sb, str);  /* No formatting. */
         else
           lj_strfmt_putfstr(sb, sf, str);
         break;
diff --git a/src/lj_ffrecord.c b/src/lj_ffrecord.c
index 97c24836..31f9b390 100644
--- a/src/lj_ffrecord.c
+++ b/src/lj_ffrecord.c
@@ -871,6 +871,84 @@ static void LJ_FASTCALL recff_string_find(jit_State *J, RecordFFData *rd)
   }
 }
 
+static void LJ_FASTCALL recff_string_format(jit_State *J, RecordFFData *rd)
+{
+  TRef trfmt = lj_ir_tostr(J, J->base[0]);
+  GCstr *fmt = argv2str(J, &rd->argv[0]);
+  int arg = 1;
+  TRef hdr, tr;
+  FormatState fs;
+  SFormat sf;
+  /* Specialize to the format string. */
+  emitir(IRTG(IR_EQ, IRT_STR), trfmt, lj_ir_kstr(J, fmt));
+  tr = hdr = recff_bufhdr(J);
+  lj_strfmt_init(&fs, strdata(fmt), fmt->len);
+  while ((sf = lj_strfmt_parse(&fs)) != STRFMT_EOF) {  /* Parse format. */
+    TRef tra = sf == STRFMT_LIT ? 0 : J->base[arg++];
+    TRef trsf = lj_ir_kint(J, (int32_t)sf);
+    IRCallID id;
+    switch (STRFMT_TYPE(sf)) {
+    case STRFMT_LIT:
+      tr = emitir(IRT(IR_BUFPUT, IRT_P32), tr,
+                  lj_ir_kstr(J, lj_str_new(J->L, fs.str, fs.len)));
+      break;
+    case STRFMT_INT:
+      id = IRCALL_lj_strfmt_putfnum_int;
+    handle_int:
+      if (!tref_isinteger(tra))
+        goto handle_num;
+      if (sf == STRFMT_INT) { /* Shortcut for plain %d. */
+        tr = emitir(IRT(IR_BUFPUT, IRT_P32), tr,
+                    emitir(IRT(IR_TOSTR, IRT_STR), tra, IRTOSTR_INT));
+      } else {
+#if LJ_HASFFI
+        tra = emitir(IRT(IR_CONV, IRT_U64), tra,
+                     (IRT_INT|(IRT_U64<<5)|IRCONV_SEXT));
+        tr = lj_ir_call(J, IRCALL_lj_strfmt_putfxint, tr, trsf, tra);
+        lj_needsplit(J);
+#else
+        recff_nyiu(J);  /* Don't bother working around this NYI. */
+#endif
+      }
+      break;
+    case STRFMT_UINT:
+      id = IRCALL_lj_strfmt_putfnum_uint;
+      goto handle_int;
+    case STRFMT_NUM:
+      id = IRCALL_lj_strfmt_putfnum;
+    handle_num:
+      tra = lj_ir_tonum(J, tra);
+      tr = lj_ir_call(J, id, tr, trsf, tra);
+      if (LJ_SOFTFP) lj_needsplit(J);
+      break;
+    case STRFMT_STR:
+      if (!tref_isstr(tra))
+        recff_nyiu(J);  /* NYI: __tostring and non-string types for %s. */
+      if (sf == STRFMT_STR)  /* Shortcut for plain %s. */
+        tr = emitir(IRT(IR_BUFPUT, IRT_P32), tr, tra);
+      else if ((sf & STRFMT_T_QUOTED))
+        tr = lj_ir_call(J, IRCALL_lj_strfmt_putquoted, tr, tra);
+      else
+        tr = lj_ir_call(J, IRCALL_lj_strfmt_putfstr, tr, trsf, tra);
+      break;
+    case STRFMT_CHAR:
+      tra = lj_opt_narrow_toint(J, tra);
+      if (sf == STRFMT_CHAR)  /* Shortcut for plain %c. */
+        tr = emitir(IRT(IR_BUFPUT, IRT_P32), tr,
+                    emitir(IRT(IR_TOSTR, IRT_STR), tra, IRTOSTR_CHAR));
+      else
+        tr = lj_ir_call(J, IRCALL_lj_strfmt_putfchar, tr, trsf, tra);
+      break;
+    case STRFMT_PTR:  /* NYI */
+    case STRFMT_ERR:
+    default:
+      recff_nyiu(J);
+      break;
+    }
+  }
+  J->base[0] = emitir(IRT(IR_BUFSTR, IRT_STR), tr, hdr);
+}
+
 /* -- Table library fast functions ---------------------------------------- */
 
 static void LJ_FASTCALL recff_table_insert(jit_State *J, RecordFFData *rd)
diff --git a/src/lj_ircall.h b/src/lj_ircall.h
index 43bd25d9..7271ceca 100644
--- a/src/lj_ircall.h
+++ b/src/lj_ircall.h
@@ -124,6 +124,13 @@ typedef struct CCallInfo {
   _(ANY,        lj_strfmt_char,         2,  FN, STR, CCI_L) \
   _(ANY,        lj_strfmt_putint,       2,  FL, P32, 0) \
   _(ANY,        lj_strfmt_putnum,       2,  FL, P32, 0) \
+  _(ANY,        lj_strfmt_putquoted,    2,  FL, P32, 0) \
+  _(ANY,        lj_strfmt_putfxint,     3,   L, P32, XA_64) \
+  _(ANY,        lj_strfmt_putfnum_int,  3,   L, P32, XA_FP) \
+  _(ANY,        lj_strfmt_putfnum_uint, 3,   L, P32, XA_FP) \
+  _(ANY,        lj_strfmt_putfnum,      3,   L, P32, XA_FP) \
+  _(ANY,        lj_strfmt_putfstr,      3,   L, P32, 0) \
+  _(ANY,        lj_strfmt_putfchar,     3,   L, P32, 0) \
   _(ANY,        lj_buf_putmem,          3,   S, P32, 0) \
   _(ANY,        lj_buf_putstr,          2,  FL, P32, 0) \
   _(ANY,        lj_buf_putchar,         2,  FL, P32, 0) \
diff --git a/src/lj_opt_fold.c b/src/lj_opt_fold.c
index e2ac28f2..f3abe8ea 100644
--- a/src/lj_opt_fold.c
+++ b/src/lj_opt_fold.c
@@ -617,6 +617,7 @@ LJFOLDF(bufstr_kfold_cse)
 LJFOLD(CALLL CARG IRCALL_lj_buf_putstr_reverse)
 LJFOLD(CALLL CARG IRCALL_lj_buf_putstr_upper)
 LJFOLD(CALLL CARG IRCALL_lj_buf_putstr_lower)
+LJFOLD(CALLL CARG IRCALL_lj_strfmt_putquoted)
 LJFOLDF(bufput_kfold_op)
 {
   if (irref_isk(fleft->op2)) {
@@ -649,6 +650,48 @@ LJFOLDF(bufput_kfold_rep)
   return EMITFOLD;  /* Always emit, CSE later. */
 }
 
+LJFOLD(CALLL CARG IRCALL_lj_strfmt_putfxint)
+LJFOLD(CALLL CARG IRCALL_lj_strfmt_putfnum_int)
+LJFOLD(CALLL CARG IRCALL_lj_strfmt_putfnum_uint)
+LJFOLD(CALLL CARG IRCALL_lj_strfmt_putfnum)
+LJFOLD(CALLL CARG IRCALL_lj_strfmt_putfstr)
+LJFOLD(CALLL CARG IRCALL_lj_strfmt_putfchar)
+LJFOLDF(bufput_kfold_fmt)
+{
+  IRIns *irc = IR(fleft->op1);
+  lua_assert(irref_isk(irc->op2));  /* SFormat must be const. */
+  if (irref_isk(fleft->op2)) {
+    SFormat sf = (SFormat)IR(irc->op2)->i;
+    IRIns *ira = IR(fleft->op2);
+    SBuf *sb = lj_buf_tmp_(J->L);
+    switch (fins->op2) {
+    case IRCALL_lj_strfmt_putfxint:
+      sb = lj_strfmt_putfxint(sb, sf, ir_k64(ira)->u64);
+      break;
+    case IRCALL_lj_strfmt_putfstr:
+      sb = lj_strfmt_putfstr(sb, sf, ir_kstr(ira));
+      break;
+    case IRCALL_lj_strfmt_putfchar:
+      sb = lj_strfmt_putfchar(sb, sf, ira->i);
+      break;
+    case IRCALL_lj_strfmt_putfnum_int:
+    case IRCALL_lj_strfmt_putfnum_uint:
+    case IRCALL_lj_strfmt_putfnum:
+    default: {
+      const CCallInfo *ci = &lj_ir_callinfo[fins->op2];
+      sb = ((SBuf * (*)(SBuf *, SFormat, lua_Number))ci->func)(sb, sf,
+                                                         ir_knum(ira)->n);
+      break;
+      }
+    }
+    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 ------------------------------ */
 
 LJFOLD(ADD KGC KINT)