From 6e216df0e8ea9aacd772c53f10518c2bffee6094 Mon Sep 17 00:00:00 2001
From: Benjamin Pollack <benjamin@bitquabit.com>
Date: Mon, 20 Mar 2017 15:25:58 -0400
Subject: [PATCH] More cleanup and some bug fixes

---
 scripts/ud_itab.py   | 12 ++++++------
 scripts/ud_opcode.py | 42 +++++++++++++++++-------------------------
 2 files changed, 23 insertions(+), 31 deletions(-)

diff --git a/scripts/ud_itab.py b/scripts/ud_itab.py
index 63de355..4b0b43e 100644
--- a/scripts/ud_itab.py
+++ b/scripts/ud_itab.py
@@ -262,16 +262,16 @@ def genInsnTable(self):
             opr_c = ["O_NONE", "O_NONE", "O_NONE", "O_NONE"]
             pfx_c = []
             opr = insn.operands
-            for i in range(len(opr)):
-                if not (opr[i] in self.OperandDict.keys()):
-                    print("error: invalid operand declaration: %s\n" % opr[i])
-                opr_c[i] = "O_" + opr[i]
+            for i, op in enumerate(opr):
+                if op not in self.OperandDict:
+                    print("error: invalid operand declaration: %s\n" % op)
+                opr_c[i] = "O_" + op
             opr = "%s %s %s %s" % (opr_c[0] + ",", opr_c[1] + ",",
                                    opr_c[2] + ",", opr_c[3])
 
             for p in insn.prefixes:
-                if not (p in self.PrefixDict.keys()):
-                    print("error: invalid prefix specification: %s \n" % pfx)
+                if p not in self.PrefixDict:
+                    print("error: invalid prefix specification: %s \n" % p)
                 pfx_c.append(self.PrefixDict[p])
             if len(insn.prefixes) == 0:
                 pfx_c.append("P_none")
diff --git a/scripts/ud_opcode.py b/scripts/ud_opcode.py
index 4495883..4a4b74c 100644
--- a/scripts/ud_opcode.py
+++ b/scripts/ud_opcode.py
@@ -25,16 +25,15 @@
 
 import os
 
-# Some compatibility stuff for supporting python 2.x as well as python 3.x
-def itemslist(dict):
-    try:
-        return dict.iteritems() # python 2.x
-    except AttributeError:
-        return list(dict.items()) # python 3.x
+
+class CollisionError(Exception):
+    pass
+
 
 class UdInsnDef:
     """An x86 instruction definition
     """
+
     def __init__(self, **insnDef):
         self.mnemonic  = insnDef['mnemonic']
         self.prefixes  = insnDef['prefixes']
@@ -77,13 +76,6 @@ class UdOpcodeTable:
        a decode field. 
     """
 
-    class CollisionError(Exception):
-        pass
-
-    class IndexError(Exception):
-        """Invalid Index Error"""
-        pass
-
     @classmethod
     def vendor2idx(cls, v):
         return (0 if v == 'amd'
@@ -174,10 +166,10 @@ def size(self):
         return self._TableInfo[self._typ]['size']
 
     def entries(self):
-        return itemslist(self._entries)
+        return self._entries.items()
 
     def numEntries(self):
-        return len(self._entries.keys())
+        return len(self._entries)
 
     def label(self):
         return self._TableInfo[self._typ]['label']
@@ -195,14 +187,14 @@ def add(self, opc, obj):
         typ = UdOpcodeTable.getOpcodeTyp(opc)
         idx = UdOpcodeTable.getOpcodeIdx(opc)
         if self._typ != typ or idx in self._entries:
-            raise CollisionError()
+            raise CollisionError
         self._entries[idx] = obj
 
     def lookup(self, opc):
         typ = UdOpcodeTable.getOpcodeTyp(opc)
         idx = UdOpcodeTable.getOpcodeIdx(opc)
         if self._typ != typ:
-            raise UdOpcodeTable.CollisionError("%s <-> %s" % (self._typ, typ))
+            raise CollisionError("%s <-> %s" % (self._typ, typ))
         return self._entries.get(idx, None)
 
     def entryAt(self, index):
@@ -212,13 +204,13 @@ def entryAt(self, index):
         """
         if index < self.size():
             return self._entries.get(index, None)
-        raise self.IndexError("index out of bounds: %s" % index)
+        raise IndexError("index out of bounds: %s" % index)
 
     def setEntryAt(self, index, obj):
         if index < self.size():
             self._entries[index] = obj
         else:
-            raise self.IndexError("index out of bounds: %s" % index)
+            raise IndexError("index out of bounds: %s" % index)
 
     @classmethod
     def getOpcodeTyp(cls, opc):
@@ -239,7 +231,7 @@ def getOpcodeIdx(cls, opc):
     @classmethod
     def getLabels(cls):
         """Returns a list of all labels"""
-        return [cls._TableInfo[k]['label'] for k in cls._TableInfo.keys()]
+        return [cls._TableInfo[k]['label'] for k in cls._TableInfo]
 
 
 class UdOpcodeTables(object):
@@ -471,8 +463,8 @@ def addSSE2AVXInsn(self, **insnDef):
         ssemnemonic = insnDef['mnemonic']
         sseopcodes = insnDef['opcodes']
         # remove vex opcode extensions
-        sseopcexts = dict([(e, v) for e, v in itemslist(insnDef['opcexts'])
-                           if not e.startswith('/vex')])
+        sseopcexts = {e: v for e, v in insnDef['opcexts'].items()
+                      if not e.startswith('/vex')}
         # strip out avx operands, preserving relative ordering
         # of remaining operands
         sseoperands = [opr for opr in insnDef['operands']
@@ -495,8 +487,8 @@ def addSSE2AVXInsn(self, **insnDef):
         vexmnemonic = 'v' + insnDef['mnemonic']
         vexprefixes = insnDef['prefixes']
         vexopcodes  = ['c4']
-        vexopcexts  = dict([(e, insnDef['opcexts'][e])
-                              for e in insnDef['opcexts'] if e != '/sse'])
+        vexopcexts  = {e: insnDef['opcexts'][e]
+                       for e in insnDef['opcexts'] if e != '/sse'}
         vexopcexts['/vex'] = insnDef['opcexts']['/sse'] + '_' + '0f'
         if insnDef['opcodes'][1] == '38' or insnDef['opcodes'][1] == '3a':
             vexopcexts['/vex'] += insnDef['opcodes'][1]
@@ -507,7 +499,7 @@ def addSSE2AVXInsn(self, **insnDef):
         for o in insnDef['operands']:
             # make the operand size explicit: x
             if o in ('V', 'W', 'H', 'U'):
-                o = o + 'x'
+                o += 'x'
             vexoperands.append(o)
         vexcpuid = [flag for flag in insnDef['cpuid']
                     if not flag.startswith('sse')]