Ruby - Subclassing array to make it randomize when flattened

Question

I'm trying to subclass Array in ruby to make it randomize its elements when flatten! is called. Looking at the source code for Array#flatten (http://ruby-doc.org/core/classes/Array.src/M002218.html), it looks like it should recursively call flatten! on any array contained within an array. So, I tried doing something like this:

class RandArray < Array
    def randomize!
        self.sort!{rand(3)-1}
    end
    def flatten!
        randomize!
        super
    end
end

However, when a normal array contains my RandArray and flatten is called on the normal array, flatten! is never called in my array. I figure ruby is just calling some other method to flatten the arrays recursively, but I can't figure out what that is. Any tips?

Answer 1

I am not an absolute expert on this but Ruby's Array is written as C code. here is the code for flatten! :

static VALUE
rb_ary_flatten_bang(ary)
    VALUE ary;
{
    long i = 0;
    int mod = 0;
    VALUE memo = Qnil;

    while (i<RARRAY(ary)->len) {
        VALUE ary2 = RARRAY(ary)->ptr[i];
        VALUE tmp;

        tmp = rb_check_array_type(ary2);
        if (!NIL_P(tmp)) {
            if (NIL_P(memo)) {
                memo = rb_ary_new();
            }
            i += flatten(ary, i, tmp, memo);
            mod = 1;
        }
        i++;
    }
    if (mod == 0) return Qnil;
    return ary;
}

As you can see on this line,

i += flatten(ary, i, tmp, memo);

and here is the implementation for this flatten C function :

static long
flatten(ary, idx, ary2, memo)
    VALUE ary;
    long idx;
    VALUE ary2, memo;
{
    VALUE id;
    long i = idx;
    long n, lim = idx + RARRAY(ary2)->len;

    id = rb_obj_id(ary2);
    if (rb_ary_includes(memo, id)) {
    rb_raise(rb_eArgError, "tried to flatten recursive array");
    }
    rb_ary_push(memo, id);
    rb_ary_splice(ary, idx, 1, ary2);
    while (i < lim) {
    VALUE tmp;

    tmp = rb_check_array_type(rb_ary_elt(ary, i));
    if (!NIL_P(tmp)) {
        n = flatten(ary, i, tmp, memo);
        i += n; lim += n;
    }
    i++;
    }
    rb_ary_pop(memo);

    return lim - idx - 1;   /* returns number of increased items */
}

The flatten! code calls directly the C flatten function for any element of the array that validates rb_check_array_type it doesn't go back to the ruby code.Instead it accesses the underlying C structure directly bypassing your overloaded implementation.

Not sure how to override this, I think one way could be to reopen Array and rewrite the flatten and flatten! function as pure ruby. You would take a performance hit, but then you would be able to overload it as you see fit. And you could always use aliasing to have a "flatten_native" and a "flatten_native!" function on your modified array, to get the perfs back on some cases.

Answer 2

Jean is correct, flatten calls a C function behind the scenes. You could patch the Array class and override the default flatten! method, while still retaining access to the original method.

class Array
  alias_method :old_flatten!, :flatten!
  def flatten!
    self.old_flatten!
    self.sort!{rand(3)-1}
  end
end

Or you could just add a flatten_with_randomize! to the Array class and use that instead and keep the original flatten! method intact.