src/soc/decoder/selectable_int.py

   1 import unittest
   2
   3
   4 class SelectableInt:
   5     def __init__(self, value, bits):
   6         self.value = value
   7         self.bits = bits
   8
   9     def __add__(self, b):
  10         assert b.bits == self.bits
  11         return SelectableInt(self.value + b.value, self.bits)
  12
  13     def __getitem__(self, key):
  14         if isinstance(key, int):
  15             assert key < self.bits
  16             assert key >= 0
  17             key = self.bits - (key + 1)
  18
  19             value = (self.value >> key) & 1
  20             return SelectableInt(value, 1)
  21         elif isinstance(key, slice):
  22             assert key.step is None or key.step == 1
  23             assert key.start < key.stop
  24             assert key.start >= 0
  25             assert key.stop <= self.bits
  26
  27             stop = self.bits - key.start
  28             start = self.bits - key.stop
  29
  30             bits = stop - start
  31             mask = (1 << bits) - 1
  32             value = (self.value >> start) & mask
  33             return SelectableInt(value, bits)
  34
  35     def __setitem__(self, key, value):
  36         if isinstance(key, int):
  37             assert key < self.bits
  38             assert key >= 0
  39             key = self.bits - (key + 1)
  40             if isinstance(value, SelectableInt):
  41                 assert value.bits == 1
  42                 value = value.value
  43
  44             value = value << key
  45             mask = 1 << key
  46             self.value = (self.value & ~mask) | (value & mask)
  47         elif isinstance(key, slice):
  48             assert key.step is None or key.step == 1
  49             assert key.start < key.stop
  50             assert key.start >= 0
  51             assert key.stop <= self.bits
  52
  53             stop = self.bits - key.start
  54             start = self.bits - key.stop
  55
  56             bits = stop - start
  57             if isinstance(value, SelectableInt):
  58                 assert value.bits == bits
  59                 value = value.value
  60             mask = ((1 << bits) - 1) << start
  61             value = value << start
  62             self.value = (self.value & ~mask) | (value & mask)
  63
  64     def __eq__(self, other):
  65         if isinstance(other, SelectableInt):
  66             return other.value == self.value and other.bits == self.bits
  67         if isinstance(other, int):
  68             return other == self.value
  69         assert False
  70
  71     def __repr__(self):
  72         return "SelectableInt(value={:x}, bits={})".format(self.value,
  73                                                            self.bits)
  74
  75
  76 class SelectableIntTestCase(unittest.TestCase):
  77     def test_add(self):
  78         a = SelectableInt(5, 8)
  79         b = SelectableInt(9, 8)
  80         c = a + b
  81         assert c.value == a.value + b.value
  82         assert c.bits == a.bits
  83
  84     def test_get(self):
  85         a = SelectableInt(0xa2, 8)
  86         # These should be big endian
  87         assert a[7] == 0
  88         assert a[0:4] == 10
  89         assert a[4:8] == 2
  90
  91     def test_set(self):
  92         a = SelectableInt(0x5, 8)
  93         a[7] = SelectableInt(0, 1)
  94         self.assertEqual(a, 4)
  95         a[4:8] = 9
  96         self.assertEqual(a, 9)
  97         a[0:4] = 3
  98         self.assertEqual(a, 0x39)
  99         a[0:4] = a[4:8]
 100         self.assertEqual(a, 0x99)
 101
 102
 103 if __name__ == "__main__":
 104     unittest.main()