core.py

   1 import sys
   2 from graph import Graph
   3 import copy
   4
   5 import utils
   6
   7
   8 class Singleton:
   9
  10     def __init__(self, n):
  11         self.n = n
  12
  13     def __repr__(self):
  14         return self.n
  15
  16 UNK = Singleton("UNK")
  17 DYN = Singleton("DYN")
  18
  19
  20 class BBlock:
  21     def __init__(self, addr):
  22         self.addr = addr
  23         self.items = []
  24         self.props = {}
  25
  26     def add(self, s):
  27         self.items.append(s)
  28
  29     def __getitem__(self, i):
  30         return self.items[i]
  31
  32     def def_addrs(self, regs_only=True):
  33         """Return all variable definitions for this basic block,
  34         as set of (var, inst_addr) pairs. Note that this includes
  35         multiple entries for the same var, if it is redefined
  36         multiple times within the basic block.
  37         """
  38         defs = set()
  39         for i in self.items:
  40             inst_defs = i.defs(regs_only)
  41             for d in inst_defs:
  42                 defs.add((d, i.addr))
  43         return defs
  44
  45     def defs(self, regs_only=True):
  46         """Return set of all variables defined in this basic block."""
  47         defs = set()
  48         for i in self.items:
  49             defs |= i.defs(regs_only)
  50         return defs
  51
  52     def uses(self):
  53         """Return set of all variables used in this basic block."""
  54         uses = set()
  55         for i in self.items:
  56             uses |= i.uses()
  57         return uses
  58
  59     def __repr__(self):
  60         return "%s(%s)" % (self.__class__.__name__, self.addr)
  61
  62     def write(self, stream, indent, s):
  63         for l in str(s).splitlines():
  64             stream.write("  " * indent)
  65             stream.write(l + "\n")
  66
  67     def dump(self, stream, indent=0, printer=str):
  68         for s in self.items:
  69             out = printer(s)
  70             if out is not None:
  71                 self.write(stream, indent, out)
  72
  73 TYPE_SORT = ("REG", "ADDR", "MEM", "EXPR", "COND", "VALUE")
  74
  75 def type_sort(t):
  76     return TYPE_SORT.index(t.__name__)
  77
  78 # Helper predicates for types below
  79
  80 def is_value(e):
  81     return isinstance(e, VALUE)
  82
  83 def is_addr(e):
  84     return isinstance(e, ADDR)
  85
  86 def is_reg(e):
  87     return isinstance(e, REG)
  88
  89 def is_mem(e):
  90     return isinstance(e, MEM)
  91
  92 def is_expr(e):
  93     return isinstance(e, EXPR)
  94
  95 def is_cond(e):
  96     return isinstance(e, COND)
  97
  98 def is_sfunc(e, name):
  99     return is_expr(e) and e.args[0] == SFUNC(name)
 100
 101
 102 class SimpleExpr:
 103     # Something which is a simple expression
 104
 105     comment = ""
 106     simple_repr = True
 107
 108     def regs(self):
 109         "Get registers referenced by the expression"
 110         return set()
 111
 112     def side_effect(self):
 113         return False
 114
 115     def __len__(self):
 116         return 1
 117
 118
 119 class REG(SimpleExpr):
 120
 121     def __init__(self, name):
 122         #assert isinstance(name, str)
 123         self.name = name
 124         self.signed = False
 125
 126     def __repr__(self):
 127         if self.simple_repr:
 128             return self.__str__()
 129         type = "REG_S" if self.signed else "REG"
 130         return self.comment + type + "(%s)" % self.name
 131
 132     def __str__(self):
 133         cast = "(i32)" if self.signed else ""
 134         return self.comment + cast + "$" + str(self.name)
 135
 136     def __eq__(self, other):
 137         return type(self) == type(other) and self.name == other.name
 138
 139     def __lt__(self, other):
 140         if type(self) != type(other):
 141             return type_sort(type(self)) < type_sort(type(other))
 142
 143         n1 = utils.natural_sort_key(self.name)
 144         n2 = utils.natural_sort_key(other.name)
 145         return n1 < n2
 146
 147     def __contains__(self, other):
 148         return self == other
 149
 150     def __hash__(self):
 151         return hash(self.name)
 152
 153     def regs(self):
 154         return {self}
 155
 156
 157 class VALUE(SimpleExpr):
 158
 159     def __init__(self, val, base=16):
 160         self.val = val
 161         self.base = base
 162
 163     def __repr__(self):
 164         if self.simple_repr:
 165             return self.__str__()
 166         return self.comment + "VALUE(%#x)" % self.val
 167
 168     def __str__(self):
 169         if isinstance(self.val, int) and self.base == 16:
 170             val = "%#x" % self.val
 171         else:
 172             val = str(self.val)
 173         return self.comment + val
 174
 175     def __eq__(self, other):
 176         return type(self) == type(other) and self.val == other.val
 177
 178     def __lt__(self, other):
 179         if type(self) != type(other):
 180             return type_sort(type(self)) < type_sort(type(other))
 181         return self.val < other.val
 182
 183     def __contains__(self, other):
 184         return self == other
 185
 186     def __hash__(self):
 187         return hash(self.val)
 188
 189
 190 class STR(SimpleExpr):
 191
 192     def __init__(self, s):
 193         self.val = s
 194
 195     def __repr__(self):
 196         if self.simple_repr:
 197             return self.__str__()
 198         return self.comment + "STR(%s)" % self.val
 199
 200     def __str__(self):
 201         return self.comment + self.val
 202
 203
 204 class ADDR(SimpleExpr):
 205
 206     resolver = staticmethod(lambda x: x)
 207
 208     def __init__(self, addr):
 209         self.addr = addr
 210
 211     def __repr__(self):
 212         return self.comment + "ADDR(%s)" % self.addr
 213
 214     def __str__(self):
 215         return self.comment + self.resolver(self.addr)
 216
 217     def __eq__(self, other):
 218         return type(self) == type(other) and self.addr == other.addr
 219
 220     def __lt__(self, other):
 221         if type(self) != type(other):
 222             return type_sort(type(self)) < type_sort(type(other))
 223         return self.addr < other.addr
 224
 225     def __contains__(self, other):
 226         return self == other
 227
 228     def __hash__(self):
 229         return hash(self.addr)
 230
 231
 232 class CVAR(SimpleExpr):
 233
 234     def __init__(self, name):
 235         self.name = name
 236
 237     def __repr__(self):
 238         return self.comment + "CVAR(%s)" % self.name
 239
 240     def __str__(self):
 241         return self.comment + self.name
 242
 243     def __eq__(self, other):
 244         return type(self) == type(other) and self.name == other.name
 245
 246     def __hash__(self):
 247         return hash(self.name)
 248
 249
 250 class MEM(SimpleExpr):
 251     def __init__(self, type, expr):
 252         self.type = type
 253         self.expr = expr
 254
 255     def __repr__(self):
 256         return self.comment + "*(%s*)%r" % (self.type, self.expr)
 257
 258     def __str__(self):
 259         if isinstance(self.expr, EXPR):
 260             return self.comment + "*(%s*)(%s)" % (self.type, self.expr)
 261         else:
 262             return self.comment + "*(%s*)%s" % (self.type, self.expr)
 263
 264     def __eq__(self, other):
 265         return type(self) == type(other) and self.type == other.type and \
 266             self.expr == other.expr
 267
 268     def __lt__(self, other):
 269         if type(self) == type(other):
 270             return self.expr < other.expr
 271         return type_sort(type(self)) < type_sort(type(other))
 272
 273     def __contains__(self, other):
 274         if other == self:
 275             return True
 276         return other in self.expr
 277
 278     def __hash__(self):
 279         return hash(self.type) ^ hash(self.expr)
 280
 281     def regs(self):
 282         return self.expr.regs()
 283
 284
 285 class SFIELD(SimpleExpr):
 286
 287     def __init__(self, type, addr, field):
 288         self.type = type
 289         self.addr = addr
 290         self.field = field
 291
 292     def __repr__(self):
 293         return self.comment + "SFIELD(%s, %s, %s)" % (self.type, self.addr, self.field)
 294
 295     def __str__(self):
 296         return self.comment + "((%s*)%s)->%s" % (self.type, self.addr, self.field)
 297
 298
 299 class SFUNC(SimpleExpr):
 300
 301     def __init__(self, name):
 302         self.name = name
 303
 304     def __repr__(self):
 305         return "(SFUNC)%s" % (self.name)
 306
 307     def __str__(self):
 308         return "%s" % self.name
 309
 310     def __eq__(self, other):
 311         return type(self) == type(other) and self.name == other.name
 312
 313     def __contains__(self, other):
 314         return False
 315
 316     def __hash__(self):
 317         return hash(self.name)
 318
 319
 320 class TYPE(SimpleExpr):
 321
 322     def __init__(self, name):
 323         self.name = name
 324
 325     def __repr__(self):
 326         return "(TYPE)%s" % (self.name)
 327
 328     def __str__(self):
 329         return "%s" % self.name
 330
 331     def __eq__(self, other):
 332         return type(self) == type(other) and self.name == other.name
 333
 334     def __contains__(self, other):
 335         return False
 336
 337     def __hash__(self):
 338         return hash(self.name)
 339
 340     def bitsize(self):
 341         assert self.name[0] in ("i", "u")
 342         return int(self.name[1:])
 343
 344
 345 class EXPR:
 346     "A recursive expression."
 347     def __init__(self, op, *args):
 348         self.op = op
 349         if isinstance(args[0], list):
 350             args = args[0]
 351         else:
 352             args = list(args)
 353         self.args = args
 354
 355     def __repr__(self):
 356         return "EXPR(%s%s)" % (self.op, self.args)
 357
 358     @staticmethod
 359     def preced(e):
 360         if is_expr(e):
 361             # See e.g. http://en.cppreference.com/w/c/language/operator_precedence
 362             return {
 363                 "+=": 14, "-=": 14, "*=": 14, "/=": 14, "%=": 14,
 364                 "<<=": 14, ">>=": 14, "&=": 14, "|=": 14, "^=": 14,
 365                 "||": 12, "&&": 11,
 366                 "|": 10, "^": 9, "&": 8,
 367                 "==": 7, "!=": 7,
 368                 ">": 6, "<": 6, ">=": 6, "<=": 6,
 369                 "<<": 5, ">>": 5,
 370                 "+": 4, "-": 4,
 371                 "*": 3, "/": 3, "%": 3,
 372                 # All the below is highest precedence
 373                 "CAST": 1, "SFUNC": 1, "NEG": 1, "!": 1,
 374             }[e.op]
 375         return 1
 376
 377     # Render this expr's arg, wrapped in parens if needed
 378     @staticmethod
 379     def strarg(expr, arg):
 380         if isinstance(arg, (set, frozenset)):
 381             s = utils.repr_stable(arg)
 382         else:
 383             s = str(arg)
 384         preced_my = EXPR.preced(expr)
 385         preced_arg = EXPR.preced(arg)
 386         full_assoc = expr.op in {"+", "*", "&", "^", "|"}
 387         if preced_arg == preced_my and full_assoc:
 388             # Render repeated fully associative operators without extra parens
 389             pass
 390         elif preced_arg > preced_my or (preced_arg == preced_my and preced_arg != 1):
 391             # Otherwise, if precedence rules require parens, render them, unless
 392             # the arg is a unary/primary term
 393             s = "(%s)" % s
 394         else:
 395             # Parens would not be required per the precedence rules, but
 396             # handle common cases of confusing precedence in C, where parens
 397             # are usually suggested.
 398             if expr.op in ("&", "^", "|") and preced_arg != 1:
 399                 # Any binary op subexpression of bitwise ops in parens
 400                 s = "(%s)" % s
 401             elif expr.op in ("<<", ">>") and preced_arg != 1:
 402                 # Any binary op subexpression of shift in parens
 403                 s = "(%s)" % s
 404         return s
 405
 406     def __str__(self):
 407         if not SimpleExpr.simple_repr:
 408             return self.__repr__()
 409
 410         if self.op == "SFUNC":
 411             return str(self.args[0]) + "(" + ", ".join([str(a) for a in self.args[1:]]) + ")"
 412         if self.op == "CAST":
 413             return "(" + str(self.args[0]) + ")" + self.strarg(self, self.args[1])
 414
 415         DICT = {
 416             "NEG": "-",
 417             "!": "!",
 418         }
 419         if self.op in DICT:
 420             assert len(self.args) == 1
 421             s = DICT [self.op]
 422             return s + self.strarg(self, self.args[0])
 423
 424         l = [self.strarg(self, self.args[0])]
 425         for a in self.args[1:]:
 426             if self.op == "+" and is_value(a) and a.val < 0:
 427                 l.append("-")
 428                 a = VALUE(-a.val, a.base)
 429             else:
 430                 l.append(self.op)
 431             l.append(self.strarg(self, a))
 432         return " ".join(l)
 433
 434     def __eq__(self, other):
 435         return type(self) == type(other) and self.op == other.op and self.args == other.args
 436
 437     def __lt__(self, other):
 438         if type(self) == type(other):
 439             return str(self) < str(other)
 440         return type_sort(type(self)) < type_sort(type(other))
 441
 442     def __contains__(self, other):
 443         if other == self:
 444             return True
 445         for a in self.args:
 446             if other in a:
 447                 return True
 448         return False
 449
 450     def __hash__(self):
 451         return hash(self.op) ^ hash(tuple(self.args))
 452
 453     def __len__(self):
 454         # One for operation itself
 455         l = 1
 456         for a in self.args:
 457             l += len(a)
 458         return l
 459
 460     def regs(self):
 461         r = set()
 462         for a in self.args:
 463             r |= set(a.regs())
 464         return r
 465
 466     uses = regs
 467
 468     def defs(self, regs_only=True):
 469         assert not self.side_effect()
 470         return set()
 471
 472     def side_effect(self):
 473         if self.op == "SFUNC":
 474             return self.args[0].name not in {
 475                 "BIT", "abs", "bitfield", "count_leading_zeroes",
 476                 "phi",
 477             }
 478         return False
 479
 480     def foreach_subexpr(self, func):
 481         # If func returned True, it means it handled entire subexpression,
 482         # so we don't recurse into it.
 483         # Note that this function recurses only within EXPR tree, it doesn't
 484         # recurse e.g. inside MEM.
 485         if func(self):
 486             return
 487         for a in self.args:
 488             if is_expr(a):
 489                 a.foreach_subexpr(func)
 490             else:
 491                 func(a)
 492
 493
 494 class Inst:
 495
 496     trail = ""
 497     show_addr = False
 498     show_comments = True
 499     annotate_calls = False
 500     comment = "//"
 501
 502     def __init__(self, dest, op, args, addr=None):
 503         self.op = op
 504         self.dest = dest
 505         self.args = args
 506         self.addr = addr
 507         self.comments = {}
 508
 509     def jump_addr(self):
 510         "If instruction may transfer control, return jump address(es), otherwise return None."
 511         if self.op in ("call", "goto"):
 512             if isinstance(self.args[0], ADDR):
 513                 return self.args[0].addr
 514         if self.op == "if":
 515             res = []
 516             for i in range(0, len(self.args), 2):
 517                 if isinstance(self.args[i + 1], ADDR):
 518                     res.append(self.args[i + 1].addr)
 519             if len(res) == 1:
 520                 return res[0]
 521             return res
 522         return None
 523
 524     def side_effect(self):
 525         if self.op == "call":
 526             return True
 527         if self.op in ("=", "SFUNC"):
 528             assert len(self.args) == 1, self.args
 529             return self.args[0].side_effect()
 530         return False
 531
 532
 533     def uses(self, cfg=None):
 534         # Avoid circular import. TODO: fix properly
 535         import arch
 536         import progdb
 537
 538         """Return set of all registers used by this instruction. Function
 539         calls (and maybe SFUNCs) require special treatment."""
 540         if self.op == "call":
 541             addr = self.args[0]
 542             uses = addr.regs()
 543             if isinstance(addr, ADDR):
 544                 # Direct call with known address
 545                 addr = addr.addr
 546                 if addr in progdb.FUNC_DB and "params" in progdb.FUNC_DB[addr]:
 547                     return uses | progdb.FUNC_DB[addr]["params"]
 548
 549             # Indirect call or not params in funcdb
 550             # TODO: need to allow saving callsite info in funcdb
 551             return uses | arch.call_params(addr)
 552
 553         if self.op == "return":
 554             if not self.args:
 555                 return arch.ret_uses(cfg)
 556
 557         uses = set()
 558         for a in self.args:
 559             for r in a.regs():
 560                 uses.add(r)
 561         if is_mem(self.dest):
 562             for r in self.dest.regs():
 563                 uses.add(r)
 564         return uses
 565
 566
 567     def defs(self, regs_only=True, cfg=None):
 568         # Avoid circular import. TODO: fix properly
 569         import arch
 570         import progdb
 571         """Return set of all registers defined by this instruction. Function
 572         calls (and maybe SFUNCs) require special treatment."""
 573         if self.op == "call":
 574             addr = self.args[0]
 575             if isinstance(addr, ADDR):
 576                 # Direct call with known address
 577                 addr = addr.addr
 578                 if addr in progdb.FUNC_DB and "modifieds" in progdb.FUNC_DB[addr]:
 579                     return progdb.FUNC_DB[addr]["modifieds"]
 580
 581             # Indirect call or not params in funcdb
 582             # TODO: need to allow saving callsite info in funcdb
 583             return arch.call_defs(addr)
 584
 585         defs = set()
 586         if self.dest:
 587             if not regs_only or isinstance(self.dest, REG):
 588                 defs.add(self.dest)
 589         return defs
 590
 591
 592     def foreach_subexpr(self, func):
 593         def do(arg):
 594             if arg:
 595                 if is_expr(arg) or is_cond(arg):
 596                     arg.foreach_subexpr(func)
 597                 else:
 598                     func(arg)
 599
 600         do(self.dest)
 601         for a in self.args:
 602             do(a)
 603
 604
 605     def __repr__(self):
 606         comments = self.comments.copy()
 607         s = ""
 608         if "org_inst" in comments:
 609             s = "// " + str(comments.pop("org_inst")) + "\n"
 610         if self.addr is not None:
 611             s += "/*%s*/ " % self.addr
 612         if self.dest is None:
 613             if self.op == "LIT":
 614                 s += self.args[0]
 615             else:
 616                 s += "%s(%s)" % (self.op, self.args)
 617         else:
 618             if self.op == "=":
 619                 # Simplify repr for assignment
 620                 s += "%s = %s" % (self.dest, self.args)
 621             else:
 622                 s += "%s = %s(%s)" % (self.dest, self.op, self.args)
 623         if comments:
 624             s += " # " + repr(comments)
 625         return s
 626
 627     def __str__(self):
 628         if not SimpleExpr.simple_repr:
 629             return self.__repr__()
 630
 631         comments = self.comments.copy()
 632
 633         addr = ""
 634         if self.show_addr:
 635             addr = "/*%s*/ " % self.addr
 636
 637         if self.op == "LIT":
 638             return addr + self.args[0]
 639
 640         s = ""
 641         if self.show_comments and "org_inst" in comments:
 642             s = self.comment + " " + str(comments.pop("org_inst")) + " "
 643
 644         s = addr + s
 645
 646         if self.op == "call" and self.annotate_calls:
 647             comments["uses"] = sorted(self.uses())
 648             comments["defs"] = sorted(self.defs())
 649
 650         tail = self.trail
 651         if self.show_comments and comments:
 652             tail += " " + self.comment + " " + utils.repr_stable_dict(comments)
 653
 654         if self.op == "return":
 655             args = ", ".join([str(a) for a in self.args])
 656             if args:
 657                 args = " " + args
 658             return s + self.op + args + tail
 659         if self.op in ("goto", "call"):
 660             return s + "%s %s" % (self.op, self.args[0]) + tail
 661         if self.op == "if":
 662             joined = ", ".join(["%s goto %s" % (self.args[i] or "else", self.args[i + 1]) for i in range(0, len(self.args), 2)])
 663             return s + "if " + joined + tail
 664
 665         if self.op == "DEAD":
 666             return s + "(dead)" + tail
 667
 668         if self.op == "SFUNC":
 669             assert self.dest is None
 670             assert len(self.args) == 1, repr(self.args)
 671             return s + str(self.args[0]) + tail
 672
 673         assert self.op == "=", repr(self.op)
 674         assert len(self.args) == 1, (self.op, repr(self.args))
 675
 676         if self.op == "=" and not is_expr(self.args[0]):
 677             s += "%s = %s" % (self.dest, self.args[0])
 678         else:
 679             e = copy.copy(self.args[0])
 680             args = e.args
 681             op = e.op
 682             if not (op == "!" or op[0].isalpha()):
 683                 # Infix operator
 684                 assert len(args) >= 2, repr(args)
 685                 if self.dest == args[0]:
 686                     s += "%s %s= " % (self.dest, op)
 687                     # Render operator as a compound statement operator
 688                     # (lowest priority, no extra parens).
 689                     e.op += "="
 690                     e.args = args[1:]
 691                 else:
 692                     s += "%s = " % self.dest
 693             else:
 694                 if self.dest is not None:
 695                     s += "%s = " % self.dest
 696             s += "%s" % e
 697
 698         s += tail
 699         return s
 700
 701
 702     def __eq__(self, other):
 703         return self.op == other.op and self.dest == other.dest and self.args == other.args
 704
 705
 706 class COND:
 707     """This class is a container of EXPR used as a condition in the
 708     'if (cond) goto' statement. It's needed because the same condition
 709     is used both in the Inst representing such a statement and a lebel
 710     of a CFG edge connecting basic blocks. If condition is updated,
 711     e.g. while transforming its Inst, the change should be mirrored
 712     to the CFG edge. Using COND class, this can be easily achieved:
 713     the same COND instance is referenced both in Inst and edge, and
 714     we can freely update or even completely replace EXPR it contains,
 715     while both users will stay up to date.
 716     """
 717
 718     NEG = {
 719         "==": "!=",
 720         "!=": "==",
 721         ">":  "<=",
 722         "<":  ">=",
 723         ">=": "<",
 724         "<=": ">",
 725         "in": "not in",
 726     }
 727
 728     SWAP = {
 729         "==": "==",
 730         "!=": "!=",
 731         ">":  "<",
 732         "<":  ">",
 733         ">=": "<=",
 734         "<=": ">=",
 735     }
 736
 737     def __init__(self, expr):
 738         self.expr = expr
 739
 740     def neg(self):
 741         op = self.expr.op
 742         if op in self.NEG:
 743             return self.__class__(EXPR(self.NEG[op], self.expr.args))
 744         elif op == "!":
 745             return self.__class__(self.expr.args[0])
 746         else:
 747             return self.__class__(EXPR("!", self.expr))
 748
 749     def swap(self):
 750         "Swap arguments in-place."
 751         self.expr.args[0], self.expr.args[1] = self.expr.args[1], self.expr.args[0]
 752         self.expr.op = self.SWAP[self.expr.op]
 753
 754     def normalize(self):
 755         if is_value(self.expr.args[0]) and not is_value(self.expr.args[1]):
 756             self.swap()
 757
 758     def is_relation(self):
 759         return is_expr(self.expr) and self.expr.op in self.NEG
 760
 761     def list(self):
 762         return [self]
 763
 764     def __str__(self):
 765         return "(%s)" % self.expr
 766
 767     def __repr__(self):
 768 #        if self.op in ("in", "not in"):
 769 #            return "COND(%r %s %s)" % (self.arg1, self.op, utils.repr_stable(self.arg2))
 770         return "COND(%r)" % self.expr
 771
 772     def __eq__(self, other):
 773         return type(self) == type(other) and self.expr == other.expr
 774
 775     def __contains__(self, other):
 776         return other in self.expr
 777
 778     def __hash__(self):
 779         return hash(self.expr) ^ hash(self.__class__)
 780
 781     def regs(self):
 782         return self.expr.regs()
 783
 784     def defs(self, regs_only=True):
 785         return self.expr.defs(regs_only)
 786
 787     def uses(self):
 788         return self.expr.uses()
 789
 790     def foreach_subexpr(self, func):
 791         self.expr.foreach_subexpr(func)
 792
 793
 794 class CompoundCond:
 795
 796     NEG = {
 797         "&&": "||",
 798         "||": "&&",
 799     }
 800
 801     def __init__(self, l):
 802         self.args = l
 803
 804     def append(self, op, arg2):
 805         self.args.extend([op, arg2])
 806
 807     def neg(self):
 808         return self.__class__([self.NEG[x] if isinstance(x, str) else x.neg() for x in self.args])
 809
 810     def list(self):
 811         return self.args
 812
 813     def __str__(self):
 814         r = " ".join([str(x) for x in self.args])
 815         return "(" + r + ")"
 816
 817     def __repr__(self):
 818         return "CCond%s" % str(self)
 819
 820 def repr_state(state):
 821     res = []
 822     unk = []
 823     for k, v in sorted(state.items()):
 824         if v == UNK:
 825             unk.append(str(k))
 826         else:
 827             res.append("%s=%s" % (k, v))
 828     res = ", ".join(res)
 829     if unk:
 830         res += " UNK: " + ",".join(unk)
 831     return "{" + res + "}"
 832
 833
 834 class CFGPrinter:
 835     """Print BBlocks in a CFG. Various printing params can be overriden
 836     via methods."""
 837
 838     header_reg_prefix = "$"
 839     addr_in_header = False
 840
 841     def __init__(self, cfg, stream=sys.stdout):
 842         self.cfg = cfg
 843         self.stream = stream
 844         # Current bblock addr
 845         self.addr = 0
 846         # Current CFG node properties
 847         self.node_props = None
 848         # Current BBlock
 849         self.bblock = None
 850         # Current BBlock properties
 851         self.bblock_props = None
 852         self.inst_printer = str
 853         self.no_dead = False
 854
 855     def bblock_order(self):
 856         "Return iterator over bblocks to be printed."
 857         return self.cfg.iter_sorted_nodes()
 858
 859     def print_graph_header(self):
 860         if self.cfg.props:
 861             print("// Graph props:", file=self.stream)
 862             for k in sorted(self.cfg.props.keys()):
 863                 v = self.cfg.props[k]
 864                 v = utils.repr_stable(v)
 865                 print("//  %s: %s" % (k, v), file=self.stream)
 866             print(file=self.stream)
 867
 868
 869     def print_header(self):
 870         if self.addr_in_header:
 871             print("// %s" % self.addr, file=self.stream)
 872         print("// Predecessors: %s" % sorted(self.cfg.pred(self.addr)), file=self.stream)
 873
 874         if self.node_props:
 875             print("// Node props:", file=self.stream)
 876             for k in sorted(self.node_props.keys()):
 877                 v = self.node_props[k]
 878                 v = utils.repr_stable(v)
 879                 v = v.replace("$", self.header_reg_prefix)
 880                 print("//  %s: %s" % (k, v), file=self.stream)
 881
 882         if self.bblock_props:
 883             print("// BBlock props:", file=self.stream)
 884
 885             for k in sorted(self.bblock_props.keys()):
 886                 v = self.bblock_props[k]
 887                 if k.startswith("state_"):
 888                     v = repr_state(v)
 889                 else:
 890                     v = utils.repr_stable(v)
 891                 v = v.replace("$", self.header_reg_prefix)
 892                 print("//  %s: %s" % (k, v), file=self.stream)
 893
 894
 895     def print_trailer(self):
 896         succ = self.cfg.succ(self.addr)
 897         exits = [(self.cfg.edge(self.addr, x).get("cond"), x) for x in succ]
 898         print("Exits:", sorted(exits, key=lambda x: utils.natural_sort_key(str(x))), file=self.stream)
 899
 900
 901     def print_label(self):
 902         print("%s:" % self.addr, file=self.stream)
 903
 904     def print_inst(self, inst):
 905         if inst.op == "DEAD" and self.no_dead:
 906             return None
 907         return self.inst_printer(inst)
 908
 909     def print_separator(self):
 910         self.stream.write("\n")
 911
 912     def print(self):
 913         self.print_graph_header()
 914         cnt = 0
 915         for self.addr, info in self.bblock_order():
 916             self.node_props = info.copy()
 917             self.bblock = self.node_props.pop("val")
 918             self.bblock_props = self.bblock.props
 919             if cnt > 0:
 920                 self.print_separator()
 921             self.print_header()
 922             self.print_label()
 923             if self.bblock is not None:
 924                 self.bblock.dump(self.stream, 0, self.print_inst)
 925             else:
 926                 print("   ", self.bblock, file=self.stream)
 927             self.print_trailer()
 928             cnt += 1
 929
 930
 931 def dump_bblocks(cfg, stream=sys.stdout, printer=str, no_graph_header=False):
 932     p = CFGPrinter(cfg, stream)
 933     p.inst_printer = printer
 934     if no_graph_header:
 935         p.print_graph_header = lambda: None
 936     p.print()