Refactor Register (#82)

* Fix register read bug

If we store a BitVec, we do NOT want to return bool() of it, which
returns True, which is totally incorrect. We do however, want to return
a symbolic Bool of it, if it is a 1 bit register (flag).

* Improve

* Raise error

* Fix register tests

Changed the interfaces:
- removed nbits param (never used)
- no longer raise AssertionError when overflowing a flag reg, just
  truncate
- rename test funcs to be more descriptive

* rm unused import

* Add symbolic tests

* Rm no longer applicable test

manticore/core/cpu/register.py

@@ -1,4 +1,4 @@
-from ..smtlib import Operators, Expression, BitVec, Bool
+from ..smtlib import Operators, BitVec, Bool
 
 class Register(object):
     '''Generic variable width register. For 1 bit registers, allows writes
@@ -9,19 +9,20 @@ class Register(object):
         self.width = width
         self.value = 0
 
-    def read(self, nbits=None):
-        if isinstance(self.value, (int, long, BitVec)):
-            val = Operators.EXTRACT(self.value, 0, nbits or self.width)
-            return bool(val) if self.width == 1 else val
-        elif isinstance(self.value, Bool):
-            return self.value
-        else:
-            raise Exception('Malformed data in a register')
+    def is_flag(self):
+        return self.width == 1
 
-    def write(self, val, nbits=None):
-        if self.width == 1:
-            assert isinstance(val, (bool, Expression)) or val in (1, 0)
-            self.value = int(val) if isinstance(val, bool) else val
-        else:
-            val = Operators.EXTRACT(val, 0, nbits or self.width)
+    def read(self):
+        return self.value
+
+    def write(self, val):
+        if isinstance(val, (Bool, bool)):
             self.value = val
+        elif isinstance(val, BitVec):
+            self.value = val.Bool() if self.is_flag() else val
+        elif isinstance(val, (int, long)):
+            self.value = Operators.EXTRACT(val, 0, self.width)
+            if self.is_flag():
+                self.value = bool(self.value)
+        else:
+            raise TypeError('Cannot store {} in Register'.format(val.__class__.__name__))

tests/test_armv7rf.py

@@ -28,10 +28,6 @@ class Armv7RF(unittest.TestCase):
         self.r.write(ARM_REG_APSR_Z, False)
         self.assertEqual(self.r.read(ARM_REG_APSR_Z), False)
 
-    def test_bad_flag_write(self):
-        with self.assertRaises(AssertionError) as e:
-            self.r.write(ARM_REG_APSR_Z, 2)
-
     def test_reg_name(self):
         self.assertEqual(self.r.reg_name(ARM_REG_R0), 'R0')
 

tests/test_register.py

@@ -1,5 +1,6 @@
 import unittest
 
+from manticore.core.smtlib import Bool, BitVecConstant
 from manticore.core.cpu.register import Register
 
 class RegisterTest(unittest.TestCase):
@@ -7,46 +8,60 @@ class RegisterTest(unittest.TestCase):
         r = Register(32)
         self.assertEqual(r.read(), 0)
 
-    def test_rd2(self):
+    def test_basic_write(self):
         r = Register(32)
         r.write(1)
         self.assertEqual(r.read(), 1)
 
-    def test_rd3(self):
-        r = Register(32)
-        r.write(2)
-        self.assertEqual(r.read(1), 0)
-
-    def test_rd4(self):
-        r = Register(32)
-        r.write(7, 2)
-        self.assertEqual(r.read(), 3)
-
-    def test_rd5(self):
+    def test_truncate(self):
         r = Register(32)
         r.write(2**32)
         self.assertEqual(r.read(), 0)
 
-    def test_rd6(self):
+    def test_largest_write(self):
         r = Register(32)
         r.write(0xffffffff)
         self.assertEqual(r.read(), 0xffffffff)
 
-    def test_rd7(self):
-        r = Register(32)
-        r.write(0xffffffff)
-        self.assertEqual(r.read(16), 0xffff)
-
-    def test_reg1(self):
+    def test_flag(self):
         r = Register(1)
         self.assertEqual(r.read(), False)
 
-    def test_reg1_write(self):
+    def test_flag_write(self):
         r = Register(1)
         r.write(True)
         self.assertEqual(r.read(), True)
 
-    def test_reg1_badwrite(self):
+    def test_flag_trunc(self):
         r = Register(1)
-        with self.assertRaises(AssertionError):
-            r.write(2)
+        r.write(3)
+        self.assertEqual(r.read(), True)
+
+    def test_bool_write_nonflag(self):
+        r = Register(32)
+        r.write(True)
+        self.assertEqual(r.read(), True)
+
+    def test_Bool(self):
+        r = Register(32)
+        b = Bool()
+        r.write(b)
+        self.assertIs(r.read(), b)
+
+    def test_bitvec_flag(self):
+        r = Register(1)
+        b = BitVecConstant(32, 0)
+        r.write(b)
+        # __nonzero__ (==) currently unimplemented for Bool
+        self.assertTrue(isinstance(r.read(), Bool))
+
+    def test_bitvec(self):
+        r = Register(32)
+        b = BitVecConstant(32, 0)
+        r.write(b)
+        self.assertIs(r.read(), b)
+
+    def test_bad_write(self):
+        r = Register(32)
+        with self.assertRaises(TypeError):
+            r.write(dict())