--- enumerator.c | 79 +++++++++++++++++++++++++++ test/ruby/test_arithmetic_sequence.rb | 17 ++++++ 2 files changed, 96 insertions(+) diff --git a/enumerator.c b/enumerator.c index 5f21455ddd..2f22afcc92 100644 --- a/enumerator.c +++ b/enumerator.c @@ -3624,6 +3624,83 @@ arith_seq_inspect(VALUE self) return str; } + +/* + * call-seq: + * arith_sequence.include?(obj) -> true or false + * arith_sequence.member?(obj) -> true or false + * + * Returns true if any member of arith_sequence + * equals obj. If obj is an integer and the beginning + * and step are integers, then this computes whether or not the + * sequence includes obj using arithmetic. If obj is a + * float, then TBD. (It might just fall back to the existing behavior + * of enumerating the sequence until the sequence ends or the number + * is found.) + * + * (1..10).include? 5 #=> true + * (1..10).include? 15 #=> false + * (1..10).member? 5 #=> true + * (1..10).member? 15 #=> false + * + */ + +static VALUE +arith_seq_member_p(VALUE self, VALUE element) +{ + VALUE b, s, offset, remainder; + + b = arith_seq_begin(self); + + if (!(CLASS_OF(b) == CLASS_OF(element))) { + return Qfalse; + } + + s = arith_seq_step(self); + + if ( (RB_FLOAT_TYPE_P(s) && FLOAT_ZERO_P(s)) || + (RB_INTEGER_TYPE_P(s) && FIXNUM_ZERO_P(s))) { + if (element == b) return Qtrue; + return Qfalse; + } + + /* TODO: Handle floats */ + + offset = rb_int_minus(element, b); + + /* Determine if the sequence will be going away from the element, + * e.g. if the sequence is 1.step then no non-positive numbers + * will be in the sequence. TODO: Also check the end of the range, + * if it exists. + */ + + if (FIXNUM_POSITIVE_P(offset)) { + if (!FIXNUM_POSITIVE_P(s)) { + return Qfalse; + } + } else { + if (FIXNUM_POSITIVE_P(s)) { + return Qfalse; + } + } + + /* Test to see if the difference between the element between + * checked and the beginning of the sequence is divisible by the + * step of the sequence. If it is, then the element is included in + * the sequence. Otherwise, it's not. + */ + + remainder = rb_int_modulo(offset, s); + + if (FIXNUM_ZERO_P(remainder)) { + return Qtrue; + } + else { + return Qfalse; + } +} + + /* * call-seq: * aseq == obj -> true or false @@ -4100,6 +4177,8 @@ InitVM_Enumerator(void) rb_define_method(rb_cArithSeq, "hash", arith_seq_hash, 0); rb_define_method(rb_cArithSeq, "each", arith_seq_each, 0); rb_define_method(rb_cArithSeq, "size", arith_seq_size, 0); + rb_define_method(rb_cArithSeq, "member?", arith_seq_member_p, 1); + rb_define_method(rb_cArithSeq, "include?", arith_seq_member_p, 1); rb_provide("enumerator.so"); /* for backward compatibility */ } diff --git a/test/ruby/test_arithmetic_sequence.rb b/test/ruby/test_arithmetic_sequence.rb index 45a0ab9222..321fa42202 100644 --- a/test/ruby/test_arithmetic_sequence.rb +++ b/test/ruby/test_arithmetic_sequence.rb @@ -486,4 +486,21 @@ def test_sum assert_equal([1.0, 2.5, 4.0, 5.5, 7.0, 8.5, 10.0].sum, (1.0..10.0).step(1.5).sum) assert_equal([1/2r, 1r, 3/2r, 2, 5/2r, 3, 7/2r, 4].sum, ((1/2r)...(9/2r)).step(1/2r).sum) end + + def test_member_p + assert_send([1.step, :member?, 10]) + assert_not_send([1.step(by: 2), :member?, 10]) + assert_send([1.step(by: 0), :member?, 1]) + assert_not_send([1.step(by: 0), :member?, 2]) + assert_send([1.step(by: -1), :member?, -5]) + assert_not_send([1.step(by: -1), :member?, 5]) + + assert_send([(1..10).step, :member?, 6]) + assert_not_send([(1..10).step(2), :member?, 6]) + assert_not_send([(1..10).step, :member?, 11]) + assert_not_send([(-5..5).step, :member? -10]) + + assert_send([1.5.step, :member?, 2.5]) + assert_not_send([1.5.step, :member?, 3]) + end end -- 2.24.0