metalua/treequery.mlua
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
local walk = require 'metalua.treequery.walk' | |
local M = { } | |
-- support for old-style modules | |
treequery = M | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
-- | |
-- multimap helper mmap: associate a key to a set of values | |
-- | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
local function mmap_add (mmap, node, x) | |
if node==nil then return false end | |
local set = mmap[node] | |
if set then set[x] = true | |
else mmap[node] = {[x]=true} end | |
end | |
-- currently unused, I throw the whole set away | |
local function mmap_remove (mmap, node, x) | |
local set = mmap[node] | |
if not set then return false | |
elseif not set[x] then return false | |
elseif next(set) then set[x]=nil | |
else mmap[node] = nil end | |
return true | |
end | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
-- | |
-- TreeQuery object. | |
-- | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
local ACTIVE_SCOPE = setmetatable({ }, {__mode="k"}) | |
-- treequery metatable | |
local Q = { }; Q.__index = Q | |
--- treequery constructor | |
-- the resultingg object will allow to filter ans operate on the AST | |
-- @param root the AST to visit | |
-- @return a treequery visitor instance | |
function M.treequery(root) | |
return setmetatable({ | |
root = root, | |
unsatisfied = 0, | |
predicates = { }, until_up = { }, | |
from_up = { }, | |
up_f = false, | |
down_f = false, | |
filters = { }, | |
}, Q) | |
end | |
-- helper to share the implementations of positional filters | |
local function add_pos_filter(self, position, inverted, inclusive, f, ...) | |
if type(f)=='string' then f = M.has_tag(f, ...) end | |
if not inverted then self.unsatisfied += 1 end | |
local x = { | |
pred = f, | |
position = position, | |
satisfied = false, | |
inverted = inverted or false, | |
inclusive = inclusive or false } | |
table.insert(self.predicates, x) | |
return self | |
end | |
-- TODO: offer an API for inclusive pos_filters | |
--- select nodes which are after one which satisfies predicate f | |
Q.after = |self, f, ...| add_pos_filter(self, 'after', false, false, f, ...) | |
--- select nodes which are not after one which satisfies predicate f | |
Q.not_after = |self, f, ...| add_pos_filter(self, 'after', true, false, f, ...) | |
--- select nodes which are under one which satisfies predicate f | |
Q.under = |self, f, ...| add_pos_filter(self, 'under', false, false, f, ...) | |
--- select nodes which are not under one which satisfies predicate f | |
Q.not_under = |self, f, ...| add_pos_filter(self, 'under', true, false, f, ...) | |
--- select nodes which satisfy predicate f | |
function Q :filter(f, ...) | |
if type(f)=='string' then f = M.has_tag(f, ...) end | |
table.insert(self.filters, f); | |
return self | |
end | |
-- private helper: apply filters and execute up/down callbacks when applicable | |
function Q :execute() | |
local cfg = { } | |
-- TODO: optimize away not_under & not_after by pruning the tree | |
function cfg.down(...) | |
--printf ("[down]\t%s\t%s", self.unsatisfied, table.tostring((...))) | |
ACTIVE_SCOPE[...] = cfg.scope | |
local satisfied = self.unsatisfied==0 | |
for _, x in ipairs(self.predicates) do | |
if not x.satisfied and x.pred(...) then | |
x.satisfied = true | |
local node, parent = ... | |
local inc = x.inverted and 1 or -1 | |
if x.position=='under' then | |
-- satisfied from after we get down this node... | |
self.unsatisfied += inc | |
-- ...until before we get up this node | |
mmap_add(self.until_up, node, x) | |
elseif x.position=='after' then | |
-- satisfied from after we get up this node... | |
mmap_add(self.from_up, node, x) | |
-- ...until before we get up this node's parent | |
mmap_add(self.until_up, parent, x) | |
elseif x.position=='under_or_after' then | |
-- satisfied from after we get down this node... | |
self.satisfied += inc | |
-- ...until before we get up this node's parent... | |
mmap_add(self.until_up, parent, x) | |
else | |
error "position not understood" | |
end -- position | |
if x.inclusive then satisfied = self.unsatisfied==0 end | |
end -- predicate passed | |
end -- for predicates | |
if satisfied then | |
for _, f in ipairs(self.filters) do | |
if not f(...) then satisfied=false; break end | |
end | |
if satisfied and self.down_f then self.down_f(...) end | |
end | |
end | |
function cfg.up(...) | |
--printf ("[up]\t%s", table.tostring((...))) | |
-- Remove predicates which are due before we go up this node | |
local preds = self.until_up[...] | |
if preds then | |
for x, _ in pairs(preds) do | |
local inc = x.inverted and -1 or 1 | |
self.unsatisfied += inc | |
x.satisfied = false | |
end | |
self.until_up[...] = nil | |
end | |
-- Execute the up callback | |
-- TODO: cache the filter passing result from the down callback | |
-- TODO: skip if there's no callback | |
local satisfied = self.unsatisfied==0 | |
if satisfied then | |
for _, f in ipairs(self.filters) do | |
if not f(self, ...) then satisfied=false; break end | |
end | |
if satisfied and self.up_f then self.up_f(...) end | |
end | |
-- Set predicate which are due after we go up this node | |
local preds = self.from_up[...] | |
if preds then | |
for p, _ in pairs(preds) do | |
local inc = p.inverted and 1 or -1 | |
self.unsatisfied += inc | |
end | |
self.from_up[...] = nil | |
end | |
ACTIVE_SCOPE[...] = nil | |
end | |
function cfg.binder(id_node, ...) | |
--printf(" >>> Binder called on %s, %s", table.tostring(id_node), | |
-- table.tostring{...}:sub(2,-2)) | |
cfg.down(id_node, ...) | |
cfg.up(id_node, ...) | |
--printf("down/up on binder done") | |
end | |
--function cfg.occurrence (binder, occ) | |
-- if binder then OCC2BIND[occ] = binder[1] end | |
--printf(" >>> %s is an occurrence of %s", occ[1], table.tostring(binder and binder[2])) | |
--end | |
--function cfg.binder(...) cfg.down(...); cfg.up(...) end | |
return walk.guess(cfg, self.root) | |
end | |
--- Execute a function on each selected node | |
-- @down: function executed when we go down a node, i.e. before its children | |
-- have been examined. | |
-- @up: function executed when we go up a node, i.e. after its children | |
-- have been examined. | |
function Q :foreach(down, up) | |
if not up and not down then | |
error "iterator not implemented" | |
end | |
self.up_f = up | |
self.down_f = down | |
return self :execute() | |
end | |
--- Return the list of nodes selected by a given treequery. | |
function Q :list() | |
local acc = { } | |
self :foreach(|x| table.insert(acc, x)) | |
return acc | |
end | |
--- Return the first matching element | |
-- TODO: dirty hack, to implement properly with a 'break' return. | |
-- Also, it won't behave correctly if a predicate causes an error, | |
-- or if coroutines are involved. | |
function Q :first() | |
local result = { } | |
local function f(...) result = {...}; error() end | |
pcall(|| self :foreach(f)) | |
return unpack(result) | |
end | |
--- Pretty printer for queries | |
function Q :__tostring() return "treequery("..table.tostring(self.root, 'nohash')..")" end | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
-- | |
-- Predicates. | |
-- | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
--- Return a predicate which is true if the tested node's tag is among the | |
-- one listed as arguments | |
-- @param ... a sequence of tag names | |
function M.has_tag(...) | |
local args = {...} | |
if #args==1 then | |
local tag = ... | |
return (|self, node| node.tag==tag) | |
--return function(self, node) printf("node %s has_tag %s?", table.tostring(node), tag); return node.tag==tag end | |
else | |
local tags = { } | |
for _, tag in ipairs(args) do tags[tag]=true end | |
return function(self, node) | |
local node_tag = node.tag | |
return node_tag and tags[node_tag] | |
end | |
end | |
end | |
--- Predicate to test whether a node represents an expression. | |
M.is_expr = M.has_tag('Nil', 'Dots', 'True', 'False', 'Number','String', | |
'Function', 'Table', 'Op', 'Paren', 'Call', 'Invoke', | |
'Id', 'Index') | |
-- helper for is_stat | |
local STAT_TAGS = { Do=1, Set=1, While=1, Repeat=1, If=1, Fornum=1, | |
Forin=1, Local=1, Localrec=1, Return=1, Break=1 } | |
--- Predicate to test whether a node represents a statement. | |
-- It is context-aware, i.e. it recognizes `Call and `Invoke nodes | |
-- used in a statement context as such. | |
function M.is_stat(self, node, parent) | |
local tag = node.tag | |
if not tag then return false | |
elseif STAT_TAGS[tag] then return true | |
elseif tag=='Call' or tag=='Invoke' then return parent.tag==nil | |
else return false end | |
end | |
--- Predicate to test whether a node represents a statements block. | |
function M.is_block(self, node) return node.tag==nil end | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
-- | |
-- Variables and scopes. | |
-- | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
local BINDER_GRAND_PARENT_TAG = { | |
Local=true, Localrec=true, Forin=true, Function=true } | |
--- Test whether a node is a binder. This is local predicate, although it | |
-- might need to inspect the parent node. | |
function M.is_binder(node, parent) | |
--printf('is_binder(%s, %s)', table.tostring(node), table.tostring(parent)) | |
if node.tag ~= 'Id' or not parent then return false end | |
if parent.tag=='Fornum' then return b[1]==a end | |
if not BINDER_GRAND_PARENT_TAG[parent.tag] then return false end | |
for _, binder in ipairs(parent[1]) do | |
if binder==node then return true end | |
end | |
return false | |
end | |
--- Retrieve the binder associated to an occurrence within root. | |
-- @param occurrence an Id node representing an occurrence in `root`. | |
-- @param root the tree in which `node` and its binder occur. | |
-- @return the binder node, and its ancestors up to root if found. | |
-- @return nil if node is global (or not an occurrence) in `root`. | |
function M.binder(occurrence, root) | |
local cfg, id_name, result = { }, occurrence[1], { } | |
function cfg.occurrence(id) | |
if id == occurrence then result = cfg.scope :get(id_name) end | |
-- TODO: break the walker | |
end | |
walk.guess(cfg, root) | |
return unpack(result) | |
end | |
--- Predicate to filter occurrences of a given binder. | |
-- Warning: it relies on internal scope book-keeping, | |
-- and for this reason, it only works as query method argument. | |
-- It won't work outside of a query. | |
-- @param binder the binder whose occurrences must be kept by predicate | |
-- @return a predicate | |
function M.is_occurrence_of(binder) | |
return function(node, ...) | |
if node.tag ~= 'Id' then return nil end | |
if M.is_binder(node, ...) then return nil end | |
local scope = ACTIVE_SCOPE[node] | |
if not scope then return nil end | |
local result = scope :get (node[1]) or { } | |
if result[1] ~= binder then return nil end | |
return unpack(result) | |
end | |
end | |
--- Transform a predicate on a node into a predicate on this node's | |
-- parent. For instance if p tests whether a node has property P, | |
-- then parent(p) tests whether this node's parent has property P. | |
-- The ancestor level is precised with n, with 1 being the node itself, | |
-- 2 its parent, 3 its grand-parent etc. | |
-- @param[optional] n the parent to examine, default=2 | |
-- @param pred the predicate to transform | |
-- @return a predicate | |
function M.parent(n, pred, ...) | |
if type(a)~='number' then n, pred = 2, n end | |
if type(pred)=='string' then pred = M.has_tag(pred, ...) end | |
return function(self, ...) | |
return select(n, ...) and pred(self, select(n, ...)) | |
end | |
end | |
--- Transform a predicate on a node into a predicate on this node's | |
-- n-th child. | |
-- @param n the child's index number | |
-- @param pred the predicate to transform | |
-- @return a predicate | |
function M.child(n, pred) | |
return function(node, ...) | |
local child = node[n] | |
return child and pred(child, node, ...) | |
end | |
end | |
--- Predicate to test the position of a node in its parent. | |
-- The predicate succeeds if the node is the n-th child of its parent, | |
-- and a <= n <= b. | |
-- nth(a) is equivalent to nth(a, a). | |
-- Negative indices are admitted, and count from the last child, | |
-- as done for instance by string.sub(). | |
-- | |
-- TODO: This is wrong, this tests the table relationship rather than the | |
-- AST node relationship. | |
-- Must build a getindex helper, based on pattern matching, then build | |
-- the predicate around it. | |
-- | |
-- @param a lower bound | |
-- @param a upper bound | |
-- @return a predicate | |
function M.is_nth(a, b) | |
b = b or a | |
return function(self, node, parent) | |
if not parent then return false end | |
local nchildren = #parent | |
local a = a<=0 and nchildren+a+1 or a | |
if a>nchildren then return false end | |
local b = b<=0 and nchildren+b+1 or b>nchildren and nchildren or b | |
for i=a,b do if parent[i]==node then return true end end | |
return false | |
end | |
end | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
-- | |
-- Comments parsing. | |
-- | |
-- ----------------------------------------------------------------------------- | |
-- ----------------------------------------------------------------------------- | |
local comment_extractor = |which_side| function (node) | |
local x = node.lineinfo | |
x = x and x[which_side] | |
x = x and x.comments | |
if not x then return nil end | |
local lines = { } | |
for _, record in ipairs(x) do | |
table.insert(lines, record[1]) | |
end | |
return table.concat(lines, '\n') | |
end | |
M.comment_prefix = comment_extractor 'first' | |
M.comment_suffix = comment_extractor 'last' | |
--- Shortcut for the query constructor | |
function M :__call(...) return self.treequery(...) end | |
setmetatable(M, M) | |
return M |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment