1 files changed, 119 insertions, 90 deletions

diff --git a/ext/io/wait/wait.c b/ext/io/wait/wait.c
index d846bba49e..a82295f72e 100644
--- a/ext/io/wait/wait.c
+++ b/ext/io/wait/wait.c
@@ -39,35 +39,6 @@
 #define FIONREAD_POSSIBLE_P(fd) ((void)(fd),Qtrue)
 #endif
 
-static VALUE io_ready_p _((VALUE io));
-static VALUE io_wait_readable _((int argc, VALUE *argv, VALUE io));
-static VALUE io_wait_writable _((int argc, VALUE *argv, VALUE io));
-void Init_wait _((void));
-
-static struct timeval *
-get_timeout(int argc, VALUE *argv, struct timeval *timerec)
-{
-    VALUE timeout = Qnil;
-    rb_check_arity(argc, 0, 1);
-    if (!argc || NIL_P(timeout = argv[0])) {
-	return NULL;
-    }
-    else {
-	*timerec = rb_time_interval(timeout);
-	return timerec;
-    }
-}
-
-static int
-wait_for_single_fd(rb_io_t *fptr, int events, struct timeval *tv)
-{
-    int i = rb_wait_for_single_fd(fptr->fd, events, tv);
-    if (i < 0)
-	rb_sys_fail(0);
-    rb_io_check_closed(fptr);
-    return (i & events);
-}
-
 /*
  * call-seq:
  *   io.nread -> int
@@ -79,13 +50,12 @@ wait_for_single_fd(rb_io_t *fptr, int events, struct timeval *tv)
 static VALUE
 io_nread(VALUE io)
 {
-    rb_io_t *fptr;
-    int len;
+    rb_io_t *fptr = NULL;
     ioctl_arg n;
 
     GetOpenFile(io, fptr);
     rb_io_check_readable(fptr);
-    len = rb_io_read_pending(fptr);
+    int len = rb_io_read_pending(fptr);
     if (len > 0) return INT2FIX(len);
     if (!FIONREAD_POSSIBLE_P(fptr->fd)) return INT2FIX(0);
     if (ioctl(fptr->fd, FIONREAD, &n)) return INT2FIX(0);
@@ -93,76 +63,113 @@ io_nread(VALUE io)
     return INT2FIX(0);
 }
 
+static VALUE
+io_wait_event(VALUE io, int event, VALUE timeout)
+{
+    VALUE result = rb_io_wait(io, RB_INT2NUM(event), timeout);
+
+    if (!RB_TEST(result)) {
+        return Qnil;
+    }
+
+    int mask = RB_NUM2INT(result);
+
+    if (mask & event) {
+        return io;
+    } else {
+        return Qfalse;
+    }
+}
+
 /*
  * call-seq:
  *   io.ready? -> true or false
  *
- * Returns true if input available without blocking, or false.
+ * Returns +true+ if input available without blocking, or +false+.
  */
 
 static VALUE
 io_ready_p(VALUE io)
 {
     rb_io_t *fptr;
-    struct timeval tv = {0, 0};
 
     GetOpenFile(io, fptr);
     rb_io_check_readable(fptr);
     if (rb_io_read_pending(fptr)) return Qtrue;
-    if (wait_for_single_fd(fptr, RB_WAITFD_IN, &tv))
-	return Qtrue;
-    return Qfalse;
+
+    return io_wait_event(io, RUBY_IO_READABLE, RB_INT2NUM(0));
 }
 
 /*
  * call-seq:
- *   io.wait_readable          -> IO, true or nil
- *   io.wait_readable(timeout) -> IO, true or nil
+ *   io.wait_readable          -> true or false
+ *   io.wait_readable(timeout) -> true or false
  *
- * Waits until IO is readable without blocking and returns +self+, or
- * +nil+ when times out.
+ * Waits until IO is readable and returns +true+, or
+ * +false+ when times out.
  * Returns +true+ immediately when buffered data is available.
  */
 
 static VALUE
 io_wait_readable(int argc, VALUE *argv, VALUE io)
 {
-    rb_io_t *fptr;
-    struct timeval timerec;
-    struct timeval *tv;
+    rb_io_t *fptr = NULL;
 
-    GetOpenFile(io, fptr);
+    RB_IO_POINTER(io, fptr);
     rb_io_check_readable(fptr);
-    tv = get_timeout(argc, argv, &timerec);
+
     if (rb_io_read_pending(fptr)) return Qtrue;
-    if (wait_for_single_fd(fptr, RB_WAITFD_IN, tv)) {
-	return io;
-    }
-    return Qnil;
+
+    rb_check_arity(argc, 0, 1);
+    VALUE timeout = (argc == 1 ? argv[0] : Qnil);
+
+    return io_wait_event(io, RUBY_IO_READABLE, timeout);
 }
 
 /*
  * call-seq:
- *   io.wait_writable          -> IO
- *   io.wait_writable(timeout) -> IO or nil
+ *   io.wait_writable          -> true or false
+ *   io.wait_writable(timeout) -> true or false
  *
- * Waits until IO is writable without blocking and returns +self+ or
- * +nil+ when times out.
+ * Waits until IO is writable and returns +true+ or
+ * +false+ when times out.
  */
 static VALUE
 io_wait_writable(int argc, VALUE *argv, VALUE io)
 {
-    rb_io_t *fptr;
-    struct timeval timerec;
-    struct timeval *tv;
+    rb_io_t *fptr = NULL;
 
-    GetOpenFile(io, fptr);
+    RB_IO_POINTER(io, fptr);
     rb_io_check_writable(fptr);
-    tv = get_timeout(argc, argv, &timerec);
-    if (wait_for_single_fd(fptr, RB_WAITFD_OUT, tv)) {
-	return io;
-    }
-    return Qnil;
+
+    rb_check_arity(argc, 0, 1);
+    VALUE timeout = (argc == 1 ? argv[0] : Qnil);
+
+    return io_wait_event(io, RUBY_IO_WRITABLE, timeout);
+}
+
+/*
+ * call-seq:
+ *   io.wait_priority          -> true or false
+ *   io.wait_priority(timeout) -> true or false
+ *
+ * Waits until IO is priority and returns +true+ or
+ * +false+ when times out.
+ */
+static VALUE
+io_wait_priority(int argc, VALUE *argv, VALUE io)
+{
+    rb_io_t *fptr = NULL;
+
+    RB_IO_POINTER(io, fptr);
+    rb_io_check_readable(fptr);
+
+    if (rb_io_read_pending(fptr)) return Qtrue;
+
+    rb_check_arity(argc, 0, 1);
+    VALUE timeout = argc == 1 ? argv[0] : Qnil;
+
+    return io_wait_event(io, RUBY_IO_PRIORITY, timeout);
 }
 
 static int
@@ -201,41 +208,60 @@ wait_mode_sym(VALUE mode)
 
 /*
  * call-seq:
- *   io.wait(timeout = nil, mode = :read) -> IO, true or nil
+ *   io.wait(events, timeout) -> event mask or false.
+ *   io.wait(timeout = nil, mode = :read) -> event mask or false (deprecated)
+ *
+ * Waits until the IO becomes ready for the specified events and returns the
+ * subset of events that become ready, or +false+ when times out.
+ *
+ * The events can be a bit mask of +IO::READABLE+, +IO::WRITABLE+ or
+ * +IO::PRIORITY+.
  *
- * Waits until IO is readable or writable without blocking and returns
- * +self+, or +nil+ when times out.
  * Returns +true+ immediately when buffered data is available.
+ *
  * Optional parameter +mode+ is one of +:read+, +:write+, or
- * +:read_write+.
+ * +:read_write+ (deprecated).
  */
 
 static VALUE
-io_wait_readwrite(int argc, VALUE *argv, VALUE io)
+io_wait(int argc, VALUE *argv, VALUE io)
 {
-    rb_io_t *fptr;
-    struct timeval timerec;
-    struct timeval *tv = NULL;
-    int event = 0;
-    int i;
+    VALUE timeout = Qnil;
+    rb_io_event_t events = 0;
 
-    GetOpenFile(io, fptr);
-    for (i = 0; i < argc; ++i) {
-	if (SYMBOL_P(argv[i])) {
-	    event |= wait_mode_sym(argv[i]);
-	}
-	else {
-	    *(tv = &timerec) = rb_time_interval(argv[i]);
-	}
+    if (argc < 2 || (argc >= 2 && RB_SYMBOL_P(argv[1]))) {
+        if (argc > 0) {
+          timeout = argv[0];
+        }
+
+        for (int i = 1; i < argc; i += 1) {
+            events |= wait_mode_sym(argv[i]);
+        }
+    } else if (argc == 2) {
+        events = RB_NUM2UINT(argv[0]);
+        
+        if (argv[1] != Qnil) {
+          timeout = argv[1];
+        }
+    } else {
+        // TODO error
+        return Qnil;
     }
-    /* rb_time_interval() and might_mode() might convert the argument */
-    rb_io_check_closed(fptr);
-    if (!event) event = RB_WAITFD_IN;
-    if ((event & RB_WAITFD_IN) && rb_io_read_pending(fptr))
-	return Qtrue;
-    if (wait_for_single_fd(fptr, event, tv))
-	return io;
-    return Qnil;
+
+    if (events == 0) {
+        events = RUBY_IO_READABLE;
+    }
+
+    if (events & RUBY_IO_READABLE) {
+        rb_io_t *fptr = NULL;
+        RB_IO_POINTER(io, fptr);
+
+        if (rb_io_read_pending(fptr)) {
+            return Qtrue;
+        }
+    }
+
+    return io_wait_event(io, events, timeout);
 }
 
 /*
@@ -247,7 +273,10 @@ Init_wait(void)
 {
     rb_define_method(rb_cIO, "nread", io_nread, 0);
     rb_define_method(rb_cIO, "ready?", io_ready_p, 0);
-    rb_define_method(rb_cIO, "wait", io_wait_readwrite, -1);
+
+    rb_define_method(rb_cIO, "wait", io_wait, -1);
+
     rb_define_method(rb_cIO, "wait_readable", io_wait_readable, -1);
     rb_define_method(rb_cIO, "wait_writable", io_wait_writable, -1);
+    rb_define_method(rb_cIO, "wait_priority", io_wait_priority, -1);
 }