Simple PNG generation example that only depends on zlib in Ruby

ngpng.rb

# Not Great PNG class.  This is a very simple example of writing a PNG.  It
# only supports colors from the color palette stored in `@palette`.  This is
# meant to be example code, but I am using it in a program for visualizing
# heap dumps from Ruby.
#
# This is free and unencumbered software released into the public domain.
#
# Anyone is free to copy, modify, publish, use, compile, sell, or
# distribute this software, either in source code form or as a compiled
# binary, for any purpose, commercial or non-commercial, and by any
# means.
#
# In jurisdictions that recognize copyright laws, the author or authors
# of this software dedicate any and all copyright interest in the
# software to the public domain. We make this dedication for the benefit
# of the public at large and to the detriment of our heirs and
# successors. We intend this dedication to be an overt act of
# relinquishment in perpetuity of all present and future rights to this
# software under copyright law.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# For more information, please refer to <https://unlicense.org>

require "zlib"

class NGPNG
  HEADER = [137, 80, 78, 71, 13, 10, 26, 10].pack("C*")

  def initialize width, height
    @width  = width
    @height = height
    # Using half a byte, so everything needs to be divided by 2
    @bits       = 4 # 4 bit depth
    @color_type = 3 # only use palette colors
    @palette = [
      [0xFF, 0xFF, 0xFF], # default color is white
      [0xFF, 0x00, 0x00],
      [0x00, 0x66, 0x00],
      [0x00, 0x00, 0x00],
    ]
    # add one for the filter. First 0 is always the filter
    @stride = (width / 2) + (width % 2).clamp(0, 1) + 1
    @buffer = "\0".b * (@stride * height)
  end

  def write_to file
    file.write HEADER
    ihdr = [@width, @height, @bits, @color_type, 0, 0, 0].pack("NNC5")
    write_chunk file, "IHDR", ihdr
    plte = @palette.flatten.pack("C*")
    write_chunk file, "PLTE", plte
    idat = Zlib::Deflate.deflate(@buffer, Zlib::DEFAULT_COMPRESSION)
    write_chunk file, "IDAT", idat
    write_chunk file, "IEND", "".b
  end

  def save filename
    File.open(filename, "wb") { |f| write_to f }
  end

  # v needs to be a palette index
  def []= x, y, v
    index = (@stride * y) + (x / 2) + 1
    mask = 0xF
    if x % 2 == 0
      mask <<= 4
      v <<= 4
    end
    @buffer.setbyte(index, (@buffer.getbyte(index) & ~mask) | v)
  end

  def write_chunk file, name, content
    file.write [content.bytesize].pack("N")
    file.write name
    file.write content
    file.write [Zlib.crc32(content, Zlib.crc32(name))].pack("N")
  end
end

# Sample usage
height = width = 100
png = NGPNG.new height, width
width.times do |x|
  height.times do |y|
    png[x, y] = rand(3)
  end
end
png.save "out.png"