Comparing my concatenation changes with main reminded me that NodeCountMapper doesn't currently accumulate node counts from functions into the overall count. It should be given a map_function_definition method that does something similar to CallSiteCountMapper:
|
@memoize_method |
|
def map_function_definition(self, /, expr: FunctionDefinition, |
|
*args: Any, **kwargs: Any) -> None: |
|
if not self.visit(expr): |
|
return |
|
|
|
new_mapper = self.clone_for_callee(expr) |
|
for subexpr in expr.returns.values(): |
|
new_mapper(subexpr, *args, **kwargs) |
|
self.count += new_mapper.count |
|
|
|
self.post_visit(expr, *args, **kwargs) |
I wonder if the default
map_function_definition implementation in
CachedWalkMapper should be disabled, to avoid bugs like this? Maybe a default implementation could be provided in a separate method, e.g.
_basic_map_function_definition, for convenience in mappers that don't need special behavior.
Comparing my concatenation changes with
mainreminded me thatNodeCountMapperdoesn't currently accumulate node counts from functions into the overall count. It should be given amap_function_definitionmethod that does something similar toCallSiteCountMapper:pytato/pytato/analysis/__init__.py
Lines 435 to 446 in 78b43c1
I wonder if the default
map_function_definitionimplementation inCachedWalkMappershould be disabled, to avoid bugs like this? Maybe a default implementation could be provided in a separate method, e.g._basic_map_function_definition, for convenience in mappers that don't need special behavior.