Replace VirtualMachine.rindex with register_to_index in utils.py

hera/utils.py

@@ -49,3 +49,20 @@ def to_u32(n):
         return 2 ** 32 + n
     else:
         return n
+
+
+def register_to_index(rname):
+    """Return the index of the register with the given name in the register array."""
+    original = rname
+    rname = rname.lower()
+    if rname == "rt":
+        return 11
+    elif rname.startswith("r"):
+        v = int(rname[1:])
+        if 0 <= v < 16:
+            return v
+    elif rname == "fp":
+        return 14
+    elif rname == "sp":
+        return 15
+    raise ValueError("{} is not a valid register".format(original))

hera/vm.py

@@ -6,7 +6,7 @@
 import functools
 
 from .preprocessor import HERA_DATA_START
-from .utils import from_u16, to_u16, to_u32
+from .utils import from_u16, to_u16, to_u32, register_to_index
 
 
 def ternary_op(f):
@@ -126,27 +126,11 @@ def exec_many(self, program):
 
     def get_register(self, name):
         """Get the contents of the register with the given name."""
-        return self.registers[self.rindex(name)]
-
-    def rindex(self, name):
-        """Return the index of the register with the given name in the register
-        array.
-        """
-        name = name.lower()
-        if name == "rt":
-            return 11
-        elif name.startswith("r"):
-            return int(name[1:])
-        elif name == "fp":
-            return 14
-        elif name == "sp":
-            return 15
-        else:
-            raise KeyError(name)
+        return self.registers[register_to_index(name)]
 
     def store_register(self, target, value):
         """Store the value in the target register (a string)."""
-        index = self.rindex(target)
+        index = register_to_index(target)
         if index != 0:
             self.registers[index] = value
 

test/test_utils.py

@@ -1,6 +1,6 @@
 import pytest
 
-from hera.utils import from_u16, to_u16, to_u32
+from hera.utils import from_u16, to_u16, to_u32, register_to_index
 
 
 def test_to_u16_with_max_negative():
@@ -125,3 +125,24 @@ def test_to_u32_with_positive_overflow():
 def test_to_u32_with_another_positive_overflow():
     with pytest.raises(OverflowError):
         to_u32(5000000000)
+
+
+def test_register_to_index_with_numbered_registers():
+    for i in range(0, 16):
+        assert register_to_index("R" + str(i)) == i
+        assert register_to_index("r" + str(i)) == i
+
+
+def test_register_to_index_with_named_registers():
+    assert register_to_index("FP") == 14
+    assert register_to_index("fp") == 14
+    assert register_to_index("SP") == 15
+    assert register_to_index("sp") == 15
+    assert register_to_index("Rt") == 11
+    assert register_to_index("rt") == 11
+
+
+def test_register_to_index_with_invalid_register():
+    with pytest.raises(ValueError) as e:
+        register_to_index("R16")
+    assert "R16" in str(e)

test/test_vm.py

@@ -12,18 +12,6 @@ def vm():
     return VirtualMachine()
 
 
-def test_rindex_with_numbered_registers(vm):
-    assert vm.rindex("R0") == 0
-    assert vm.rindex("R7") == 7
-    assert vm.rindex("R13") == 13
-
-
-def test_rindex_with_special_registers(vm):
-    assert vm.rindex("SP") == 15
-    assert vm.rindex("FP") == 14
-    assert vm.rindex("Rt") == 11
-
-
 def test_exec_one_delegates_to_setlo(vm):
     with patch("hera.vm.VirtualMachine.exec_setlo") as mock_exec_setlo:
         vm.exec_one(Op("SETLO", ["R1", 47]))