lua cairo railroad diagram

gistfile1.txt

#!/usr/bin/env lua
-- "Quick" hack to generate rail road diagrams.
--
local cairo = require("lcairo")

-- reverse ipairs():
local function ripairs(t)
    local function ripairs_it(t,i)
        local v=t[i]
        if v==nil then return v end
        return i-1,v
    end
    return ripairs_it, t, #t
end


local w = 600
local h = 240
local outfile = "rr.png"

local cs = cairo.ImageSurface (cairo.FORMAT_RGB24, w, h)
local cr = cairo.Context (cs)

-- Normalize to an image of width 1:
cr:scale(w, w)
h = h/w
w = 1

-- Set background to white:
cr:set_source_rgb (1, 1, 1)
cr:paint()

-- Line object:
local Line = {}
Line.__index = Line
Line.line_width = 1/400

local Line_axis = {
    east  = "x",
    west  = "x",
    north = "y",
    south = "y",
}

function Line:new(direction, name)
    local self = {
        axis = Line_axis[direction] or
            error("InvalidDirection: '" .. direction ..
                  "' must be north, south, east, or west."),
        direction = direction,
        name      = name,
    }
    setmetatable(self, Line)
    return self
end
setmetatable(Line, { __call = Line.new })
-- Bubble object:
local Bubble = {}
Bubble.__index = Bubble
Bubble.font_size  = 1/30
Bubble.line_width = 1/400
Bubble.margin     = 1/120 + Bubble.line_width/2
do
    -- get the size of the text.  We will use the height of X so we
    -- get a consistent height.
    local extents = cairo.TextExtents()
    cr:select_font_face("sans", cairo.FONT_SLANT_NORMAL,
                        cairo.FONT_WEIGHT_BOLD)
    cr:set_font_size(Bubble.font_size)
    cr:text_extents("X", extents)

    Bubble.text_height = extents.height

    -- Now calculate the radius to be half the height:
    Bubble.radius = (Bubble.text_height+2*Bubble.margin)/2
end

function Bubble:new(str, color)
    local self = {text=str, color=color}
    setmetatable(self, Bubble)
    return self
end
function Bubble:pack(cr)
    local extents = cairo.TextExtents()
    cr:text_extents(self.text, extents)

    local text_width = math.max(extents.x_advance, self.radius)
    return text_width + 2*self.margin
end
function Bubble:leadup(line)
    return 0
end
function Bubble:render(cr)
    -- Various attributes:
    local color       = self.color
    local x           = self.x
    local y           = self.y
    local text        = self.text
    local font_size   = self.font_size
    local line_width  = self.line_width
    local margin      = self.margin
    local text_height = self.text_height
    local radius      = self.radius

    -- Save state, start a new path:
    cr:save()
    cr:new_path()

    -- Set-up drawing context:
    cr:set_line_width(line_width)
    cr:select_font_face("sans", cairo.FONT_SLANT_NORMAL,
                        cairo.FONT_WEIGHT_BOLD)
    cr:set_font_size(font_size)

    -- Now to calculate the width, we want the x_advance of the text,
    -- but if the x_advance is less than the radius, then we need to
    -- use the radius since the rounded corners would overlap.
    local extents = cairo.TextExtents()
    cr:text_extents(text, extents)

    local text_width = math.max(extents.x_advance, radius)

    local top = y - text_height/2 - margin
    local left = x
    local bottom = top + text_height + 2*margin
    local right = left + text_width  + 2*margin

    local begin_text_x = left + ((right - left) - extents.x_advance)/2
    local begin_text_y = bottom - margin

    cr:arc(left+radius, top+radius, radius, 2*(math.pi/2), 3*(math.pi/2))
    cr:arc(right-radius, top+radius, radius, 3*(math.pi/2), 0)
    cr:arc(right-radius, bottom-radius, radius, 0, math.pi/2)
    cr:arc(left+radius, bottom-radius, radius, math.pi/2, 2*(math.pi/2))
    cr:close_path()

    cr:set_source_rgb(unpack(color))
    cr:fill_preserve()

    cr:set_source_rgb(0.114, 0.29, 0.49)
    cr:stroke()

    cr:set_source_rgb(0, 0, 0)
    cr:move_to(begin_text_x, begin_text_y)
    cr:text_path(text)
    cr:fill()

    print("rendering '" .. text .. "' location top=" .. top ..
          " bottom=" .. bottom .. " left=" .. left .. " right=" .. right)

    cr:restore()
end
setmetatable(Bubble, { __call = Bubble.new })
-- Elbow object:
local Elbow = {}
Elbow.__index = Elbow
Elbow.radius  = 1/50
Elbow.line_width = 1/400
function Elbow:new(from_line, to_line, quadrant)
    local self = {
        quadrant = quadrant,
        from = from_line,
        to   = to_line
    }
    setmetatable(self, Elbow)
    return self
end
function Elbow:__eq(other)
    return
        self.quadrant == other.quadrant and
        self.from == other.from and
        self.to == other.to
end

function Elbow:pack(cr, line)
    local other_line =
        ( self.from == line ) and self.to or self.from

    other_line:pack(cr, self)

    -- TODO: Remove this crutch!
    return self.radius
end
function Elbow:leadup(line)
    --   _____
    --  /3   4\
    -- |       |
    --  \2   1/
    --   -----

    local choose = {
        self.radius,
        (line.axis == "x") and 0 or self.radius,
        0,
        (line.axis == "x") and self.radius or 0,
    }

    return choose[self.quadrant]
end
function Elbow:render(cr, line)
    cr:save()
    cr:new_path()

    --   _____
    --  /3   4\
    -- |       |
    --  \2   1/
    --   -----

    local x, y
    local fixup = {
        function()
            x = self.x - self.radius
            y = self.y - self.radius
        end,
        function()
            x = self.x + self.radius
            y = self.y - self.radius
        end,
        function()
            x = self.x + self.radius
            y = self.y + self.radius
        end,
        function()
            x = self.x - self.radius
            y = self.y + self.radius
        end,
    }
    fixup[self.quadrant]()

    cr:set_line_width(self.line_width)
    cr:set_source_rgb(0.114, 0.29, 0.49)
    cr:arc(x, y, self.radius,
           (self.quadrant-1)*(math.pi/2),
           self.quadrant*(math.pi/2))

    cr:stroke()
    cr:restore()

    if ( self.from == line ) then
        self.to:render(cr)
    else
        self.from:render(cr)
    end
end
setmetatable(Elbow, { __call = Elbow.new })
-- Circle object:
local Circle = {}
Circle.__index = Circle
Circle.radius     = 1/150
Circle.line_width = 1/200
function Circle:new()
    local self = {}
    setmetatable(self, Circle)
    return self
end
function Circle:pack(cr)
    return 2*self.radius
end
function Circle:leadup(line)
    return 0
end
function Circle:render(cr)
    cr:save()
    cr:new_path()

    cr:set_line_width(self.line_width)
    cr:set_source_rgb(0.114, 0.29, 0.49)
    -- TODO: Fix this "hack" that assumes the circle is rendered on
    -- the X axis.
    cr:arc(self.x + self.radius, self.y, self.radius, 0, 2*math.pi)
    cr:stroke()

    cr:restore()
end
setmetatable(Circle, { __call = Circle.new })

function Line:circle()
    self[#self+1] = Circle()
    return self
end
local r_quadrant = {
    southwest = 1,
    westnorth = 2,
    northeast = 3,
    eastsouth = 4,
}
function Line:r_exit(other)
    local quadrant = r_quadrant[self.direction .. other.direction] or
        error("Invalid r_exit: " .. self.direction .. other.direction)
    self[#self+1] = Elbow(self, other, quadrant)
    return self
end
function Line:r_enter(other)
    local quadrant = r_quadrant[other.direction .. self.direction] or
        error("Invalid r_enter: " .. self.direction .. other.direction)
    self[#self+1] = Elbow(other, self, quadrant)
    return self
end
local l_quadrant = {
    eastnorth = 1,
    southeast = 2,
    westsouth = 3,
    northwest = 4,
}
function Line:l_exit(other)
    local quadrant = l_quadrant[self.direction .. other.direction] or
        error("Invalid l_exit: " .. self.direction .. other.direction)
    self[#self+1] = Elbow(self, other, quadrant)
    return self
end
function Line:l_enter(other)
    local quadrant = l_quadrant[other.direction .. self.direction] or
        error("Invalid l_enter: " .. self.direction .. other.direction)
    self[#self+1] = Elbow(other, self, quadrant)
    return self
end
function Line:terminal(str)
    self[#self+1] = Bubble(str, {1,1,0})
    return self
end
function Line:non_terminal(str)
    self[#self+1] = Bubble(str, {1,1,1})
    return self
end

local function increment(i, amount)
    amount = amount or 1
    return i + amount
end
local function decrement(i, amount)
    amount = amount or 1
    return i - amount
end

-- Works with nil
local function max(...)
    local result
    for n=1, select('#', ...) do
        local elm = select(n, ...)
        if ( not result or
             elm and elm > result )
        then
            result = elm
        end
    end
    return result
end

-- Works with nil
local function min(...)
    local result
    for n=1, select('#', ...) do
        local elm = select(n, ...)
        if ( not result or
             elm and elm < result )
        then
            result = elm
        end
    end
    return result
end

-- Set the (X, Y) of each element.
function Line:pack(cr, elbow)
    -- Find the matching elbow.
    local elbow_i
    if ( getmetatable(elbow) ~= Elbow ) then
        elbow_i = 1
    else
        for i,node in ipairs(self) do
            if ( node == elbow ) then
                elbow_i = i
                break
            end
        end
    end
    assert(elbow_i, "InvalidElbow not on this line!")

    -- Determine if elbow requires adjustment and if so, which way do
    -- we need to ripple this adjustment?  Note that we only ripple to
    -- the SouthEast.
    local new_coords = {
        x = max(self[elbow_i].x, elbow.x),
        y = max(self[elbow_i].y, elbow.y),
    }

    -- Adjust the entire line in this new direction:
    local other_axis = ( self.axis == "x" ) and "y" or "x"
    local full_repack = false
    if ( self[elbow_i][other_axis] ~= new_coords[other_axis] ) then
        for i,node in ipairs(self) do
            print(self.name .. "[" .. i .. "]." .. other_axis .. " = " .. new_coords[other_axis])

            node[other_axis] = new_coords[other_axis]
        end
        full_repack = true
    end

    local choose = {
        north = decrement,
        east  = increment,
        south = increment,
        west  = decrement,
    }
    local next_i = choose[self.direction] or error("UnexpectedDirection")

    if ( self[elbow_i][self.axis] ~= new_coords[self.axis] ) then
        -- Shift everything down:
        local i = elbow_i
        while self[i] do
            print(self.name .. "[" .. i .. "]." .. self.axis .. " = " .. new_coords[self.axis])

            new_coords[self.axis] = max(self[i][self.axis], new_coords[self.axis])

            self[i][self.axis] = new_coords[self.axis]

            local width = self[i]:pack(cr, self)

            self[i].width = width

            new_coords[self.axis] =
                increment(new_coords[self.axis], width)

            local i_next = next_i(i)
            if self[i_next] then
                new_coords[self.axis] = new_coords[self.axis] + self[i_next]:leadup(self)
            end
            i = i_next
        end
    end

    if not full_repack then
        return
    end

    local choose = {
        north = #self,
        east  = 1,
        south = 1,
        west  = #self,
    }
    local i = choose[self.direction] or erro("UnexpectedDirection")
    while ( i > 0 and i <= #self ) do
        if ( i == elbow_i ) then
            return
        end
        if self[i].x and self[i].y then
            self[i]:pack(cr, self)
        end
        i = next_i(i)
    end

    -- return value ignored by Elbow
    return
end
function Line:render(cr)
    if self.rendered then
        return
    end
    self.rendered = true

    cr:save()
    cr:new_path()

    cr:set_line_width(self.line_width)
    cr:set_source_rgb(0.114, 0.29, 0.49)

    local choose = {
        north = decrement,
        east  = increment,
        south = increment,
        west  = decrement,
    }
    local advance = choose[self.direction] or error("InvalidDirection")
    local coords = {
        x = self[1].x,
        y = self[1].y,
    }
    coords[self.axis] = advance(coords[self.axis], self[1].width)
    cr:move_to(coords.x, coords.y)

    local coords = {
        x = self[#self].x,
        y = self[#self].y,
    }
    coords[self.axis] = coords[self.axis] - self[#self].width


    local choose = {
        north = increment,
        east  = decrement,
        south = decrement,
        west  = increment,
    }
    local retreat = choose[self.direction] or error("InvalidDirection")
    local coords = {
        x = self[#self].x,
        y = self[#self].y,
    }
    -- TODO: This is a hack, think of a better way :-(.
    if getmetatable(self[#self]) ~= Circle then
        coords[self.axis] = retreat(coords[self.axis], self[#self].width)
    end
    cr:line_to(coords.x, coords.y)
    cr:stroke()

    cr:restore()

    for _,node in ipairs(self) do
        node:render(cr)
    end
end

-- Declare all lines:
local east1  = Line("east", "east1")
local east2  = Line("east", "east2")
local east3  = Line("east", "east3")
local west4  = Line("west", "west4")
local north1 = Line("north", "north1")
local south2 = Line("south", "south2")
local north3 = Line("north", "north3")
local south4 = Line("south", "south4")
local north5 = Line("north", "north5")
local south6 = Line("south", "south6")
local south7 = Line("south", "south7")
local north8 = Line("north", "north8")
local north9 = Line("north", "north9")

-- Horizontal lines:
east1
:circle()
:r_enter(north1)
:r_exit(south2)
:r_enter(north3)
:r_enter(north5)
:r_exit(south7)
:r_enter(north8)
:r_enter(north9)
:circle()

east2
:l_enter(south2)
:non_terminal("stat")
:l_exit(north3)
:r_exit(south4)
:terminal(";")
:l_exit(north5)
:r_exit(south6)

east3
:l_enter(south7)
:non_terminal("laststat")
:l_exit(north8)
:terminal(";")
:l_exit(north9)

west4:r_enter(south6):r_enter(south4):r_exit(north1)

-- Vertical lines:
north1:r_enter(west4):r_exit(east1)
south2:r_enter(east1):l_exit(east2)
north3:l_enter(east2):r_exit(east1)
south4:r_enter(east2):r_exit(west4)
north5:l_enter(east2):r_exit(east1)
south6:r_enter(east2):r_exit(west4)
south7:r_enter(east1):l_exit(east3)
north8:l_enter(east3):r_exit(east1)
north9:l_enter(east3):r_exit(east1)

east1:pack(cr, {x=1/40, y=1/10})
east1:render(cr)

-- I'm not sure why this doesn't work:
--   cr:get_target():write_to_png(outfile)
cairo.surface_write_to_png(cr:get_target(), outfile)