# Redirect memchr imported by the Ruby DLL to another implementation.
require 'fiddle'

module PatchMemchr
  PAGE_READWRITE = 0x04

  module_function

  def read(addr, len) = Fiddle::Pointer.new(addr)[0, len]

  def apply(new_addr, dll_name = 'x64-ucrt-ruby400.dll')
    k32 = Fiddle.dlopen('kernel32.dll')
    get_module = Fiddle::Function.new(k32['GetModuleHandleA'], [Fiddle::TYPE_VOIDP], Fiddle::TYPE_VOIDP)
    protect = Fiddle::Function.new(k32['VirtualProtect'], [Fiddle::TYPE_VOIDP, Fiddle::TYPE_SIZE_T, Fiddle::TYPE_INT, Fiddle::TYPE_VOIDP], Fiddle::TYPE_INT)
    base = get_module.call(dll_name).to_i
    raise "#{dll_name} not loaded" if base.zero?
    nt = base + read(base + 0x3C, 4).unpack1('l<')
    opt = nt + 24
    raise 'not PE32+' unless read(opt, 2).unpack1('S<') == 0x20b
    import_rva = read(opt + 112 + 8, 4).unpack1('L<')
    desc = base + import_rva
    patched = []
    loop do
      oft, _ts, _fc, name_rva, ft = read(desc, 20).unpack('L<5')
      break if oft.zero? && name_rva.zero? && ft.zero?
      i = 0
      loop do
        thunk = read(base + oft + i * 8, 8).unpack1('Q<')
        break if thunk.zero?
        if (thunk >> 63).zero? && Fiddle::Pointer.new(base + thunk + 2).to_s == 'memchr'
          slot = base + ft + i * 8
          old = Fiddle::Pointer.malloc(4, Fiddle::RUBY_FREE)
          # VirtualProtect returns 0 on failure, and 0 is truthy in Ruby.
          raise 'VirtualProtect failed' if protect.call(slot, 8, PAGE_READWRITE, old).zero?
          Fiddle::Pointer.new(slot)[0, 8] = [new_addr].pack('Q<')
          raise 'VirtualProtect failed to restore' if protect.call(slot, 8, old[0, 4].unpack1('L<'), old).zero?
          raise 'import slot was not updated' unless read(slot, 8).unpack1('Q<') == new_addr
          patched << Fiddle::Pointer.new(base + name_rva).to_s
        end
        i += 1
      end
      desc += 20
    end
    patched
  end
end

variant = ENV['MEMCHR'] || 'ucrt'
unless variant == 'ucrt'
  addr = case variant
         when 'swar' then Fiddle.dlopen(File.join(__dir__, 'fastmemchr.dll'))['swar_memchr']
         when 'sse2' then Fiddle.dlopen(File.join(__dir__, 'fastmemchr.dll'))['sse2_memchr']
         else raise "unknown MEMCHR=#{variant}"
         end
  dlls = PatchMemchr.apply(addr)
  raise 'memchr import was not found, nothing was patched' if dlls.empty?
  warn "memchr -> #{variant} (patched import from #{dlls.join(', ')})"
end
s = ("a" * 1023 + "\n").b * 1000
raise 'index broken' unless s.index("\n") == 1023 && s.index("\n", 1024) == 2047 && s.each_line.count == 1000