1 #!/usr/bin/env python 2 3 """ 4 Simplify AST structures for easier type propagation and analysis. 5 6 Copyright (C) 2006 Paul Boddie <paul@boddie.org.uk> 7 8 This software is free software; you can redistribute it and/or 9 modify it under the terms of the GNU General Public License as 10 published by the Free Software Foundation; either version 2 of 11 the License, or (at your option) any later version. 12 13 This software is distributed in the hope that it will be useful, 14 but WITHOUT ANY WARRANTY; without even the implied warranty of 15 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 16 GNU General Public License for more details. 17 18 You should have received a copy of the GNU General Public 19 License along with this library; see the file LICENCE.txt 20 If not, write to the Free Software Foundation, Inc., 21 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA 22 """ 23 24 from compiler.visitor import ASTVisitor 25 import compiler.ast 26 from simplified import * 27 28 class Simplifier(ASTVisitor): 29 30 """ 31 A simplifying visitor for AST nodes. 32 33 Covered: AssAttr, AssList, AssName, AssTuple, Assign, AugAssign, Break, 34 CallFunc, Class, Const, Continue, Dict, Discard, For, Function, 35 Getattr, Global, If, Keyword, Lambda, List, Module, Name, Pass, 36 Return, Stmt, TryExcept, TryFinally, Tuple, While. 37 38 Missing: Add, And, Assert, Backquote, Bitand, Bitor, Bitxor, Compare, 39 Decorators, Div, Ellipsis, Exec, FloorDiv, From, Import, Invert, 40 LeftShift, ListComp, ListCompFor, ListCompIf, Mod, Mul, Not, Or, 41 Power, Print, Printnl, Raise, RightShift, Slice, Sliceobj, Sub, 42 Subscript, UnaryAdd, UnarySub, Yield. 43 """ 44 45 def __init__(self): 46 ASTVisitor.__init__(self) 47 self.result = None # The resulting tree. 48 self.subprograms = [] # Subprograms outside the tree. 49 self.current_subprograms = [] # Current subprograms being processed. 50 51 # Generic visitor methods. 52 53 def default(self, node, *args): 54 raise ValueError, node.__class__ 55 56 def dispatch(self, node, *args): 57 return ASTVisitor.dispatch(self, node, *args) 58 59 def dispatches(self, nodes, *args): 60 results = [] 61 for node in nodes: 62 results.append(self.dispatch(node, *args)) 63 return results 64 65 # Placeholder or deletion transformations. 66 67 def visitStmt(self, stmt): 68 return self.dispatches(stmt.nodes) 69 70 def visitPass(self, pass_): 71 return Pass(pass_) 72 73 def visitDiscard(self, discard): 74 return self.dispatch(discard.expr) 75 76 # Relatively trivial transformations. 77 78 def visitModule(self, module): 79 self.result = Module(module) 80 self.result.code = self.dispatch(module.node) 81 return self.result 82 83 def visitClass(self, class_): 84 result = Class(class_, name=class_.name, bases=class_.bases) 85 result.code = self.dispatch(class_.code) 86 return result 87 88 def visitGetattr(self, getattr): 89 result = LoadAttr(getattr, name=getattr.attrname) 90 result.expr = self.dispatch(getattr.expr) 91 return result 92 93 def visitKeyword(self, keyword): 94 result = Keyword(keyword, name=keyword.name) 95 result.expr = self.dispatch(keyword.expr) 96 return result 97 98 def visitGlobal(self, global_): 99 result = Global(global_, names=global_.names) 100 return result 101 102 def visitName(self, name): 103 result = LoadName(name, name=name.name) 104 return result 105 106 def visitConst(self, const): 107 result = LoadConst(const, value=const.value) 108 return result 109 110 def visitReturn(self, return_): 111 result = Return(return_) 112 result.expr = self.dispatch(return_.value) 113 return result 114 115 def visitBreak(self, break_): 116 result = Return(break_) 117 return result 118 119 def visitContinue(self, continue_): 120 result = Invoke(continue_, same_frame=1, star=None, dstar=None, args=[]) 121 result.expr = LoadRef(ref=self.current_subprograms[-1]) 122 return result 123 124 def visitIf(self, if_): 125 result = If(if_, else_=[]) 126 tests = [] 127 for compare, stmt in if_.tests: 128 # Produce something like... 129 # expr.__true__() ? body 130 test = Conditional(else_=[], test=Invoke( 131 expr=LoadAttr(expr=self.dispatch(compare), name="__true__"), 132 params=[], star=None, dstar=None)) 133 test.body = self.dispatch(stmt) 134 tests.append(test) 135 result.tests = tests 136 if if_.else_ is not None: 137 result.else_ = self.dispatch(if_.else_) 138 return result 139 140 def _visitBuiltin(self, builtin, name): 141 result = Invoke(builtin, expr=LoadName(name=name)) 142 result.args = self.dispatches(builtin.nodes) 143 return result 144 145 def visitTuple(self, tuple): 146 return self._visitBuiltin(tuple, "Tuple") 147 148 def visitList(self, list): 149 return self._visitBuiltin(list, "List") 150 151 def visitDict(self, dict): 152 result = Invoke(expr=LoadName(name="Dict")) 153 args = [] 154 for key, value in dict.items: 155 tuple = Invoke(expr=LoadName(name="Tuple"), star=None, dstar=None) 156 tuple.args = [self.dispatch(key), self.dispatch(value)] 157 args.append(tuple) 158 result.args = args 159 return result 160 161 # Logical operators. 162 163 def visitAnd(self, and_): 164 result = And(and_) 165 nodes = [] 166 for node in and_.nodes: 167 nodes.append(Invoke(expr=LoadAttr(expr=self.dispatch(node), name="__true__"), 168 params=[], star=None, dstar=None)) 169 result.nodes = nodes 170 return result 171 172 def visitOr(self, or_): 173 result = Or(or_) 174 nodes = [] 175 for node in or_.nodes: 176 nodes.append(Invoke(expr=LoadAttr(expr=self.dispatch(node), name="__true__"), 177 params=[], star=None, dstar=None)) 178 result.nodes = nodes 179 return result 180 181 def visitNot(self, not_): 182 result = Not(not_, expr=Invoke(expr=LoadAttr(expr=self.dispatch(not_.expr), name="__true__"), 183 params=[], star=None, dstar=None)) 184 return result 185 186 # Assignments. 187 188 augassign_methods = { 189 "+=" : "__iadd__", "-=" : "__isub__", "*=" : "__imul__", "/=" : "__idiv__" 190 } 191 192 def visitAugAssign(self, augassign): 193 result = Assign(augassign) 194 expr = self.dispatch(augassign.expr) 195 196 # Simple augmented assignment: name += expr 197 198 if isinstance(augassign.node, compiler.ast.Name): 199 node = self.dispatch(augassign.node) 200 get_incremented = StoreTemp( 201 expr=Invoke(expr=LoadAttr(expr=node, name=self.augassign_methods[augassign.op]), args=[expr]) 202 ) 203 store = StoreName(expr=LoadTemp(), name=augassign.node.name) 204 result.code = [get_incremented, store, ReleaseTemp()] 205 206 # Complicated augmented assignment: expr.attr += expr 207 208 elif isinstance(augassign.node, compiler.ast.Getattr): 209 store_expr = StoreTemp(index=1, expr=self.dispatch(augassign.node.expr)) 210 node_attr = LoadAttr(expr=LoadTemp(index=1), name=augassign.node.attrname) 211 get_incremented = StoreTemp( 212 expr=Invoke(expr=LoadAttr(expr=node_attr, name=self.augassign_methods[augassign.op]), args=[expr]) 213 ) 214 store = StoreAttr(expr=LoadTemp(), lvalue=LoadTemp(index=1), name=augassign.node.attrname) 215 result.code = [store_expr, get_incremented, store, ReleaseTemp(index=1), ReleaseTemp()] 216 217 else: 218 raise NotImplementedError, augassign.node.__class__ 219 220 return result 221 222 def visitAssign(self, assign): 223 result = Assign(assign) 224 store = StoreTemp(expr=self.dispatch(assign.expr)) 225 release = ReleaseTemp() 226 result.code = [store] + self.dispatches(assign.nodes, 0) + [release] 227 return result 228 229 def visitAssList(self, asslist, in_sequence=0): 230 if not in_sequence: 231 expr = LoadTemp(asslist) 232 else: 233 expr = Invoke(asslist, expr=LoadAttr(expr=LoadTemp(), name="next")) 234 result = Assign(asslist) 235 store = StoreTemp(expr=Invoke(expr=LoadAttr(name="__iter__", expr=expr))) 236 release = ReleaseTemp() 237 result.code = [store] + self.dispatches(asslist.nodes, 1) + [release] 238 return result 239 240 visitAssTuple = visitAssList 241 242 def _visitAssNameOrAttr(self, node, in_sequence): 243 if not in_sequence: 244 return LoadTemp(node) 245 else: 246 return Invoke(node, expr=LoadAttr(expr=LoadTemp(), name="next")) 247 248 def visitAssName(self, assname, in_sequence=0): 249 expr = self._visitAssNameOrAttr(assname, in_sequence) 250 result = StoreName(assname, name=assname.name, expr=expr) 251 return result 252 253 def visitAssAttr(self, assattr, in_sequence=0): 254 expr = self._visitAssNameOrAttr(assattr, in_sequence) 255 lvalue = self.dispatch(assattr.expr) 256 result = StoreAttr(assattr, name=assattr.attrname, lvalue=lvalue, expr=expr) 257 return result 258 259 # Invocation and subprogram transformations. 260 261 def _visitFunction(self, function, subprogram): 262 if function.flags & 4 != 0: has_star = 1 263 else: has_star = 0 264 if function.flags & 8 != 0: has_dstar = 1 265 else: has_dstar = 0 266 ndefaults = len(function.defaults) 267 npositional = len(function.argnames) - has_star - has_dstar 268 if has_star: star = function.argnames[npositional] 269 else: star = None 270 if has_dstar: dstar = function.argnames[npositional + has_star] 271 else: dstar = None 272 273 params = [] 274 for i in range(0, npositional - ndefaults): 275 params.append((function.argnames[i], None)) 276 277 # NOTE: Fix/process defaults. 278 279 for i in range(0, ndefaults): 280 default = function.defaults[i] 281 if default is not None: 282 params.append((function.argnames[npositional - ndefaults + i], self.dispatch(default))) 283 else: 284 params.append((function.argnames[npositional - ndefaults + i], default)) 285 286 subprogram.params = params 287 subprogram.star = star 288 subprogram.dstar = dstar 289 self.subprograms.append(subprogram) 290 291 def visitFunction(self, function): 292 293 # Make a subprogram for the function and record it outside the main 294 # tree. 295 296 subprogram = Subprogram(function, name=function.name, star=None, dstar=None) 297 self.current_subprograms.append(subprogram) 298 subprogram.code = self.dispatch(function.code) 299 self.current_subprograms.pop() 300 self._visitFunction(function, subprogram) 301 302 # Make a definition of the function associating it with a name. 303 304 result = Assign(function) 305 load = LoadRef(ref=subprogram) 306 store = StoreName(name=function.name) 307 result.code = [load, store] 308 return result 309 310 def visitLambda(self, lambda_): 311 312 # Make a subprogram for the function and record it outside the main 313 # tree. 314 315 subprogram = Subprogram(lambda_, name=hex(id(lambda_)), star=None, dstar=None) 316 self.current_subprograms.append(subprogram) 317 subprogram.code = [Return(expr=self.dispatch(lambda_.code))] 318 self.current_subprograms.pop() 319 self._visitFunction(lambda_, subprogram) 320 321 # Get the subprogram reference to the lambda. 322 323 return LoadRef(ref=subprogram) 324 325 def visitCallFunc(self, callfunc): 326 result = Invoke(callfunc, same_frame=0, star=None, dstar=None) 327 result.args = self.dispatches(callfunc.args) 328 if callfunc.star_args is not None: 329 result.star = self.dispatch(callfunc.star_args) 330 if callfunc.dstar_args is not None: 331 result.dstar = self.dispatch(callfunc.dstar_args) 332 result.expr = self.dispatch(callfunc.node) 333 return result 334 335 def visitWhile(self, while_): 336 337 # Make a subprogram for the block and record it outside the main tree. 338 339 subprogram = Subprogram(while_, name=hex(id(while_)), acquire_locals=1, params=[], star=None, dstar=None) 340 self.current_subprograms.append(subprogram) 341 342 # Include a conditional statement in the subprogram. 343 344 test = Conditional(else_=[]) 345 test.test = Invoke(expr=LoadAttr(expr=self.dispatch(while_.test), name="__true__"), 346 params=[], star=None, dstar=None) 347 348 # Inside the conditional, add a recursive invocation to the subprogram 349 # if the test condition was satisfied. 350 351 continuation = Invoke(same_frame=1, star=None, dstar=None, args=[]) 352 continuation.expr = LoadRef(ref=subprogram) 353 test.body = self.dispatch(while_.body) + [continuation] 354 if while_.else_ is not None: 355 test.else_ = self.dispatch(while_.else_) 356 subprogram.code = [test] 357 358 self.current_subprograms.pop() 359 self.subprograms.append(subprogram) 360 361 # Make an invocation of the subprogram. 362 363 result = Invoke(while_, same_frame=1, star=None, dstar=None, args=[]) 364 result.expr = LoadRef(ref=subprogram) 365 return result 366 367 def visitFor(self, for_): 368 369 # Make a subprogram for the block and record it outside the main tree. 370 371 subprogram = Subprogram(for_, name=hex(id(for_)), acquire_locals=1, params=[], star=None, dstar=None) 372 self.current_subprograms.append(subprogram) 373 374 # Wrap the assignment in a try...except statement. 375 376 try_except = Try(body=[], handlers=[], else_=[], finally_=[]) 377 except_spec = Invoke(expr=LoadName(name="Tuple"), params=[LoadName(name="StopIteration")]) 378 stopiteration = Except(spec=except_spec) 379 stopiteration.code = self.dispatch(for_.else_) 380 try_except.handlers = [stopiteration] 381 382 assign = Assign() 383 assign.code = [ 384 StoreTemp(expr=Invoke(expr=LoadAttr(expr=LoadTemp(), name="next"))), 385 self.dispatch(for_.assign), 386 ReleaseTemp() 387 ] 388 389 # Inside the conditional, add a recursive invocation to the subprogram 390 # if the test condition was satisfied. 391 392 continuation = Invoke(same_frame=1, star=None, dstar=None, args=[]) 393 continuation.expr = LoadRef(ref=subprogram) 394 try_except.body = [assign] + self.dispatch(for_.body) + [continuation] 395 subprogram.code = [try_except] 396 397 self.subprograms.append(subprogram) 398 self.current_subprograms.pop() 399 400 # Obtain an iterator for the sequence involved. 401 # Then, make an invocation of the subprogram. 402 403 result = Assign(for_) 404 result.code = [ 405 StoreTemp(expr=Invoke(expr=LoadAttr(name="__iter__", expr=self.dispatch(for_.list)))), 406 Invoke(expr=LoadRef(ref=subprogram), same_frame=1, star=None, dstar=None, args=[]), 407 ReleaseTemp() 408 ] 409 return result 410 411 # Exception node transformations. 412 413 def visitTryExcept(self, tryexcept): 414 result = Try(tryexcept, body=[], handlers=[], else_=[], finally_=[]) 415 if tryexcept.body is not None: 416 result.body = self.dispatch(tryexcept.body) 417 if tryexcept.else_ is not None: 418 result.else_ = self.dispatch(tryexcept.else_) 419 handlers = [] 420 for spec, assign, stmt in tryexcept.handlers: 421 get_exc = Assign() 422 get_exc.code = [StoreTemp(expr=LoadExc()), self.dispatch(assign), ReleaseTemp()] 423 handler = Except(spec=self.dispatch(spec)) 424 handler.code = [get_exc] + self.dispatch(stmt) 425 handlers.append(handler) 426 result.handlers = handlers 427 return result 428 429 def visitTryFinally(self, tryfinally): 430 result = Try(tryfinally, body=[], handlers=[], else_=[], finally_=[]) 431 if tryfinally.body is not None: 432 result.body = self.dispatch(tryfinally.body) 433 if tryfinally.final is not None: 434 result.finally_ = self.dispatch(tryfinally.final) 435 return result 436 437 # vim: tabstop=4 expandtab shiftwidth=4