Multikey Map

multikey_map.rb

# frozen_string_literal: true

class MultikeyMap
  include Enumerable

  # Creates multikey map from a dict { value => alias | [aliases...] },
  # i.e. both key and values will be mapped to key.
  # @param dict [Hash]
  # @raise [TypeError] if input can't be coerced as hash
  # @return [MultikeyMap]
  def self.build(dict, &block)
    Types::Coercible::Hash.try(dict) { |result| raise result.error.cause }

    mm = new

    if block_given?
      mm.default_proc = block.to_proc
    elsif dict.default_proc
      mm.default_proc = dict.default_proc
    elsif dict.default
      mm.default = dict.default
    end

    dict.each do |key, als|
      mm.let(key, *Array.wrap(als), to: key)
    end

    mm
  end

  # @param hash [Hash] unfolded dictionary
  # @return [MultikeyMap]
  def self.fold(hash)
    build(hash.each_key.group_by(&hash))
  end

  # @!method self.[]
  #   @see build
  singleton_class.alias_method :[], :build

  attr_reader :values, :mapping

  attr_accessor :default, :default_proc

  # @!visibility private
  def default_proc=(proc)
    Types::StrictProc.try(proc) { |result| raise result.error.cause }
    @default_proc = proc
  end

  delegate :size, :include?, to: :values
  alias_method :key?, :include?

  # @!attribute [r] values
  #   @return [Array] Sorted list of values

  # @!attribute [r] mapping
  #   @return [Hash { Integer => Array<Object> }] Map index of each {#values} to keys

  # @!attribute [rw] default
  #   @return [Any, nil] default value if key didn't exist in mapping

  # @!attribute [rw] default_proc
  #   @return [Proc] to be executed on each failed key lookup

  # @!visibility private
  def initialize
    @values = []
    @mapping = {}
  end

  # @overload each_pair { ... }
  #   @yieldparam [Object] key
  #   @yieldparam [Object, nil] value
  # @return [Enumerator]
  def each_pair
    return enum_for(__method__){size} unless block_given?

    @values.each_with_index do |value, index|
      yield(@mapping[index], value)
    end
  end

  alias_method :each, :each_pair

  # @param key [Object]
  # @return [Object]
  def [](key)
    if n = index(key)
      @values[n]
    else
      @default_proc&.call(self, key) || @default
    end
  end

  # @see #let
  def []=(*keys, to)
    let(*keys, to: to)
  end

  # Returns a value from the hash for the given key.
  #
  # If the key can't be found, there are several options:
  #
  #   * With no other arguments, it will raise a KeyError exception;
  #   * if default is given, then that will be returned;
  #   * if the optional code block is specified, then that will be run and its result returned.
  #
  # @param key [Any]
  #
  # @overload fetch(key, default)
  #   @param default [Any] fallback value
  #
  # @overload fetch(key) { |key| ... }
  #   @yield [key] Fallback proc
  #   @yieldparam [Any] key
  #   @yieldreturn [Object] fallback value
  #
  # @raise [KeyError] if key can't be found and default is not set and no other arguments.
  #
  # @return [Object, nil]
  def fetch(key, *args)
    if args.size > 1
      raise ArgumentError,
        "wrong number of arguments (given #{args.size.next}, expected 1 or 2)"
    end

    n = index(key)
    return @values[n] if n
    return yield(key) if block_given?
    return args[0] if args.size == 1
    raise KeyError.new(receiver: self, key: key)
  end

  # Set multiple keys to the given value.
  # @param keys [Array<Object>]
  # @param to [Object, nil]
  # @return [Array<Object>] keys
  def let(*keys, to:)
    @values << to unless @values.include?(to)
    index = @values.index(to)
    @mapping[index] ||= []
    @mapping[index].concat(keys)
    @mapping[index].uniq!
    keys
  end

  # @param key [Object]
  # @return [Object] deleted object
  def delete(key)
    n = index(key)
    @mapping[n].delete(key)
    @values.delete_at(n)
  end

  # @param key [Object]
  # @return [Object] first mapped key
  def index(key)
    @mapping.each_key.detect { |k| @mapping[k].include?(key) }
  end

  # @return [Hash{ Array<keys...> => value }]
  def to_hash
    each_pair.to_h
  end

  # Inverted {#to_hash}, reverse to {build}.
  # @return [Hash{ value => Array<keys...> }]
  def as_dict
    to_hash.invert
  end

  # Inverted {fold}
  # @return [Hash{ key => value }]
  def unfold
    hsh = Hash[@mapping.flat_map { |i, keys| keys.product([@values[i]]) }]
    hsh.default = @default if defined?(@default)
    hsh.default_proc = @default_proc if defined?(@default_proc)
    hsh
  end

  alias_method :to_h, :unfold

  # @!method ==(other)
  #   @see Hash#==
  #   @return [true, false]
  # @!method as_json(*)
  #   @see Hash#as_json
  #   @return [Hash{ String => as_json }]
  delegate :as_json, :==, to: :to_h
end