Implement File#size and File#truncate

mruby · Dec 14, 2019 · 6d9ac89 · 6d9ac89
1 parent 6c5ee8f
commit 6d9ac89
Show file tree

Hide file tree

Showing 2 changed files with 116 additions and 1 deletion.
diff --git a/mrbgems/mruby-io/src/file.c b/mrbgems/mruby-io/src/file.c
@@ -70,7 +70,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
@@ -384,6 +392,95 @@ mrb_file_flock(mrb_state *mrb, mrb_value self)
   return mrb_fixnum_value(0);
 }
 
+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)
 {
@@ -490,6 +587,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));

diff --git a/mrbgems/mruby-io/test/file.rb b/mrbgems/mruby-io/test/file.rb
@@ -80,6 +80,22 @@
   end
 end
 
+assert('File#size and File#truncate') do
+  fname = "#{$mrbtest_io_wfname}.resize"
+  begin
+    File.open(fname, 'w') do |f|
+      assert_equal 0, f.size
+      assert_equal 0, f.truncate(100)
+      assert_equal 100, f.size
+      assert_equal 0, f.pos
+      assert_equal 0, f.truncate(5)
+      assert_equal 5, f.size
+    end
+  ensure
+    File.delete(fname)
+  end
+end
+
 assert('File.join') do
   assert_equal "", File.join()
   assert_equal "a", File.join("a")