DecodePNG.lua

local PNG_IDENTIFIER = "\137PNG\13\10\26\10"
local BYTES_PER_PIXEL = { [0] = 1, [2] = 3, [3] = 1, [4] = 2, [6] = 4 }

-- https://github.com/jiwonz/zlib91.lua/blob/main/rbx/zlib91/compression.luau
local zlib = require(script.zlib).Zlib

local chunkParsers = {
	IHDR = function(b, offset, output, length)
		output.width = bit32.byteswap(buffer.readu32(b, offset))
		output.height = bit32.byteswap(buffer.readu32(b, offset + 4))
		output.byteDepth = buffer.readu8(b, offset + 8) // 8
		output.colorType = buffer.readu8(b, offset + 9)
		return offset + 13
	end,
	IDAT = function(b, offset, output, length)
		table.insert(output.imageData, buffer.readstring(b, offset, length))
		return offset + length
	end,
	PLTE = function(b, offset, output, length)		
		local palette = buffer.create(length)
		output.palette = palette
		
		buffer.copy(palette, 0, b, offset, length)
		
		return offset + length
	end,
}

local colorCalculations = {
	[0] = function(x: number, scanLine: buffer, palette: buffer, output: buffer, offset: number) -- Greyscale
		buffer.fill(output, offset, buffer.readu8(scanLine, x), 3)
		buffer.writeu8(output, offset + 3, 255)
	end,
	[2] = function(x: number, scanLine: buffer, palette: buffer, output: buffer, offset: number) -- RGB
		buffer.copy(output, offset, scanLine, x, 3)
		buffer.writeu8(output, offset + 3, 255)
	end,
	[3] = function(x: number, scanLine: buffer, palette: buffer, output: buffer, offset: number) -- Palette
		local paletteIndex = buffer.readu8(scanLine, x)
		buffer.copy(output, offset, palette, paletteIndex, 3)
		buffer.writeu8(output, offset + 3, 255)
	end,
	[4] = function(x: number, scanLine: buffer, palette: buffer, output: buffer, offset: number) -- Greyscale & Alpha
		buffer.fill(output, offset, buffer.readu8(scanLine, x), 3)
		buffer.copy(output, offset + 3, scanLine, x + 1, 1)
	end,
	[6] = function(x: number, scanLine: buffer, palette: buffer, output: buffer, offset: number) -- RGBA
		buffer.copy(output, offset, scanLine, x, 4)
	end,
}

local unfilterScanLine; unfilterScanLine = {
	[1] = function(scanLine: buffer, previousRow: buffer?, bytesPerPixel: number) -- Sub
		for i = bytesPerPixel, buffer.len(scanLine) - 1 do
			local x = buffer.readu8(scanLine, i)
			local left = buffer.readu8(scanLine, i - bytesPerPixel)
			buffer.writeu8(scanLine, i, x + left)
		end
	end,
	[2] = function(scanLine: buffer, previousRow: buffer?, bytesPerPixel: number) -- Up
		if not previousRow then return end

		for i = 0, buffer.len(scanLine) - 1 do
			local x = buffer.readu8(scanLine, i)
			local above = buffer.readu8(previousRow, i)
			buffer.writeu8(scanLine, i, x + above)
		end
	end,
	[3] = function(scanLine: buffer, previousRow: buffer?, bytesPerPixel: number) -- Average
		if previousRow then
			for i = 0, bytesPerPixel - 1 do
				local x = buffer.readu8(scanLine, i)
				local left = buffer.readu8(previousRow, i)
				buffer.writeu8(scanLine, i, x + left // 2)
			end

			for i = bytesPerPixel, buffer.len(scanLine) - 1 do
				local x = buffer.readu8(scanLine, i)
				local left = buffer.readu8(scanLine, i - bytesPerPixel)
				local above = buffer.readu8(previousRow, i)
				buffer.writeu8(scanLine, i, x + ( above + left ) // 2)
			end
		else
			for i = bytesPerPixel, buffer.len(scanLine) - 1 do
				local x = buffer.readu8(scanLine, i)
				local left = buffer.readu8(scanLine, i - bytesPerPixel)
				buffer.writeu8(scanLine, i, x + left // 2)
			end
		end
	end,
	[4] = function(scanLine: buffer, previousRow: buffer, bytesPerPixel: number) -- Paeth
		if previousRow then
			for i = 0, bytesPerPixel - 1 do
				local x = buffer.readu8(scanLine, i)
				local above = buffer.readu8(previousRow, i)
				buffer.writeu8(scanLine, i, x + above)
			end

			for i = bytesPerPixel, buffer.len(scanLine) - 1 do
				local x = buffer.readu8(scanLine, i)
				local left = buffer.readu8(scanLine, i - bytesPerPixel)
				local above = buffer.readu8(previousRow, i)
				local leftAbove = buffer.readu8(previousRow, i - bytesPerPixel)

				local p = left + above - leftAbove
				local pa = math.abs(p - left)
				local pb = math.abs(p - above)
				local pc = math.abs(p - leftAbove)

				local pr = if pa <= pb and pa <= pc then left
					elseif pb <= pc then above
					else leftAbove

				buffer.writeu8(scanLine, i, x + pr)
			end
		else
			-- Sub filter
			unfilterScanLine[1](scanLine, previousRow, bytesPerPixel)
		end
	end
}

local u24Buffer = buffer.create(4)
local byteReaders = {
	[1] = buffer.readu8,
	[2] = buffer.readu16,
	[3] = function(b: buffer, offset: number)
		buffer.copy(u24Buffer, 0, b, offset, 3)
		return buffer.readu32(u24Buffer, 0)
	end,
	[4] = buffer.readu32,
}

local function extractChunkData(b, offset)
	local output = { imageData = {} }

	repeat
		local length = bit32.byteswap(buffer.readu32(b, offset))
		local chunkType = buffer.readstring(b, offset + 4, 4)
		offset += 8

		local parser = chunkParsers[chunkType]
		if parser then
			offset = parser(b, offset, output, length)
		else
			offset += length
		end
		
		offset += 4
	until chunkType == "IEND"

	return output, offset
end

local function DecodePNG(imageString: string): (buffer, number, number)
	local imageData = buffer.fromstring(imageString)
	
	local identifier = buffer.readstring(imageData, 0, 8)
	assert(identifier == PNG_IDENTIFIER, "The provided image link is not a PNG")

	local chunkData, offset = extractChunkData(imageData, 8)
	
	local width = chunkData.width
	local height = chunkData.height
	local byteDepth = chunkData.byteDepth
	local colorType = chunkData.colorType
	local palette = chunkData.palette
	
	local calculateColour = colorCalculations[colorType]
	local pixelReader = byteReaders[byteDepth]
	local bytesPerPixel = BYTES_PER_PIXEL[colorType]

	local output = buffer.create(width * height * 4)
	local scanLines = table.create(height)

	local decompressedData = buffer.fromstring(
		zlib.Decompress(table.concat(chunkData.imageData))
	)

	local outputIndex = 0
	local dataOffset = 0

	for y = 1, height do
		local filterType = buffer.readu8(decompressedData, dataOffset)
		dataOffset += 1

		local scanLine = buffer.create(width * bytesPerPixel)
		scanLines[y] = scanLine

		for x = 0, width - 1 do
			local index = x * bytesPerPixel
			for i = 0, bytesPerPixel - 1 do
				buffer.writeu8(scanLine, index + i, pixelReader(decompressedData, dataOffset))
				dataOffset += byteDepth
			end
		end

		local unfilter = unfilterScanLine[filterType]
		if unfilter then
			unfilter(scanLine, scanLines[y - 1], bytesPerPixel)
		end

		for x = 0, width * bytesPerPixel - 1, bytesPerPixel do
			calculateColour(x, scanLine, palette, output, outputIndex)
			outputIndex += 4
		end
	end
	
	return output, width, height
end

return DecodePNG