Simple pickler

[Luis Carvalho <carvalho@...>]

Hi,

I've just found out about Roberto's struct lib, and so I thought it would not
be hard to come up with a simple pickle/unpickle module. I'm sharing it here,
since it's small, in case anyone is interested. :)

Taking a (slightly adapted) example from PiL2:

require "pickler"

a = {x = 1, y = 2; {3, 4, 5}}
a[2] = a -- cycle
a.z = a[1] -- shared subtable
a.f = function(n) return n+1 end

-- cross refs
b = {k = a[1]}
a[3] = b

pickler.tofile(a, "a.test")
pickler.tofile(b, "b.test")

We have in another session:

Lua 5.1.2  Copyright (C) 1994-2007 Lua.org, PUC-Rio
> require"pickler"
> a = pickler.fromfile"a.test"
> b = pickler.fromfile"b.test"
> print(a)
table: 0x309de0
> table.foreach(a, print)
1       table: 0x309ee0
2       table: 0x309de0
3       table: 0x30a010
y       2
x       1
z       table: 0x309ee0
f       function: 0x30a310
> print(b)
table: 0x30a010
> table.foreach(b, print)
k       table: 0x309ee0
> table.foreach(a[1], print)
1       3
2       4
3       5

It shouldn't be hard to extend the module to handle metatables or even add
a compression layer on top of the packing.

Cheers,
Luis.

-- 
A mathematician is a device for turning coffee into theorems.
        -- P. Erdos 

-- 
Luis Carvalho
Applied Math PhD Student - Brown University
PGP Key: E820854A <carvalho@dam.brown.edu>

-- pickler.lua
-- Simple pickle library based on Roberto's struct library at
--    http://www.inf.puc-rio.br/~roberto/struct
-- All Lua types but thread and userdata are supported. Function pickling is
-- based on string.dump/loadstring, and so closures are not an (easy) option.
-- Table with references to other tables are ok (including cyclic references).
-- This code is in public domain.

require "struct"
local pack = struct.pack
local unpack = struct.unpack

local type = type
local dump = string.dump
local loadstring = loadstring
local setmetatable = setmetatable
local pairs = pairs
local concat = table.concat
local error = error
local assert = assert
local open = io.open
local function taddress(t) -- table address: simple hash
  return tostring(t):sub(10)
end

module(...)

-- from lua.h
local TNIL = 0
local TBOOLEAN = 1
local TNUMBER = 3
local TSTRING = 4
local TTABLE = 5
local TFUNCTION = 6
-- extra
local TREF = 9
local TGUARD = 10

-- keep references to tables
local ref = setmetatable({}, {__mode = "v"})

pickle = function(o)
  local t = type(o)
  if t == "nil" then
    return pack("B", TNIL)
  elseif t == "boolean" then
    return pack("BB", TBOOLEAN, o and 1 or 0)
  elseif t == "number" then
    return pack("Bd", TNUMBER, o)
  elseif t == "string" then
    return pack("BBc0", TSTRING, #o, o)
  elseif t == "table" then
    local a = taddress(o)
    if ref[a] == nil then -- not interned?
      ref[a] = o
      local p = {} -- hold packs
      p[1] = pack("B", TTABLE) -- type
      p[2] = pack("Bc0", #a, a) -- ref key
      for k, v in pairs(o) do
        p[#p + 1] = pickle(k)
        p[#p + 1] = pickle(v)
      end
      p[#p + 1] = pack("B", TGUARD) -- table end
      return concat(p)
    end
    return pack("BBc0", TREF, #a, a)
  elseif t == "function" then
    local f = dump(o)
    return pack("BBc0", TFUNCTION, #f, f)
  else
    error("type not supported: " .. t)
  end
end

unpickle = function(o, start)
  local p = start or 1
  local t, p = unpack("B", o, p)
  if t == TNIL then
    return nil, 2
  elseif t == TBOOLEAN then
    return unpack("B", o, p) == 1, 3
  elseif t == TNUMBER then
    return unpack("d", o, p)
  elseif t == TSTRING then
    return unpack("Bc0", o, p)
  elseif t == TTABLE then
    local a, p = unpack("Bc0", o, p)
    -- intern table
    local r = {}
    ref[taddress(r)] = r
    ref[a] = r
    local k, v
    while (unpack("B", o, p)) ~= TGUARD do
      k, p = unpickle(o, p)
      v, p = unpickle(o, p)
      r[k] = v
    end
    p = p + 1 -- after TGUARD
    return r, p
  elseif t == TFUNCTION then
    local u, p = unpack("Bc0", o, p)
    return loadstring(u), p
  else -- TREF
    local a, p = unpack("Bc0", o, p)
    assert(ref[a] ~= nil, "table not interned")
    return ref[a], p
  end
end

tofile = function(o, fname)
  local f = assert(open(fname, "w"), "cannot open file: `" .. fname .."'")
  local p = pickle(o)
  f:write(p)
  f:close()
end

fromfile = function(fname)
  local f = assert(open(fname), "cannot open file: `" .. fname .."'")
  local p = f:read"*a"
  local o = unpickle(p)
  f:close()
  return o
end