From 6bedbf462544a7917fdc8d8c44276079a6e156cf Mon Sep 17 00:00:00 2001 From: Yusuke Endoh Date: Sun, 28 Apr 2019 23:24:09 +0900 Subject: [PATCH] numeric.c: Extend Integer#[] to support range arguments ```` 0b01001101[2, 4] #=> 0b0011 0b01001100[2..5] #=> 0b0011 0b01001100[2...6] #=> 0b0011 ^^^^ ```` [Feature #8842] --- numeric.c | 148 ++++++++++++++++++++++++++++++++------ test/ruby/test_integer.rb | 25 +++++++ 2 files changed, 150 insertions(+), 23 deletions(-) diff --git a/numeric.c b/numeric.c index 4af829742a..b5b62ae469 100644 --- a/numeric.c +++ b/numeric.c @@ -4630,24 +4630,6 @@ rb_int_rshift(VALUE x, VALUE y) return Qnil; } -/* - * Document-method: Integer#[] - * call-seq: - * int[n] -> 0, 1 - * - * Bit Reference---Returns the nth bit in the - * binary representation of +int+, where int[0] - * is the least significant bit. - * - * a = 0b11001100101010 - * 30.downto(0) {|n| print a[n] } - * #=> 0000000000000000011001100101010 - * - * a = 9**15 - * 50.downto(0) {|n| print a[n] } - * #=> 000101110110100000111000011110010100111100010111001 - */ - static VALUE fix_aref(VALUE fix, VALUE idx) { @@ -4675,18 +4657,138 @@ fix_aref(VALUE fix, VALUE idx) return INT2FIX(0); } -static VALUE -int_aref(VALUE num, VALUE idx) + +/* copied from "r_less" in range.c */ +/* compares _a_ and _b_ and returns: + * < 0: a < b + * = 0: a = b + * > 0: a > b or non-comparable + */ +static int +compare_indexes(VALUE a, VALUE b) { + VALUE r = rb_funcall(a, id_cmp, 1, b); + + if (NIL_P(r)) + return INT_MAX; + return rb_cmpint(r, a, b); +} + +static VALUE +generate_mask(VALUE len) { + return rb_int_minus(rb_int_lshift(INT2FIX(1), len), INT2FIX(1)); +} + +static VALUE +int_aref1(VALUE num, VALUE arg) +{ + VALUE orig_num = num, beg, end; + int excl; + + if (rb_range_values(arg, &beg, &end, &excl)) { + if (NIL_P(beg)) { + /* beginless range */ + if (!RTEST(num_negative_p(end))) { + if (!excl) end = rb_int_plus(end, INT2FIX(1)); + VALUE mask = generate_mask(end); + if (RTEST(num_zero_p(rb_int_and(num, mask)))) { + return INT2FIX(0); + } + else { + rb_raise(rb_eArgError, "The beginless range for Integer#[] results in infinity"); + } + } + else { + return INT2FIX(0); + } + } + num = rb_int_rshift(num, beg); + + int cmp = compare_indexes(beg, end); + if (!NIL_P(end) && cmp < 0) { + VALUE len = rb_int_minus(end, beg); + if (!excl) len = rb_int_plus(len, INT2FIX(1)); + VALUE mask = generate_mask(len); + num = rb_int_and(num, mask); + } + else if (cmp == 0) { + if (excl) return INT2FIX(0); + num = orig_num; + arg = beg; + goto one_bit; + } + return num; + } + +one_bit: if (FIXNUM_P(num)) { - return fix_aref(num, idx); + return fix_aref(num, arg); } else if (RB_TYPE_P(num, T_BIGNUM)) { - return rb_big_aref(num, idx); + return rb_big_aref(num, arg); } return Qnil; } +static VALUE +int_aref2(VALUE num, VALUE beg, VALUE len) +{ + num = rb_int_rshift(num, beg); + VALUE mask = generate_mask(len); + num = rb_int_and(num, mask); + return num; +} + +/* + * Document-method: Integer#[] + * call-seq: + * int[n] -> 0, 1 + * int[n, m] -> num + * int[range] -> num + * + * Bit Reference---Returns the nth bit in the + * binary representation of +int+, where int[0] + * is the least significant bit. + * + * a = 0b11001100101010 + * 30.downto(0) {|n| print a[n] } + * #=> 0000000000000000011001100101010 + * + * a = 9**15 + * 50.downto(0) {|n| print a[n] } + * #=> 000101110110100000111000011110010100111100010111001 + * + * In principle, n[i] is equivalent to (n >> i) & 1. + * Thus, any negative index always returns zero: + * + * p 255[-1] #=> 0 + * + * Range operations n[i, len] and n[i..j] + * are naturally extended. + * + * * n[i, len] equals to (n >> i) & ((1 << len) - 1). + * * n[i..j] equals to (n >> i) & ((1 << (j - i + 1)) - 1). + * * n[i...j] equals to (n >> i) & ((1 << (j - i)) - 1). + * * n[i..] equals to (n >> i). + * * n[..j] is zero if n & ((1 << (j + 1)) - 1) is zero. Otherwise, raises an ArgumentError. + * * n[...j] is zero if n & ((1 << j) - 1) is zero. Otherwise, raises an ArgumentError. + * + * Note that range operation may exhaust memory. + * For example, -1[0, 1000000000000] will raise NoMemoryError. + */ + +static VALUE +int_aref(int const argc, VALUE * const argv, VALUE const num) +{ + rb_check_arity(argc, 1, 2); + if (argc == 2) { + return int_aref2(num, argv[0], argv[1]); + } + return int_aref1(num, argv[0]); + + return Qnil; +} + /* * Document-method: Integer#to_f * call-seq: @@ -5555,7 +5657,7 @@ Init_Numeric(void) rb_define_method(rb_cInteger, "&", rb_int_and, 1); rb_define_method(rb_cInteger, "|", int_or, 1); rb_define_method(rb_cInteger, "^", int_xor, 1); - rb_define_method(rb_cInteger, "[]", int_aref, 1); + rb_define_method(rb_cInteger, "[]", int_aref, -1); rb_define_method(rb_cInteger, "<<", rb_int_lshift, 1); rb_define_method(rb_cInteger, ">>", rb_int_rshift, 1); diff --git a/test/ruby/test_integer.rb b/test/ruby/test_integer.rb index ad088aa72f..2334cab50a 100644 --- a/test/ruby/test_integer.rb +++ b/test/ruby/test_integer.rb @@ -23,6 +23,31 @@ class TestInteger < Test::Unit::TestCase (-64..64).each do |idx| assert_equal((n >> idx) & 1, n[idx]) end + [*-66..-62, *-34..-30, *-5..5, *30..34, *62..66].each do |idx| + (0..100).each do |len| + assert_equal((n >> idx) & ((1 << len) - 1), n[idx, len], "#{ n }[#{ idx }, #{ len }]") + end + (0..100).each do |len| + assert_equal((n >> idx) & ((1 << (len + 1)) - 1), n[idx..idx+len], "#{ n }[#{ idx }..#{ idx+len }]") + assert_equal((n >> idx) & ((1 << len) - 1), n[idx...idx+len], "#{ n }[#{ idx }...#{ idx+len }]") + end + + # endless + assert_equal((n >> idx), n[idx..], "#{ n }[#{ idx }..]") + assert_equal((n >> idx), n[idx...], "#{ n }[#{ idx }...#]") + + # beginless + if idx >= 0 && n & ((1 << (idx + 1)) - 1) != 0 + assert_raise(ArgumentError, "#{ n }[..#{ idx }]") { n[..idx] } + else + assert_equal(0, n[..idx], "#{ n }[..#{ idx }]") + end + if idx >= 0 && n & ((1 << idx) - 1) != 0 + assert_raise(ArgumentError, "#{ n }[...#{ idx }]") { n[...idx] } + else + assert_equal(0, n[...idx], "#{ n }[...#{ idx }]") + end + end end # assert_equal(1, (1 << 0x40000000)[0x40000000], "[ruby-dev:31271]")