Do not choke on invalid data statements

iafisher · Dec 14, 2018 · 9be83cc · 9be83cc
1 parent 09727f4
commit 9be83cc
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/hera/symtab.py b/hera/symtab.py
@@ -13,17 +13,19 @@ def get_symtab(program):
     dc = HERA_DATA_START
     for op in program:
         odc = dc
-        if op.name == "LABEL":
+        if op.name == "LABEL" and len(op.args) == 1:
             labels[op.args[0]] = pc
-        elif op.name == "DLABEL":
+        elif op.name == "DLABEL" and len(op.args) == 1:
             labels[op.args[0]] = dc
-        elif op.name == "CONSTANT":
+        elif op.name == "CONSTANT" and len(op.args) == 2:
             labels[op.args[0]] = op.args[1]
         elif op.name == "INTEGER":
             dc += 1
-        elif op.name == "DSKIP":
+        elif op.name == "DSKIP" and len(op.args) == 1 and isinstance(op.args[0], int):
             dc += op.args[0]
-        elif op.name == "LP_STRING":
+        elif (
+            op.name == "LP_STRING" and len(op.args) == 1 and isinstance(op.args[0], str)
+        ):
             dc += len(op.args[0]) + 1
         else:
             pc += len(convert(op))

diff --git a/test/test_symtab.py b/test/test_symtab.py
@@ -63,3 +63,9 @@ def test_get_symtab_with_empty_lp_string():
     assert len(labels) == 2
     assert labels["S"] == HERA_DATA_START
     assert labels["X"] == HERA_DATA_START + 1
+
+
+def test_get_symtab_with_invalid_instructions():
+    labels = get_symtab([Op("CONSTANT", ["N"]), Op("CONSTANT", ["X", 42])])
+    assert len(labels) == 1
+    assert labels["X"] == 42