Merge branch 'file-size-truncate' of https://github.com/dearblue/mruby into dearblue-file-size-truncate

1 files changed, 100 insertions, 1 deletions

diff --git a/mrbgems/mruby-io/src/file.c b/mrbgems/mruby-io/src/file.c
index 673decc20..9d5bc8ca2 100644
--- a/mrbgems/mruby-io/src/file.c
+++ b/mrbgems/mruby-io/src/file.c
@@ -65,7 +65,15 @@
 #define LOCK_UN 8
 #endif
 
-#define STAT(p, s)        stat(p, s)
+#ifndef _WIN32
+typedef struct stat         mrb_stat;
+# define mrb_stat(path, sb) stat(path, sb)
+# define mrb_fstat(fd, sb)  fstat(fd, sb)
+#else
+typedef struct __stat64     mrb_stat;
+# define mrb_stat(path, sb) _stat64(path, sb)
+# define mrb_fstat(fd, sb)  _fstat64(fd, sb)
+#endif
 
 #ifdef _WIN32
 static int
@@ -380,6 +388,95 @@ mrb_file_flock(mrb_state *mrb, mrb_value self)
 }
 
 static mrb_value
+mrb_file_size(mrb_state *mrb, mrb_value self)
+{
+  mrb_stat st;
+  int fd;
+
+  fd = (int)mrb_fixnum(mrb_io_fileno(mrb, self));
+  if (mrb_fstat(fd, &st) == -1) {
+    mrb_raise(mrb, E_RUNTIME_ERROR, "fstat failed");
+  }
+
+  if (st.st_size > MRB_INT_MAX) {
+#ifdef MRB_WITHOUT_FLOAT
+    mrb_raise(mrb, E_RUNTIME_ERROR, "File#size too large for MRB_WITHOUT_FLOAT");
+#else
+    return mrb_float_value(mrb, st.st_size);
+#endif
+  }
+
+  return mrb_fixnum_value(st.st_size);
+}
+
+static int
+mrb_ftruncate(int fd, int64_t length)
+{
+#ifndef _WIN32
+  return ftruncate(fd, (off_t)length);
+#else
+  HANDLE file;
+  __int64 cur;
+
+  file = (HANDLE)_get_osfhandle(fd);
+  if (file == INVALID_HANDLE_VALUE) {
+    return -1;
+  }
+
+  cur = _lseeki64(fd, 0, SEEK_CUR);
+  if (cur == -1) return -1;
+
+  if (_lseeki64(fd, (__int64)length, SEEK_SET) == -1) return -1;
+
+  if (!SetEndOfFile(file)) {
+    errno = EINVAL; /* TODO: GetLastError to errno */
+    return -1;
+  }
+
+  if (_lseeki64(fd, cur, SEEK_SET) == -1) return -1;
+
+  return 0;
+#endif /* _WIN32 */
+}
+
+static mrb_value
+mrb_file_truncate(mrb_state *mrb, mrb_value self)
+{
+  int fd;
+  int64_t length;
+  mrb_value lenv;
+
+  fd = (int)mrb_fixnum(mrb_io_fileno(mrb, self));
+  mrb_get_args(mrb, "o", &lenv);
+  switch (mrb_type(lenv)) {
+#ifndef MRB_WITHOUT_FLOAT
+    case MRB_TT_FLOAT:
+      {
+        mrb_float lenf = mrb_float(lenv);
+        if (lenf > INT64_MAX) {
+          mrb_raise(mrb, E_ARGUMENT_ERROR, "length too large");
+        }
+        length = (int64_t)lenf;
+      }
+      break;
+#endif
+    case MRB_TT_FIXNUM:
+    default:
+      {
+        mrb_int leni = mrb_int(mrb, lenv);
+        length = (int64_t)leni;
+      }
+      break;
+  }
+
+  if (mrb_ftruncate(fd, length) != 0) {
+    mrb_raise(mrb, E_IO_ERROR, "ftruncate failed");
+  }
+
+  return mrb_fixnum_value(0);
+}
+
+static mrb_value
 mrb_file_s_symlink(mrb_state *mrb, mrb_value klass)
 {
 #if defined(_WIN32) || defined(_WIN64)
@@ -485,6 +582,8 @@ mrb_init_file(mrb_state *mrb)
 
   mrb_define_method(mrb, file, "flock", mrb_file_flock, MRB_ARGS_REQ(1));
   mrb_define_method(mrb, file, "mtime", mrb_file_mtime, MRB_ARGS_NONE());
+  mrb_define_method(mrb, file, "size", mrb_file_size, MRB_ARGS_NONE());
+  mrb_define_method(mrb, file, "truncate", mrb_file_truncate, MRB_ARGS_REQ(1));
 
   cnst = mrb_define_module_under(mrb, file, "Constants");
   mrb_define_const(mrb, cnst, "LOCK_SH", mrb_fixnum_value(LOCK_SH));