src/lib_io.c

   1 /*
   2 ** I/O library.
   3 ** Copyright (C) 2005-2012 Mike Pall. See Copyright Notice in luajit.h
   4 **
   5 ** Major portions taken verbatim or adapted from the Lua interpreter.
   6 ** Copyright (C) 1994-2011 Lua.org, PUC-Rio. See Copyright Notice in lua.h
   7 */
   8
   9 #include <errno.h>
  10 #include <stdio.h>
  11
  12 #define lib_io_c
  13 #define LUA_LIB
  14
  15 #include "lua.h"
  16 #include "lauxlib.h"
  17 #include "lualib.h"
  18
  19 #include "lj_obj.h"
  20 #include "lj_err.h"
  21 #include "lj_str.h"
  22 #include "lj_ff.h"
  23 #include "lj_lib.h"
  24
  25 /* Userdata payload for I/O file. */
  26 typedef struct IOFileUD {
  27   FILE *fp;             /* File handle. */
  28   uint32_t type;        /* File type. */
  29 } IOFileUD;
  30
  31 #define IOFILE_TYPE_FILE        0       /* Regular file. */
  32 #define IOFILE_TYPE_PIPE        1       /* Pipe. */
  33 #define IOFILE_TYPE_STDF        2       /* Standard file handle. */
  34 #define IOFILE_TYPE_MASK        3
  35
  36 #define IOFILE_FLAG_CLOSE       4       /* Close after io.lines() iterator. */
  37
  38 #define IOSTDF_UD(L, id)        (&gcref(G(L)->gcroot[(id)])->ud)
  39 #define IOSTDF_IOF(L, id)       ((IOFileUD *)uddata(IOSTDF_UD(L, (id))))
  40
  41 /* -- Open/close helpers -------------------------------------------------- */
  42
  43 static IOFileUD *io_tofilep(lua_State *L)
  44 {
  45   if (!(L->base < L->top && tvisudata(L->base) &&
  46         udataV(L->base)->udtype == UDTYPE_IO_FILE))
  47     lj_err_argtype(L, 1, "FILE*");
  48   return (IOFileUD *)uddata(udataV(L->base));
  49 }
  50
  51 static IOFileUD *io_tofile(lua_State *L)
  52 {
  53   IOFileUD *iof = io_tofilep(L);
  54   if (iof->fp == NULL)
  55     lj_err_caller(L, LJ_ERR_IOCLFL);
  56   return iof;
  57 }
  58
  59 static FILE *io_stdfile(lua_State *L, ptrdiff_t id)
  60 {
  61   IOFileUD *iof = IOSTDF_IOF(L, id);
  62   if (iof->fp == NULL)
  63     lj_err_caller(L, LJ_ERR_IOSTDCL);
  64   return iof->fp;
  65 }
  66
  67 static IOFileUD *io_file_new(lua_State *L)
  68 {
  69   IOFileUD *iof = (IOFileUD *)lua_newuserdata(L, sizeof(IOFileUD));
  70   GCudata *ud = udataV(L->top-1);
  71   ud->udtype = UDTYPE_IO_FILE;
  72   /* NOBARRIER: The GCudata is new (marked white). */
  73   setgcrefr(ud->metatable, curr_func(L)->c.env);
  74   iof->fp = NULL;
  75   iof->type = IOFILE_TYPE_FILE;
  76   return iof;
  77 }
  78
  79 static IOFileUD *io_file_open(lua_State *L, const char *mode)
  80 {
  81   const char *fname = strdata(lj_lib_checkstr(L, 1));
  82   IOFileUD *iof = io_file_new(L);
  83   iof->fp = fopen(fname, mode);
  84   if (iof->fp == NULL)
  85     luaL_argerror(L, 1, lj_str_pushf(L, "%s: %s", fname, strerror(errno)));
  86   return iof;
  87 }
  88
  89 static int io_file_close(lua_State *L, IOFileUD *iof)
  90 {
  91   int ok;
  92   if ((iof->type & IOFILE_TYPE_MASK) == IOFILE_TYPE_FILE) {
  93     ok = (fclose(iof->fp) == 0);
  94   } else if ((iof->type & IOFILE_TYPE_MASK) == IOFILE_TYPE_PIPE) {
  95     int stat = -1;
  96 #if LJ_TARGET_POSIX
  97     stat = pclose(iof->fp);
  98 #elif LJ_TARGET_WINDOWS
  99     stat = _pclose(iof->fp);
 100 #else
 101     lua_assert(0);
 102     return 0;
 103 #endif
 104 #if LJ_52
 105     iof->fp = NULL;
 106     return luaL_execresult(L, stat);
 107 #else
 108     ok = (stat != -1);
 109 #endif
 110   } else {
 111     lua_assert((iof->type & IOFILE_TYPE_MASK) == IOFILE_TYPE_STDF);
 112     setnilV(L->top++);
 113     lua_pushliteral(L, "cannot close standard file");
 114     return 2;
 115   }
 116   iof->fp = NULL;
 117   return luaL_fileresult(L, ok, NULL);
 118 }
 119
 120 /* -- Read/write helpers -------------------------------------------------- */
 121
 122 static int io_file_readnum(lua_State *L, FILE *fp)
 123 {
 124   lua_Number d;
 125   if (fscanf(fp, LUA_NUMBER_SCAN, &d) == 1) {
 126     if (LJ_DUALNUM) {
 127       int32_t i = lj_num2int(d);
 128       if (d == (lua_Number)i && !tvismzero((cTValue *)&d)) {
 129         setintV(L->top++, i);
 130         return 1;
 131       }
 132     }
 133     setnumV(L->top++, d);
 134     return 1;
 135   } else {
 136     setnilV(L->top++);
 137     return 0;
 138   }
 139 }
 140
 141 static int io_file_testeof(lua_State *L, FILE *fp)
 142 {
 143   int c = getc(fp);
 144   ungetc(c, fp);
 145   lua_pushlstring(L, NULL, 0);
 146   return (c != EOF);
 147 }
 148
 149 static int io_file_readline(lua_State *L, FILE *fp, size_t chop)
 150 {
 151   luaL_Buffer b;
 152   luaL_buffinit(L, &b);
 153   for (;;) {
 154     size_t len;
 155     char *p = luaL_prepbuffer(&b);
 156     if (fgets(p, LUAL_BUFFERSIZE, fp) == NULL) {  /* EOF? */
 157       luaL_pushresult(&b);
 158       return (strV(L->top-1)->len > 0);  /* Anything read? */
 159     }
 160     len = strlen(p);
 161     if (len == 0 || p[len-1] != '\n') {  /* Partial line? */
 162       luaL_addsize(&b, len);
 163     } else {
 164       luaL_addsize(&b, len - chop);  /* Keep or remove EOL. */
 165       luaL_pushresult(&b);
 166       return 1;  /* Got at least an EOL. */
 167     }
 168   }
 169 }
 170
 171 static int io_file_readchars(lua_State *L, FILE *fp, size_t n)
 172 {
 173   size_t rlen;  /* how much to read */
 174   size_t nr;  /* number of chars actually read */
 175   luaL_Buffer b;
 176   luaL_buffinit(L, &b);
 177   rlen = LUAL_BUFFERSIZE;  /* try to read that much each time */
 178   do {
 179     char *p = luaL_prepbuffer(&b);
 180     if (rlen > n) rlen = n;  /* cannot read more than asked */
 181     nr = fread(p, 1, rlen, fp);
 182     luaL_addsize(&b, nr);
 183     n -= nr;  /* still have to read `n' chars */
 184   } while (n > 0 && nr == rlen);  /* until end of count or eof */
 185   luaL_pushresult(&b);  /* close buffer */
 186   return (n == 0 || strV(L->top-1)->len > 0);
 187 }
 188
 189 static int io_file_read(lua_State *L, FILE *fp, int start)
 190 {
 191   int ok, n, nargs = (int)(L->top - L->base) - start;
 192   clearerr(fp);
 193   if (nargs == 0) {
 194     ok = io_file_readline(L, fp, 1);
 195     n = start+1;  /* Return 1 result. */
 196   } else {
 197     /* The results plus the buffers go on top of the args. */
 198     luaL_checkstack(L, nargs+LUA_MINSTACK, "too many arguments");
 199     ok = 1;
 200     for (n = start; nargs-- && ok; n++) {
 201       if (tvisstr(L->base+n)) {
 202         const char *p = strVdata(L->base+n);
 203         if (p[0] != '*')
 204           lj_err_arg(L, n+1, LJ_ERR_INVOPT);
 205         if (p[1] == 'n')
 206           ok = io_file_readnum(L, fp);
 207         else if ((p[1] & ~0x20) == 'L')
 208           ok = io_file_readline(L, fp, (p[1] == 'l'));
 209         else if (p[1] == 'a')
 210           io_file_readchars(L, fp, ~((size_t)0));
 211         else
 212           lj_err_arg(L, n+1, LJ_ERR_INVFMT);
 213       } else if (tvisnumber(L->base+n)) {
 214         size_t len = (size_t)lj_lib_checkint(L, n+1);
 215         ok = len ? io_file_readchars(L, fp, len) : io_file_testeof(L, fp);
 216       } else {
 217         lj_err_arg(L, n+1, LJ_ERR_INVOPT);
 218       }
 219     }
 220   }
 221   if (ferror(fp))
 222     return luaL_fileresult(L, 0, NULL);
 223   if (!ok)
 224     setnilV(L->top-1);  /* Replace last result with nil. */
 225   return n - start;
 226 }
 227
 228 static int io_file_write(lua_State *L, FILE *fp, int start)
 229 {
 230   cTValue *tv;
 231   int status = 1;
 232   for (tv = L->base+start; tv < L->top; tv++) {
 233     if (tvisstr(tv)) {
 234       MSize len = strV(tv)->len;
 235       status = status && (fwrite(strVdata(tv), 1, len, fp) == len);
 236     } else if (tvisint(tv)) {
 237       char buf[LJ_STR_INTBUF];
 238       char *p = lj_str_bufint(buf, intV(tv));
 239       size_t len = (size_t)(buf+LJ_STR_INTBUF-p);
 240       status = status && (fwrite(p, 1, len, fp) == len);
 241     } else if (tvisnum(tv)) {
 242       status = status && (fprintf(fp, LUA_NUMBER_FMT, numV(tv)) > 0);
 243     } else {
 244       lj_err_argt(L, (int)(tv - L->base) + 1, LUA_TSTRING);
 245     }
 246   }
 247   if (LJ_52 && status) {
 248     L->top = L->base+1;
 249     if (start == 0)
 250       setudataV(L, L->base, IOSTDF_UD(L, GCROOT_IO_OUTPUT));
 251     return 1;
 252   }
 253   return luaL_fileresult(L, status, NULL);
 254 }
 255
 256 /* -- I/O file methods ---------------------------------------------------- */
 257
 258 #define LJLIB_MODULE_io_method
 259
 260 LJLIB_CF(io_method_close)
 261 {
 262   IOFileUD *iof = L->base < L->top ? io_tofile(L) :
 263                   IOSTDF_IOF(L, GCROOT_IO_OUTPUT);
 264   return io_file_close(L, iof);
 265 }
 266
 267 LJLIB_CF(io_method_read)
 268 {
 269   return io_file_read(L, io_tofile(L)->fp, 1);
 270 }
 271
 272 LJLIB_CF(io_method_write)               LJLIB_REC(io_write 0)
 273 {
 274   return io_file_write(L, io_tofile(L)->fp, 1);
 275 }
 276
 277 LJLIB_CF(io_method_flush)               LJLIB_REC(io_flush 0)
 278 {
 279   return luaL_fileresult(L, fflush(io_tofile(L)->fp) == 0, NULL);
 280 }
 281
 282 LJLIB_CF(io_method_seek)
 283 {
 284   FILE *fp = io_tofile(L)->fp;
 285   int opt = lj_lib_checkopt(L, 2, 1, "\3set\3cur\3end");
 286   int64_t ofs = 0;
 287   cTValue *o;
 288   int res;
 289   if (opt == 0) opt = SEEK_SET;
 290   else if (opt == 1) opt = SEEK_CUR;
 291   else if (opt == 2) opt = SEEK_END;
 292   o = L->base+2;
 293   if (o < L->top) {
 294     if (tvisint(o))
 295       ofs = (int64_t)intV(o);
 296     else if (tvisnum(o))
 297       ofs = (int64_t)numV(o);
 298     else if (!tvisnil(o))
 299       lj_err_argt(L, 3, LUA_TNUMBER);
 300   }
 301 #if LJ_TARGET_POSIX
 302   res = fseeko(fp, ofs, opt);
 303 #elif _MSC_VER >= 1400
 304   res = _fseeki64(fp, ofs, opt);
 305 #elif defined(__MINGW32__)
 306   res = fseeko64(fp, ofs, opt);
 307 #else
 308   res = fseek(fp, (long)ofs, opt);
 309 #endif
 310   if (res)
 311     return luaL_fileresult(L, 0, NULL);
 312 #if LJ_TARGET_POSIX
 313   ofs = ftello(fp);
 314 #elif _MSC_VER >= 1400
 315   ofs = _ftelli64(fp);
 316 #elif defined(__MINGW32__)
 317   ofs = ftello64(fp);
 318 #else
 319   ofs = (int64_t)ftell(fp);
 320 #endif
 321   setint64V(L->top-1, ofs);
 322   return 1;
 323 }
 324
 325 LJLIB_CF(io_method_setvbuf)
 326 {
 327   FILE *fp = io_tofile(L)->fp;
 328   int opt = lj_lib_checkopt(L, 2, -1, "\4full\4line\2no");
 329   size_t sz = (size_t)lj_lib_optint(L, 3, LUAL_BUFFERSIZE);
 330   if (opt == 0) opt = _IOFBF;
 331   else if (opt == 1) opt = _IOLBF;
 332   else if (opt == 2) opt = _IONBF;
 333   return luaL_fileresult(L, setvbuf(fp, NULL, opt, sz) == 0, NULL);
 334 }
 335
 336 LJLIB_PUSH(top-2)  /* io_lines_iter */
 337 LJLIB_CF(io_method_lines)
 338 {
 339   io_tofile(L);
 340   setfuncV(L, L->top, funcV(lj_lib_upvalue(L, 1)));
 341   setudataV(L, L->top+1, udataV(L->base));
 342   L->top += 2;
 343   return 2;
 344 }
 345
 346 LJLIB_CF(io_method___gc)
 347 {
 348   IOFileUD *iof = io_tofilep(L);
 349   if (iof->fp != NULL && (iof->type & IOFILE_TYPE_MASK) != IOFILE_TYPE_STDF)
 350     io_file_close(L, iof);
 351   return 0;
 352 }
 353
 354 LJLIB_CF(io_method___tostring)
 355 {
 356   IOFileUD *iof = io_tofilep(L);
 357   if (iof->fp != NULL)
 358     lua_pushfstring(L, "file (%p)", iof->fp);
 359   else
 360     lua_pushliteral(L, "file (closed)");
 361   return 1;
 362 }
 363
 364 LJLIB_PUSH(top-1) LJLIB_SET(__index)
 365
 366 #include "lj_libdef.h"
 367
 368 /* -- I/O library functions ----------------------------------------------- */
 369
 370 #define LJLIB_MODULE_io
 371
 372 LJLIB_PUSH(top-2) LJLIB_SET(!)  /* Set environment. */
 373
 374 LJLIB_CF(io_open)
 375 {
 376   const char *fname = strdata(lj_lib_checkstr(L, 1));
 377   GCstr *s = lj_lib_optstr(L, 2);
 378   const char *mode = s ? strdata(s) : "r";
 379   IOFileUD *iof = io_file_new(L);
 380   iof->fp = fopen(fname, mode);
 381   return iof->fp != NULL ? 1 : luaL_fileresult(L, 0, fname);
 382 }
 383
 384 LJLIB_CF(io_popen)
 385 {
 386 #if LJ_TARGET_POSIX || LJ_TARGET_WINDOWS
 387   const char *fname = strdata(lj_lib_checkstr(L, 1));
 388   GCstr *s = lj_lib_optstr(L, 2);
 389   const char *mode = s ? strdata(s) : "r";
 390   IOFileUD *iof = io_file_new(L);
 391   iof->type = IOFILE_TYPE_PIPE;
 392 #if LJ_TARGET_POSIX
 393   fflush(NULL);
 394   iof->fp = popen(fname, mode);
 395 #else
 396   iof->fp = _popen(fname, mode);
 397 #endif
 398   return iof->fp != NULL ? 1 : luaL_fileresult(L, 0, fname);
 399 #else
 400   return luaL_error(L, LUA_QL("popen") " not supported");
 401 #endif
 402 }
 403
 404 LJLIB_CF(io_tmpfile)
 405 {
 406   IOFileUD *iof = io_file_new(L);
 407 #if LJ_TARGET_PS3
 408   iof->fp = NULL; errno = ENOSYS;
 409 #else
 410   iof->fp = tmpfile();
 411 #endif
 412   return iof->fp != NULL ? 1 : luaL_fileresult(L, 0, NULL);
 413 }
 414
 415 LJLIB_CF(io_close)
 416 {
 417   return lj_cf_io_method_close(L);
 418 }
 419
 420 LJLIB_CF(io_read)
 421 {
 422   return io_file_read(L, io_stdfile(L, GCROOT_IO_INPUT), 0);
 423 }
 424
 425 LJLIB_CF(io_write)              LJLIB_REC(io_write GCROOT_IO_OUTPUT)
 426 {
 427   return io_file_write(L, io_stdfile(L, GCROOT_IO_OUTPUT), 0);
 428 }
 429
 430 LJLIB_CF(io_flush)              LJLIB_REC(io_flush GCROOT_IO_OUTPUT)
 431 {
 432   return luaL_fileresult(L, fflush(io_stdfile(L, GCROOT_IO_OUTPUT)) == 0, NULL);
 433 }
 434
 435 static int io_std_getset(lua_State *L, ptrdiff_t id, const char *mode)
 436 {
 437   if (L->base < L->top && !tvisnil(L->base)) {
 438     if (tvisudata(L->base)) {
 439       io_tofile(L);
 440       L->top = L->base+1;
 441     } else {
 442       io_file_open(L, mode);
 443     }
 444     /* NOBARRIER: The standard I/O handles are GC roots. */
 445     setgcref(G(L)->gcroot[id], gcV(L->top-1));
 446   } else {
 447     setudataV(L, L->top++, IOSTDF_UD(L, id));
 448   }
 449   return 1;
 450 }
 451
 452 LJLIB_CF(io_input)
 453 {
 454   return io_std_getset(L, GCROOT_IO_INPUT, "r");
 455 }
 456
 457 LJLIB_CF(io_output)
 458 {
 459   return io_std_getset(L, GCROOT_IO_OUTPUT, "w");
 460 }
 461
 462 LJLIB_NOREG LJLIB_CF(io_lines_iter)
 463 {
 464   IOFileUD *iof = io_tofile(L);
 465   int ok = io_file_readline(L, iof->fp, 1);
 466   if (ferror(iof->fp))
 467     lj_err_callermsg(L, strerror(errno));
 468   if (!ok && (iof->type & IOFILE_FLAG_CLOSE))
 469     io_file_close(L, iof);  /* Return values are ignored (ok is 0). */
 470   return ok;
 471 }
 472
 473 LJLIB_PUSH(top-3)  /* io_lines_iter */
 474 LJLIB_CF(io_lines)
 475 {
 476   if (L->base < L->top && !tvisnil(L->base)) {  /* io.lines(fname) */
 477     IOFileUD *iof = io_file_open(L, "r");
 478     iof->type = IOFILE_TYPE_FILE|IOFILE_FLAG_CLOSE;
 479     setfuncV(L, L->top-2, funcV(lj_lib_upvalue(L, 1)));
 480   } else {  /* io.lines() iterates over stdin. */
 481     setfuncV(L, L->top, funcV(lj_lib_upvalue(L, 1)));
 482     setudataV(L, L->top+1, IOSTDF_UD(L, GCROOT_IO_INPUT));
 483     L->top += 2;
 484   }
 485   return 2;
 486 }
 487
 488 LJLIB_CF(io_type)
 489 {
 490   cTValue *o = lj_lib_checkany(L, 1);
 491   if (!(tvisudata(o) && udataV(o)->udtype == UDTYPE_IO_FILE))
 492     setnilV(L->top++);
 493   else if (((IOFileUD *)uddata(udataV(o)))->fp != NULL)
 494     lua_pushliteral(L, "file");
 495   else
 496     lua_pushliteral(L, "closed file");
 497   return 1;
 498 }
 499
 500 #include "lj_libdef.h"
 501
 502 /* ------------------------------------------------------------------------ */
 503
 504 static GCobj *io_std_new(lua_State *L, FILE *fp, const char *name)
 505 {
 506   IOFileUD *iof = (IOFileUD *)lua_newuserdata(L, sizeof(IOFileUD));
 507   GCudata *ud = udataV(L->top-1);
 508   ud->udtype = UDTYPE_IO_FILE;
 509   /* NOBARRIER: The GCudata is new (marked white). */
 510   setgcref(ud->metatable, gcV(L->top-3));
 511   iof->fp = fp;
 512   iof->type = IOFILE_TYPE_STDF;
 513   lua_setfield(L, -2, name);
 514   return obj2gco(ud);
 515 }
 516
 517 LUALIB_API int luaopen_io(lua_State *L)
 518 {
 519   lj_lib_pushcf(L, lj_cf_io_lines_iter, FF_io_lines_iter);
 520   LJ_LIB_REG(L, NULL, io_method);
 521   copyTV(L, L->top, L->top-1); L->top++;
 522   lua_setfield(L, LUA_REGISTRYINDEX, LUA_FILEHANDLE);
 523   LJ_LIB_REG(L, LUA_IOLIBNAME, io);
 524   setgcref(G(L)->gcroot[GCROOT_IO_INPUT], io_std_new(L, stdin, "stdin"));
 525   setgcref(G(L)->gcroot[GCROOT_IO_OUTPUT], io_std_new(L, stdout, "stdout"));
 526   io_std_new(L, stderr, "stderr");
 527   return 1;
 528 }
 529