PNGlib

-- libPNGimage by TehSomeLuigi
-- Revision: 1
--
-- A library to load, edit and save PNGs for OpenComputers
--

--[[

    Feel free to use however you wish.
    This header must however be preserved should this be redistributed, even
    if in a modified form.

    This software comes with no warranties whatsoever.

    2014 TehSomeLuigi

]]--
local PNGImage = {}                                    ----1
local PNGImagemetatable = {}
PNGImagemetatable.__index = PNGImage

local function __unpack_msb_uint32(s)                  ----56
    local a,b,c,d = s:byte(1,#s)
    local num = (((a*256) + b) * 256 + c) * 256 + d
    return num
end

-- Read 32-bit unsigned integer (most-significant-byte, MSB, first) from file.
local function __read_msb_uint32(fh)                                               ----110
    return __unpack_msb_uint32(fh:read(4))
end

-- Read unsigned byte (integer) from file
local function __read_byte(fh)                       ----115
    return fh:read(1):byte()
end

local function getBitWidthPerPixel(ihdr)             ----121
    if ihdr.color_type == 0 then -- Greyscale
        return ihdr.bit_depth
    end
    if ihdr.color_type == 2 then -- Truecolour
        return ihdr.bit_depth * 3
    end
    if ihdr.color_type == 3 then -- Indexed-colour
        return ihdr.bit_depth
    end
    if ihdr.color_type == 4 then -- Greyscale + Alpha
        return ihdr.bit_depth * 2
    end
    if ihdr.color_type == 6 then -- Truecolour + Alpha
        return ihdr.bit_depth * 4
    end
end

local function getByteWidthPerScanline(ihdr)         ----139
    return math.ceil((ihdr.width * getBitWidthPerPixel(ihdr)) / 8)
end

local outssmt = {}                                     ----143

function outssmt:__call(write)
    self.str = self.str .. string.char(write)
end

function outssmt.OutStringStream()
    local outss = {str=""}
    setmetatable(outss, outssmt)
    return outss
end

local function __parse_IHDR(fh, len)                 ----158
    if len ~= 13 then
        error("PNG IHDR Corrupt - should be 13 bytes long")
    end
    local ihdr = {}
    ihdr.width = __read_msb_uint32(fh)
    ihdr.height = __read_msb_uint32(fh)
    ihdr.bit_depth = __read_byte(fh)
    ihdr.color_type = __read_byte(fh)
    ihdr.compression_method = __read_byte(fh)
    ihdr.filter_method = __read_byte(fh)
    ihdr.interlace_method = __read_byte(fh)
    return ihdr
end

local function __parse_IDAT(fh, len, commeth, outss)   ----206
    if commeth ~= 0 then
        error("Only zlib/DEFLATE compression supported")
    end
    local input = fh:read(len)
    local cfg = {input=input, output=outss, disable_crc=true}
    inflate_zlib(cfg)
    return true
end

local function getPNGStdByteAtXY(ihdr, oss, x, y)               ----237
    local bpsl = getByteWidthPerScanline(ihdr) -- don't include filterType byte -- we don't store that after it has been read
    if (x <= 0) or (y <= 0) then
        return 0 -- this is what the spec says we should return when the coordinate is out of bounds -- in this part of the code, the coordinates are ONE-BASED like in good Lua
    end
    local offset_by_y = (y - 1) * bpsl
    -- now read it!
    local idx = offset_by_y + x
    return oss.str:sub(idx, idx):byte()
end

local function __paeth_predictor(a, b, c) --249
    local p = a + b - c
    local pa = math.abs(p - a)
    local pb = math.abs(p - b)
    local pc = math.abs(p - c)
    if pa <= pb and pa <= pc then
        return a
    elseif pb <= pc then
        return b
    else
        return c
    end
end

local function __parse_IDAT_effective_bytes(outss, ihdr)    ----265
    local bpsl = getByteWidthPerScanline(ihdr)
    local bypsl = math.ceil(getBitWidthPerPixel(ihdr) / 8)
    if outss.str:len() == 0 then
        error("Empty string: outss")
    end
    local bys = bytestream_from_string(outss.str)

    if not bys then
        error("Did not get a bytestream from string", bys, outss)
    end
    local out2 = outssmt.OutStringStream() -- __callable table with metatable that stores what you give it
    local y = 0
    -- x the byte being filtered;
    -- a the byte corresponding to x in the pixel immediately before the pixel containing x (or the byte immediately before x, when the bit depth is less than 8);
    -- b the byte corresponding to x in the previous scanline;
    -- c the byte corresponding to b in the pixel immediately before the pixel containing b (or the byte immediately before b, when the bit depth is less than 8).
    while true do
        local filterType = bys:read()
        if filterType == nil then
            break
        end
        y = y + 1
        for x = 1, bpsl do
            local a = getPNGStdByteAtXY(ihdr, out2, x - bypsl, y)
            local b = getPNGStdByteAtXY(ihdr, out2, x, y - 1)
            local c = getPNGStdByteAtXY(ihdr, out2, x - bypsl, y - 1)
            local outVal = 0
            if filterType == 0 then outVal = bys:read()
            elseif filterType == 1 then outVal = bys:read() + a
            elseif filterType == 2 then outVal = bys:read() + b
            elseif filterType == 3 then outVal = bys:read() + math.floor((a + b) / 2)
            elseif filterType == 4 then outVal = bys:read() + __paeth_predictor(a, b, c)
            else
                error("Unsupported Filter Type: " .. tostring(filterType))
            end
            outVal = outVal % 256
            out2(outVal)
        end
    end
    return out2
end

-- Warning: Co-ordinates are Zero-based but strings are 1-based
function PNGImage:getByteOffsetForPixel(x, y)      ----498
    return (((y * self.ihdr.width) + x) * 4) + 1
end

function PNGImage:getPixel(x, y)
    local off = self:getByteOffsetForPixel(x, y)
    return self.data:byte(off, off + 3)
end
--[[
function PNGImage:setPixel(x, y, col)
    local off = self:getByteOffsetForPixel(x, y)
    self.data = table.concat({self.data:sub(1, off - 1), string.char(col[1], col[2], col[3], col[4]), self.data:sub(off + 4)})
end
function PNGImage:lineXAB(ax, y, bx, col)
    for x=ax, bx do
        self:setPixel(x, y, col)
    end
end
function PNGImage:lineYAB(x, ay, by, col)
    for y=ay, by do
        self:setPixel(x, y, col)
    end
end
function PNGImage:lineRectangleAB(ax, ay, bx, by, col)
    self:lineXAB(ax, ay, bx, col)
    self:lineXAB(ax, by, bx, col)
    self:lineYAB(ax, ay, by, col)
    self:lineYAB(bx, ay, by, col)
end
function PNGImage:fillRectangleAB(ax, ay, bx, by, col)
    for x=ax, bx do
        for y=ay, by do
            self:setPixel(x, y, col)
        end
    end
end
function PNGImage:saveToFile(fn)
    local fh = io.open(fn, 'wb')
    if not fh then
        error("Could not open for writing: " .. fn)
    end
    self:saveToFileHandle(fh)
    fh:close()
end
function PNGImage:generateRawIDATData(outbuf)
    for y = 0, self.ihdr.height - 1 do
        outbuf(0) -- filter type is 0 (Filt(x) = Orig(x))
        for x = 0, self.ihdr.width - 1 do
            local r, g, b, a = self:getPixel(x, y)
            outbuf(r)
            outbuf(g)
            outbuf(b)
            outbuf(a)
        end
    end
end
]]
function PNGImage:getSize()
    return self.ihdr.width, self.ihdr.height
end

local function newFromFile(fh)
    local fh = io.open(fh, 'rb')
    if not fh then
        error("Could not open PNG file")
    end
    local pngi = {}
    setmetatable(pngi, PNGImagemetatable)
    local expecting = "\137\080\078\071\013\010\026\010"
    if fh:read(8) ~= expecting then -- check the 8-byte PNG header exists
        error("Not a PNG file")
    end
    local ihdr
    local outss = outssmt.OutStringStream()
    while true do
        local len = __read_msb_uint32(fh)
        local stype = fh:read(4)
        if stype == 'IHDR' then
            ihdr, msg = __parse_IHDR(fh, len)
        elseif stype == 'IDAT' then
            local res, msg = __parse_IDAT(fh, len, ihdr.compression_method, outss)
        else
            fh:read(len) -- dummy read
        end
        local crc = __read_msb_uint32(fh)
        if stype == 'IEND' then
            break
        end
    end
    fh:close()
    if ihdr.filter_method ~= 0 then
        error("Unsupported Filter Method: " .. ihdr.filter_method)
    end
    if ihdr.interlace_method ~= 0 then
        error("Unsupported Interlace Method (Interlacing is currently unsupported): " .. ihdr.interlace_method)
    end
    if ihdr.color_type ~= 6 --[[TruecolourAlpha]] and ihdr.color_type ~= 2 --[[Truecolour]] then
        error("Currently, only Truecolour and Truecolour+Alpha images are supported.")
    end
    if ihdr.bit_depth ~= 8 then
        error("Currently, only images with a bit depth of 8 are supported.")
    end
    -- now parse the IDAT chunks
    local out2 = __parse_IDAT_effective_bytes(outss, ihdr)
    if ihdr.color_type == 2 --[[Truecolour]] then
        -- add an alpha layer so it effectively becomes RGBA, not RGB
        local inp = out2.str
        out2 = outssmt.OutStringStream()
        for i=1, ihdr.width*ihdr.height do
            local b = ((i - 1)*3) + 1
            out2(inp:byte(b)) -- R
            out2(inp:byte(b + 1)) -- G
            out2(inp:byte(b + 2)) -- B
            out2(255) -- A
        end
    end
    pngi.ihdr = ihdr
    pngi.data = out2.str
    return pngi
end