From 5384fa47b14e604d8fe458f34ba0114483b8d19d Mon Sep 17 00:00:00 2001
From: Ken Lauer <klauer@slac.stanford.edu>
Date: Fri, 14 Jan 2022 16:28:14 -0800
Subject: [PATCH] ENH: add union support

---
 blark/iec.lark                  |  9 +++--
 blark/summary.py                |  1 +
 blark/tests/test_transformer.py | 22 ++++++++++-
 blark/transform.py              | 68 ++++++++++++++++++++++++++++++++-
 4 files changed, 95 insertions(+), 5 deletions(-)

diff --git a/blark/iec.lark b/blark/iec.lark
index eed799e..90bb7bf 100644
--- a/blark/iec.lark
+++ b/blark/iec.lark
@@ -242,6 +242,7 @@ data_type_declaration: "TYPE"i [ _type_declaration ] ";"* "END_TYPE"i ";"*
 
 _type_declaration: array_type_declaration
                  | structure_type_declaration
+                 | union_type_declaration
                  | string_type_declaration
                  | simple_type_declaration
                  | subrange_type_declaration
@@ -305,16 +306,18 @@ _array_initial_element: constant
                       | structure_initialization
                       | enumerated_value
 
-structure_type_declaration: structure_type_name [ extends ] ":" [ indirection_type ] _structure_declaration
+structure_type_declaration: structure_type_name [ extends ] ":" [ indirection_type ] "STRUCT"i ( structure_element_declaration ";"+ )* "END_STRUCT"i
 
 initialized_structure_type_declaration: structure_type_name [ extends ] ":" initialized_structure
 
 initialized_structure: structure_type_name ":=" structure_initialization
 
-_structure_declaration: "STRUCT"i ( structure_element_declaration ";"+ )* "END_STRUCT"i ";"*
-
 structure_element_declaration: structure_element_name [ incomplete_location ] ":" ( initialized_structure | array_spec_init | simple_spec_init | subrange_spec_init | enumerated_spec_init )
 
+union_element_declaration: structure_element_name ":" ( array_specification | simple_specification | subrange_specification | enumerated_specification )
+
+union_type_declaration: structure_type_name ":" "UNION"i ( union_element_declaration ";"+ )* "END_UNION"i ";"*
+
 structure_initialization: "(" structure_element_initialization ( "," structure_element_initialization )* ")"
 
 structure_element_initialization: constant
diff --git a/blark/summary.py b/blark/summary.py
index d9216ed..7562587 100644
--- a/blark/summary.py
+++ b/blark/summary.py
@@ -325,6 +325,7 @@ def from_function_block(
 @dataclass
 class DataTypeSummary(Summary):
     """Summary representation of a single data type."""
+    # Note: structures only for now.
     name: str
     source_code: str
     type: str
diff --git a/blark/tests/test_transformer.py b/blark/tests/test_transformer.py
index feba3a9..721a86a 100644
--- a/blark/tests/test_transformer.py
+++ b/blark/tests/test_transformer.py
@@ -990,7 +990,7 @@ def test_incomplete_located_var_decls(rule_name, value):
                 STRUCT
                     xPLC_CnBitsValid : BOOL;
                     xPLC_CnBits : ARRAY [0..20] OF BYTE;
-                END_STRUCT;
+                END_STRUCT
             END_TYPE
             """
         )),
@@ -1030,6 +1030,26 @@ def test_incomplete_located_var_decls(rule_name, value):
             END_TYPE
             """
         )),
+        param("data_type_declaration", tf.multiline_code_block(
+            """
+            TYPE TypeName :
+                UNION
+                    intval : INT;
+                    as_bytes : ARRAY [0..2] OF BYTE;
+                END_UNION
+            END_TYPE
+            """
+        )),
+        param("data_type_declaration", tf.multiline_code_block(
+            """
+            TYPE TypeName :
+                UNION
+                    intval : INT;
+                    enum : (iValue := 1, iValue2 := 2) INT;
+                END_UNION
+            END_TYPE
+            """
+        )),
     ],
 )
 def test_data_type_declaration(rule_name, value):
diff --git a/blark/transform.py b/blark/transform.py
index 8ad6264..1089ac1 100644
--- a/blark/transform.py
+++ b/blark/transform.py
@@ -1053,6 +1053,65 @@ def __str__(self) -> str:
         return f"{name_and_location} : {self.init};"
 
 
+UnionElementSpecification = Union[
+    ArraySpecification,
+    lark.Token,  # simple_specification
+    SubrangeSpecification,
+    EnumeratedSpecification,
+]
+
+
+@dataclass
+@_rule_handler("union_element_declaration", comments=True)
+class UnionElementDeclaration:
+    name: lark.Token
+    spec: UnionElementSpecification
+    meta: Optional[Meta] = meta_field()
+
+    @property
+    def variables(self) -> List[str]:
+        """API compat"""
+        return [self.name]
+
+    def __str__(self) -> str:
+        return f"{self.name} : {self.spec};"
+
+
+@dataclass
+@_rule_handler("union_type_declaration", comments=True)
+class UnionTypeDeclaration:
+    name: lark.Token
+    declarations: List[UnionElementDeclaration]
+    meta: Optional[Meta] = meta_field()
+
+    @staticmethod
+    def from_lark(name: lark.Token, *decls: UnionElementDeclaration):
+        return UnionTypeDeclaration(
+            name=name,
+            declarations=list(decls),
+        )
+
+    def __str__(self) -> str:
+        if not self.declarations:
+            decls = None
+        else:
+            decls = indent(
+                "\n".join(
+                    str(decl) for decl in self.declarations
+                )
+            )
+
+        return "\n".join(
+            line for line in (
+                f"{self.name} :",
+                "UNION",
+                decls,
+                "END_UNION",
+            )
+            if line is not None
+        )
+
+
 @dataclass
 @_rule_handler("initialized_structure")
 class InitializedStructure:
@@ -3141,9 +3200,16 @@ def __str__(self) -> str:
         if not self.declaration:
             return "TYPE\nEND_TYPE"
 
+        decl = indent(self.declaration).lstrip()
+        if not isinstance(
+            self.declaration, (StructureTypeDeclaration, UnionTypeDeclaration)
+        ):
+            # note: END_STRUCT; END_UNION; result in "END_TYPE expected not ;"
+            decl = decl + ";"
+
         return "\n".join(
             (
-                "TYPE " + indent(f"{self.declaration};").lstrip(),
+                f"TYPE {decl}",
                 "END_TYPE",
             )
         )