string_methods.rb

require 'mock_redis/assertions'

class MockRedis
  module StringMethods
    include Assertions
    include UtilityMethods

    def append(key, value)
      assert_stringy(key)
      data[key] ||= ''
      data[key] << value
      data[key].length
    end

    def bitfield(*args)
      if args.length < 4
        raise Redis::CommandError, 'ERR wrong number of arguments for BITFIELD'
      end

      key = args.shift
      output = []
      overflow_method = 'wrap'

      until args.empty?
        command = args.shift.to_s

        if command == 'overflow'
          new_overflow_method = args.shift.to_s.downcase

          unless %w[wrap sat fail].include? new_overflow_method
            raise Redis::CommandError, 'ERR Invalid OVERFLOW type specified'
          end

          overflow_method = new_overflow_method
          next
        end

        type, offset = args.shift(2)

        is_signed = type.slice(0) == 'i'
        type_size = type[1..-1].to_i

        if (type_size > 64 && is_signed) || (type_size >= 64 && !is_signed)
          raise Redis::CommandError,
            'ERR Invalid bitfield type. Use something like i16 u8. ' \
            'Note that u64 is not supported but i64 is.'
        end

        if offset.to_s[0] == '#'
          offset = offset[1..-1].to_i * type_size
        end

        bits = []

        type_size.times do |i|
          bits.push(getbit(key, offset + i))
        end

        val = is_signed ? twos_complement_decode(bits) : bits.join('').to_i(2)

        case command
        when 'get'
          output.push(val)
        when 'set'
          output.push(val)

          set_bitfield(key, args.shift.to_i, is_signed, type_size, offset)
        when 'incrby'
          new_val = incr_bitfield(val, args.shift.to_i, is_signed, type_size, overflow_method)

          set_bitfield(key, new_val, is_signed, type_size, offset) if new_val
          output.push(new_val)
        end
      end

      output
    end

    def decr(key)
      decrby(key, 1)
    end

    def decrby(key, n)
      incrby(key, -n)
    end

    def get(key)
      key = key.to_s
      assert_stringy(key)
      data[key]
    end

    def getbit(key, offset)
      assert_stringy(key)

      offset_of_byte = offset / 8
      offset_within_byte = offset % 8

      # String#getbyte would be lovely, but it's not in 1.8.7.
      byte = (data[key] || '').each_byte.drop(offset_of_byte).first

      if byte
        (byte & (2**7 >> offset_within_byte)) > 0 ? 1 : 0
      else
        0
      end
    end

    def getrange(key, start, stop)
      assert_stringy(key)
      (data[key] || '')[start..stop]
    end

    def getset(key, value)
      retval = get(key)
      set(key, value)
      retval
    end

    def incr(key)
      incrby(key, 1)
    end

    def incrby(key, n)
      assert_stringy(key)
      unless can_incr?(data[key])
        raise Redis::CommandError, 'ERR value is not an integer or out of range'
      end

      unless looks_like_integer?(n.to_s)
        raise Redis::CommandError, 'ERR value is not an integer or out of range'
      end

      new_value = data[key].to_i + n.to_i
      data[key] = new_value.to_s
      # for some reason, redis-rb doesn't return this as a string.
      new_value
    end

    def incrbyfloat(key, n)
      assert_stringy(key)
      unless can_incr_float?(data[key])
        raise Redis::CommandError, 'ERR value is not a valid float'
      end

      unless looks_like_float?(n.to_s)
        raise Redis::CommandError, 'ERR value is not a valid float'
      end

      new_value = data[key].to_f + n.to_f
      data[key] = new_value.to_s
      # for some reason, redis-rb doesn't return this as a string.
      new_value
    end

    def mget(*keys)
      assert_has_args(keys, 'mget')

      keys.flatten!
      keys.map do |key|
        get(key) if stringy?(key)
      end
    end

    def mapped_mget(*keys)
      Hash[keys.zip(mget(*keys))]
    end

    def mset(*kvpairs)
      assert_has_args(kvpairs, 'mset')
      kvpairs = kvpairs.first if kvpairs.size == 1 && kvpairs.first.is_a?(Enumerable)

      if kvpairs.length.odd?
        raise Redis::CommandError, 'ERR wrong number of arguments for MSET'
      end

      kvpairs.each_slice(2) do |(k, v)|
        set(k, v)
      end

      'OK'
    end

    def mapped_mset(hash)
      mset(*hash.to_a.flatten)
    end

    def msetnx(*kvpairs)
      assert_has_args(kvpairs, 'msetnx')

      if kvpairs.each_slice(2).any? { |(k, _)| exists?(k) }
        false
      else
        mset(*kvpairs)
        true
      end
    end

    def mapped_msetnx(hash)
      msetnx(*hash.to_a.flatten)
    end

    def set(key, value, options = {})
      key = key.to_s
      return_true = false
      options = options.dup
      if options.delete(:nx)
        if exists?(key)
          return false
        else
          return_true = true
        end
      end
      if options.delete(:xx)
        if exists?(key)
          return_true = true
        else
          return false
        end
      end
      data[key] = value.to_s

      duration = options.delete(:ex)
      if duration
        if duration == 0
          raise Redis::CommandError, 'ERR invalid expire time in set'
        end
        expire(key, duration)
      end

      duration = options.delete(:px)
      if duration
        if duration == 0
          raise Redis::CommandError, 'ERR invalid expire time in set'
        end
        pexpire(key, duration)
      end
      unless options.empty?
        raise ArgumentError, "unknown keyword: #{options.keys[0]}"
      end

      return_true ? true : 'OK'
    end

    def setbit(key, offset, value)
      assert_stringy(key, 'ERR bit is not an integer or out of range')
      retval = getbit(key, offset)

      str = data[key] || ''

      offset_of_byte = offset / 8
      offset_within_byte = offset % 8

      if offset_of_byte >= str.bytesize
        str = zero_pad(str, offset_of_byte + 1)
      end

      char_index = byte_index = offset_within_char = 0
      str.each_char do |c|
        if byte_index < offset_of_byte
          char_index += 1
          byte_index += c.bytesize
        else
          offset_within_char = byte_index - offset_of_byte
          break
        end
      end

      char = str[char_index]
      char = char.chr if char.respond_to?(:chr) # ruby 1.8 vs 1.9
      char_as_number = char.each_byte.reduce(0) do |a, byte|
        (a << 8) + byte
      end

      bitmask_length = (char.bytesize * 8 - offset_within_char * 8 - offset_within_byte - 1)
      bitmask = 1 << bitmask_length

      if value.zero?
        bitmask ^= 2**(char.bytesize * 8) - 1
        char_as_number &= bitmask
      else
        char_as_number |= bitmask
      end

      str[char_index] = char_as_number.chr

      data[key] = str
      retval
    end

    def bitcount(key, start = 0, stop = -1)
      assert_stringy(key)

      str   = data[key] || ''
      count = 0
      m1    = 0x5555555555555555
      m2    = 0x3333333333333333
      m4    = 0x0f0f0f0f0f0f0f0f
      m8    = 0x00ff00ff00ff00ff
      m16   = 0x0000ffff0000ffff
      m32   = 0x00000000ffffffff

      str.bytes.to_a[start..stop].each do |byte|
        # Naive Hamming weight
        c = byte
        c = (c & m1) + ((c >> 1) & m1)
        c = (c & m2) + ((c >> 2) & m2)
        c = (c & m4) + ((c >> 4) & m4)
        c = (c & m8) + ((c >> 8) & m8)
        c = (c & m16) + ((c >> 16) & m16)
        c = (c & m32) + ((c >> 32) & m32)
        count += c
      end

      count
    end

    def setex(key, seconds, value)
      if seconds <= 0
        raise Redis::CommandError, 'ERR invalid expire time in setex'
      else
        set(key, value)
        expire(key, seconds)
        'OK'
      end
    end

    def setnx(key, value)
      if exists?(key)
        false
      else
        set(key, value)
        true
      end
    end

    def setrange(key, offset, value)
      assert_stringy(key)
      value = value.to_s
      old_value = (data[key] || '')

      prefix = zero_pad(old_value[0...offset], offset)
      data[key] = prefix + value + (old_value[(offset + value.length)..-1] || '')
      data[key].length
    end

    def strlen(key)
      assert_stringy(key)
      (data[key] || '').bytesize
    end

    private

    def stringy?(key)
      data[key].nil? || data[key].is_a?(String)
    end

    def assert_stringy(key,
        message = 'WRONGTYPE Operation against a key holding the wrong kind of value')
      unless stringy?(key)
        raise Redis::CommandError, message
      end
    end

    def set_bitfield(key, value, is_signed, type_size, offset)
      if is_signed
        val_array = twos_complement_encode(value, type_size)
      else
        str = left_pad(value.to_i.abs.to_s(2), type_size)
        val_array = str.split('').map(&:to_i)
      end

      val_array.each_with_index do |bit, i|
        setbit(key, offset + i, bit)
      end
    end

    def incr_bitfield(val, incrby, is_signed, type_size, overflow_method)
      new_val = val + incrby

      max = is_signed ? (2**(type_size - 1)) - 1 : (2**type_size) - 1
      min = is_signed ? (-2**(type_size - 1)) : 0
      size = 2**type_size

      return new_val if (min..max).cover?(new_val)

      case overflow_method
      when 'fail'
        new_val = nil
      when 'sat'
        new_val = new_val > max ? max : min
      when 'wrap'
        if is_signed
          if new_val > max
            remainder = new_val - (max + 1)
            new_val = min + remainder.abs
          else
            remainder = new_val - (min - 1)
            new_val = max - remainder.abs
          end
        else
          new_val = new_val > max ? new_val % size : size - new_val.abs
        end
      end

      new_val
    end
  end
end