177 lines · python
1# Automatically formatted with yapf (https://github.com/google/yapf)2"""Utility functions for creating and manipulating LLVM 'opt' NPM pipeline objects."""3 4 5def fromStr(pipeStr):6 """Create pipeline object from string representation."""7 stack = []8 curr = []9 tok = ""10 kind = ""11 for c in pipeStr:12 if c == ",":13 if tok != "":14 curr.append([None, tok])15 tok = ""16 elif c == "(":17 stack.append([kind, curr])18 kind = tok19 curr = []20 tok = ""21 elif c == ")":22 if tok != "":23 curr.append([None, tok])24 tok = ""25 oldKind = kind26 oldCurr = curr27 [kind, curr] = stack.pop()28 curr.append([oldKind, oldCurr])29 else:30 tok += c31 if tok != "":32 curr.append([None, tok])33 return curr34 35 36def toStr(pipeObj):37 """Create string representation of pipeline object."""38 res = ""39 lastIdx = len(pipeObj) - 140 for i, c in enumerate(pipeObj):41 if c[0]:42 res += c[0] + "("43 res += toStr(c[1])44 res += ")"45 else:46 res += c[1]47 if i != lastIdx:48 res += ","49 return res50 51 52def count(pipeObj):53 """Count number of passes (pass-managers excluded) in pipeline object."""54 cnt = 055 for c in pipeObj:56 if c[0]:57 cnt += count(c[1])58 else:59 cnt += 160 return cnt61 62 63def split(pipeObj, splitIndex):64 """Create two new pipeline objects by splitting pipeObj in two directly after pass with index splitIndex."""65 66 def splitInt(src, splitIndex, dstA, dstB, idx):67 for s in src:68 if s[0]:69 dstA2 = []70 dstB2 = []71 idx = splitInt(s[1], splitIndex, dstA2, dstB2, idx)72 dstA.append([s[0], dstA2])73 dstB.append([s[0], dstB2])74 else:75 if idx <= splitIndex:76 dstA.append([None, s[1]])77 else:78 dstB.append([None, s[1]])79 idx += 180 return idx81 82 listA = []83 listB = []84 splitInt(pipeObj, splitIndex, listA, listB, 0)85 return [listA, listB]86 87 88def remove(pipeObj, removeIndex):89 """Create new pipeline object by removing pass with index removeIndex from pipeObj."""90 91 def removeInt(src, removeIndex, dst, idx):92 for s in src:93 if s[0]:94 dst2 = []95 idx = removeInt(s[1], removeIndex, dst2, idx)96 dst.append([s[0], dst2])97 else:98 if idx != removeIndex:99 dst.append([None, s[1]])100 idx += 1101 return idx102 103 dst = []104 removeInt(pipeObj, removeIndex, dst, 0)105 return dst106 107 108def copy(srcPipeObj):109 """Create copy of pipeline object srcPipeObj."""110 111 def copyInt(dst, src):112 for s in src:113 if s[0]:114 dst2 = []115 copyInt(dst2, s[1])116 dst.append([s[0], dst2])117 else:118 dst.append([None, s[1]])119 120 dstPipeObj = []121 copyInt(dstPipeObj, srcPipeObj)122 return dstPipeObj123 124 125def prune(srcPipeObj):126 """Create new pipeline object by removing empty pass-managers (those with count = 0) from srcPipeObj."""127 128 def pruneInt(dst, src):129 for s in src:130 if s[0]:131 if count(s[1]):132 dst2 = []133 pruneInt(dst2, s[1])134 dst.append([s[0], dst2])135 else:136 dst.append([None, s[1]])137 138 dstPipeObj = []139 pruneInt(dstPipeObj, srcPipeObj)140 return dstPipeObj141 142 143if __name__ == "__main__":144 import unittest145 146 class Test(unittest.TestCase):147 def test_0(self):148 pipeStr = "a,b,A(c,B(d,e),f),g"149 pipeObj = fromStr(pipeStr)150 151 self.assertEqual(7, count(pipeObj))152 153 self.assertEqual(pipeObj, pipeObj)154 self.assertEqual(pipeObj, prune(pipeObj))155 self.assertEqual(pipeObj, copy(pipeObj))156 157 self.assertEqual(pipeStr, toStr(pipeObj))158 self.assertEqual(pipeStr, toStr(prune(pipeObj)))159 self.assertEqual(pipeStr, toStr(copy(pipeObj)))160 161 [pipeObjA, pipeObjB] = split(pipeObj, 3)162 self.assertEqual("a,b,A(c,B(d))", toStr(pipeObjA))163 self.assertEqual("A(B(e),f),g", toStr(pipeObjB))164 165 self.assertEqual("b,A(c,B(d,e),f),g", toStr(remove(pipeObj, 0)))166 self.assertEqual("a,b,A(c,B(d,e),f)", toStr(remove(pipeObj, 6)))167 168 pipeObjC = remove(pipeObj, 4)169 self.assertEqual("a,b,A(c,B(d),f),g", toStr(pipeObjC))170 pipeObjC = remove(pipeObjC, 3)171 self.assertEqual("a,b,A(c,B(),f),g", toStr(pipeObjC))172 pipeObjC = prune(pipeObjC)173 self.assertEqual("a,b,A(c,f),g", toStr(pipeObjC))174 175 unittest.main()176 exit(0)177