diff --git a/src/gboml/ast/__init__.py b/src/gboml/ast/__init__.py index 916c97dddf4c19c09481219648ad21e7d96df727..2289b63892b753d6db255962d36f928362138cb4 100644 --- a/src/gboml/ast/__init__.py +++ b/src/gboml/ast/__init__.py @@ -1,7 +1,7 @@ __all__ = [ "Meta", "GBOMLObject", "VarScope", "VarType", "SOSType", "ObjType", "Operator", "ExpressionObj", "Expression", "BoolExpression", "VarOrParamLeaf", - "VarOrParam", "Array", "Loop", "BaseLoop", "Function", "BoolExpressionOp", + "VarOrParam", "Array", "Loop", "BaseLoop", "EqLoop", "Function", "BoolExpressionOp", "BoolExpressionComparison", "ScopeChange", "ImportFile", "Definition", "Constraint", "StdConstraint", "SOSConstraint", "Objective", "VariableDefinition", "Node", "HyperEdge", "NodeDefinition", "HyperEdgeDefinition", @@ -9,7 +9,7 @@ __all__ = [ "Range", "MultiLoop", "DictEntry", "Dictionary", "NodeGenerator", "HyperEdgeGenerator", "DefinitionType", "FunctionDefinition", "ConstantDefinition", "ExpressionDefinition", "CtrActivation", "ObjActivation", "ActivationType", "Activation", "Extends", "NamedGBOMLObject", - "ExpressionUseGenScope", "AnyGBOMLObject" + "ExpressionUseGenScope", "AnyGBOMLObject", "IndexingParameterDefinition" ] from gboml.ast.arrays import * diff --git a/src/gboml/ast/loops.py b/src/gboml/ast/loops.py index adb3e31bc26f051ccea370ad257300ba7b4b9c0a..ce17bdaff35f27dff15c07c61309846b075d47a4 100644 --- a/src/gboml/ast/loops.py +++ b/src/gboml/ast/loops.py @@ -21,6 +21,13 @@ class BaseLoop(Loop): condition: Optional[BoolExpression] +@dataclass +class EqLoop(Loop): + varid: str + on: VarOrParam + condition: Optional[BoolExpression] + + @dataclass class ImplicitLoop(BaseLoop): varid: str = field(default="t", init=False) diff --git a/src/gboml/ast/variables.py b/src/gboml/ast/variables.py index d4e1e59987ffbfddb2bf68dbef4519286916a254..1ca9e80e6b63e442ea6f7fdf29255c5f83828b20 100644 --- a/src/gboml/ast/variables.py +++ b/src/gboml/ast/variables.py @@ -2,6 +2,8 @@ from dataclasses import dataclass, field from enum import Enum from typing import Optional +from gboml.ast.functions import Function +from gboml.ast.arrays import Array, Range from gboml.ast.base import GBOMLObject, NamedGBOMLObject from gboml.ast.path import VarOrParam from gboml.ast.rvalue import RValue @@ -45,6 +47,9 @@ class FunctionDefinition(Definition): value: RValue tags: set[str] = field(default_factory=set) +@dataclass +class IndexingParameterDefinition(Definition): + value: Function | Array | Range | VarOrParam @dataclass class VariableDefinition(NamedGBOMLObject): diff --git a/src/gboml/gboml.lark b/src/gboml/gboml.lark index ce0716688f690ce496768354c04e86ebb42b033f..bf7ff85a91c2dbda265cd0e494197654f47bcede 100644 --- a/src/gboml/gboml.lark +++ b/src/gboml/gboml.lark @@ -83,8 +83,9 @@ OBJ_TYPE: "min" | "max" // LOOPS ?loop: multi_loop | implicit_loop -?multi_loop: base_loop+ +?multi_loop: (base_loop | eq_loop)+ base_loop: "for" ID "in" iterable ["where" bool_expression] +eq_loop: "for" ID "=" var_or_param ["where" bool_expression] implicit_loop: "where" bool_expression // BOOLEAN EXPRESSIONS @@ -98,7 +99,9 @@ bool_expression_comparison: expression (COMPARISON_OPERATOR | CTR_OPERATOR) expr COMPARISON_OPERATOR: "<" | ">" | "!=" // DEFINITIONS -definition: ID ["(" separated_list{ID, ","} ")"] DEF_TYPE rvalue tags ";" +?definition: definition_std_param | definition_indexing_param +definition_indexing_param: ID "in" iterable ";" +definition_std_param: ID ["(" separated_list{ID, ","} ")"] DEF_TYPE rvalue tags ";" DEF_TYPE: "=" | "<-" // ARRAYS diff --git a/src/gboml/parsing.py b/src/gboml/parsing.py index ac9be844d4dea7f453442a255bd82c4601b71a36..0d1f0dfd3b007d22a75f953f2e4d0b4236a83cf3 100644 --- a/src/gboml/parsing.py +++ b/src/gboml/parsing.py @@ -72,6 +72,7 @@ def _lark_to_gboml(tree: Tree, filename: Optional[str] = None) -> GBOMLGraph: "constraint_sos": SOSConstraint, "objective": Objective, "base_loop": BaseLoop, + "eq_loop": EqLoop, "implicit_loop": ImplicitLoop, "subtraction": op_transform(Operator.minus), "sum": op_transform(Operator.plus), @@ -92,6 +93,7 @@ def _lark_to_gboml(tree: Tree, filename: Optional[str] = None) -> GBOMLGraph: "dict_entry": DictEntry, "array": Array, "dict": Dictionary, + "definition_indexing_param": IndexingParameterDefinition, "ctr_activate": lambda *x, meta: CtrActivation(ActivationType.activate, *x, meta=meta), "ctr_deactivate": lambda *x, meta: CtrActivation(ActivationType.deactivate, *x, meta=meta), "obj_activate": lambda *x, meta: ObjActivation(ActivationType.activate, *x, meta=meta), @@ -211,7 +213,7 @@ def _lark_to_gboml(tree: Tree, filename: Optional[str] = None) -> GBOMLGraph: return Array(entries, meta=meta) raise Exception("An array cannot contain dictionary entries (and conversely)") - def definition(self, meta: Meta, name: str, args: Optional[list[str]], typ: DefinitionType, val: RValue, tags: set[str]): + def definition_std_param(self, meta: Meta, name: str, args: Optional[list[str]], typ: DefinitionType, val: RValue, tags: set[str]): if args is not None: if typ != DefinitionType.expression: raise Exception("Functions can only be defined as expressions (use `<-` instead of `=`)")