def con_lhs_vars_from_comp(self, comp): """Return a tuple of constrained variables. An variable is constrained if its first occurrence in the query is in a constrained position. For cyclic object queries like {(x, y) for y in x for x in y}, after translating into clauses over M, there are two possible sets of constrained vars: {x} and {y}. This function processes clauses left-to-right, so {y} will be chosen. """ uncon = OrderedSet() con = OrderedSet() for cl in comp.clauses: # The new unconstrained vars are the ones that occur in # unconstrained positions and are not already known to be # constrained. uncon.update(v for v in self.uncon_vars(cl) if v not in con) # Vice versa for the constrained vars. Uncons are processed # first so that "x in x" makes x unconstrained if it hasn't # been seen before. con.update(v for v in self.con_lhs_vars(cl) if v not in uncon) return tuple(con)
class IdentFinder(L.NodeVisitor): """Return an OrderedSet of all identifiers in the specified contexts. """ fun_ctxs = ('Fun.name', 'Call.func') query_ctxs = ('Query.name', 'ResetDemand.names') rel_ctxs = ('RelUpdate.rel', 'RelClear.rel', 'RelMember.rel') @classmethod def find_functions(cls, tree): return cls().run(tree, contexts=cls.fun_ctxs) @classmethod def find_vars(cls, tree): ctxs = (cls.fun_ctxs + cls.query_ctxs) return cls().run(tree, contexts=ctxs, invert=True) @classmethod def find_non_rel_uses(cls, tree): ctxs = (cls.fun_ctxs + cls.query_ctxs + cls.rel_ctxs) return cls().run(tree, contexts=ctxs, invert=True) def __init__(self, contexts=None, invert=False): if contexts is not None: for c in contexts: node_name, field_name = c.split('.') if not field_name in L.ident_fields.get(node_name, []): raise ValueError( 'Unknown identifier context "{}"'.format(c)) self.contexts = contexts """Collection of contexts to include/exclude. Each context is a string of the form '<node type name>.<field name>'. A value of None is equivalent to specifying all contexts. """ self.invert = bool(invert) """If True, find identifiers that occur in any context besides the ones given. """ def process(self, tree): self.names = OrderedSet() super().process(tree) return self.names def generic_visit(self, node): super().generic_visit(node) clsname = node.__class__.__name__ id_fields = L.ident_fields.get(clsname, []) for f in id_fields: inctx = (self.contexts is None or clsname + '.' + f in self.contexts) if inctx != self.invert: # Normalize for either one id or a sequence of ids. ids = getattr(node, f) if isinstance(ids, str): ids = [ids] if ids is not None: self.names.update(ids)
def con_lhs_vars_from_comp(self, comp): """Return a tuple of constrained variables. An variable is constrained if its first occurrence in the query is in a constrained position. For cyclic object queries like {(x, y) for y in x for x in y}, after translating into clauses over M, there are two possible sets of constrained vars: {x} and {y}. This function processes clauses left-to-right, so {y} will be chosen. """ uncon = OrderedSet() con = OrderedSet() for cl in comp.clauses: # The new unconstrained vars are the ones that occur in # unconstrained positions and are not already known to be # constrained. uncon.update(v for v in self.uncon_vars(cl) if v not in con) # Vice versa for the constrained vars. Uncons are processed # first so that "x in x" makes x unconstrained if it hasn't # been seen before. con.update(v for v in self.con_lhs_vars(cl) if v not in uncon) return tuple(con)
def lhs_vars_from_clauses(self, clauses): """Return a tuple of all LHS vars appearing in the given clauses, in order, without duplicates. """ vars = OrderedSet() for cl in clauses: vars.update(self.lhs_vars(cl)) return tuple(vars)
def lhs_vars_from_clauses(self, clauses): """Return a tuple of all LHS vars appearing in the given clauses, in order, without duplicates. """ vars = OrderedSet() for cl in clauses: vars.update(self.lhs_vars(cl)) return tuple(vars)
class AggrMaintainer(L.NodeTransformer): def __init__(self, fresh_vars, aggrinv): super().__init__() self.fresh_vars = fresh_vars self.aggrinv = aggrinv self.maint_funcs = OrderedSet() def visit_Module(self, node): node = self.generic_visit(node) fv = self.fresh_vars ops = [L.SetAdd(), L.SetRemove()] funcs = [] for op in ops: func = make_aggr_oper_maint_func(fv, self.aggrinv, op) funcs.append(func) if self.aggrinv.uses_demand: for op in ops: func = make_aggr_restr_maint_func(fv, self.aggrinv, op) funcs.append(func) func_names = L.get_defined_functions(tuple(funcs)) self.maint_funcs.update(func_names) node = node._replace(body=tuple(funcs) + node.body) return node def visit_RelUpdate(self, node): if not isinstance(node.op, (L.SetAdd, L.SetRemove)): return node if node.rel == self.aggrinv.rel: func = self.aggrinv.get_oper_maint_func_name(node.op) code = L.insert_rel_maint_call(node, func) elif self.aggrinv.uses_demand and node.rel == self.aggrinv.restr: func = self.aggrinv.get_restr_maint_func_name(node.op) code = L.insert_rel_maint_call(node, func) else: code = node return code def visit_RelClear(self, node): # We should clear if we are not using demand and our operand is # being cleared, or if we are using demand and our demand set is # being cleared. aggrinv = self.aggrinv uses_demand = aggrinv.uses_demand if uses_demand: should_clear = node.rel == aggrinv.restr else: should_clear = node.rel == aggrinv.rel if not should_clear: return node clear_code = (L.MapClear(self.aggrinv.map),) code = L.insert_rel_maint((node,), clear_code, L.SetRemove()) return code
class AggrMaintainer(L.NodeTransformer): def __init__(self, fresh_vars, aggrinv): super().__init__() self.fresh_vars = fresh_vars self.aggrinv = aggrinv self.maint_funcs = OrderedSet() def visit_Module(self, node): node = self.generic_visit(node) fv = self.fresh_vars ops = [L.SetAdd(), L.SetRemove()] funcs = [] for op in ops: func = make_aggr_oper_maint_func(fv, self.aggrinv, op) funcs.append(func) if self.aggrinv.uses_demand: for op in ops: func = make_aggr_restr_maint_func(fv, self.aggrinv, op) funcs.append(func) func_names = L.get_defined_functions(tuple(funcs)) self.maint_funcs.update(func_names) node = node._replace(body=tuple(funcs) + node.body) return node def visit_RelUpdate(self, node): if not isinstance(node.op, (L.SetAdd, L.SetRemove)): return node if node.rel == self.aggrinv.rel: func = self.aggrinv.get_oper_maint_func_name(node.op) code = L.insert_rel_maint_call(node, func) elif self.aggrinv.uses_demand and node.rel == self.aggrinv.restr: func = self.aggrinv.get_restr_maint_func_name(node.op) code = L.insert_rel_maint_call(node, func) else: code = node return code def visit_RelClear(self, node): # We should clear if we are not using demand and our operand is # being cleared, or if we are using demand and our demand set is # being cleared. aggrinv = self.aggrinv uses_demand = aggrinv.uses_demand if uses_demand: should_clear = node.rel == aggrinv.restr else: should_clear = node.rel == aggrinv.rel if not should_clear: return node clear_code = (L.MapClear(self.aggrinv.map), ) code = L.insert_rel_maint((node, ), clear_code, L.SetRemove()) return code
def prefix_locals(tree, prefix, extra_boundvars): """Rename the local variables in a block of code with the given prefix. The extra_boundvars are treated as additional variables known to be bound even if they don't appear in a binding occurrence within tree. """ localvars = OrderedSet(extra_boundvars) localvars.update(BindingFinder.run(tree)) tree = Templater.run(tree, {v: prefix + v for v in localvars}) return tree
def prefix_locals(tree, prefix, extra_boundvars): """Rename the local variables in a block of code with the given prefix. The extra_boundvars are treated as additional variables known to be bound even if they don't appear in a binding occurrence within tree. """ localvars = OrderedSet(extra_boundvars) localvars.update(BindingFinder.run(tree)) tree = Templater.run(tree, {v: prefix + v for v in localvars}) return tree
class Flattener(L.QueryMapper): def process(self, tree): self.trels = OrderedSet() tree = super().process(tree) return tree, self.trels def map_Comp(self, node): new_comp, new_trels = flatten_tuples_comp(node) self.trels.update(new_trels) return new_comp
class Flattener(L.QueryMapper): def process(self, tree): self.trels = OrderedSet() tree = super().process(tree) return tree, self.trels def map_Comp(self, node): new_comp, new_trels = flatten_tuples_comp(node) self.trels.update(new_trels) return new_comp
class Flattener(L.QueryMapper): def process(self, tree): self.use_mset = False self.fields = OrderedSet() self.use_mapset = False tree = super().process(tree) return tree, self.use_mset, self.fields, self.use_mapset def map_Comp(self, node): new_comp, new_mset, new_fields, new_mapset = \ flatten_comp(node, input_rels) self.use_mset |= new_mset self.fields.update(new_fields) self.use_mapset |= new_mapset return new_comp
class AggrMapCompRewriter(S.QueryRewriter): """Rewrite comprehension queries so that map lookups from aggregates are flattened into SetFromMap clauses. Return a pair of the new tree and a set of SetFromMap invariants that need to be transformed. """ def process(self, tree): self.sfm_invs = OrderedSet() tree = super().process(tree) return tree, self.sfm_invs def rewrite_comp(self, symbol, name, comp): rewriter = AggrMapReplacer(self.symtab.fresh_names.vars) comp = L.rewrite_comp(comp, rewriter.process) self.sfm_invs.update(rewriter.sfm_invs) return comp
class Flattener(L.QueryMapper): def process(self, tree): self.use_mset = False self.fields = OrderedSet() self.use_mapset = False tree = super().process(tree) return tree, self.use_mset, self.fields, self.use_mapset def map_Comp(self, node): new_comp, new_mset, new_fields, new_mapset = \ flatten_comp(node, input_rels) self.use_mset |= new_mset self.fields.update(new_fields) self.use_mapset |= new_mapset return new_comp
class AggrMapCompRewriter(S.QueryRewriter): """Rewrite comprehension queries so that map lookups from aggregates are flattened into SetFromMap clauses. Return a pair of the new tree and a set of SetFromMap invariants that need to be transformed. """ def process(self, tree): self.sfm_invs = OrderedSet() tree = super().process(tree) return tree, self.sfm_invs def rewrite_comp(self, symbol, name, comp): rewriter = AggrMapReplacer(self.symtab.fresh_names.vars) comp = L.rewrite_comp(comp, rewriter.process) self.sfm_invs.update(rewriter.sfm_invs) return comp
class BindingFinder(L.NodeVisitor): """Return names of variables that appear in a binding context, i.e., that would have a Store context in Python's AST. """ def process(self, tree): self.vars = OrderedSet() self.write_ctx = False super().process(tree) return self.vars def visit_Fun(self, node): self.vars.update(node.args) self.generic_visit(node) def visit_For(self, node): self.vars.add(node.target) self.generic_visit(node) def visit_DecompFor(self, node): self.vars.update(node.vars) self.generic_visit(node) def visit_Assign(self, node): self.vars.add(node.target) self.generic_visit(node) def visit_DecompAssign(self, node): self.vars.update(node.vars) self.generic_visit(node)
class BindingFinder(L.NodeVisitor): """Return names of variables that appear in a binding context, i.e., that would have a Store context in Python's AST. """ def process(self, tree): self.vars = OrderedSet() self.write_ctx = False super().process(tree) return self.vars def visit_Fun(self, node): self.vars.update(node.args) self.generic_visit(node) def visit_For(self, node): self.vars.add(node.target) self.generic_visit(node) def visit_DecompFor(self, node): self.vars.update(node.vars) self.generic_visit(node) def visit_Assign(self, node): self.vars.add(node.target) self.generic_visit(node) def visit_DecompAssign(self, node): self.vars.update(node.vars) self.generic_visit(node)
class InvariantTransformer(L.AdvNodeTransformer): """Insert maintenance functions, insert calls to these functions at updates, and replace expressions with uses of stored results. """ # There can be at most one SetFromMap mask for a given map. # Multiple distinct masks would have to differ on their arity, # which would be a type error. def __init__(self, fresh_vars, auxmaps, setfrommaps, wraps): super().__init__() self.fresh_vars = fresh_vars self.maint_funcs = OrderedSet() # Index over auxmaps for fast retrieval. self.auxmaps_by_rel = OrderedDict() self.auxmaps_by_relmask = OrderedDict() for auxmap in auxmaps: self.auxmaps_by_rel.setdefault(auxmap.rel, []).append(auxmap) # Index by relation, mask, and whether value is a singleton. # Not indexed by whether key is a singleton; if there are # two separate invariants that differ only on that, we may # end up only using one but maintaining both. stats = (auxmap.rel, auxmap.mask, auxmap.unwrap_value) self.auxmaps_by_relmask[stats] = auxmap # Index over setfrommaps. self.setfrommaps_by_map = sfm_by_map = OrderedDict() for sfm in setfrommaps: if sfm.map in sfm_by_map: raise L.ProgramError('Multiple SetFromMap invariants on ' 'same map {}'.format(sfm.map)) sfm_by_map[sfm.map] = sfm self.wraps_by_rel = OrderedDict() for wrap in wraps: self.wraps_by_rel.setdefault(wrap.oper, []).append(wrap) def in_loop_rhs(self): """Return True if the current node is the iter of a For node. Note that this checks the _visit_stack, and is sensitive to where in the recursive visiting we are. """ stack = self._visit_stack if len(stack) < 2: return False _current, field, _index = stack[-1] parent, _field, _index = stack[-2] return isinstance(parent, L.For) and field == 'iter' def visit_Module(self, node): node = self.generic_visit(node) funcs = [] for auxmaps in self.auxmaps_by_rel.values(): for auxmap in auxmaps: for op in [L.SetAdd(), L.SetRemove()]: func = make_auxmap_maint_func(self.fresh_vars, auxmap, op) funcs.append(func) for sfm in self.setfrommaps_by_map.values(): for op in ['assign', 'delete']: func = make_setfrommap_maint_func(self.fresh_vars, sfm, op) funcs.append(func) for wraps in self.wraps_by_rel.values(): for wrap in wraps: for op in [L.SetAdd(), L.SetRemove()]: func = make_wrap_maint_func(self.fresh_vars, wrap, op) funcs.append(func) func_names = L.get_defined_functions(tuple(funcs)) self.maint_funcs.update(func_names) node = node._replace(body=tuple(funcs) + node.body) return node def visit_RelUpdate(self, node): if not isinstance(node.op, (L.SetAdd, L.SetRemove)): return node code = (node,) auxmaps = self.auxmaps_by_rel.get(node.rel, set()) for auxmap in auxmaps: func_name = auxmap.get_maint_func_name(node.op) call_code = (L.Expr(L.Call(func_name, [L.Name(node.elem)])),) code = L.insert_rel_maint(code, call_code, node.op) wraps = self.wraps_by_rel.get(node.rel, set()) for wrap in wraps: func_name = wrap.get_maint_func_name(node.op) call_code = (L.Expr(L.Call(func_name, [L.Name(node.elem)])),) code = L.insert_rel_maint(code, call_code, node.op) return code def visit_RelClear(self, node): code = (node,) auxmaps = self.auxmaps_by_rel.get(node.rel, set()) for auxmap in auxmaps: clear_code = (L.MapClear(auxmap.map),) code = L.insert_rel_maint(code, clear_code, L.SetRemove()) wraps = self.wraps_by_rel.get(node.rel, set()) for wrap in wraps: clear_code = (L.RelClear(wrap.rel),) code = L.insert_rel_maint(code, clear_code, L.SetRemove()) return code def visit_MapAssign(self, node): sfm = self.setfrommaps_by_map.get(node.map, None) if sfm is None: return node code = (node,) func_name = sfm.get_maint_func_name('assign') call_code = (L.Expr(L.Call(func_name, [L.Name(node.key), L.Name(node.value)])),) code = L.insert_rel_maint(code, call_code, L.SetAdd()) return code def visit_MapDelete(self, node): sfm = self.setfrommaps_by_map.get(node.map, None) if sfm is None: return node code = (node,) func_name = sfm.get_maint_func_name('delete') call_code = (L.Expr(L.Call(func_name, [L.Name(node.key)])),) code = L.insert_rel_maint(code, call_code, L.SetRemove()) return code def visit_MapClear(self, node): sfm = self.setfrommaps_by_map.get(node.map, None) if sfm is None: return node code = (node,) clear_code = (L.RelClear(sfm.rel),) code = L.insert_rel_maint(code, clear_code, L.SetRemove()) return code def imglookup_helper(self, node, *, in_unwrap, on_loop_rhs): """Return the replacement for an ImgLookup node, or None if no replacement is applicable. """ if not isinstance(node.set, L.Name): return None rel = node.set.id stats = (rel, node.mask, in_unwrap) auxmap = self.auxmaps_by_relmask.get(stats, None) if auxmap is None: return None key = L.tuplify(node.bounds, unwrap=auxmap.unwrap_key) empty = L.Parser.pe('()') if on_loop_rhs else L.Parser.pe('Set()') return L.Parser.pe('_MAP[_KEY] if _KEY in _MAP else _EMPTY', subst={'_MAP': auxmap.map, '_KEY': key, '_EMPTY': empty}) def visit_ImgLookup(self, node): node = self.generic_visit(node) repl = self.imglookup_helper(node, in_unwrap=False, on_loop_rhs=self.in_loop_rhs()) if repl is not None: node = repl return node def visit_SetFromMap(self, node): node = self.generic_visit(node) if not isinstance(node.map, L.Name): return node map = node.map.id sfm = self.setfrommaps_by_map.get(map, None) if sfm is None: return node if not sfm.mask == node.mask: raise L.ProgramError('Multiple SetFromMap expressions on ' 'same map {}'.format(map)) return L.Name(sfm.rel) def wrap_helper(self, node): """Process a wrap invariant at a Wrap or Unwrap node. Don't recurse. """ if not isinstance(node.value, L.Name): return node rel = node.value.id wraps = self.wraps_by_rel.get(rel, []) for wrap in wraps: if ((isinstance(node, L.Wrap) and not wrap.unwrap) or (isinstance(node, L.Unwrap) and wrap.unwrap)): return L.Name(wrap.rel) return node def visit_Unwrap(self, node): # Handle special case of an auxmap with the unwrap_value flag. if isinstance(node.value, L.ImgLookup): # Recurse over children below the ImgLookup. value = self.generic_visit(node.value) node = node._replace(value=value) # See if an auxmap invariant applies. repl = self.imglookup_helper(value, in_unwrap=True, on_loop_rhs=self.in_loop_rhs()) if repl is not None: return repl else: # Don't run in case we already did generic_visit() above # and didn't return. node = self.generic_visit(node) return self.wrap_helper(node) def visit_Wrap(self, node): node = self.generic_visit(node) node = self.wrap_helper(node) return node
class CompMaintainer(L.NodeTransformer): """Insert comprehension maintenance functions and calls to these functions at relevant updates. """ def __init__(self, clausetools, fresh_vars, fresh_join_names, comp, result_var, *, counted): super().__init__() self.clausetools = clausetools self.fresh_vars = fresh_vars self.fresh_join_names = fresh_join_names self.comp = comp self.result_var = result_var self.counted = counted self.rels = self.clausetools.rhs_rels_from_comp(self.comp) self.maint_funcs = OrderedSet() def visit_Module(self, node): ct = self.clausetools node = self.generic_visit(node) funcs = [] for rel in self.rels: for op in [L.SetAdd(), L.SetRemove()]: fresh_var_prefix = next(self.fresh_vars) func = make_comp_maint_func( ct, fresh_var_prefix, self.fresh_join_names, self.comp, self.result_var, rel, op, counted=self.counted) funcs.append(func) func_names = L.get_defined_functions(tuple(funcs)) self.maint_funcs.update(func_names) node = node._replace(body=tuple(funcs) + node.body) return node def visit_RelUpdate(self, node): if not isinstance(node.op, (L.SetAdd, L.SetRemove)): return node if node.rel not in self.rels: return node op_name = L.set_update_name(node.op) func_name = N.get_maint_func_name(self.result_var, node.rel, op_name) code = (node,) call_code = (L.Expr(L.Call(func_name, [L.Name(node.elem)])),) code = L.insert_rel_maint(code, call_code, node.op) return code def visit_RelClear(self, node): if node.rel not in self.rels: return node code = (node,) clear_code = (L.RelClear(self.result_var),) code = L.insert_rel_maint(code, clear_code, L.SetRemove()) return code
class InvariantTransformer(L.AdvNodeTransformer): """Insert maintenance functions, insert calls to these functions at updates, and replace expressions with uses of stored results. """ # There can be at most one SetFromMap mask for a given map. # Multiple distinct masks would have to differ on their arity, # which would be a type error. def __init__(self, fresh_vars, auxmaps, setfrommaps, wraps): super().__init__() self.fresh_vars = fresh_vars self.maint_funcs = OrderedSet() # Index over auxmaps for fast retrieval. self.auxmaps_by_rel = OrderedDict() self.auxmaps_by_relmask = OrderedDict() for auxmap in auxmaps: self.auxmaps_by_rel.setdefault(auxmap.rel, []).append(auxmap) # Index by relation, mask, and whether value is a singleton. # Not indexed by whether key is a singleton; if there are # two separate invariants that differ only on that, we may # end up only using one but maintaining both. stats = (auxmap.rel, auxmap.mask, auxmap.unwrap_value) self.auxmaps_by_relmask[stats] = auxmap # Index over setfrommaps. self.setfrommaps_by_map = sfm_by_map = OrderedDict() for sfm in setfrommaps: if sfm.map in sfm_by_map: raise L.ProgramError('Multiple SetFromMap invariants on ' 'same map {}'.format(sfm.map)) sfm_by_map[sfm.map] = sfm self.wraps_by_rel = OrderedDict() for wrap in wraps: self.wraps_by_rel.setdefault(wrap.oper, []).append(wrap) def in_loop_rhs(self): """Return True if the current node is the iter of a For node. Note that this checks the _visit_stack, and is sensitive to where in the recursive visiting we are. """ stack = self._visit_stack if len(stack) < 2: return False _current, field, _index = stack[-1] parent, _field, _index = stack[-2] return isinstance(parent, L.For) and field == 'iter' def visit_Module(self, node): node = self.generic_visit(node) funcs = [] for auxmaps in self.auxmaps_by_rel.values(): for auxmap in auxmaps: for op in [L.SetAdd(), L.SetRemove()]: func = make_auxmap_maint_func(self.fresh_vars, auxmap, op) funcs.append(func) for sfm in self.setfrommaps_by_map.values(): for op in ['assign', 'delete']: func = make_setfrommap_maint_func(self.fresh_vars, sfm, op) funcs.append(func) for wraps in self.wraps_by_rel.values(): for wrap in wraps: for op in [L.SetAdd(), L.SetRemove()]: func = make_wrap_maint_func(self.fresh_vars, wrap, op) funcs.append(func) func_names = L.get_defined_functions(tuple(funcs)) self.maint_funcs.update(func_names) node = node._replace(body=tuple(funcs) + node.body) return node def visit_RelUpdate(self, node): if not isinstance(node.op, (L.SetAdd, L.SetRemove)): return node code = (node, ) auxmaps = self.auxmaps_by_rel.get(node.rel, set()) for auxmap in auxmaps: func_name = auxmap.get_maint_func_name(node.op) call_code = (L.Expr(L.Call(func_name, [L.Name(node.elem)])), ) code = L.insert_rel_maint(code, call_code, node.op) wraps = self.wraps_by_rel.get(node.rel, set()) for wrap in wraps: func_name = wrap.get_maint_func_name(node.op) call_code = (L.Expr(L.Call(func_name, [L.Name(node.elem)])), ) code = L.insert_rel_maint(code, call_code, node.op) return code def visit_RelClear(self, node): code = (node, ) auxmaps = self.auxmaps_by_rel.get(node.rel, set()) for auxmap in auxmaps: clear_code = (L.MapClear(auxmap.map), ) code = L.insert_rel_maint(code, clear_code, L.SetRemove()) wraps = self.wraps_by_rel.get(node.rel, set()) for wrap in wraps: clear_code = (L.RelClear(wrap.rel), ) code = L.insert_rel_maint(code, clear_code, L.SetRemove()) return code def visit_MapAssign(self, node): sfm = self.setfrommaps_by_map.get(node.map, None) if sfm is None: return node code = (node, ) func_name = sfm.get_maint_func_name('assign') call_code = (L.Expr( L.Call(func_name, [L.Name(node.key), L.Name(node.value)])), ) code = L.insert_rel_maint(code, call_code, L.SetAdd()) return code def visit_MapDelete(self, node): sfm = self.setfrommaps_by_map.get(node.map, None) if sfm is None: return node code = (node, ) func_name = sfm.get_maint_func_name('delete') call_code = (L.Expr(L.Call(func_name, [L.Name(node.key)])), ) code = L.insert_rel_maint(code, call_code, L.SetRemove()) return code def visit_MapClear(self, node): sfm = self.setfrommaps_by_map.get(node.map, None) if sfm is None: return node code = (node, ) clear_code = (L.RelClear(sfm.rel), ) code = L.insert_rel_maint(code, clear_code, L.SetRemove()) return code def imglookup_helper(self, node, *, in_unwrap, on_loop_rhs): """Return the replacement for an ImgLookup node, or None if no replacement is applicable. """ if not isinstance(node.set, L.Name): return None rel = node.set.id stats = (rel, node.mask, in_unwrap) auxmap = self.auxmaps_by_relmask.get(stats, None) if auxmap is None: return None key = L.tuplify(node.bounds, unwrap=auxmap.unwrap_key) empty = L.Parser.pe('()') if on_loop_rhs else L.Parser.pe('Set()') return L.Parser.pe('_MAP[_KEY] if _KEY in _MAP else _EMPTY', subst={ '_MAP': auxmap.map, '_KEY': key, '_EMPTY': empty }) def visit_ImgLookup(self, node): node = self.generic_visit(node) repl = self.imglookup_helper(node, in_unwrap=False, on_loop_rhs=self.in_loop_rhs()) if repl is not None: node = repl return node def visit_SetFromMap(self, node): node = self.generic_visit(node) if not isinstance(node.map, L.Name): return node map = node.map.id sfm = self.setfrommaps_by_map.get(map, None) if sfm is None: return node if not sfm.mask == node.mask: raise L.ProgramError('Multiple SetFromMap expressions on ' 'same map {}'.format(map)) return L.Name(sfm.rel) def wrap_helper(self, node): """Process a wrap invariant at a Wrap or Unwrap node. Don't recurse. """ if not isinstance(node.value, L.Name): return node rel = node.value.id wraps = self.wraps_by_rel.get(rel, []) for wrap in wraps: if ((isinstance(node, L.Wrap) and not wrap.unwrap) or (isinstance(node, L.Unwrap) and wrap.unwrap)): return L.Name(wrap.rel) return node def visit_Unwrap(self, node): # Handle special case of an auxmap with the unwrap_value flag. if isinstance(node.value, L.ImgLookup): # Recurse over children below the ImgLookup. value = self.generic_visit(node.value) node = node._replace(value=value) # See if an auxmap invariant applies. repl = self.imglookup_helper(value, in_unwrap=True, on_loop_rhs=self.in_loop_rhs()) if repl is not None: return repl else: # Don't run in case we already did generic_visit() above # and didn't return. node = self.generic_visit(node) return self.wrap_helper(node) def visit_Wrap(self, node): node = self.generic_visit(node) node = self.wrap_helper(node) return node
class VarsFinder(NodeVisitor): """Simple finder of variables (Name nodes). Flags: ignore_store: Name nodes with Store context are ignored, as are update operations (e.g. SetUpdate). As an exception, Name nodes on the LHS of Enumerators are not ignored. This is to ensure safety under pattern matching semantics. ignore_functions: Name nodes that appear to be functions are ignored. ignore_rels: Names that appear to be relations are ignored. The builtins None, True, and False are always excluded, as they are NameConstants, not variables. """ def __init__(self, *, ignore_store=False, ignore_functions=False, ignore_rels=False): super().__init__() self.ignore_store = ignore_store self.ignore_functions = ignore_functions self.ignore_rels = ignore_rels def process(self, tree): self.usedvars = OrderedSet() super().process(tree) return self.usedvars def visit_Name(self, node): self.generic_visit(node) if not (self.ignore_store and isinstance(node.ctx, Store)): self.usedvars.add(node.id) def visit_Call(self, node): class IGNORE(iast.AST): _meta = True if isinstance(node.func, Name) and self.ignore_functions: self.generic_visit(node._replace(func=IGNORE())) else: self.generic_visit(node) def visit_Enumerator(self, node): if is_vartuple(node.target): # Bypass ignore_store. vars = get_vartuple(node.target) self.usedvars.update(vars) else: self.visit(node.target) if not (self.ignore_rels and isinstance(node.iter, Name)): self.visit(node.iter) def visit_Comp(self, node): self.visit(node.resexp) # Hack to ensure we don't grab rels on RHS of # membership conditions. for i in node.clauses: if (self.ignore_rels and isinstance(i, Compare) and len(i.ops) == len(i.comparators) == 1 and isinstance(i.ops[0], In) and isinstance(i.comparators[0], Name)): self.visit(i.left) else: self.visit(i) def visit_Aggregate(self, node): if isinstance(node.value, Name) and self.ignore_rels: return else: self.generic_visit(node) def visit_SetMatch(self, node): if isinstance(node.target, Name) and self.ignore_rels: self.visit(node.key) else: self.generic_visit(node) def update_helper(self, node): IGNORE = object() if isinstance(node.target, Name) and self.ignore_store: self.generic_visit(node._replace(target=IGNORE)) else: self.generic_visit(node) visit_SetUpdate = update_helper visit_RCSetRefUpdate = update_helper visit_AssignKey = update_helper visit_DelKey = update_helper
class CompMaintainer(L.NodeTransformer): """Insert comprehension maintenance functions and calls to these functions at relevant updates. """ def __init__(self, clausetools, fresh_vars, fresh_join_names, comp, result_var, *, counted): super().__init__() self.clausetools = clausetools self.fresh_vars = fresh_vars self.fresh_join_names = fresh_join_names self.comp = comp self.result_var = result_var self.counted = counted self.rels = self.clausetools.rhs_rels_from_comp(self.comp) self.maint_funcs = OrderedSet() def visit_Module(self, node): ct = self.clausetools node = self.generic_visit(node) funcs = [] for rel in self.rels: for op in [L.SetAdd(), L.SetRemove()]: fresh_var_prefix = next(self.fresh_vars) func = make_comp_maint_func(ct, fresh_var_prefix, self.fresh_join_names, self.comp, self.result_var, rel, op, counted=self.counted) funcs.append(func) func_names = L.get_defined_functions(tuple(funcs)) self.maint_funcs.update(func_names) node = node._replace(body=tuple(funcs) + node.body) return node def visit_RelUpdate(self, node): if not isinstance(node.op, (L.SetAdd, L.SetRemove)): return node if node.rel not in self.rels: return node op_name = L.set_update_name(node.op) func_name = N.get_maint_func_name(self.result_var, node.rel, op_name) code = (node, ) call_code = (L.Expr(L.Call(func_name, [L.Name(node.elem)])), ) code = L.insert_rel_maint(code, call_code, node.op) return code def visit_RelClear(self, node): if node.rel not in self.rels: return node code = (node, ) clear_code = (L.RelClear(self.result_var), ) code = L.insert_rel_maint(code, clear_code, L.SetRemove()) return code
class VarsFinder(NodeVisitor): """Simple finder of variables (Name nodes). Flags: ignore_store: Name nodes with Store context are ignored, as are update operations (e.g. SetUpdate). As an exception, Name nodes on the LHS of Enumerators are not ignored. This is to ensure safety under pattern matching semantics. ignore_functions: Name nodes that appear to be functions are ignored. ignore_rels: Names that appear to be relations are ignored. The builtins None, True, and False are always excluded, as they are NameConstants, not variables. """ def __init__(self, *, ignore_store=False, ignore_functions=False, ignore_rels=False): super().__init__() self.ignore_store = ignore_store self.ignore_functions = ignore_functions self.ignore_rels = ignore_rels def process(self, tree): self.usedvars = OrderedSet() super().process(tree) return self.usedvars def visit_Name(self, node): self.generic_visit(node) if not (self.ignore_store and isinstance(node.ctx, Store)): self.usedvars.add(node.id) def visit_Call(self, node): class IGNORE(iast.AST): _meta = True if isinstance(node.func, Name) and self.ignore_functions: self.generic_visit(node._replace(func=IGNORE())) else: self.generic_visit(node) def visit_Enumerator(self, node): if is_vartuple(node.target): # Bypass ignore_store. vars = get_vartuple(node.target) self.usedvars.update(vars) else: self.visit(node.target) if not (self.ignore_rels and isinstance(node.iter, Name)): self.visit(node.iter) def visit_Comp(self, node): self.visit(node.resexp) # Hack to ensure we don't grab rels on RHS of # membership conditions. for i in node.clauses: if (self.ignore_rels and isinstance(i, Compare) and len(i.ops) == len(i.comparators) == 1 and isinstance(i.ops[0], In) and isinstance(i.comparators[0], Name)): self.visit(i.left) else: self.visit(i) def visit_Aggregate(self, node): if isinstance(node.value, Name) and self.ignore_rels: return else: self.generic_visit(node) def visit_SetMatch(self, node): if isinstance(node.target, Name) and self.ignore_rels: self.visit(node.key) else: self.generic_visit(node) def update_helper(self, node): IGNORE = object() if isinstance(node.target, Name) and self.ignore_store: self.generic_visit(node._replace(target=IGNORE)) else: self.generic_visit(node) visit_SetUpdate = update_helper visit_RCSetRefUpdate = update_helper visit_AssignKey = update_helper visit_DelKey = update_helper
class IdentFinder(L.NodeVisitor): """Return an OrderedSet of all identifiers in the specified contexts. """ fun_ctxs = ('Fun.name', 'Call.func') query_ctxs = ('Query.name', 'ResetDemand.names') rel_ctxs = ('RelUpdate.rel', 'RelClear.rel', 'RelMember.rel') @classmethod def find_functions(cls, tree): return cls().run(tree, contexts=cls.fun_ctxs) @classmethod def find_vars(cls, tree): ctxs = (cls.fun_ctxs + cls.query_ctxs) return cls().run(tree, contexts=ctxs, invert=True) @classmethod def find_non_rel_uses(cls, tree): ctxs = (cls.fun_ctxs + cls.query_ctxs + cls.rel_ctxs) return cls().run(tree, contexts=ctxs, invert=True) def __init__(self, contexts=None, invert=False): if contexts is not None: for c in contexts: node_name, field_name = c.split('.') if not field_name in L.ident_fields.get(node_name, []): raise ValueError('Unknown identifier context "{}"' .format(c)) self.contexts = contexts """Collection of contexts to include/exclude. Each context is a string of the form '<node type name>.<field name>'. A value of None is equivalent to specifying all contexts. """ self.invert = bool(invert) """If True, find identifiers that occur in any context besides the ones given. """ def process(self, tree): self.names = OrderedSet() super().process(tree) return self.names def generic_visit(self, node): super().generic_visit(node) clsname = node.__class__.__name__ id_fields = L.ident_fields.get(clsname, []) for f in id_fields: inctx = (self.contexts is None or clsname + '.' + f in self.contexts) if inctx != self.invert: # Normalize for either one id or a sequence of ids. ids = getattr(node, f) if isinstance(ids, str): ids = [ids] if ids is not None: self.names.update(ids)