Added possibility to add argument to enum/repeat/select in STA's grammar

LLNL · Mar 5, 2024 · 26ba7f9 · 26ba7f9
1 parent bc4a9dc
commit 26ba7f9
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 10 deletions.
diff --git a/autocog/sta/ast.py b/autocog/sta/ast.py
@@ -175,10 +175,13 @@ def gvlbl(self):
 
 class EnumType(ASTNode):
     kind: str
+    arguments: List[Argument] = []
     source: Union[Path,List[Expression]]
 
     def gvtree(self):
         yield from super().gvtree()
+        for (i,arg) in enumerate(self.arguments):
+            yield ("argument",i,arg)
         if isinstance(self.source, Path):
             yield ("path",None,self.source)
         else:

diff --git a/autocog/sta/frontend.py b/autocog/sta/frontend.py
@@ -410,18 +410,42 @@ def visit_string_expr(self, node, visited_children):
         return visited_children[0]
 
     def visit_repeat_def(self, node, visited_children):
-        assert len(visited_children) == 5
-        return EnumType(kind='repeat', source=visited_children[3])
+        assert len(visited_children) == 7
+        arguments = visited_children[2]
+        if 'children' in arguments:
+            assert len(arguments['children']) == 1
+            arguments = arguments['children'][0]
+            assert arguments['kind'] == 'param_list'
+            arguments = arguments['children']
+        else:
+            arguments = []
+        return EnumType(kind='repeat', source=visited_children[5], arguments=arguments)
 
     def visit_select_def(self, node, visited_children):
-        assert len(visited_children) == 5
-        return EnumType(kind='select', source=visited_children[3])
+        assert len(visited_children) == 7
+        arguments = visited_children[2]
+        if 'children' in arguments:
+            assert len(arguments['children']) == 1
+            arguments = arguments['children'][0]
+            assert arguments['kind'] == 'param_list'
+            arguments = arguments['children']
+        else:
+            arguments = []
+        return EnumType(kind='select', source=visited_children[5], arguments=arguments)
 
     def visit_enum_def(self, node, visited_children):
-        assert len(visited_children) == 5
-        source = visited_children[3]
+        assert len(visited_children) == 7
+        arguments = visited_children[2]
+        if 'children' in arguments:
+            assert len(arguments['children']) == 1
+            arguments = arguments['children'][0]
+            assert arguments['kind'] == 'param_list'
+            arguments = arguments['children']
+        else:
+            arguments = []
+        source = visited_children[5]
         assert source['kind'] == 'string_expr_list'
-        return EnumType(kind='enum', source=source['children'])
+        return EnumType(kind='enum', source=source['children'], arguments=arguments)
 
     def visit_string_expr_list__(self, node, visited_children):
         if len(visited_children) == 1:

diff --git a/autocog/sta/grammar.py b/autocog/sta/grammar.py
@@ -83,9 +83,9 @@
 array_slice_cont = array_slice__?
 array_slice__    = WS COLON WS expression
 
-repeat_def       = REPEAT WS LPAR path_expr RPAR
-select_def       = SELECT WS LPAR path_expr RPAR
-enum_def         = ENUM   WS LPAR string_expr_list RPAR
+repeat_def       = REPEAT WS type_ref_param? WS LPAR path_expr RPAR
+select_def       = SELECT WS type_ref_param? WS LPAR path_expr RPAR
+enum_def         = ENUM   WS type_ref_param? WS LPAR string_expr_list RPAR
 
 type_ref         = identifier WS type_ref_param?
 type_ref_param   = LT WS param_list WS GT