arti-pack/tests/deflate.lua

-- A minimal DEFLATE encoder, so compressed fixtures are built here rather than
-- shelled out to a compressor that may be absent and silently skip the tests.
local M = {}

local MIN_MATCH, MAX_MATCH, WINDOW = 3, 258, 32768
local END_OF_BLOCK = 256
local MAX_LIT_SYMBOLS, MAX_DIST_SYMBOLS, MAX_CL_SYMBOLS = 286, 30, 19

local function nums(spec)
  local out = {}
  for n in spec:gmatch('%d+') do
    table.insert(out, tonumber(n))
  end
  return out
end

-- RFC 1951 section 3.2.5: length codes 257..285 and distance codes 0..29.
local LENGTH_BASE = nums('3 4 5 6 7 8 9 10 11 13 15 17 19 23 27 31 35 43 51 59 67 83 99 115 131 163 195 227 258')
local LENGTH_EXTRA = nums('0 0 0 0 0 0 0 0 1 1 1 1 2 2 2 2 3 3 3 3 4 4 4 4 5 5 5 5 0')
local DIST_BASE = nums(
  '1 2 3 4 5 7 9 13 17 25 33 49 65 97 129 193 257 385 513 769 1025 1537 2049 3073 4097 6145 8193 12289 16385 24577'
)
local DIST_EXTRA = nums('0 0 0 0 1 1 2 2 3 3 4 4 5 5 6 6 7 7 8 8 9 9 10 10 11 11 12 12 13 13')
local CL_ORDER = { 16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15 }

local Writer = {}
Writer.__index = Writer

local function writer()
  return setmetatable({ out = {}, acc = 0, n = 0 }, Writer)
end

-- Plain values go out least significant bit first.
function Writer:bits(value, count)
  for _ = 1, count do
    self.acc = self.acc + (value % 2) * 2 ^ self.n
    value = math.floor(value / 2)
    self.n = self.n + 1
    if self.n == 8 then
      table.insert(self.out, string.char(self.acc))
      self.acc, self.n = 0, 0
    end
  end
end

-- Huffman codes go out most significant bit first.
function Writer:code(code, len)
  for i = len - 1, 0, -1 do
    self:bits(math.floor(code / 2 ^ i) % 2, 1)
  end
end

function Writer:finish()
  if self.n > 0 then
    table.insert(self.out, string.char(self.acc))
    self.acc, self.n = 0, 0
  end
  return table.concat(self.out)
end

local function length_code(len)
  for i = #LENGTH_BASE, 1, -1 do
    if len >= LENGTH_BASE[i] then
      return 256 + i, len - LENGTH_BASE[i], LENGTH_EXTRA[i]
    end
  end
  error('length out of range: ' .. len)
end

local function dist_code(dist)
  for i = #DIST_BASE, 1, -1 do
    if dist >= DIST_BASE[i] then
      return i - 1, dist - DIST_BASE[i], DIST_EXTRA[i]
    end
  end
  error('distance out of range: ' .. dist)
end

-- Greedy LZ77 over a single-slot hash of the next three bytes. Good enough to
-- put real matches, extra bits and overlapping copies into the stream.
local function tokenize(data)
  local tokens, head, pos, n = {}, {}, 1, #data
  while pos <= n do
    local best, dist = 0, 0
    if pos + MIN_MATCH - 1 <= n then
      local key = data:sub(pos, pos + MIN_MATCH - 1)
      local prev = head[key]
      if prev and pos - prev <= WINDOW then
        local len = 0
        while len < MAX_MATCH and pos + len <= n and data:byte(prev + len) == data:byte(pos + len) do
          len = len + 1
        end
        if len >= MIN_MATCH then
          best, dist = len, pos - prev
        end
      end
      head[key] = pos
    end
    if best > 0 then
      table.insert(tokens, { len = best, dist = dist })
      for i = pos + 1, math.min(pos + best - 1, n - MIN_MATCH + 1) do
        head[data:sub(i, i + MIN_MATCH - 1)] = i
      end
      pos = pos + best
    else
      table.insert(tokens, data:byte(pos))
      pos = pos + 1
    end
  end
  return tokens
end

local function fixed_lit(sym)
  if sym <= 143 then
    return 0x30 + sym, 8
  elseif sym <= 255 then
    return 0x190 + sym - 144, 9
  elseif sym <= 279 then
    return sym - 256, 7
  end
  return 0xc0 + sym - 280, 8
end

local function emit(w, tokens, litcode, distcode)
  for _, tk in ipairs(tokens) do
    if type(tk) == 'number' then
      w:code(litcode(tk))
    else
      local lc, lextra, lbits = length_code(tk.len)
      w:code(litcode(lc))
      if lbits > 0 then
        w:bits(lextra, lbits)
      end
      local dc, dextra, dbits = dist_code(tk.dist)
      w:code(distcode(dc))
      if dbits > 0 then
        w:bits(dextra, dbits)
      end
    end
  end
  w:code(litcode(END_OF_BLOCK))
end

-- A one-leaf alphabet cannot produce a complete prefix code, so pad it.
local function ensure_two(freq, limit)
  local used = 0
  for _ in pairs(freq) do
    used = used + 1
  end
  for sym = 0, limit - 1 do
    if used >= 2 then
      return
    end
    if not freq[sym] then
      freq[sym] = 1
      used = used + 1
    end
  end
end

local function build_lengths(weights)
  local nodes, lens = {}, {}
  for sym, weight in pairs(weights) do
    table.insert(nodes, { w = weight, syms = { sym } })
    lens[sym] = 0
  end
  assert(#nodes >= 2, 'huffman needs at least two symbols')
  while #nodes > 1 do
    table.sort(nodes, function(a, b)
      if a.w ~= b.w then
        return a.w < b.w
      end
      return a.syms[1] < b.syms[1]
    end)
    local a, b = table.remove(nodes, 1), table.remove(nodes, 1)
    local syms = {}
    for _, list in ipairs({ a.syms, b.syms }) do
      for _, sym in ipairs(list) do
        lens[sym] = lens[sym] + 1
        table.insert(syms, sym)
      end
    end
    table.insert(nodes, { w = a.w + b.w, syms = syms })
  end
  return lens
end

-- Flattening the weights shortens the deepest codes; a flat alphabet of at most
-- 286 symbols never needs more than nine bits, so this always settles.
local function huffman(freq, limit)
  local weights = {}
  for sym, weight in pairs(freq) do
    weights[sym] = weight
  end
  while true do
    local lens, worst = build_lengths(weights), 0
    for _, len in pairs(lens) do
      worst = math.max(worst, len)
    end
    if worst <= limit then
      return lens
    end
    for sym, weight in pairs(weights) do
      weights[sym] = math.ceil(weight / 2)
    end
  end
end

local function canonical(lens)
  local counts, maxlen, syms = {}, 0, {}
  for sym, len in pairs(lens) do
    counts[len] = (counts[len] or 0) + 1
    maxlen = math.max(maxlen, len)
    table.insert(syms, sym)
  end
  local next_code, code = {}, 0
  for bits = 1, maxlen do
    code = (code + (counts[bits - 1] or 0)) * 2
    next_code[bits] = code
  end
  table.sort(syms)
  local codes = {}
  for _, sym in ipairs(syms) do
    codes[sym] = next_code[lens[sym]]
    next_code[lens[sym]] = next_code[lens[sym]] + 1
  end
  return codes
end

local function highest(lens)
  local top = 0
  for sym in pairs(lens) do
    top = math.max(top, sym)
  end
  return top
end

local function fixed_block(w, tokens, final)
  w:bits(final, 1)
  w:bits(1, 2)
  emit(w, tokens, fixed_lit, function(sym)
    return sym, 5
  end)
end

local function dynamic_block(w, tokens, final)
  local litfreq, distfreq = { [END_OF_BLOCK] = 1 }, {}
  for _, tk in ipairs(tokens) do
    if type(tk) == 'number' then
      litfreq[tk] = (litfreq[tk] or 0) + 1
    else
      local lc = length_code(tk.len)
      litfreq[lc] = (litfreq[lc] or 0) + 1
      local dc = dist_code(tk.dist)
      distfreq[dc] = (distfreq[dc] or 0) + 1
    end
  end
  ensure_two(litfreq, MAX_LIT_SYMBOLS)
  ensure_two(distfreq, MAX_DIST_SYMBOLS)

  local litlens, distlens = huffman(litfreq, 15), huffman(distfreq, 15)
  local hlit = math.max(257, highest(litlens) + 1)
  local hdist = highest(distlens) + 1

  local seq = {}
  for sym = 0, hlit - 1 do
    table.insert(seq, litlens[sym] or 0)
  end
  for sym = 0, hdist - 1 do
    table.insert(seq, distlens[sym] or 0)
  end

  -- Code lengths are written one per symbol; the run-length codes 16..18 are
  -- left unused so the encoder stays small.
  local clfreq = {}
  for _, len in ipairs(seq) do
    clfreq[len] = (clfreq[len] or 0) + 1
  end
  ensure_two(clfreq, 16)
  local cllens = huffman(clfreq, 7)
  local clcodes = canonical(cllens)

  w:bits(final, 1)
  w:bits(2, 2)
  w:bits(hlit - 257, 5)
  w:bits(hdist - 1, 5)
  w:bits(MAX_CL_SYMBOLS - 4, 4)
  for _, sym in ipairs(CL_ORDER) do
    w:bits(cllens[sym] or 0, 3)
  end
  for _, len in ipairs(seq) do
    w:code(clcodes[len], cllens[len])
  end

  local litcodes, distcodes = canonical(litlens), canonical(distlens)
  emit(w, tokens, function(sym)
    return litcodes[sym], litlens[sym]
  end, function(sym)
    return distcodes[sym], distlens[sym]
  end)
end

-- Returns a raw DEFLATE stream. mode is 'dynamic' (default) or 'fixed'.
function M.compress(data, mode)
  local w = writer()
  local tokens = tokenize(data)
  if mode == 'fixed' then
    fixed_block(w, tokens, 1)
  else
    dynamic_block(w, tokens, 1)
  end
  return w:finish()
end

return M