Fix register allocation function

1 files changed, 12 insertions, 8 deletions

diff --git a/pp2cc.py b/pp2cc.py
index dda8c2f..ca50191 100755
--- a/pp2cc.py
+++ b/pp2cc.py
@@ -44,13 +44,14 @@ class Registers(object):
             register = register[1]
             if self.registers[register]:
                 raise RuntimeError("Register 'R{}' is already in use".format(register))
-        for register in range(0, 6):
-            register = str(register)
-            # find a free register
-            if not self.registers[register]:
-                break
         else:
-            raise RuntimeError("No free registers are available")
+            for register in range(0, 6):
+                register = str(register)
+                # find a free register
+                if not self.registers[register]:
+                    break
+            else:
+                raise RuntimeError("No free registers are available")
         self.registers[register] = True
         return "R" + register
     def free(self, register):
@@ -63,9 +64,12 @@ class Registers(object):
         # remove leading R
         register = register[1:]
         if register in self.registers:
-            self.registers[register] = False;
+            if self.registers[register]:
+                self.registers[register] = False;
+            else:
+                raise RuntimeError("free() of unused register")
         else:
-            raise RuntimeError("free() of unused register")
+            raise RuntimeError("free() of invalid register")
     def get_register(self, text):
         """Retrieve the first register from a binary instruction"""
         text = text.strip()