def visitSubscript(self,
ctx: CCL_Parser.CCL.SubscriptContext) -> ast.Subscript:
name = ast.Name(self.get_pos(ctx.name), ctx.name.text)
indices = []
for idx in ctx.indices:
indices.append(ast.Name(self.get_pos(idx), idx.text))
return ast.Subscript(self.get_pos(ctx), name, tuple(indices))
def visitFor_loop(self, ctx: CCL_Parser.CCL.For_loopContext) -> ast.For:
name = ast.Name(self.get_pos(ctx.identifier), ctx.identifier.text)
value_from = self.visit(ctx.value_from)
value_to = self.visit(ctx.value_to)
body = []
for statement in ctx.body:
body.append(self.visit(statement))
return ast.For(self.get_pos(ctx), name, value_from, value_to, body)
def visit_ForEach(self, node: ast.ForEach) -> None:
s = self.current_table.resolve(node.name.val)
if s is not None:
raise CCLSymbolError(
node.name, f'Loop variable {node.name.val} already defined.')
table = SymbolTable(self.current_table)
atom_indices: Set[str] = set()
if node.atom_indices is not None:
i1, i2 = node.atom_indices
s1 = self.current_table.resolve(i1)
s2 = self.current_table.resolve(i2)
if s1 is not None or s2 is not None:
raise CCLSymbolError(
node,
f'Decomposition of bond symbol {node.name.val} used already defined names.'
)
bonded_constraint = ast.Predicate(
(node.line, node.column), 'bonded', (
ast.Name((node.line, node.column), i1),
ast.Name((node.line, node.column), i2),
))
table.define(
ObjectSymbol(i1, node, ObjectType.ATOM, bonded_constraint))
table.define(
ObjectSymbol(i2, node, ObjectType.ATOM, bonded_constraint))
atom_indices = {i1, i2}
self._iterating_over |= atom_indices | {node.name.val}
node.symbol_table = table
table.define(
ObjectSymbol(node.name.val, node, node.type, node.constraints))
self.current_table = table
for statement in node.body:
self.visit(statement)
self._iterating_over -= atom_indices | {node.name.val}
assert self.current_table.parent is not None
self.current_table = self.current_table.parent
def visitFor_each(self,
ctx: CCL_Parser.CCL.For_eachContext) -> ast.ForEach:
name = ast.Name(self.get_pos(ctx.identifier), ctx.identifier.text)
object_type = ast.ObjectType(self.visit(ctx.abtype).capitalize())
if ctx.constraint() is not None:
constraint = self.visit(ctx.constraint())
else:
constraint = None
if ctx.bond_decomp():
if object_type != ast.ObjectType.BOND:
raise CCLSyntaxError(ast.ASTNode(self.get_pos(ctx.abtype)),
f'Only bonds can be decomposed.')
atom_indices = tuple(i.text for i in ctx.bond_decomp().indices)
else:
atom_indices = None
body = []
for statement in ctx.body:
body.append(self.visit(statement))
return ast.ForEach(self.get_pos(ctx), name, object_type, atom_indices,
constraint, body)
def visitSumOp(self, ctx: CCL_Parser.CCL.SumOpContext) -> ast.Sum:
name = ast.Name(self.get_pos(ctx.identifier), ctx.identifier.text)
expr = self.visit(ctx.expr())
return ast.Sum(self.get_pos(ctx), name, expr)
def visitBasename(self, ctx: CCL_Parser.CCL.BasenameContext) -> ast.Name:
return ast.Name(self.get_pos(ctx), ctx.name.text)
def visit_Subscript(self, node: ast.Subscript) -> None:
s = self.current_table.resolve(node.name.val)
if not isinstance(s, SubstitutionSymbol):
self.visit(node.name)
symbol_type = node.name.result_type
index_types_list = []
for idx in node.indices:
self.visit(idx)
index_types_list.append(idx.result_type)
mapped_val = self._indices_mapping.get(idx.val, idx.val)
if isinstance(idx.result_type, ObjectType) and \
mapped_val not in self.iterating_over:
raise CCLSymbolError(
idx,
f'Object {mapped_val} not bound to any For/ForEach/Sum.')
index_types = tuple(index_types_list)
index_types_str = ', '.join(str(i) for i in index_types)
if isinstance(s, ParameterSymbol):
if symbol_type == ParameterType.ATOM and index_types != (
ObjectType.ATOM, ):
raise CCLTypeError(
node,
f'Cannot index atom parameter with {index_types_str}.')
if symbol_type == ParameterType.BOND:
if index_types not in ((ObjectType.BOND, ), (ObjectType.ATOM,
ObjectType.ATOM)):
raise CCLTypeError(
node,
f'Cannot index bond parameter with {index_types_str}.')
# Check whether two atoms are actually bonded, i.e., 'bonded' predicate exists
if index_types == (ObjectType.ATOM, ObjectType.ATOM):
s1 = self.current_table.resolve(node.indices[0].val)
s2 = self.current_table.resolve(node.indices[1].val)
assert s1 is not None and isinstance(s1, ObjectSymbol)
assert s2 is not None and isinstance(s2, ObjectSymbol)
idx1 = node.indices[0].val
idx2 = node.indices[1].val
results: Set[Optional[ast.ASTNode]] = {None}
for c in (c for c in (s1.constraints, s2.constraints)
if c is not None):
# Search for either bonded(i, j) or bonded(j, i)
results.add(
ast.search_ast_element(
c,
ast.Predicate((-1, -1), 'bonded', (ast.Name(
(-1, -1), idx1), ast.Name(
(-1, -1), idx2)))))
results.add(
ast.search_ast_element(
c,
ast.Predicate((-1, -1), 'bonded', (ast.Name(
(-1, -1), idx2), ast.Name(
(-1, -1), idx1)))))
if results == {None}:
raise CCLSymbolError(
node,
f'Cannot index bond parameter by two non-bonded atoms.'
)
if symbol_type == NumericType.FLOAT:
raise CCLTypeError(node, f'Cannot index common parameter.')
elif isinstance(s, VariableSymbol) and isinstance(
symbol_type, ArrayType):
if symbol_type.indices != index_types:
raise CCLTypeError(
node, f'Cannot index Array of type {symbol_type} '
f'using index/indices of type(s) {index_types_str}.')
elif isinstance(s, FunctionSymbol):
assert isinstance(symbol_type, FunctionType)
if symbol_type.args != index_types:
raise CCLTypeError(
node,
f'Cannot use function {s.function.name}: {s.function.type} '
f'with arguments of type(s) {index_types_str}')
assert isinstance(s.function.type.return_type,
(NumericType, ArrayType))
node.result_type = s.function.type.return_type
return
elif isinstance(s, SubstitutionSymbol):
if len(s.indices) != len(index_types):
raise CCLTypeError(
node,
f'Bad number of indices for {s.name}, got {len(index_types)}, '
f'expected {len(s.indices)}.')
if not all(isinstance(t, ObjectType) for t in index_types):
raise CCLTypeError(
node,
f'Substitution indices for symbol {s.name} must have type Atom or Bond.'
)
self._indices_mapping.update(
{si.val: ni.val
for si, ni in zip(s.indices, node.indices)})
types = set()
for constraint, expr in s.rules.items():
if constraint is not None:
self.visit(constraint)
self.visit(expr)
types.add(expr.result_type)
if len(types) > 1:
raise CCLTypeError(
node,
f'All expressions within a substitution symbol {s.name} must have same type.'
)
for si in s.indices:
self._indices_mapping.pop(si.val)
else:
raise CCLTypeError(
node,
f'Cannot index type {symbol_type} with indices of type(s) {index_types_str}'
)
# Return Float if not assigned already
node.result_type = NumericType.FLOAT