TreeQuery: a DSL for syntax tree exploration

My previous post teased about IDE support through Metalua, and among others, about a robust and readable DSL to describe AST visitors. I have hacked together a prototype of this library. This post will start describing what it can do.

This post will showcase TreeQuery's API, the interfaces which allow to deal with trees. Methods and functions fall in two categories:

• Those which allow to describe an interesting subset of nodes;
• Those which allow to perform action on those selected nodes.

Source code

The implementation is also cut in two parts:

• a low-level walker metalua.treequery.walk, which goes mechanically through an AST; it is a variation upon the existing metalua.walk.
• metalua.treequery, which uses it to expose a declarative user interface.

Part I: selecting nodes in an AST

Queries

TreeQuery works on query objects, which represent sets of AST nodes (those sets are computed lazily). Queries are created by calling treequery() on an AST, and the resulting query initially represents the set of all nodes in the AST. By calling dedicated methods on the query, we'll eliminate irrelevant nodes, until only the wanted ones are left.

In the whole post, we'll alias "treequery" to "Q", to make things a bit terser; consider it a tribute to jQuery's $() shortcut. The following snippet loads the library, create an AST which we will use as example, then a query which represents every node in the AST:

local Q = require 'metalua.treequery'
ast = +{ block:
    local x=1
    for y=1,10 do
        print (x+i)
    end
    return math.cos(x) }
all_nodes = Q(ast)

The +{...} notation is a Metalua syntax which lets write an AST using regular syntax; it's the equivalent to Lisp's quasi-quoting forms (the anti-quote is written -{...}). Without the syntax sugar, ast would be written:

ast = {
    `Local{ { `Id "x" }, { `Number 1 } },
    `Fornum{ `Id "y", `Number 1, `Number 10, {
        `Call{ `Id "print", `Op{ "add", `Id "x", `Id "i" } }
    }
    `Return{ `Call{ `Index{ `Id "math", `String "cos" }, `Id "x" } } }

(the `Foo{bar} notation is itself a standard Metalua shortcut for {tag='Foo', bar}, which improves AST readability).

In the example above, Q(ast) represents all nodes and sub-nodes in ast: the block, the local statement, its "x" binder declaration, the "1" value, the "for" statement, etc. We will now intoduce the different methods which allow to filter those nodes and only keep the ones we're interested in.

:filter(pred), selecting a node according to its properties

The simplest operation consists of only keeping nodes which have a given property. This is done by method :filter(), which takes a predicate (a function taking a node and returning true or false). For instance, the following denotes the set of all nodes in the ast which have the 'Call' tag:

call_nodes = Q(ast) :filter (function(node) return node.tag=='Call' end)

Given the previous definition of ast, this query denotes the nodes +{print(x+i)} and +{math.cos(x)}.

:filter() allows some more sophisticated operations: in addition to receiving the tested node as first parameter, the predicate receives the node's parent as its second parameter, its grand-parent as third, etc. up to the AST root.

For instance, let's say we want to filter function-calls-as-statements (as opposed to function-calls-within-sub-expressions). These are the nodes which (1) have their tag equal to "Call" and (2) are in a block, i.e. their parent's tag is nil.

call_stats = Q(ast) :filter (function(node, parent)
    return node.tag=='Call' and parent and parent.tag==nil
end)

Moreover, we don't want users to write all of their predicates by hand. It's tedious, error-prone, and hard to read back. So we provide a library of standard predicate and predicate generators in TreeQuery.

Q.has_tag(tag_1, tag_2, ... tag_n) will return a predicate, which itself returns true if the node's tag is one of those listed as arguments. The first example, filtering all "Call" nodes, can therefore be rewritten:

call_nodes = Q(ast) :filter(Q.has_tag 'Call')

To further enhance readability, TreeQuery supports some shortcuts for the most common operations. Among others, when it expects a predicate but receives a string (or a sequence of strings), it assumes that the user meant to use the Q.has_tag() predicate generator. The last example can therefore be simplified into:

call_nodes = Q(ast) :filter 'Call'

Testing the parents

As mentioned above, predicates receive not only the node to test, but also its parent, grand-parent etc. up to the root node. We can therefore transform any predicate on a node into a predicate on this node's parent (it's just a matter of removing the predicate's first argument). Q.parent() does just that: for instance, predicate Q.is_block filters nodes which are blocks (i.e. whose tag equals nil); Q.parent(Q.is_block) therefore filters nodes whose parent is a block.

The example about "Call' nodes within blocks can therefore be rewritten by chaining two filters together:

call_stats = Q(ast) :filter 'Call' :filter (Q.parent(Q.is_block))

Positional filtering methods

:filter() acts locally. We also want to filter node according to their relative position to other node. To do that, we offer a series of methods :after(), :not_after(), :under(), not_under(), :under_or_after(), :not_under_or_after(). Each of these take a predicate, and only keep the nodes which are in the specified position relative to a node which passed the predicate.

For instance, if we want to eliminate all nodes under a "Function" node, we can write any of the following:

not_in_func = Q(ast) :not_under (function(x) x.tag=='Function' end)
not_in_func = Q(ast) :not_under (Q.has_tag 'Function')
not_in_func = Q(ast) :not_under 'Function'

Let's turn this into a more useful query: we want to find all return statements in a given function body. Those are the nodes with a 'Return' tag; however, returns which are in a nested 'Function' node must be ignored: they return from the nested function body, not from the body currently considered. For instance, in the following AST:

ast = +{block:
    if foo then return a end
    local function bar()
         return b
    end}

The first "return a" must be selected, but not the second "return b", which belongs to function bar. This is done as follows:

my_returns = Q(ast)
    :filter(Q.has_tag('Return'))
    :not_under(Q.has_tag('Function'))

Or with the string shortcuts:

my_returns = Q(ast) :filter 'Return' :not_under 'Function'

These positional filtering methods are still a work in progress. Today, they are to be interpreted strictly (i.e. a node is not considered to be under nor after itself). The option should also be offered to interpret them inclusively, and it's trivial to implement, but I can't think of a naming scheme or calling convention which makes this inclusiveness choice clear and readable. I'm OK with doubling the number of functions in the API, but I really want the resulting queries to read naturally; suggestions are welcome.

Variables and scopes

Many typical queries have to deal with identifiers: variable renaming, closing and moving a code block, finding rogue globals... For all of these jobs, we must parse variables correctly: detect potential variable captures, recognize occurrences of a same variable without being fooled by homonyms, etc. TreeQuery integrates a native support for these tasks.

First, TreeQuery makes a difference between binders, which create a local variable in a given scope ("for x=... end", "local x", "function (x)... end" etc.), and occurrences, i.e. cases when the variable is used as an expression (e.g. "print(x)").

To do so, it offers predicates to tell apart binders from occurrences, to retrieve the binder associated to a given occurrence (or return nil if it is an occurrence of a global varriable), and conversely to filter occurrences of a given binder. This is a pretty rich matter, which deserves to be explored in a post of its own. Let's just list the available APIs for now:

• Q.binder(occurrence, ast_root) returns the binder node associated to the occurrence, or nil if the variable is open (i.e. global) in this AST.
• Q.is_occurrence_of(binder) is a predicate which filters variables which are occurrences of the binder argument.
• Q.is_binder(...) tests whether an 'Id' node is a binder or an occurrence.

Miscelaneous predicates

• Q.is_nth(n)(node, ...) is true if node is the child number n of its parent
• Q.is_nth(a, b)(node, ...) is true if node is the child number #n of its parent, and a <= n <= b
• Q.is_stat(), Q.is_expr(), Q.is.block() test for statements, expressions, blocks.

Q.child() predicate transformer

This is the converse of Q.parent(): predicate Q.child(n, P) will return true if, when given the n-th child of the node, P returns true. It could be generalized to arbitrary descendants: for instance, Q.child(n1, n2, P) tests P on the n2-th child of the n1-th child of the tested node.

Part II: Acting on selected nodes

We've seen various ways to describe a set of nodes, within an AST root, which we're interested in; we haven't done anything with them yet. This part describes methods intended to act on this nodes set.

Extracting a list with :list()

The simplest way to do whatever you want on the nodes is to get a list of them. This is what's returned by the method :list(). The nodes in the list are ordered according to the depth-first traversal order, i.e. all the nodes in "a(b1(c11,c12),b2(c21,c22))" will be listed in order "{a, b1, c11, c12, b2, c21, c22}".

Extracting the first node with :first()

Sometimes you know you're only looking for one node, and there's no point traversing the rest of the tree once you've found it. :first() will stop as soon as it has found the first matching node (still in depth-first order), and returrn it. Moreover, it will return a multiple value: the node, its parent, grand-parent, etc. up to the root.

For instance, "Q(+{print(1+2*3)}) :filter 'Op' :first()" will return "+{1+2*3}, +{print(1+2*3)}".

Iterating on nodes with :foreach()

the :foreach() method takes one (or two, cf. below) function, and applies it on every node selected by the query. As usual with treequery, the parents of the node are passed as extra parameters to the callback.

(Historic note) In the old metalua.walk library, there were to visitor callbacks, down() and up(). down() corresponds to the top-to-bottom, depth first traversal order, whereas up() was called when going back up the tree. They guaranteed the following invariants:

• when down() is called on a node, down() has already been called on all of its parents;
• when up() is called on a node, up() has already been called on all of its children;
• on a given node, down() is called before up().

This ways of controlling the traversal order sometimes remains important; therefore, when :foreach() receives two callbacks, the first is used as the down() visitor, and the second is used as the up() visitor.

(Not implemented) Mapping transformations on nodes with :map()

This method allows to transform nodes into other things before they're passed to :foreach() callback, :list() or :first(). But maybe more importantly, they allow these transformations to be conditional, to only apply on nodes which pass some predicates. This allows to perform several, partially specific operations in a single pass. :map(p1, p2, ..., pn, f) won't eliminate any node from the query: it will replace the node N which pass all predicates p1,...pn with f(N). Several mappings can be put in a single request, which will either transform different nodes, or transform the same node more than once.

As usual, f receives the node's ancestors. By returning more than one node, f() can transform not only the node itself, but also its ancestry.

A question still to be determined: when several :map() are chained, and a first map already transformed a node, should the predicates of the second map receive the transformed nodes, or the original ones?

(Not implemented) for loop iterators

It would be nicer, and more idiomatic, to let write:

for x, x_parent in Q(ast) :filter 'Call' do
     various_stuff_on_call_nodes(x, x_parent)
end

Instead of:

Q(ast) :filter 'Call' :foreach (function (x, x_parent)
     various_stuff_on_call_nodes(x, x_parent)
end)

This is tricky to do correctly, due to Lua 5.1's "don't yield across C/Lua boundaries" limitation. I keep that in mind nonetheless.

Enough already for this post. Next time I'll show some interesting uses of the query language described above.

The other facet of Metalua: making IDEs smarter

Metalua is basically a tool to manipulate Lua programs in Lua. There are two main possible applications of it:

• using it as a self-extensible language, a.k.a. "Lisp without the parenthetical speech impediment";
• using it to analyze, and possibly modify, plain Lua source code.

I find the first application to be the funnier one, by far. It also raises plenty of software engineering open problems, about dynamic parsers, macro composition, hygiene, multiple views of a same object, etc. However, it remains a niche within a niche. Yet the later application, static source analysis and transformation, has a much wider potential audience. Java IDEs have transformed the expectations of many developers; we now expect lot of intelligence and assistance from and IDE, and it requires a deep static understanding of the programs being written. This is very tricky with dynamic languages such as Lua: with Java, you spend tremendous amounts of time making your program intelligible to the type system, with declarations, adapters, interface implementations etc. All this tedious bookkeeping is reused by the IDE to understand your programs. Dynamic languages free you from all this, but that leaves the IDE mostly clueless. So without statically checked types, either your IDE can know and do very little about your programs, or it has to make wild guesses based on heuristics ("heuristics" being an euphemism for "algorithmic-like thingy which sometimes fails, without even realizing how badly it failed"). The incentive to have those heuristics for Lua manipulations written in Lua, and easily modified, is huge: you want Lua hackers to tweak their own heuristics, in the language they all know--Lua. If they need to learn Eclipse, DLTK, XText, your own Lua framework, and possibly Java to describe their peculiar way of declaring a module, they simply won't: they'll keep using the IDE as a Notepad with syntax highlight. Providing a better IDE support Here comes the Metalua-based solution: interface your IDE with Metalua; let the latter tell the former what the code means, and perform any refactoring asked by the user. If Metalua is usable enough for your above-average-but-not-quite-a-wizard Lua hacker, then you can expect interesting things to happen. I believe such interesting things are indeed about to happen:

• my company is about to release a Lua support plugin for Eclipse, based on Metalua, as part of the wider Eclipse project Koneki.
• I'm working on making Metalua more accessible for source-to-source analysis and transformation.

Metalua makes Lua source file manipulations easier, by turning them into AST. unfortunately, visiting and modifying these trees is still a bit tedious. Moreover, such trees are easy to compile into bytecode, but not to convert back into proper source code. Of course, transforming an AST into a syntactically valid source code is trivial, and Metalua does this. But when an AST has been generated from a source file, modified, and redumped, you want the resulting file to keep as much of the original formatting as possible: comments, spaces, line jumps, syntax sugar etc. You'd also like this happen automatically: AST are supposed to be easy to parse; if keeping all the formatting details accurate is a chore, AST lose most of their appeal.

The two needs identified for IDE support are therefore:

• a robust, readable, intuitive tree visitor library;
• a robust "source -> AST -> refactored AST -> back to source" round-trip converter.

I'm currently focusing on the first points, although I think I've also got a rather elegant solution fo the second one. It's called TreeQuery, it's inspired by declarative tree visitors such as XPath or jQuery, and it defines a robust and readable Domain-Specific Language in Lua. It will be the subject of my next post.

Refactoring Lua with Metalua: source generation

In a previous post, I've begun to write a system to keep the AST / source code of a program in sync together; as a result, we were able to associate a source skeleton (string with holes in it for children nodes) to any given AST node, provided that this node was generated from sources. To make this useful, we must be able to generate source code from modified and/or generated AST parts; this post will describe how to do so. Be warned, there's nothing really fancy about it, it's pretty much brute force, systematic code walking. The resulting code is available in the Metalua repository, there: src/samples/synth.mlua First, I won't bother to use the walk library that comes with Metalua. This lib is a support for the visitor pattern, which is useful if you want to leave alone most of the AST nodes, or if you want to do pretty much the same thing on each node. In short, it saves its users from writing the tree traversal code for each kind of node, but this won't save much work if you have to do something different for each kind of node anyway. So everything will be embedded in a synth class: this class will:

• keep a generated source code accumulator;
• keep track of the current indentation level;
• Support some basic code creation methods: adding strings (:acc()), properly indented newlines (:nl()), handling indentation level (:nlindent() and :nldedent()),
• and actually walk the AST nodes (:node() and :list()).

I won't go into a painstaking description of each of these methods, and the source file above is extensively commented anyway. The only noteworthy point is the way :node() works: its job is to properly generate source from an AST node, and it does so by delegating to a method specialized for the given node's tag. That's one of those cases where Lua's "Everything is a table" paradygm is useful: an AST whose tag is "Foo" is going to be rendered by a method synth:Foo(ast, ...). There's no need to establish an explicit tag->method association, the method's name is the association. So :node() dispatches nodes to those specialized methods, which perform series of calls to :acc(), :nl(), :nlindent(), :nldedent()), :node() and :list(). Now Lua supports some syntax sugar, and there are some patterns which should be recognized and printed in a nicer way:

• "foo['bar']" should be printed "foo.bar" when bar is a valid identifier name;
• "x.y.z = function() ... end" should be printed "function x.y.z() ... end";
• "x.y.z.methodname = function (self, ...) end" should be printed "function x.y.z:methodname (...) end";
• "{ ['foo'] = bar }" should be printed "{ foo = bar }" if foo is a valid identifier name;
• "foo('bar')", "foo({bar})", "foo:bar('gnat')" and "foo:bar({gnat})" should be respectively printed "foo 'bar'", "foo{ bar }", "foo:bar 'gnat'" and "foo:bar{ gnat }";
• Unnecessary parentheses around operators should be avoided.

All of these patterns, plus a couple of trivial ones, are easily detected with structural pattern matching: This is the main reason why synth's code is significantly nicer in Metalua than it would have been in plain Lua. As a result, it's a fairly good example of advanced patterns usage, for a purpose that will be intuitive to all Lua programmer. This example will now have, in a future post, to be mixed with the "weaver" from a previous post. This weaver caused an error when it hadn't enough information to retrieve an AST's code skeleton. Now instead of causing an error, it should ask the synthetizer to produce decent looking source code for such AST parts. We'll see that it still hides a couple of non-trivial issues :) Finally, I'd like to state that this program is useful by itself, as opposed to the previous one: the source it produces is readable, significantly more so than big chunks of ASTs, and as such it is a nice way to debug code generated by hairy macros; I've written it 3 days ago and I'm already dependent on it! It will certainly be added, under a form or another, to the compiler's optional outputs.

Metalua moves to Github

I've cloned the Metalua repository from http://repo.or.cz/w/metalua.git to Github, and plan to use the later as the primary repository. The rationale, technical reason behind this migration is that Github's pages look nicer :) More seriously, I'm eager to see whether "something interesting" happens thanks to Github's social networking dimension. Next commits will be done in priority on Github, although I'll try to keep repo.or.cz updated regularly.

Refactoring Lua code with Metalua

Metalua is primarily designed as a self-extensible language, close to Scheme through Lua's heritage, but with a more opinionated syntax than lisps (I've written something about the importance of opinionated syntax over there). But its static introspection features and its backward compatibility with Lua make it a potentially powerful tool for plain Lua code analysis and refactoring. This post will be about the first steps toward support for Lua source refactoring in Metalua. The main service offered by Metalua, in this domain, is to let users manipulate AST rather than ASCII text. However, AST are much easier to analyze and modify than source code, but they don't carry as much information as sources do: the later have some presentation features (indentation, skipped lines, comments, uses of syntax sugar, even code generating macros in Metalua) which are not present in an AST, yet must be preserved when refactoring sources. The most natural reflex would be to enrich the AST grammar to the point where it contains as much information as the source code, but that would be a ridiculously complex structure, which would become unpractical to define, learn and handle. So we take a different approach: every node of an AST generated from sources will remember, in an optional "lineinfo" field, the first and last characters which defined it in the source code. For instance, let's consider this code and its AST:

"for i=1,10 do print(i) end"
`For{ `Id "i", `Number 1, `Number 10, 
      { `Call{ `Id "print", `Id "i" } } }

The `For node will have a lineinfo range 1-27; the first `Id "i" will have 5-5, the second 21-21; `Call will have 15-22, and in it, `Id "print", 15-19; etc. Now, for each node we know both its original textual representation, and its AST structure. The idea is that when we refactor a piece of code:

• we'll keep every untouched parts of the code under their original form;
• we'll resynthetize code from AST for the parts we modified, and only those.

In this example, let's say that I decide to change "print" into "_G.print". In AST terms, that would be replacing `Id "print" with `Index{ `Id "_G", `String "print" }. The replacement in AST is trivial (ast [4] [1] = +{_G.print}). When regenerating the source code of the for-loop, we need to take its "skeleton", its source with holes in the places occupied by its children nodes: for [ ] = [ ], [ ] do [ ] end Then in this source skeleton, which keeps the original formatting, we replace each of the children-holes by the corresponding child's regenerated source:

• the one taken straight out of the original sources for the first 3 children;
• and a synthesized source, generated from the lineinfo-free +{_G.print} AST node, for the last one that we modified.

The first step happened a while ago: it has been to properly implement and fix the lineinfo generation system, thanks to users' friendly pressure :). That feature had been planned for a long time but remained quite sketchy until then. Step #2 now: ability to regenerate a source file from its AST+source. I won't insist on the ability to synthesize source out of raw AST, only on the ability to weave together the bits of sources, recursively, from AST leaves to top-level node; since I don't plan to mix those recomposed parts with generated ones yet, the initial goal is simply to retrieve the original source file! The program is basically a a code walker, which will, on its way down,

• check which AST nodes can be weaved: to be able to weave it, we need to be able to extract its skeleton form the original sources, so we need the node and all of its direct children to have lineinfo fields; The list of conforming nodes will be kept in a weaveable node->Boolean table.
• generate an easy to read parent / child relation. This is not the same as table inclusion in the AST: some tables in the AST don't represent an expression, a statement or a block, and we're only interested into these element kinds. The relation will be kept in an ast_children parent->children list table.

Here's the code of this first walker:

require 'walk'

function weave_ast (src, ast, name)
   local ast_children, weaveable, node_cfg = { }, { }, { }

   function node_cfg.down(ast, parent)
      if   not ast.lineinfo 
      then weaveable [parent] = nil
      else weaveable [ast] = true end
      ast_children [ast] = { }
      if parent then table.insert (ast_children [parent], ast) end
   end

   local cfg = { expr=node_cfg; stat=node_cfg; block=node_cfg }
   walk.block (cfg, ast)
end

local src = ???
local ast = mlc.luastring_to_ast (src, name)
weave_ast (src, ast)

Now, this information can be used on our way "back up": each weaveable node will be generated by stitching its children's sources in its own source skeleton. This has to be done children first, parents later, of course. Moreover, we'll need a tiny trick here: the code walker is imperative, not functional, it's designed to modify ASTs rather than returning a value. To "return" a node's source, we'll create a gen_source AST->source string table, which will remember all of the translatations for the duration of the operations.

function weave_ast (src, ast, name)
   local gen_source, ast_children, weaveable, node_cfg = { }, { }, { }, { }

   function node_cfg.down(ast, parent) ... end

   function node_cfg.up(ast)
      local function weave_skeleton (ast)
         local buffer = { }
         local a, d = ast.lineinfo.first[3], ast.lineinfo.last[3]
         for child in ivalues (ast_children [ast]) do
            local b, c = child.lineinfo.first[3], child.lineinfo.last[3]
            table.insert (buffer, acc (src:sub (a, b-1)) .. acc (gen_source [child]))
            a = c+1
         end
         gen_source [ast] = table.concat (buffer) .. src:sub (a, d)
      end
      if weaveable [ast] then weave (ast)
      else error "Can't handle nodes without proper lineinfo yet" end
   end

   local cfg = { expr=node_cfg; stat=node_cfg; block=node_cfg }
   walk.block (cfg, ast)
   return gen_source [ast]
end

This will be enough for today. Next steps will be:

• step #3: ability to generate decent sources from an AST without any lineinfo nor sources.
• step #4: ability to mix both approaches, generate decent sources from an AST that's only partially lineinfo'd.

Statements as expressions

As most Lua users and every Metalua users know, Lua makes a difference between expressions and statements, and it's sometimes cumbersome to put a statement where the context expects an expression. This can be addressed in Lua by using a function closure, but at a considerable runtime cost. Metalua introduces a `Stat{ } AST node which lets metaprogrammers introduce a statement in an expression's context, but no concrete syntax was provided, partly to discourage the use of stats-as-exprs (which often is of poor taste), partly because I couldn't come up with a satisfying syntax. In this post I'll propose a syntax that seems better, if trickier to implement, than anything I've come by until now. `Stat{ block, expr } is an expression node, which evaluates the block, then evaluates the expression in the block's context (therefore having access to its local variables; that's also how "repeat foo until bar" works in Lua, where condition "bar" has access to "foo"'s local variables). The node's expression value is the second child's value. For instance, the following AST, when executed, prints 42:

print( -{ `Stat{ +{block: local x = 21}, +{expr: 2*x} })

I propose to present stat node between keyword "stat" ... "end". In that block, the last statement should be a return, and its value is the returned expression. The example above would be written: print(stat local x=21; return 2*x end). This is easily done by analyzing the content of the block and extracting the content of the return statement:

mlp.lexer:add{ "stat" }
mlp.expr:add{ "stat", mlp.block, "end", builder = stat_builder }
function stat_builder(x)
  local returned_expr
  local walker_cfg = { ... } -- Here the magic happens
  walk.block (walker_cfg, x[1])
  return `Stat{ x[1], expr }
end

Walking a block to extract returns (which are not inside a nested function) has already been addressed multiple times, e.g. here. Now what happens when no return or multiple returns are found? If there's no return, it's easy, just return nil. So above, we simply initialize returned_expr at +{nil}. In case of multiple returns, we could fail at compile time, but we can actually do better, and rewrite this as follows (supposing that goto and label are supported by the syntax):

print(stat 
   local x=21; 
   if foo then return foo else return 2*x end
end)

is rewritten as:

print(stat
   local $1=nil
   local x=21
   if foo then $1=foo; goto 'retval'
   else $1=2*x; goto 'retval' end
   label 'retval'
   return $1
end)

This way, multiple returns are handled gracefully. Moreover, this code has a semantically equivalent (although much less efficient) Lua implementation: "stat block end" can be translated into the closure "((function() block end)())". The complete implementation looks like this:

-{ extension 'xmatch' }

require 'walk'

local function stat_builder (x)
   local body, result = x[1]

   -- Accumulate direct return statements in `return_stats' list
   local return_stats = { }
   local walker_cfg = { expr={ }, stat={ } }
   match function walker_cfg.stat.up
   | `Return{ _ } == x -> table.insert (return_stats, x) 
   | `Return{   } == x -> x[1] = +{nil}; table.insert (return_stats, x) 
   | `Return{...} -> error "No multi-returns in stat blocks!"
   | _ -> 
   end
   match function walker_cfg.expr.down
   | `Function{...} -> return 'break' | _ -> 
   end
   walk.block (walker_cfg, x[1])

   match return_stats with
   -- No return statement -> resulting value is nil, don't change the body
   | { } -> result = +{ nil }
   -- One return statement at the end -> just extract it
   | { ret } if body[#body] == ret -> result = table.remove (body) [1]
   -- Several returns -> use variable and goto statements
   | _ -> 
      local retvar = mlp.gensym 'stat_value' 
      local label  = mlp.gensym 'end_of_stat'
      table.insert (body, 1, `Local{ {retvar}, { } })
      table.insert (body, `Label{ label })
      for ret in ivalues (return_stats) do
         ret.tag = nil
         ret <- { `Set{ {retvar}, {ret[1]} }, `Goto{ label } }
      end
      result = retvar
   end
   return `Stat{ body, result }
end

mlp.lexer:add{ "stat" }
mlp.expr:add{ "stat", mlp.block, "end", builder = stat_builder }

(I've used the xmatch extension described in a previous post). Notice, if you find that implementation too verbose, that it performs some operations more complex than most macros found in other languages: code analysis, different replacements depending on the code structure, (manual) hygiene within a Lisp-1 context... The simplest possible implementation, not supporting advanced code walking, would have been:

function stat_builder (x)
   local body, result = x[1], +{ nil }
   if body[#body].tag == 'Return' then 
      result = table.remove (body) [1]
   end
   return `Stat{ body, result }
end

mlp.lexer:add{ "stat" }
mlp.expr:add{ "stat", mlp.block, "end", builder = stat_builder }

The only limitation of the complete version is that multiple returns inside a stat blocks aren't supported. I'd rather keep it that way: the implementation of `Stat{ } is already a hack, since Lua really dislikes statements-as-expressions, at low level as well as at the syntax level. Making it multiple-values compatible would be horribly messy for little added value. However, it could be handled through rewriting: when at least one return in the block might have multiple values (that means, all returns with a comma, and all returns of a function call), we can:

• Change all "return foo" statements into "return {foo}".
• Change the default return value from +{nil} to +{{}}
• Surround the `Stat{ } node with a call to unpack().

Definitely not worth the kludge, nor the performance hit when treating a statement of the form "return f()", in my opinion. However, just because I can, here is the code:

-{ extension 'xmatch' }

require 'walk'

local multireturn_tags = table.transpose{ 'Dots', 'Call', 'Invoke' }

local function stat_builder (x)
   local body, result = x[1]

   -- Accumulate direct return statements in `return_stats' list
   local multireturn  = false
   local return_stats = { }
   local walker_cfg = { expr={ }, stat={ } }
   match function walker_cfg.stat.up
   | `Return{...} == x ->
      if #x>1 or multireturn_tags[x[#x].tag] then multireturn = true end
      table.insert (return_stats, x) 
   | _ -> 
   end
   match function walker_cfg.expr.down
   | `Function{...} -> return 'break' | _ -> 
   end
   walk.block (walker_cfg, x[1])

   match return_stats with
   -- No return statement -> resulting value is nil, don't change the body
   | { } -> result = +{ nil }; assert(not multireturn)
   -- One return statement at the end -> just extract it
   | { ret } if body[#body] == ret -> 
      if multireturn then
         result     = table.remove (body) 
         result.tag = 'Table' -- change return stat into list
      else
         result     = table.remove (body) [1]
      end
   -- Several returns -> use variable and goto statements
   | _ -> 
      local retvar = mlp.gensym 'stat_value' 
      local label  = mlp.gensym 'end_of_stat'
      table.insert (body, 1, `Local{ {retvar}, { multireturn and `Table } })
      table.insert (body, `Label{ label })
      for ret in ivalues (return_stats) do
         local this_result
         if multireturn then 
            this_result = table.shallow_copy(ret)
            this_result.tag = 'Table'
         else 
            this_result = ret[1] 
         end 
         ret.tag = nil
         ret <- { `Set{ {retvar}, {this_result} }, `Goto{ label } }
      end
      result = retvar
   end
   if multireturn then
      return `Call{ `Id "unpack", `Stat{ body, result } }
   else
      return `Stat{ body, result }
   end
end

mlp.lexer:add{ "stat" }
mlp.expr:add{ "stat", mlp.block, "end", builder = stat_builder }

One more catch

(This post might only be interesting to people with fairly good Lua proficiency). Some time ago, I wrote a post about designing a try...catch extension. The point was to show that even seemingly simple extensions require care to many corner cases, if one wants them to work reliably; and the implied consequence was that lower level approaches to meta-programming, akin to source preprocessing, are not suitable to write robust language features. I unwillingly demonstrated this by leaving an corner case unattended. The extension wraps portions of code that might fail into a "pcall(function()...end)"; this breaks return statements: they would return from the newly created function, not from the surrounding user's function as he would expect. Replacing return statements with something suitable was one of the extension's tasks (and one that would be hard to achieve with a code preprocessor). However, I forgot that "...", the extra arguments of a function call, can't be passed across function boundaries either. If they're used in a pcall+function wrapper, they have to be stored in a temporary table and unpacked when used. Here's the piece of code that has shown the issue:

-{ extension 'withdo' }

function run (fmt, ...)
   with h = io.popen(string.format(fmt, ...), 'r') do
      return h:read '*a'
   end
end

It uses "with...do...end", which limits the lifetime of a resource (here a file handle) to a lexical scope: the "h" declared in "with h = ... do ... end" will be closed once "end" is reached, even if returns and/or errors happen in the block; of course, it is implemented using the try/catch extension, and expands to:

function run (fmt, ...)
   do
      local h
      try 
         h = io.popen(string.format(fmt, ...), 'r') do
         return h:read '*a'
      finally
         h:close()
      end
   end
end

This in turns is equivalent to the following plain Lua code:

function run (fmt, ...)
   do
      local h
      local args = {...}
      local caught_return
      local user_success, user_error = pcall( function()
         h = io.popen(string.format(fmt, unpack(args)), 'r')
         caught_return =  { h:read '*a' }
      end)
      h:close()
      if user_success then
         if caught_return then return unpack(caught_return) end
      else
         error (user_error)
      end
   end
end

The (repeated) moral of the story is: language extensions are even harder than regular code to make reusable. If your language encourages writing quick and dirty macros, it will have a hard time growing extensive libraries. And if your meta-programming framework hasn't got a deep understanding of the code it manipulates, it's simply pointless.