Add indexed color, other png improvements

4 files changed, 126 insertions, 26 deletions

diff --git a/src/img/png.nim b/src/img/png.nim
index 9bcbd348..50b0e980 100644
--- a/src/img/png.nim
+++ b/src/img/png.nim
@@ -3,8 +3,7 @@ import math
 import bindings/zlib
 import img/bitmap
 import types/color
-
-import lib/endians2
+import utils/endians
 
 type PNGWriter = object
   buf: pointer
@@ -50,24 +49,21 @@ type PNGColorType {.size: sizeof(uint8).} = enum
   GRAYSCALE_WITH_ALPHA = 4
   TRUECOLOR_WITH_ALPHA = 6
 
-func u8toc(x: openArray[uint8]): string =
-  #TODO ew
-  var s = newString(x.len)
-  copyMem(addr s[0], unsafeAddr x[0], x.len)
-  return s
-
 const PNGSignature = "\x89PNG\r\n\x1A\n"
 proc writeIHDR(writer: var PNGWriter, width, height: uint32,
     bitDepth: uint8, colorType: PNGColorType,
     compressionMethod, filterMethod, interlaceMethod: uint8) =
   writer.writeStr(PNGSignature)
-  let ihdr = u8toc(pngInt(width)) &
-    u8toc(pngInt(height)) &
-    char(bitDepth) &
-    char(uint8(colorType)) &
-    char(compressionMethod) &
-    char(filterMethod) &
-    char(interlaceMethod)
+  var ihdr {.noInit.}: array[13, uint8]
+  var pw = pngInt(width)
+  var ph = pngInt(height)
+  copyMem(addr ihdr[0], addr pw[0], 4)
+  copyMem(addr ihdr[4], addr ph[0], 4)
+  ihdr[8] = bitDepth
+  ihdr[9] = uint8(colorType)
+  ihdr[10] = compressionMethod
+  ihdr[11] = filterMethod
+  ihdr[12] = interlaceMethod
   writer.writeChunk("IHDR", ihdr)
 
 proc writeIDAT(writer: var PNGWriter, bmp: Bitmap) =
@@ -121,6 +117,9 @@ type PNGReader = object
   strm: z_stream
   strmend: bool
   atline: int
+  plteseen: bool
+  palette: seq[RGBAColor]
+  trns: RGBAColor
 
 func width(reader: PNGReader): int {.inline.} = int(reader.bmp.width)
 
@@ -165,7 +164,7 @@ template readU8(reader: var PNGReader): uint8 =
 template readU32(reader: var PNGReader): uint32 =
   if reader.i + 4 > reader.limit:
     reader.err "too short"
-  let x = fromBytesBE(uint32, toOpenArray(reader.iq, reader.i, reader.i + 3))
+  let x = fromBytesBEu32(toOpenArray(reader.iq, reader.i, reader.i + 3))
   reader.i += 4
   x
 
@@ -222,10 +221,11 @@ proc readIHDR(reader: var PNGReader) =
 proc readbKGD(reader: var PNGReader) =
   case reader.colorType
   of GRAYSCALE, GRAYSCALE_WITH_ALPHA:
-    discard reader.readU8() #TODO bit depth > 8
+    # We can't really use bit depth > 8
+    discard reader.readU8()
     reader.background = gray(reader.readU8())
   of TRUECOLOR, TRUECOLOR_WITH_ALPHA:
-    discard reader.readU8() #TODO bit depth > 8
+    discard reader.readU8()
     let r = reader.readU8()
     discard reader.readU8()
     let g = reader.readU8()
@@ -233,7 +233,30 @@ proc readbKGD(reader: var PNGReader) =
     let b = reader.readU8()
     reader.background = rgb(r, g, b)
   of INDEXED_COLOR:
-    discard #TODO
+    let i = int(reader.readU8())
+    if i >= reader.palette.len:
+      reader.err "invalid palette index"
+    reader.background = reader.palette[i]
+
+proc readtRNS(reader: var PNGReader) =
+  case reader.colorType
+  of GRAYSCALE, GRAYSCALE_WITH_ALPHA:
+    # We can't really use bit depth > 8
+    discard reader.readU8()
+    reader.trns = gray(reader.readU8())
+  of TRUECOLOR, TRUECOLOR_WITH_ALPHA:
+    discard reader.readU8()
+    let r = reader.readU8()
+    discard reader.readU8()
+    let g = reader.readU8()
+    discard reader.readU8()
+    let b = reader.readU8()
+    reader.trns = rgb(r, g, b)
+  of INDEXED_COLOR:
+    if reader.limit - reader.i > reader.palette.len:
+      reader.err "too many trns values"
+    for i in 0 ..< reader.palette.len:
+      reader.palette[i].a = reader.readU8()
 
 proc unfilter(reader: var PNGReader, irow: openArray[uint8], bpp: int) =
   # none, sub, up -> replace uprow directly
@@ -290,7 +313,12 @@ proc writepxs(reader: var PNGReader, crow: var openArray[RGBAColor]) =
       let b = reader.uprow[i]
       i += step
       crow[x] = rgba(r, g, b, 255u8)
-  of INDEXED_COLOR: discard #TODO
+  of INDEXED_COLOR:
+    for x in 0 ..< crow.len:
+      let i = int(reader.uprow[x])
+      if unlikely(i >= reader.palette.len):
+        reader.err "invalid palette index"
+      crow[x] = reader.palette[i]
   of GRAYSCALE_WITH_ALPHA:
     let step = int(reader.bitDepth) div 8
     var i = 0
@@ -309,11 +337,32 @@ proc writepxs(reader: var PNGReader, crow: var openArray[RGBAColor]) =
       let a = reader.uprow[(x + 3) * step]
       crow[x] = rgba(r, g, b, a)
 
+proc readPLTE(reader: var PNGReader) =
+  # For non-indexed-color, palette is just a suggestion for quantization.
+  #TODO support this in term
+  const CanHavePLTE = {TRUECOLOR, INDEXED_COLOR, TRUECOLOR_WITH_ALPHA}
+  if reader.plteseen:
+    reader.err "too many PLTE chunks"
+  if reader.colorType notin CanHavePLTE:
+    reader.err "unexpected PLTE chunk for color type"
+  let len = reader.limit - reader.i
+  if len mod 3 != 0:
+    reader.err "palette length not divisible by 3"
+  reader.palette = newSeq[RGBAColor](len)
+  for i in 0 ..< len div 3:
+    let r = reader.readU8()
+    let g = reader.readU8()
+    let b = reader.readU8()
+    reader.palette[i] = rgba(r, g, b, 255)
+  reader.plteseen = true
+
 proc readIDAT(reader: var PNGReader) =
   if reader.idatAt == reader.idatBuf.len:
     reader.err "idat buffer already filled"
   if reader.strmend:
     reader.err "stream already ended"
+  if reader.colorType == INDEXED_COLOR and not reader.plteseen:
+    reader.err "palette expected for indexed color"
   reader.strm.avail_in = cuint(reader.limit - reader.i)
   reader.strm.next_in = addr reader.iq[reader.i]
   let olen = reader.idatBuf.len - reader.idatAt
@@ -386,8 +435,6 @@ proc fromPNG*(iq: openArray[uint8]): Bitmap =
   if reader.bmp == nil: return
   if reader.width == 0 or reader.height == 0:
     reader.err "invalid zero sized png"
-  if reader.colorType == INDEXED_COLOR:
-    reader.err "indexed color not implemented yet"
   reader.initZStream()
   while reader.i < iq.len and not reader.isend:
     let len = int(reader.readPNGInt())
@@ -398,9 +445,11 @@ proc fromPNG*(iq: openArray[uint8]): Bitmap =
     reader.limit = reader.i + len
     case t
     of "IHDR": reader.err "IHDR expected to be first chunk"
+    of "PLTE": reader.readPLTE()
     of "IDAT": reader.readIDAT()
     of "IEND": reader.readIEND()
     of "bKGD": reader.readbKGD()
+    of "tRNS": reader.readtRNS()
     else: reader.readUnknown(t)
     if reader.bmp == nil: return
     let crc = crc32(0, unsafeAddr iq[j], cuint(len + 4))
diff --git a/src/io/request.nim b/src/io/request.nim
index da3437fb..4ae461df 100644
--- a/src/io/request.nim
+++ b/src/io/request.nim
@@ -317,7 +317,7 @@ proc text*(response: Response): string {.jsfunc.} =
   result = response.body.readAll()
   response.close()
 
-proc json(response: Response, ctx: JSContext): JSValue =
+proc json(response: Response, ctx: JSContext): JSValue {.jsfunc.} =
   var s = response.text()
   return JS_ParseJSON(ctx, cstring(s), cast[csize_t](s.len), cstring"<input>")
 
diff --git a/src/types/color.nim b/src/types/color.nim
index 25c93d9b..b9431340 100644
--- a/src/types/color.nim
+++ b/src/types/color.nim
@@ -212,16 +212,28 @@ const ColorsRGB* = {
 }.map((a) => (a[0], RGBColor(a[1]))).toTable()
 
 func r*(c: RGBAColor): int =
-  return int(uint32(c) shr 16 and 0xff)
+  return int((uint32(c) shr 16) and 0xff)
 
 func g*(c: RGBAColor): int =
-  return int(uint32(c) shr 8 and 0xff)
+  return int((uint32(c) shr 8) and 0xff)
 
 func b*(c: RGBAColor): int =
   return int(uint32(c) and 0xff)
 
 func a*(c: RGBAColor): int =
-  return int(uint32(c) shr 24 and 0xff)
+  return int((uint32(c) shr 24) and 0xff)
+
+proc `r=`*(c: var RGBAColor, r: uint8) =
+  c = RGBAColor(uint32(c) or (uint32(r) shl 16))
+
+proc `g=`*(c: var RGBAColor, g: uint8) =
+  c = RGBAColor(uint32(c) or (uint32(g) shl 8))
+
+proc `b=`*(c: var RGBAColor, b: uint8) =
+  c = RGBAColor(uint32(c) or uint32(b))
+
+proc `a=`*(c: var RGBAColor, a: uint8) =
+  c = RGBAColor(uint32(c) or (uint32(a) shl 24))
 
 # https://html.spec.whatwg.org/#serialisation-of-a-color
 func serialize*(color: RGBAColor): string =
diff --git a/src/utils/endians.nim b/src/utils/endians.nim
new file mode 100644
index 00000000..ea34efa7
--- /dev/null
+++ b/src/utils/endians.nim
@@ -0,0 +1,39 @@
+when defined(gcc) or defined(llvm_gcc) or defined(clang):
+  const useBuiltinSwap = true
+  proc builtin_bswap32(a: uint32): uint32 {.
+      importc: "__builtin_bswap32", nodecl, noSideEffect.}
+elif defined(icc):
+  const useBuiltinSwap = true
+  proc builtin_bswap32(a: uint32): uint32 {.
+      importc: "_bswap", nodecl, noSideEffect.}
+elif defined(vcc):
+  const useBuiltinSwap = true
+  proc builtin_bswap32(a: uint32): uint32 {.
+      importc: "_byteswap_ulong", nodecl, header: "<intrin.h>", noSideEffect.}
+else:
+  const useBuiltinSwap = false
+
+when useBuiltinSwap:
+  proc swap(u: uint32): uint32 {.inline.} =
+    return builtin_bswap32(u)
+else:
+  proc swap(u: uint32): uint32 {.inline.} =
+    return ((u and 0xFF000000) shr 24) or
+      ((u and 0x00FF0000) shr 8) or
+      ((u and 0x0000FF00) shl 8) or
+      (u shl 24)
+
+proc fromBytesBEu32*(x: openArray[uint8]): uint32 {.inline.} =
+  var u {.noInit.}: uint32
+  copyMem(addr u, unsafeAddr x[0], sizeof(uint32))
+  when system.cpuEndian == littleEndian:
+    return swap(u)
+  else:
+    return u
+
+proc toBytesBE*(u: uint32): array[sizeof(uint32), uint8] {.inline.} =
+  when system.cpuEndian == littleEndian:
+    var u = swap(u)
+    copyMem(addr result[0], addr u, sizeof(uint32))
+  else:
+    copyMem(addr result[0], unsafeAddr u, sizeof(uint32))