gen1recomp/src/import/Rom.lua

local Rom = {}
Rom.__index = Rom

local BANK_SIZE = 0x4000

function Rom.new(data)
  assert(type(data) == "string", "ROM data must be a string")
  return setmetatable({ data = data }, Rom)
end

function Rom.offset(bank, address)
  if bank == 0 then
    assert(address >= 0 and address < BANK_SIZE,
      ("ROM0 address out of range: $%04X"):format(address))
    return address
  end
  assert(address >= BANK_SIZE and address < BANK_SIZE * 2,
    ("bank %02X address out of range: $%04X"):format(bank, address))
  return bank * BANK_SIZE + address - BANK_SIZE
end

function Rom:byte(bank, address)
  local value = self.data:byte(Rom.offset(bank, address) + 1)
  assert(value, ("ROM read past end at %02X:%04X"):format(bank, address))
  return value
end

function Rom:word(bank, address)
  return self:byte(bank, address) + self:byte(bank, address + 1) * 0x100
end

function Rom:bytes(bank, address, length)
  local first = Rom.offset(bank, address) + 1
  local last = first + length - 1
  assert(last <= #self.data,
    ("ROM read past end at %02X:%04X + %d"):format(bank, address, length))
  local out = {}
  for index = 1, length do
    out[index] = self.data:byte(first + index - 1)
  end
  return out
end

function Rom:decodeText(raw, charmap, stop)
  local out = {}
  stop = stop or 0x50
  for _, value in ipairs(raw) do
    if value == stop then break end
    out[#out + 1] = charmap[tostring(value)]
      or ("{BYTE:%02X}"):format(value)
  end
  return table.concat(out)
end

function Rom:readString(bank, address, charmap, stop, maxLength)
  local out = {}
  stop = stop or 0x50
  maxLength = maxLength or 4096
  for offset = 0, maxLength - 1 do
    local value = self:byte(bank, address + offset)
    if value == stop then return table.concat(out), offset + 1 end
    out[#out + 1] = charmap[tostring(value)]
      or ("{BYTE:%02X}"):format(value)
  end
  error(("unterminated string at %02X:%04X"):format(bank, address))
end

function Rom.bcd(raw)
  local value = 0
  for _, byte in ipairs(raw) do
    value = value * 100 + math.floor(byte / 16) * 10 + byte % 16
  end
  return value
end

local BitReader = {}
BitReader.__index = BitReader

function BitReader.new(data)
  return setmetatable({ data = data, byte = 1, bit = 7 }, BitReader)
end

function BitReader:read(count)
  local value = 0
  for _ = 1, count or 1 do
    local byte = self.data[self.byte]
    if not byte then error("compressed picture ended unexpectedly") end
    value = value * 2 + math.floor(byte / 2 ^ self.bit) % 2
    self.bit = self.bit - 1
    if self.bit < 0 then
      self.byte = self.byte + 1
      self.bit = 7
    end
  end
  return value
end

local function fillPicPlane(reader, width)
  local mode = reader:read()
  local groupCount = width * width * 0x20
  local groups = {}
  while #groups < groupCount do
    if mode ~= 0 then
      while #groups < groupCount do
        local group = reader:read(2)
        if group == 0 then break end
        groups[#groups + 1] = group
      end
    else
      local prefix = 0
      while reader:read() ~= 0 do
        prefix = prefix + 1
        if prefix >= 16 then error("invalid compressed picture zero run") end
      end
      local zeroCount = 2 ^ (prefix + 1) - 1 + reader:read(prefix + 1)
      for _ = 1, math.min(zeroCount, groupCount - #groups) do
        groups[#groups + 1] = 0
      end
    end
    mode = 1 - mode
  end

  local reordered = {}
  for y = 0, width - 1 do
    for x = 0, width * 8 - 1 do
      for group = 0, 3 do
        local source = (y * 4 + group) * width * 8 + x
        reordered[#reordered + 1] = groups[source + 1]
      end
    end
  end
  local packed = {}
  for index = 0, width * width * 8 - 1 do
    local start = index * 4
    packed[index + 1] = reordered[start + 1] * 0x40
      + reordered[start + 2] * 0x10
      + reordered[start + 3] * 4
      + reordered[start + 4]
  end
  return packed
end

local PIC_CODES = {
  { 0x0, 0x1, 0x3, 0x2, 0x7, 0x6, 0x4, 0x5,
    0xF, 0xE, 0xC, 0xD, 0x8, 0x9, 0xB, 0xA },
  { 0xF, 0xE, 0xC, 0xD, 0x8, 0x9, 0xB, 0xA,
    0x0, 0x1, 0x3, 0x2, 0x7, 0x6, 0x4, 0x5 },
}

local function unfilterPicPlane(plane, width)
  for x = 0, width * 8 - 1 do
    local bit = 0
    for y = 0, width - 1 do
      local index = y * width * 8 + x + 1
      local high = PIC_CODES[bit + 1][math.floor(plane[index] / 16) + 1]
      bit = high % 2
      local low = PIC_CODES[bit + 1][plane[index] % 16 + 1]
      bit = low % 2
      plane[index] = high * 16 + low
    end
  end
end

local function transposePicTiles(data, width)
  local tileCount = width * width
  for index = 0, tileCount - 1 do
    local other = (index * width + math.floor(index / width)) % tileCount
    if index < other then
      for offset = 1, 16 do
        local left = index * 16 + offset
        local right = other * 16 + offset
        data[left], data[right] = data[right], data[left]
      end
    end
  end
end

function Rom.decompressPic(data)
  local reader = BitReader.new(data)
  local width, height = reader:read(4), reader:read(4)
  if width == 0 or width ~= height then
    error(("compressed picture is not a non-empty square (%dx%d)")
      :format(width, height))
  end

  local order = reader:read()
  local planes = {}
  planes[order + 1] = fillPicPlane(reader, width)
  local mode = reader:read()
  if mode ~= 0 then mode = mode + reader:read() end
  planes[(1 - order) + 1] = fillPicPlane(reader, width)

  unfilterPicPlane(planes[order + 1], width)
  if mode ~= 1 then unfilterPicPlane(planes[(1 - order) + 1], width) end
  if mode ~= 0 then
    for index = 1, width * width * 8 do
      planes[(1 - order) + 1][index] =
        bit.bxor(planes[(1 - order) + 1][index], planes[order + 1][index])
    end
  end

  local output = {}
  for index = 1, width * width * 8 do
    output[#output + 1] = planes[1][index]
    output[#output + 1] = planes[2][index]
  end
  transposePicTiles(output, width)
  return output, width
end

-- pokegold's "lz3" compression (home/decompress.asm), used for Gen 2
-- graphics (tilesets, Pokemon pics, title screen art, ...).  `data` is a Lua
-- array of bytes (as returned by Rom:bytes) or a raw string; returns a Lua
-- array of decompressed bytes.  Ported instruction-for-instruction against
-- pokegold's Decompress routine and cross-checked against tools/lzcompress.c
-- (--uncompress path), which is the canonical reference for this format.
local LZ_END = 0xFF
local LZ_LITERAL   = 0
local LZ_ITERATE   = 1
local LZ_ALTERNATE = 2
local LZ_ZERO      = 3
local LZ_FLIP      = 5
local LZ_REVERSE   = 6
local LZ_LONG      = 7

local function flipBits(value)
  local flipped = 0
  for bitIndex = 0, 7 do
    flipped = flipped + math.floor(value / 2 ^ bitIndex) % 2 * 2 ^ (7 - bitIndex)
  end
  return flipped
end

function Rom.decompressLz3(data)
  local bytes = data
  if type(data) == "string" then
    bytes = {}
    for index = 1, #data do bytes[index] = data:byte(index) end
  end

  local pos = 1
  local function nextByte()
    local value = bytes[pos]
    if not value then error("lz3 stream ended unexpectedly") end
    pos = pos + 1
    return value
  end

  local out = {}
  while true do
    local first = bytes[pos]
    if not first then error("lz3 stream ended without a terminator") end
    pos = pos + 1
    if first == LZ_END then break end

    local command, length
    if math.floor(first / 0x20) == LZ_LONG then
      -- 111xxxyy yyyyyyyy: xxx is the real command, yy.. is a 10-bit length
      command = math.floor(first / 4) % 8
      local high = first % 4
      length = high * 0x100 + nextByte() + 1
    else
      command = math.floor(first / 0x20)
      length = first % 0x20 + 1
    end

    if command == LZ_LITERAL then
      for _ = 1, length do out[#out + 1] = nextByte() end
    elseif command == LZ_ITERATE then
      local value = nextByte()
      for _ = 1, length do out[#out + 1] = value end
    elseif command == LZ_ALTERNATE then
      local a, b = nextByte(), nextByte()
      for index = 1, length do
        out[#out + 1] = (index % 2 == 1) and a or b
      end
    elseif command == LZ_ZERO then
      for _ = 1, length do out[#out + 1] = 0 end
    else
      -- Lookback commands (LZ_REPEAT/LZ_FLIP/LZ_REVERSE, and the unused id
      -- 7 which the hardware routine falls through to LZ_REPEAT for).
      -- A high-bit offset byte is a 7-bit negative lookback from the
      -- current output position; otherwise it is a 15-bit absolute offset
      -- from the start of the output buffer (two bytes, big-endian).
      local offsetByte = nextByte()
      local from
      if offsetByte >= 0x80 then
        from = #out - (offsetByte % 0x80)
      else
        from = offsetByte * 0x100 + nextByte() + 1
      end
      if command == LZ_FLIP then
        for index = 0, length - 1 do
          out[#out + 1] = flipBits(out[from + index])
        end
      elseif command == LZ_REVERSE then
        for index = 0, length - 1 do
          out[#out + 1] = out[from - index]
        end
      else
        for index = 0, length - 1 do
          out[#out + 1] = out[from + index]
        end
      end
    end
  end
  return out
end

return Rom