diff --git a/src/gboml/ast/__init__.py b/src/gboml/ast/__init__.py index 7a1e59d9ce03e5084997cc769fc3fc94d98895d8..550739ba0eb6a3bac35ab448aeac7327fdc01765 100644 --- a/src/gboml/ast/__init__.py +++ b/src/gboml/ast/__init__.py @@ -6,7 +6,7 @@ __all__ = [ "StdConstraint", "SOSConstraint", "Objective", "VariableDefinition", "Node", "HyperEdge", "NodeDefinition", "NodeImport", "HyperEdgeDefinition", "HyperEdgeImport", "ExpressionOp", "GBOMLGraph", "ImplicitLoop", "RValue", "RValueWithGen", "GeneratedRValue", - "Range", "MultiLoop", "DictEntry", "Dictionary" + "Range", "MultiLoop", "DictEntry", "Dictionary", "NodeGenerator", "HyperEdgeGenerator" ] from gboml.ast.arrays import * diff --git a/src/gboml/ast/hyperedges.py b/src/gboml/ast/hyperedges.py index c918750db061029c2cbdc97267f314a0a4d4dcfc..2f894647f133f53da9f6c542b08eb0c206592a4d 100644 --- a/src/gboml/ast/hyperedges.py +++ b/src/gboml/ast/hyperedges.py @@ -1,5 +1,6 @@ from dataclasses import dataclass +from gboml.ast import Loop from gboml.ast.base import GBOMLObject from gboml.ast.constraints import Constraint from gboml.ast.path import VarOrParam @@ -18,6 +19,14 @@ class HyperEdgeDefinition(HyperEdge): constraints: list[Constraint] +@dataclass +class HyperEdgeGenerator(HyperEdge): + name: VarOrParam + loop: Loop + parameters: list[Definition] + constraints: list[Constraint] + + @dataclass class HyperEdgeImport(HyperEdge): name: str diff --git a/src/gboml/ast/nodes.py b/src/gboml/ast/nodes.py index d93ae5f3129d9ffbc1810123cdd45127b97ab58f..3b6a7d7cac40174e80fbfd4cc46692d29b9664c6 100644 --- a/src/gboml/ast/nodes.py +++ b/src/gboml/ast/nodes.py @@ -1,5 +1,6 @@ from dataclasses import dataclass +from gboml.ast import Loop from gboml.ast.base import GBOMLObject from gboml.ast.constraints import Constraint from gboml.ast.path import VarOrParam @@ -24,6 +25,18 @@ class NodeDefinition(Node): objectives: list[Objective] +@dataclass +class NodeGenerator(Node): + name: VarOrParam + loop: Loop + parameters: list[Definition] + nodes: list[Node] + hyperedges: list[HyperEdge] + variables: list[VariableDefinition] + constraints: list[Constraint] + objectives: list[Objective] + + @dataclass class NodeImport(Node): name: str diff --git a/src/gboml/gboml.lark b/src/gboml/gboml.lark index f8ff7585c545949640073e10f6b3b077016b8c50..a5864f761f72ba0895a873e80e096c426dd23463 100644 --- a/src/gboml/gboml.lark +++ b/src/gboml/gboml.lark @@ -25,7 +25,7 @@ _program: node | hyperedge // NODES ?node: node_definition | node_import -node_definition: "#NODE" ID \ +node_definition: "#NODE" var_or_param [loop] \ parameters_block \ program_block \ variables_block \ @@ -45,7 +45,7 @@ variable_scope_change: ID SCOPE ";" // HYPEREDGES ?hyperedge: hyperedge_definition | hyperedge_import -hyperedge_definition: "#HYPEREDGE" ID \ +hyperedge_definition: "#HYPEREDGE" var_or_param [loop] \ parameters_block \ constraints_block diff --git a/src/gboml/parsing.py b/src/gboml/parsing.py index 19e99d75ac24aaca0953bda5daa55abb28317bac..83fd7cfde1e75ecaf0bec656ca0dd45ab20c3e76 100644 --- a/src/gboml/parsing.py +++ b/src/gboml/parsing.py @@ -61,7 +61,6 @@ def _lark_to_gboml(tree: Tree, filename: Optional[str] = None) -> GBOMLGraph: # obj(*children, meta=meta) # to_obj = { - "hyperedge_definition": HyperEdgeDefinition, "hyperedge_import": HyperEdgeImport, "definition": Definition, "var_or_param_leaf": VarOrParamLeaf, @@ -119,13 +118,31 @@ def _lark_to_gboml(tree: Tree, filename: Optional[str] = None) -> GBOMLGraph: def program_block(self, meta: Meta, *childrens: list[Node | HyperEdge]) -> NodesAndHyperEdges: return self.NodesAndHyperEdges([x for x in childrens if isinstance(x, Node)], [x for x in childrens if isinstance(x, HyperEdge)]) - def node_definition(self, meta: Meta, name: str, param_block: list[Definition], subprogram_block: NodesAndHyperEdges, + def hyperedge_definition(self, meta: Meta, name: VarOrParam, loop: Optional[Loop], param_block: list[Definition], constraint_block: list[Constraint]): + if loop is None: + if len(name.path) != 1 or len(name.path[0].indices) != 0: + raise Exception(f"Invalid name for node: {name}") + return HyperEdgeDefinition(name.path[0].name, param_block, constraint_block, meta=meta) + else: + return HyperEdgeGenerator(name, loop, param_block, constraint_block, meta=meta) + + def node_definition(self, meta: Meta, name: VarOrParam, loop: Optional[Loop], param_block: list[Definition], subprogram_block: NodesAndHyperEdges, variable_block: list[VariableDefinition], constraint_block: list[Constraint], objectives_block: list[Objective]): - return NodeDefinition(name, param_block, - subprogram_block.nodes, subprogram_block.hyperedges, - variable_block, constraint_block, - objectives_block, meta=meta) + if loop is None: + if len(name.path) != 1 or len(name.path[0].indices) != 0: + raise Exception(f"Invalid name for node: {name}") + return NodeDefinition(name.path[0].name, + param_block, + subprogram_block.nodes, subprogram_block.hyperedges, + variable_block, constraint_block, + objectives_block, meta=meta) + else: + return NodeGenerator(name, loop, + param_block, + subprogram_block.nodes, subprogram_block.hyperedges, + variable_block, constraint_block, + objectives_block, meta=meta) def node_import(self, meta: Meta, name: str, imported_name: VarOrParam, imported_from: str, redef: list[ScopeChange | Definition]): return NodeImport(name, imported_name, imported_from, diff --git a/tests/instances/ok/complex_parsing.txt b/tests/instances/ok/complex_parsing.txt index 4549ea1633748521cf2e12758c66d0a7d3d11928..fed325b24bd164514610fae54de247fe052065b1 100644 --- a/tests/instances/ok/complex_parsing.txt +++ b/tests/instances/ok/complex_parsing.txt @@ -13,6 +13,10 @@ a = 2; b external; c internal; +#NODE nodeL[i] for i in [0:10] + #VARIABLES +#HYPEREDGE E[i] for i in [0:10] + #CONSTRAINTS #NODE node1 #PARAMETERS a = 2;