metalua/treequery.mlua

   1 -------------------------------------------------------------------------------
   2 -- Copyright (c) 2006-2013 Fabien Fleutot and others.
   3 --
   4 -- All rights reserved.
   5 --
   6 -- This program and the accompanying materials are made available
   7 -- under the terms of the Eclipse Public License v1.0 which
   8 -- accompanies this distribution, and is available at
   9 -- http://www.eclipse.org/legal/epl-v10.html
  10 --
  11 -- This program and the accompanying materials are also made available
  12 -- under the terms of the MIT public license which accompanies this
  13 -- distribution, and is available at http://www.lua.org/license.html
  14 --
  15 -- Contributors:
  16 --     Fabien Fleutot - API and implementation
  17 --
  18 -------------------------------------------------------------------------------
  19
  20 local walk = require 'metalua.treequery.walk'
  21
  22 local M = { }
  23 -- support for old-style modules
  24 treequery = M
  25
  26 -- -----------------------------------------------------------------------------
  27 -- -----------------------------------------------------------------------------
  28 --
  29 -- multimap helper mmap: associate a key to a set of values
  30 --
  31 -- -----------------------------------------------------------------------------
  32 -- -----------------------------------------------------------------------------
  33
  34 local function mmap_add (mmap, node, x)
  35     if node==nil then return false end
  36     local set = mmap[node]
  37     if set then set[x] = true
  38     else mmap[node] = {[x]=true} end
  39 end
  40
  41 -- currently unused, I throw the whole set away
  42 local function mmap_remove (mmap, node, x)
  43     local set = mmap[node]
  44     if not set then return false
  45     elseif not set[x] then return false
  46     elseif next(set) then set[x]=nil
  47     else mmap[node] = nil end
  48     return true
  49 end
  50
  51 -- -----------------------------------------------------------------------------
  52 -- -----------------------------------------------------------------------------
  53 --
  54 -- TreeQuery object.
  55 --
  56 -- -----------------------------------------------------------------------------
  57 -- -----------------------------------------------------------------------------
  58
  59 local ACTIVE_SCOPE = setmetatable({ }, {__mode="k"})
  60
  61 -- treequery metatable
  62 local Q = { }; Q.__index = Q
  63
  64 --- treequery constructor
  65 --  the resultingg object will allow to filter ans operate on the AST
  66 --  @param root the AST to visit
  67 --  @return a treequery visitor instance
  68 function M.treequery(root)
  69     return setmetatable({
  70         root = root,
  71         unsatisfied = 0,
  72         predicates  = { },
  73         until_up    = { },
  74         from_up     = { },
  75         up_f        = false,
  76         down_f      = false,
  77         filters     = { },
  78     }, Q)
  79 end
  80
  81 -- helper to share the implementations of positional filters
  82 local function add_pos_filter(self, position, inverted, inclusive, f, ...)
  83     if type(f)=='string' then f = M.has_tag(f, ...) end
  84     if not inverted then self.unsatisfied += 1 end
  85     local x = {
  86         pred      = f,
  87         position  = position,
  88         satisfied = false,
  89         inverted  = inverted  or false,
  90         inclusive = inclusive or false }
  91     table.insert(self.predicates, x)
  92     return self
  93 end
  94
  95 function Q :if_unknown(f)
  96     self.unknown_handler = f or (||nil)
  97     return self
  98 end
  99
 100 -- TODO: offer an API for inclusive pos_filters
 101
 102 --- select nodes which are after one which satisfies predicate f
 103 Q.after     = |self, f, ...| add_pos_filter(self, 'after', false, false, f, ...)
 104 --- select nodes which are not after one which satisfies predicate f
 105 Q.not_after = |self, f, ...| add_pos_filter(self, 'after', true,  false, f, ...)
 106 --- select nodes which are under one which satisfies predicate f
 107 Q.under     = |self, f, ...| add_pos_filter(self, 'under', false, false, f, ...)
 108 --- select nodes which are not under one which satisfies predicate f
 109 Q.not_under = |self, f, ...| add_pos_filter(self, 'under', true,  false, f, ...)
 110
 111 --- select nodes which satisfy predicate f
 112 function Q :filter(f, ...)
 113     if type(f)=='string' then f = M.has_tag(f, ...) end
 114     table.insert(self.filters, f);
 115     return self
 116 end
 117
 118 --- select nodes which satisfy predicate f
 119 function Q :filter_not(f, ...)
 120     if type(f)=='string' then f = M.has_tag(f, ...) end
 121     table.insert(self.filters, |...| not f(...))
 122     return self
 123 end
 124
 125 -- private helper: apply filters and execute up/down callbacks when applicable
 126 function Q :execute()
 127     local cfg = { }
 128     -- TODO: optimize away not_under & not_after by pruning the tree
 129     function cfg.down(...)
 130         --printf ("[down]\t%s\t%s", self.unsatisfied, table.tostring((...)))
 131         ACTIVE_SCOPE[...] = cfg.scope
 132         local satisfied = self.unsatisfied==0
 133         for _, x in ipairs(self.predicates) do
 134             if not x.satisfied and x.pred(...) then
 135                 x.satisfied = true
 136                 local node, parent = ...
 137                 local inc = x.inverted and 1 or -1
 138                 if x.position=='under' then
 139                     -- satisfied from after we get down this node...
 140                     self.unsatisfied += inc
 141                     -- ...until before we get up this node
 142                     mmap_add(self.until_up, node, x)
 143                 elseif x.position=='after' then
 144                     -- satisfied from after we get up this node...
 145                     mmap_add(self.from_up, node, x)
 146                     -- ...until before we get up this node's parent
 147                     mmap_add(self.until_up, parent, x)
 148                 elseif x.position=='under_or_after' then
 149                     -- satisfied from after we get down this node...
 150                     self.satisfied += inc
 151                     -- ...until before we get up this node's parent...
 152                     mmap_add(self.until_up, parent, x)
 153                 else
 154                     error "position not understood"
 155                 end -- position
 156                 if x.inclusive then satisfied = self.unsatisfied==0 end
 157             end -- predicate passed
 158         end -- for predicates
 159
 160         if satisfied then
 161             for _, f in ipairs(self.filters) do
 162                 if not f(...) then satisfied=false; break end
 163             end
 164             if satisfied and self.down_f then self.down_f(...) end
 165         end
 166     end
 167
 168     function cfg.up(...)
 169         --printf ("[up]\t%s", table.tostring((...)))
 170
 171         -- Remove predicates which are due before we go up this node
 172         local preds = self.until_up[...]
 173         if preds then
 174             for x, _ in pairs(preds) do
 175                 local inc = x.inverted and -1 or 1
 176                 self.unsatisfied += inc
 177                 x.satisfied = false
 178             end
 179             self.until_up[...] = nil
 180         end
 181
 182         -- Execute the up callback
 183         -- TODO: cache the filter passing result from the down callback
 184         -- TODO: skip if there's no callback
 185         local satisfied = self.unsatisfied==0
 186         if satisfied then
 187             for _, f in ipairs(self.filters) do
 188                 if not f(self, ...) then satisfied=false; break end
 189             end
 190             if satisfied and self.up_f then self.up_f(...) end
 191         end
 192
 193         -- Set predicate which are due after we go up this node
 194         local preds = self.from_up[...]
 195         if preds then
 196             for p, _ in pairs(preds) do
 197                 local inc = p.inverted and 1 or -1
 198                 self.unsatisfied += inc
 199             end
 200             self.from_up[...] = nil
 201         end
 202         ACTIVE_SCOPE[...] = nil
 203     end
 204
 205     function cfg.binder(id_node, ...)
 206         --printf(" >>> Binder called on %s, %s", table.tostring(id_node),
 207         --      table.tostring{...}:sub(2,-2))
 208         cfg.down(id_node, ...)
 209         cfg.up(id_node, ...)
 210         --printf("down/up on binder done")
 211     end
 212
 213     cfg.unknown = self.unknown_handler
 214
 215     --function cfg.occurrence (binder, occ)
 216     --   if binder then OCC2BIND[occ] = binder[1] end
 217        --printf(" >>> %s is an occurrence of %s", occ[1], table.tostring(binder and binder[2]))
 218     --end
 219
 220     --function cfg.binder(...) cfg.down(...); cfg.up(...) end
 221     return walk.guess(cfg, self.root)
 222 end
 223
 224 --- Execute a function on each selected node
 225 --  @down: function executed when we go down a node, i.e. before its children
 226 --         have been examined.
 227 --  @up: function executed when we go up a node, i.e. after its children
 228 --       have been examined.
 229 function Q :foreach(down, up)
 230     if not up and not down then
 231         error "iterator missing"
 232     end
 233     self.up_f = up
 234     self.down_f = down
 235     return self :execute()
 236 end
 237
 238 --- Return the list of nodes selected by a given treequery.
 239 function Q :list()
 240     local acc = { }
 241     self :foreach(|x| table.insert(acc, x))
 242     return acc
 243 end
 244
 245 --- Return the first matching element
 246 --  TODO:  dirty hack, to implement properly with a 'break' return.
 247 --  Also, it won't behave correctly if a predicate causes an error,
 248 --  or if coroutines are involved.
 249 function Q :first()
 250    local result = { }
 251    local function f(...) result = {...}; error() end
 252    pcall(|| self :foreach(f))
 253    return unpack(result)
 254 end
 255
 256 --- Pretty printer for queries
 257 function Q :__tostring() return "<treequery>" end
 258
 259 -- -----------------------------------------------------------------------------
 260 -- -----------------------------------------------------------------------------
 261 --
 262 -- Predicates.
 263 --
 264 -- -----------------------------------------------------------------------------
 265 -- -----------------------------------------------------------------------------
 266
 267 --- Return a predicate which is true if the tested node's tag is among the
 268 --  one listed as arguments
 269 -- @param ... a sequence of tag names
 270 function M.has_tag(...)
 271     local args = {...}
 272     if #args==1 then
 273         local tag = ...
 274         return (|node| node.tag==tag)
 275         --return function(self, node) printf("node %s has_tag %s?", table.tostring(node), tag); return node.tag==tag end
 276     else
 277         local tags = { }
 278         for _, tag in ipairs(args) do tags[tag]=true end
 279         return function(node)
 280             local node_tag = node.tag
 281             return node_tag and tags[node_tag]
 282         end
 283     end
 284 end
 285
 286 --- Predicate to test whether a node represents an expression.
 287 M.is_expr = M.has_tag('Nil', 'Dots', 'True', 'False', 'Number','String',
 288                   'Function', 'Table', 'Op', 'Paren', 'Call', 'Invoke',
 289                   'Id', 'Index')
 290
 291 -- helper for is_stat
 292 local STAT_TAGS = { Do=1, Set=1, While=1, Repeat=1, If=1, Fornum=1,
 293                     Forin=1, Local=1, Localrec=1, Return=1, Break=1 }
 294
 295 --- Predicate to test whether a node represents a statement.
 296 --  It is context-aware, i.e. it recognizes `Call and `Invoke nodes
 297 --  used in a statement context as such.
 298 function M.is_stat(node, parent)
 299     local tag = node.tag
 300     if not tag then return false
 301     elseif STAT_TAGS[tag] then return true
 302     elseif tag=='Call' or tag=='Invoke' then return parent.tag==nil
 303     else return false end
 304 end
 305
 306 --- Predicate to test whether a node represents a statements block.
 307 function M.is_block(node) return node.tag==nil end
 308
 309 -- -----------------------------------------------------------------------------
 310 -- -----------------------------------------------------------------------------
 311 --
 312 -- Variables and scopes.
 313 --
 314 -- -----------------------------------------------------------------------------
 315 -- -----------------------------------------------------------------------------
 316
 317 local BINDER_PARENT_TAG = {
 318    Local=true, Localrec=true, Forin=true, Function=true }
 319
 320 --- Test whether a node is a binder. This is local predicate, although it
 321 --  might need to inspect the parent node.
 322 function M.is_binder(node, parent)
 323    --printf('is_binder(%s, %s)', table.tostring(node), table.tostring(parent))
 324    if node.tag ~= 'Id' or not parent then return false end
 325    if parent.tag=='Fornum' then  return parent[1]==node end
 326    if not BINDER_PARENT_TAG[parent.tag] then return false end
 327    for _, binder in ipairs(parent[1]) do
 328        if binder==node then return true end
 329    end
 330    return false
 331 end
 332
 333 --- Retrieve the binder associated to an occurrence within root.
 334 --  @param occurrence an Id node representing an occurrence in `root`.
 335 --  @param root the tree in which `node` and its binder occur.
 336 --  @return the binder node, and its ancestors up to root if found.
 337 --  @return nil if node is global (or not an occurrence) in `root`.
 338 function M.binder(occurrence, root)
 339     local cfg, id_name, result = { }, occurrence[1], { }
 340     function cfg.occurrence(id)
 341         if id == occurrence then result = cfg.scope :get(id_name) end
 342         -- TODO: break the walker
 343     end
 344     walk.guess(cfg, root)
 345     return unpack(result)
 346 end
 347
 348 --- Predicate to filter occurrences of a given binder.
 349 --  Warning: it relies on internal scope book-keeping,
 350 --  and for this reason, it only works as query method argument.
 351 --  It won't work outside of a query.
 352 --  @param binder the binder whose occurrences must be kept by predicate
 353 --  @return a predicate
 354
 355 -- function M.is_occurrence_of(binder)
 356 --     return function(node, ...)
 357 --         if node.tag ~= 'Id' then return nil end
 358 --         if M.is_binder(node, ...) then return nil end
 359 --         local scope = ACTIVE_SCOPE[node]
 360 --         if not scope then return nil end
 361 --         local result = scope :get (node[1]) or { }
 362 --         if result[1] ~= binder then return nil end
 363 --         return unpack(result)
 364 --     end
 365 -- end
 366
 367 function M.is_occurrence_of(binder)
 368     return function(node, ...)
 369         local b = M.get_binder(node)
 370         return b and b==binder
 371     end
 372 end
 373
 374 function M.get_binder(occurrence, ...)
 375     if occurrence.tag ~= 'Id' then return nil end
 376     if M.is_binder(occurrence, ...) then return nil end
 377     local scope = ACTIVE_SCOPE[occurrence]
 378     local binder_hierarchy = scope :get(occurrence[1])
 379     return unpack (binder_hierarchy or { })
 380 end
 381
 382 --- Transform a predicate on a node into a predicate on this node's
 383 --  parent. For instance if p tests whether a node has property P,
 384 --  then parent(p) tests whether this node's parent has property P.
 385 --  The ancestor level is precised with n, with 1 being the node itself,
 386 --  2 its parent, 3 its grand-parent etc.
 387 --  @param[optional] n the parent to examine, default=2
 388 --  @param pred the predicate to transform
 389 --  @return a predicate
 390 function M.parent(n, pred, ...)
 391     if type(n)~='number' then n, pred = 2, n end
 392     if type(pred)=='string' then pred = M.has_tag(pred, ...) end
 393     return function(self, ...)
 394         return select(n, ...) and pred(self, select(n, ...))
 395     end
 396 end
 397
 398 --- Transform a predicate on a node into a predicate on this node's
 399 --  n-th child.
 400 --  @param n the child's index number
 401 --  @param pred the predicate to transform
 402 --  @return a predicate
 403 function M.child(n, pred)
 404     return function(node, ...)
 405         local child = node[n]
 406         return child and pred(child, node, ...)
 407     end
 408 end
 409
 410 --- Predicate to test the position of a node in its parent.
 411 --  The predicate succeeds if the node is the n-th child of its parent,
 412 --  and a <= n <= b.
 413 --  nth(a) is equivalent to nth(a, a).
 414 --  Negative indices are admitted, and count from the last child,
 415 --  as done for instance by string.sub().
 416 --
 417 --  TODO: This is wrong, this tests the table relationship rather than the
 418 --  AST node relationship.
 419 --  Must build a getindex helper, based on pattern matching, then build
 420 --  the predicate around it.
 421 --
 422 --  @param a lower bound
 423 --  @param a upper bound
 424 --  @return a predicate
 425 function M.is_nth(a, b)
 426     b = b or a
 427     return function(self, node, parent)
 428         if not parent then return false end
 429         local nchildren = #parent
 430         local a = a<=0 and nchildren+a+1 or a
 431         if a>nchildren then return false end
 432         local b = b<=0 and nchildren+b+1 or b>nchildren and nchildren or b
 433         for i=a,b do if parent[i]==node then return true end end
 434         return false
 435     end
 436 end
 437
 438
 439 -- -----------------------------------------------------------------------------
 440 -- -----------------------------------------------------------------------------
 441 --
 442 -- Comments parsing.
 443 --
 444 -- -----------------------------------------------------------------------------
 445 -- -----------------------------------------------------------------------------
 446
 447 local comment_extractor = |which_side| function (node)
 448     local x = node.lineinfo
 449     x = x and x[which_side]
 450     x = x and x.comments
 451     if not x then return nil end
 452     local lines = { }
 453     for _, record in ipairs(x) do
 454         table.insert(lines, record[1])
 455     end
 456     return table.concat(lines, '\n')
 457 end
 458
 459 M.comment_prefix = comment_extractor 'first'
 460 M.comment_suffix = comment_extractor 'last'
 461
 462
 463 --- Shortcut for the query constructor
 464 function M :__call(...) return self.treequery(...) end
 465 setmetatable(M, M)
 466
 467 return M