ENH: first pass at JMP and labeled statements

Author: klauer

blark/iec.lark

@@ -727,7 +727,7 @@ parenthesized_expression: "(" expression ")"
                    | constant
 
 // B.3.2
-statement_list: _statement+
+statement_list: _statement+ end_of_statement_list_label?
 
 _statement: ";"
           | assignment_statement
@@ -745,6 +745,8 @@ _statement: ";"
           | repeat_statement
           | exit_statement
           | continue_statement
+          | jmp_statement
+          | labeled_statement
 
 
 // B.3.2.1
@@ -779,7 +781,29 @@ case_statement: "CASE"i expression "OF"i case_elements [ else_clause ] "END_CASE
 
 case_elements: case_element+
 
-case_element: case_list ":" [ statement_list ]
+case_element_statement_list: _case_element_statement+
+
+// Specifically exclude labeled statements as they are aliased to case
+// statements themselves:
+_case_element_statement: ";"
+                       | assignment_statement
+                       | no_op_statement
+                       | set_statement
+                       | reset_statement
+                       | reference_assignment_statement
+                       | return_statement
+                       | chained_function_call_statement
+                       | function_call_statement
+                       | if_statement
+                       | case_statement
+                       | for_statement
+                       | while_statement
+                       | repeat_statement
+                       | exit_statement
+                       | continue_statement
+                       | jmp_statement
+
+case_element: case_list ":" [ case_element_statement_list ]
 
 case_list: case_list_element ( "," case_list_element )*
 
@@ -804,3 +828,11 @@ repeat_statement: "REPEAT"i statement_list "UNTIL"i expression "END_REPEAT"i ";"
 exit_statement.1: "EXIT"i ";"+
 
 continue_statement.1: "CONTINUE"i ";"+
+
+LABEL: IDENTIFIER
+labeled_statement.1: LABEL ":" _statement
+
+// End-of-statement list may have a label associated with it:
+end_of_statement_list_label: LABEL ":"
+
+jmp_statement: "JMP"i LABEL ";"+

blark/tests/test_transformer.py

@@ -1724,3 +1724,60 @@ def test_statement_priority(statement: str, cls: type):
     assert isinstance(transformed, tf.StatementList)
     transformed_statement, = transformed.statements
     assert isinstance(transformed_statement, cls)
+
+
+@pytest.mark.parametrize(
+    "statements, labels",
+    [
+        param(
+            tf.multiline_code_block(
+                """
+                JMP _label;
+                _label :
+                """,
+            ),
+            [],
+            id="simple_jump_no_statement"
+        ),
+        param(
+            tf.multiline_code_block(
+                """
+                JMP _label;
+                _label : A := 5;
+                """,
+            ),
+            [],
+            id="simple_jump_with_assignment"
+        ),
+        param(
+            tf.multiline_code_block(
+                """
+                JMP _label;
+                _label :
+                IF a = 2 THEN
+                    b := 3;
+                END_IF
+                """,
+            ),
+            [],
+            id="simple_jump_with_assignment"
+        ),
+        param(
+            tf.multiline_code_block(
+                """
+                JMP _label1;
+                _label1 :
+                _label2 :
+                IF a = 2 THEN
+                    b := 3;
+                END_IF
+                """,
+            ),
+            [],
+            id="multi_label"
+        ),
+    ],
+)
+def test_labeled_statements_roundtrip(statements: str, labels: List[str]):
+    transformed = roundtrip_rule("statement_list", statements)
+    assert isinstance(transformed, tf.StatementList)

blark/transform.py

@@ -3055,7 +3055,40 @@ def __str__(self) -> str:
 
 
 @dataclass
-@_rule_handler("statement_list")
+@_rule_handler(
+    "labeled_statement",
+    "end_of_statement_list_label",
+    comments=True,
+)
+class LabeledStatement(Statement):
+    label: lark.Token
+    statement: Optional[Statement] = None
+    meta: Optional[Meta] = meta_field()
+
+    def __str__(self) -> str:
+        if self.statement is None:
+            return f"{self.label} :"
+
+        statement = str(self.statement)
+        if statement.count("\n") > 1:
+            # Multiline statement after label - put it on the next line
+            return f"{self.label} :\n{statement}"
+        # Single line statement after label - put it on the same line
+        return f"{self.label} : {statement}"
+
+
+@dataclass
+@_rule_handler("jmp_statement", comments=True)
+class JumpStatement(Statement):
+    label: lark.Token
+    meta: Optional[Meta] = meta_field()
+
+    def __str__(self) -> str:
+        return f"JMP {self.label};"
+
+
+@dataclass
+@_rule_handler("statement_list", "case_element_statement_list")
 class StatementList:
     statements: List[Statement]
     meta: Optional[Meta] = meta_field()