D3.js's nest operator in Ruby

nest.rb

# D3.js's nest operator in Ruby

class Nest
  def initialize
    #   var nest = {}, keys = [], sortKeys = [], sortValues, rollup;
    @nest = {}
    @keys = []
    @sort_keys = []
  end

  #   nest.map = function(array) {
  #     return map(array, 0);
  #   };
  def map(array)
    _map(array, 0)
  end

  #   nest.entries = function(array) {
  #     return entries(map(array, 0), 0);
  #   };
  def entries(array)
    _entries(_map(array, 0), 0)
  end

  #   nest.key = function(d) {
  #     keys.push(d);
  #     return nest;
  #   };
  def key(&f)
    @keys << f
    self
  end

  #   nest.sortKeys = function(order) {
  #     sortKeys[keys.length - 1] = order;
  #     return nest;
  #   };
  def sort_keys(&order)
    @sort_keys[@keys.size - 1] = order
    self
  end

  #   nest.sortValues = function(order) {
  #     sortValues = order;
  #     return nest;
  #   };
  def sort_values(&order)
    @sort_values = order
    self
  end

  #   nest.rollup = function(f) {
  #     rollup = f;
  #     return nest;
  #   };
  def rollup(&f)
    @rollup = f
    self
  end

  private

  #   function map(array, depth) {
  #     if (depth >= keys.length) return rollup ? rollup.call(nest, array) : sortValues ? array.sort(sortValues) : array;
  #     var i = -1, n = array.length, key = keys[depth++], keyValue, object, valuesByKey = new d3_Map, values, o = {};
  #     while (++i < n) {
  #       if (values = valuesByKey.get(keyValue = key(object = array[i]))) {
  #         values.push(object);
  #       } else {
  #         valuesByKey.set(keyValue, [ object ]);
  #       }
  #     }
  #     valuesByKey.forEach(function(keyValue, values) {
  #       o[keyValue] = map(values, depth);
  #     });
  #     return o;
  #   }
  def _map(array, depth)
    if depth >= @keys.size
      return @rollup.call(array) if @rollup
      return array.sort { |a, b| @sort_values.call(a, b) } if @sort_values
      return array
    end

    key = @keys[depth]
    depth += 1
    values_by_key = {}

    array.each_with_index do |object, i| 
      key_value = key.call(object)
      values_by_key[key_value] ||= []
      values_by_key[key_value] << object
    end

    o = {}
    values_by_key.each do |key_value, values|
      o[key_value] = _map(values, depth)
    end
    o
  end

  #   function entries(map, depth) {
  #     if (depth >= keys.length) return map;
  #     var a = [], sortKey = sortKeys[depth++], key;
  #     for (key in map) {
  #       a.push({
  #         key: key,
  #         values: entries(map[key], depth)
  #       });
  #     }
  #     if (sortKey) a.sort(function(a, b) {
  #       return sortKey(a.key, b.key);
  #     });
  #     return a;
  #   }
  def _entries(map, depth)
    return map if depth >= @keys.size

    a = []
    sort_key = @sort_keys[depth]
    depth += 1

    map.each do |key, values|
      a << {
        key: key,
        values: _entries(values, depth)
      }
    end

    if sort_key
      a.sort { |a, b| sort_key.call(a[:key], b[:key]) }
    else
      a
    end
  end
end


# TEST

# require 'pp'

# data = [
#   {year: "2013", category: "Cars", value: 1000},
#   {year: "2011", category: "Cars", value: 1200},
#   {year: "2012", category: "Cars", value: 1300},
#   {year: "2011", category: "Planes", value: 1000},
#   {year: "2011", category: "Planes", value: 1100},
#   {year: "2012", category: "Planes", value: 1200},
#   {year: "2011", category: "Planes", value: 1300},
#   {year: "2012", category: "Bikes", value: 1000},
#   {year: "2011", category: "Bikes", value: 1100},
#   {year: "2011", category: "Bikes", value: 1200},
# ]

# nest = Nest.new
# nest.key { |d| d[:year] }
# nest.sort_keys { |a,b| a <=> b }
# nest.key { |d| d[:category] }
# # nest.sort_values { |a,b| b[:value] <=> a[:value] }
# nest.rollup do |values| 
#   {
#     count: values.size,
#     total_value: values.map { |d| d[:value] }.reduce(:+),
#     raw: values
#   } 
# end

# # pp nest.entries(data)
# pp nest.map(data)