def visit_IfExp(self, node):
r'''
We'd like to be able to use the "?" operator in C++, but the
problem is lazy evaluation. It could be when we look at one or the
other item, a bunch of prep work has to be done - and that will
show up in separate statements. So we have to use if/then/else with
a result value.
'''
# The result we'll store everything in.
result = crep.cpp_variable(unique_name("if_else_result"),
self._gc.current_scope(),
cpp_type=ctyp.terminal("double"))
self._gc.declare_variable(result)
# We always have to evaluate the test.
current_scope = self._gc.current_scope()
test_expr = self.get_rep(node.test)
self._gc.add_statement(statement.iftest(test_expr))
if_scope = self._gc.current_scope()
# Next, we do the true and false if statement.
self._gc.add_statement(
statement.set_var(result, self.get_rep(node.body)))
self._gc.set_scope(if_scope)
self._gc.pop_scope()
self._gc.add_statement(statement.elsephrase())
self._gc.add_statement(
statement.set_var(result, self.get_rep(node.orelse)))
self._gc.set_scope(current_scope)
# Done, the result is the rep of this node!
node.rep = result
self._result = result
def visit_Call_Aggregate_only(self, node: ast.Call):
'''
- (acc lambda): the accumulator is set to the first element, and the lambda is called to
update it after that. This is called `agg_only`.
'''
agg_lambda = node.args[0]
# Get the sequence we are calling against and the accumulator
if not isinstance(node.func, ast.Attribute):
raise BaseException("Wrong type of function")
seq = self.as_sequence(node.func.value)
accumulator, accumulator_scope = self.create_accumulator(seq)
# We have to do a simple if statement here so that the first time through we can set the
# accumulator, and the second time we can add to it.
is_first_iter = crep.cpp_variable(unique_name("is_first"),
self._gc.current_scope(),
cpp_type=ctyp.terminal('bool'),
initial_value=crep.cpp_value(
'true', self._gc.current_scope(),
ctyp.terminal('bool')))
accumulator_scope.declare_variable(is_first_iter)
# Set the scope where we will be doing the accumulation
sv = seq.sequence_value()
if isinstance(sv, crep.cpp_sequence):
self._gc.set_scope(sv.iterator_value().scope()[-1])
else:
self._gc.set_scope(sv.scope())
# Code up if statement to select out the first element.
if_first = statement.iftest(is_first_iter)
self._gc.add_statement(if_first)
self._gc.add_statement(
statement.set_var(
is_first_iter,
crep.cpp_value("false", self._gc.current_scope(),
ctyp.terminal('bool'))))
# Set the accumulator
self._gc.add_statement(
statement.set_var(accumulator, seq.sequence_value()))
self._gc.pop_scope()
# Now do the if statement and make the call to calculate the accumulation.
self._gc.add_statement(statement.elsephrase())
call = ast.Call(
func=agg_lambda,
args=[accumulator.as_ast(),
seq.sequence_value().as_ast()])
self._gc.add_statement(
statement.set_var(accumulator, self.get_rep(call)))
# Finally, since this is a terminal, we need to pop off the top.
self._gc.set_scope(accumulator_scope)
# Cache the results in our result in case we are skipping nodes in the AST.
node.rep = accumulator
self._result = accumulator
def visit_BoolOp(self, node):
'''A bool op like And or Or on a set of values
This is a bit more complex than just "anding" things as we want to make sure to short-circuit the
evaluation if we need to.
'''
# The result of this test
result = crep.cpp_variable(unique_name('bool_op'),
self._gc.current_scope(),
cpp_type='bool')
self._gc.declare_variable(result)
# How we check and short-circuit depends on if we are doing and or or.
check_expr = result.as_cpp() if type(
node.op) == ast.And else '!{0}'.format(result.as_cpp())
check = crep.cpp_value(check_expr,
self._gc.current_scope(),
cpp_type='bool')
first = True
scope = self._gc.current_scope()
for v in node.values:
if not first:
self._gc.add_statement(statement.iftest(check))
rep_v = self.get_rep(v)
self._gc.add_statement(statement.set_var(result, rep_v))
if not first:
self._gc.set_scope(scope)
first = False
# Cache result variable so those above us have something to use.
self._result = result
node.rep = result
def visit_call_Aggregate_initial(self, node: ast.Call):
'''
- (const, acc lambda): the accumulator is set to the value, and then the lambda is called to
update it on every single element. This is called `agg_initial`
'''
agg_lambda = node.args[1]
init_val = self.get_rep(node.args[0])
# Get the sequence we are calling against and the accumulator
if not isinstance(node.func, ast.Attribute):
raise BaseException("Wrong type of function")
seq = self.as_sequence(node.func.value)
accumulator, accumulator_scope = self.create_accumulator(
seq, initial_value=init_val, acc_type=init_val.cpp_type())
# Now do the accumulation. This happens at the current iterator scope.
sv = seq.sequence_value()
if isinstance(sv, crep.cpp_sequence):
self._gc.set_scope(sv.iterator_value().scope()[-1])
else:
self._gc.set_scope(sv.scope())
call = ast.Call(
func=agg_lambda,
args=[accumulator.as_ast(),
seq.sequence_value().as_ast()])
self._gc.add_statement(
statement.set_var(accumulator, self.get_rep(call)))
# Finally, since this is a terminal, we need to pop off the top.
self._gc.set_scope(accumulator_scope)
# Cache the results in our result in case we are skipping nodes in the AST.
node.rep = accumulator
self._result = accumulator
def test_insert_two_levels():
s1 = statement.iftest("true")
s2 = statement.set_var("v1", "true")
g = generated_code()
g.add_statement(s1)
g.add_statement(s2)
assert 1 == len(s1._statements)
def process_ast_node(visitor, gc, call_node: ast.Call):
r'''Inject the proper code into the output stream to deal with this C++ code.
We expect this to be run on the back-end of the system.
visitor - The node visitor that is converting the code into C++
gc - the generated code object that we fill with actual code
call_node - a Call ast node, with func being a CPPCodeValue.
Result:
representation - A value that represents the output
'''
# We write everything into a new scope to prevent conflicts. So we have to declare the result ahead of time.
cpp_ast_node = call_node.func
result_rep = cpp_ast_node.result_rep(gc.current_scope())
gc.declare_variable(result_rep)
# Include files
for i in cpp_ast_node.include_files:
gc.add_include(i)
# Build the dictionary for replacement for the object we are calling
# against, if any.
repl_list = []
if cpp_ast_node.replacement_instance_obj is not None:
repl_list += [
(cpp_ast_node.replacement_instance_obj[0],
visitor.resolve_id(
cpp_ast_node.replacement_instance_obj[1]).rep.as_cpp())
]
# Process the arguments that are getting passed to the function
for arg, dest in zip(cpp_ast_node.args, call_node.args):
rep = visitor.get_rep(dest)
repl_list += [(arg, rep.as_cpp())]
# Emit the statements.
blk = statements.block()
visitor._gc.add_statement(blk)
for s in cpp_ast_node.running_code:
l = s
for src, dest in repl_list:
l = l.replace(src, str(dest))
blk.add_statement(statements.arbitrary_statement(l))
# Set the result and close the scope
blk.add_statement(
statements.set_var(
result_rep,
cpp_value(cpp_ast_node.result, gc.current_scope(),
result_rep.cpp_type())))
gc.pop_scope()
return result_rep
def test_deepest_scope_equal():
g = generated_code()
s1 = statement.iftest("true")
s2 = statement.set_var("v1", "true")
g.add_statement(s1)
scope_1 = g.current_scope()
v1 = crep.cpp_value("v1", scope_1, ctyp.terminal('int'))
v2 = crep.cpp_value("v2", scope_1, ctyp.terminal('int'))
assert v1 == deepest_scope(v1, v2)
assert v2 == deepest_scope(v2, v1)
def test_insert_in_middle():
s1 = statement.iftest("true")
s2 = statement.set_var("v1", "true")
g = generated_code()
g.add_statement(s1)
g.add_statement(s2)
s3 = statement.iftest("fork")
g.add_statement(s3, below=s1)
assert 1 == len(s1._statements)
assert 1 == len(s3._statements)
assert s1._statements[0] is s3
assert s3._statements[0] is s2
def visit_First(self, node):
'We are in a sequence. Take the first element of the sequence and use that for future things.'
# Make sure we are in a loop.
seq = self.as_sequence(node.source)
# The First terminal works by protecting the code with a if (first_time) {} block.
# We need to declare the first_time variable outside the block where the thing we are
# looping over here is defined. This is a little tricky, so we delegate to another method.
loop_scope = seq.iterator_value().scope()
outside_block_scope = loop_scope[-1]
# Define the variable to track this outside that block.
is_first = crep.cpp_variable(unique_name('is_first'),
outside_block_scope,
cpp_type=ctyp.terminal('bool'),
initial_value=crep.cpp_value(
'true', self._gc.current_scope(),
ctyp.terminal('bool')))
outside_block_scope.declare_variable(is_first)
# Now, as long as is_first is true, we can execute things inside this statement.
# The trick is putting the if statement in the right place. We need to locate it just one level
# below where we defined the scope above.
s = statement.iftest(is_first)
s.add_statement(
statement.set_var(
is_first,
crep.cpp_value('false',
top_level_scope(),
cpp_type=ctyp.terminal('bool'))))
sv = seq.sequence_value()
if isinstance(sv, crep.cpp_sequence):
self._gc.set_scope(sv.iterator_value().scope()[-1])
else:
self._gc.set_scope(sv.scope())
self._gc.add_statement(s)
# If we just found the first sequence in a sequence, return that.
# Otherwise return a new version of the value.
first_value = sv if isinstance(
sv, crep.cpp_sequence) else sv.copy_with_new_scope(
self._gc.current_scope())
node.rep = first_value
self._result = first_value
def visit_ResultTTree(self, node: query_result_asts.ResultTTree):
'''This AST means we are taking an iterable and converting it to a ROOT file.
'''
# Get the representations for each variable. We expect some sort of structure
# for the variables - or perhaps a single variable.
self.generic_visit(node)
v_rep_not_norm = self.as_sequence(node.source)
# What we have is a sequence of the data values we want to fill. The iterator at play
# here is the scope we want to use to run our Fill() calls to the TTree.
scope_fill = v_rep_not_norm.iterator_value().scope()
# Clean the data up so it is uniform and the next bit can proceed smoothly.
# If we don't have a tuple of data to log, turn it into a tuple.
seq_values = v_rep_not_norm.sequence_value()
if not isinstance(seq_values, crep.cpp_tuple):
seq_values = crep.cpp_tuple((v_rep_not_norm.sequence_value(), ),
scope_fill)
# Make sure the number of items is the same as the number of columns specified.
if len(seq_values.values()) != len(node.column_names):
raise BaseException(
"Number of columns ({0}) is not the same as labels ({1}) in TTree creation"
.format(len(seq_values.values()), len(node.column_names)))
# Next, look at each on in turn to decide if it is a vector or a simple variable.
# Create a variable that we will fill for each one.
var_names = [
(name,
crep.cpp_variable(unique_name(name, is_class_var=True),
self._gc.current_scope(),
cpp_type=get_ttree_type(rep)))
for name, rep in zip(node.column_names, seq_values.values())
]
# For each incoming variable, we need to declare something we are going to write.
for cv in var_names:
self._gc.declare_class_variable(cv[1])
# Next, emit the booking code
tree_name = unique_name(node.tree_name)
self._gc.add_book_statement(statement.book_ttree(tree_name, var_names))
# Note that the output file and tree are what we are going to return.
# The output filename is fixed - the hose code in AnalysisBase has that hard coded.
# To allow it to be different we have to modify that template too, and pass the
# information there. If more than one tree is written, the current code would
# lead to a bug.
node.rep = rh.cpp_ttree_rep("ANALYSIS.root", tree_name,
self._gc.current_scope())
# For each varable we need to save, cache it or push it back, depending.
# Make sure that it happens at the proper scope, where what we are after is defined!
s_orig = self._gc.current_scope()
for e_rep, e_name in zip(seq_values.values(), var_names):
# Set the scope. Normally we want to do it where the variable was calculated
# (think of cases when you have to calculate something with a `push_back`),
# but if the variable was already calculated, we want to make sure we are at least
# in the same scope as the tree fill.
e_rep_scope = e_rep.scope() if not isinstance(
e_rep, crep.cpp_sequence) else e_rep.sequence_value().scope()
if e_rep_scope.starts_with(scope_fill):
self._gc.set_scope(e_rep_scope)
else:
self._gc.set_scope(scope_fill)
# If the variable is something we are iterating over, then fill it, otherwise,
# just set it.
if rep_is_collection(e_rep):
self._gc.add_statement(
statement.push_back(e_name[1], e_rep.sequence_value()))
else:
self._gc.add_statement(statement.set_var(e_name[1], e_rep))
cs = self._gc.current_scope()
if cs.starts_with(scope_fill):
scope_fill = cs
# The fill statement. This should happen at the scope where the tuple was defined.
# The scope where this should be done is a bit tricky (note the update above):
# - If a sequence, you want it where the sequence iterator is defined - or outside that scope
# - If a value, you want it at the level where the value is set.
self._gc.set_scope(scope_fill)
self._gc.add_statement(statement.ttree_fill(tree_name))
for e in zip(seq_values.values(), var_names):
if rep_is_collection(e[0]):
self._gc.add_statement(statement.container_clear(e[1][1]))
# And we are a terminal, so pop off the block.
self._gc.set_scope(s_orig)
self._gc.pop_scope()