relationalai 0.12.13__py3-none-any.whl → 0.13.0.dev0__py3-none-any.whl

relationalai/semantics/metamodel/helpers.py

@@ -1,361 +0,0 @@
-"""
-Helpers to analyze the metamodel IR.
-"""
-from __future__ import annotations
-
-import re
-from dataclasses import fields
-from typing import cast, Tuple, Iterable, Optional, TypeVar
-from relationalai.semantics.metamodel import ir, visitor, builtins, types, factory as f
-from relationalai.semantics.metamodel.util import NameCache, OrderedSet, FrozenOrderedSet, flatten_tuple, ordered_set
-
-
-
-#--------------------------------------------------
-# Name helpers
-#--------------------------------------------------
-
-def sanitize(name:str) -> str:
-    """ Cleanup the name to make it more palatable to names. """
-    x = re.sub(r"[ ,\.\(\)\|]", "_", name)
-    return x[0:-1] if x[-1] == "_" else x
-
-#--------------------------------------------------
-# Checks
-#--------------------------------------------------
-
-def is_concept_lookup(node: ir.Lookup|ir.Relation):
-    """ Whether this task is a concept lookup. """
-    if isinstance(node, ir.Lookup) and is_concept_lookup(node.relation):
-        return True
-    return builtins.concept_relation_annotation in node.annotations
-
-def is_external(relation: ir.Relation):
-   """ Whether this relation is external, by being marked with the external annotation. """
-   return builtins.external_annotation in relation.annotations
-
-def is_from_cast(node: ir.Lookup|ir.Relation):
-    """ Whether this relation is from cast, by being marked with the from_cast_annotation annotation. """
-    if isinstance(node, ir.Lookup) and is_from_cast(node.relation):
-        return True
-    return builtins.from_cast_annotation in node.annotations
-
-def is_aggregate_input(var: ir.Var, agg: ir.Aggregate):
-    """ Whether this var is an input to this aggregation. """
-    return (var in agg.args and agg.aggregation.fields[agg.args.index(var)].input)
-
-def is_effective_logical(n: ir.Task):
-    """ Whether this task is a Logical and contains an Update child, recursively. """
-    return isinstance(n, ir.Logical) and len(visitor.collect_by_type(ir.Update, n)) > 0
-
-def is_nullable_logical(n: ir.Task):
-    """ Whether this task is a Logical that contains a hoisted variable with a None default. """
-    return isinstance(n, ir.Logical) and any(isinstance(v, ir.Default) and v.value is None for v in n.hoisted)
-
-def relation_is_subtype(r1: ir.Relation, r2: ir.Relation):
-    if r1 is r2:
-        return True
-    if len(r1.fields) != len(r2.fields):
-        return False
-    return all(types.is_subtype(f1.type, f2.type) for f1, f2 in zip(r1.fields, r2.fields))
-
-def relation_is_proper_subtype(r1: ir.Relation, r2: ir.Relation):
-    if r1 is r2:
-        return False
-    if relation_is_subtype(r1, r2):
-        return any(types.is_proper_subtype(f1.type, f2.type) for f1, f2 in zip(r1.fields, r2.fields))
-    else:
-        return False
-
-def relation_name_prefix(relation: ir.Relation):
-    prefix = ""
-    if len(relation.fields) > 0 and not is_concept_lookup(relation):
-        main_type = relation.fields[0].type
-        if (isinstance(main_type, ir.ScalarType) and
-            main_type not in types.builtin_types and
-            not relation.name.startswith(main_type.name) and
-            not relation.name.startswith("_")):
-                prefix = f"{main_type.name.lower()}_"
-    return prefix
-
-def get_outputs(lookup: ir.Lookup):
-    """
-        Return an array with the arguments of this lookup that are referring to output fields
-        on the relation being looked up.
-    """
-    if builtins.is_eq(lookup.relation):
-        # special case eq because it can be input or output
-        x, y = lookup.args[0], lookup.args[1]
-        if isinstance(x, ir.Var) and not isinstance(y, ir.Var):
-            return [x]
-        elif not isinstance(x, ir.Var) and isinstance(y, ir.Var):
-            return [y]
-        # both are inputs
-        return []
-    else:
-        outputs = []
-        # register variables depending on the input flag of the relation bound to the lookup
-        for idx, fld in enumerate(lookup.relation.fields):
-            arg = lookup.args[idx]
-            if isinstance(arg, Iterable):
-                # deal with ListType fields that pack arguments in a tuple
-                for element in arg:
-                    if isinstance(element, ir.Var) and not fld.input:
-                        outputs.append(element)
-            else:
-                if isinstance(arg, ir.Var) and not fld.input:
-                    outputs.append(arg)
-        return outputs
-
-def get_agg_outputs(lookup: ir.Aggregate):
-    """
-        Return an array with the arguments of this aggregate that are referring to output fields
-        on the relation being looked up.
-    """
-    outputs = []
-    # register variables depending on the input flag of the relation bound to the lookup
-    for idx, fld in enumerate(lookup.aggregation.fields):
-        arg = lookup.args[idx]
-        if isinstance(arg, Iterable):
-            # deal with ListType fields that pack arguments in a tuple
-            for element in arg:
-                if isinstance(element, ir.Var) and not fld.input:
-                    outputs.append(element)
-        else:
-            if isinstance(arg, ir.Var) and not fld.input:
-                outputs.append(arg)
-    return outputs
-
-#--------------------------------------------------
-# Filters
-#--------------------------------------------------
-
-def aggregate_outputs(agg: ir.Aggregate) -> list[ir.Var]:
-    """ Get the list of vars bound to the outputs of this aggregation. """
-    return list(filter(lambda arg: isinstance(arg, ir.Var) and not is_aggregate_input(arg, agg), agg.args)) # type: ignore
-
-def aggregate_inputs(agg: ir.Aggregate) -> list[ir.Var]:
-    """ Get the list of vars bound to the args that are inputs of this aggregation. """
-    return list(filter(lambda arg: isinstance(arg, ir.Var) and is_aggregate_input(arg, agg), agg.args)) # type: ignore
-
-def effective_logicals(tasks: OrderedSet[ir.Task]) -> OrderedSet[ir.Logical]:
-    """ Filter tasks to return only the Logical tasks that are effective. """
-    return OrderedSet.from_iterable(filter(lambda t: is_effective_logical(t), tasks))
-
-def nullable_logicals(tasks: OrderedSet[ir.Task]) -> OrderedSet[ir.Logical]:
-    """ Filter tasks to return only the Logical tasks that are nullable. """
-    return OrderedSet.from_iterable(filter(lambda t: is_nullable_logical(t), tasks))
-
-def hoisted_vars(hoisted: Iterable[ir.VarOrDefault]) -> list[ir.Var]:
-    """ Extract the vars from defaults in the hoisted list, returning just Vars. """
-    return [hoisted_var(v) for v in hoisted]
-
-def hoisted_var(hoisted: ir.VarOrDefault) -> ir.Var:
-    """ Extract the var from VarOrDefault, returning just Var. """
-    return hoisted.var if isinstance(hoisted, ir.Default) else hoisted
-
-def vars(args: Tuple[ir.Value, ...]) -> list[ir.Var]:
-    """ Filter this list of values, keeping only Vars. """
-    return cast(list[ir.Var], list(filter(lambda v: isinstance(v, ir.Var), flatten_tuple(args, ir.Value))))
-
-def output_vars(aliases: FrozenOrderedSet[Tuple[str, ir.Value]]) -> list[ir.Var]:
-    return [alias[1] for alias in aliases if isinstance(alias[1], ir.Var)]
-
-def output_values(aliases: FrozenOrderedSet[Tuple[str, ir.Value]]) -> list[ir.Value]:
-    return [alias[1] for alias in aliases]
-
-def output_alias_names(aliases: FrozenOrderedSet[Tuple[str, ir.Value]]) -> list[str]:
-    return [alias[0] for alias in aliases]
-
-#--------------------------------------------------
-# Visitors/Collectors
-#--------------------------------------------------
-
-def collect_vars(*nodes: ir.Node) -> OrderedSet[ir.Var]:
-    """ Collect all Vars starting at this node. """
-    return cast(OrderedSet[ir.Var],
-        visitor.collect_by_type(ir.Var, *nodes)
-    )
-
-def collect_quantified_vars(*nodes: ir.Node) -> OrderedSet[ir.Var]:
-    """ Collect all Vars that are children of Exists and ForAll. """
-    return cast(OrderedSet[ir.Var],
-        visitor.collect(
-            lambda n, parent: isinstance(n, ir.Var) and isinstance(parent, (ir.Exists, ir.ForAll)),
-            *nodes)
-    )
-
-def collect_aggregate_vars(*nodes: ir.Node) -> OrderedSet[ir.Var]:
-    """ Collect vars that are declared by aggregates in Rel (projection + over). """
-    return cast(OrderedSet[ir.Var],
-        # TODO - when dealing with multiple aggregations we will need to consider groupbys
-        visitor.collect(
-            lambda n, parent:
-                # parent of var is an aggregate and either
-                isinstance(parent, ir.Aggregate) and isinstance(n, ir.Var) and (
-                # var is in the projection
-                n in parent.projection or
-                # var is an input to the aggregation's relation
-                is_aggregate_input(n, parent)
-                ),
-            *nodes)
-    )
-
-def collect_rank_vars(*nodes: ir.Node) -> OrderedSet[ir.Var]:
-    """ Collect vars that are declared by ranks in Rel (projection + over). """
-    return cast(OrderedSet[ir.Var],
-        # TODO - when dealing with multiple aggregations we will need to consider groupbys
-        visitor.collect(
-            lambda n, parent:
-                # parent of var is an aggregate and either
-                isinstance(parent, ir.Rank) and isinstance(n, ir.Var) and (
-                # var is in the projection
-                n in parent.projection or
-                n in parent.args
-                ),
-            *nodes)
-    )
-
-def collect_implicit_vars(*nodes: ir.Node) -> OrderedSet[ir.Var]:
-    """ Collect vars except the quantified and aggregate vars. """
-    if not nodes:
-        return ordered_set()
-    return collect_vars(*nodes) - collect_quantified_vars(*nodes) - collect_aggregate_vars(*nodes) - collect_rank_vars(*nodes)
-
-#--------------------------------------------------
-# Useful node categories
-#--------------------------------------------------
-
-BINDERS = (ir.Lookup, ir.Construct, ir.Aggregate, ir.Exists, ir.Data, ir.Not)
-COMPOSITES = (ir.Logical, ir.Sequence, ir.Union, ir.Match, ir.Until, ir.Wait)
-EFFECTS = (ir.Update, ir.Output)
-
-#--------------------------------------------------
-# Helper classes
-#--------------------------------------------------
-
-class RewriteContext():
-    """
-    Container of information collected during a rewrite pass.
-    """
-    def __init__(self):
-        # the logicals that will be at the top level at the end of the rewrite
-        self.top_level: list[ir.Logical] = []
-        # new relations created during the pass
-        self.relations: list[ir.Relation] = []
-
-
-#--------------------------------------------------
-# Rewrite helpers
-#--------------------------------------------------
-
-def extract(task: ir.Task, body: OrderedSet[ir.Task], exposed_vars: list[ir.Var], ctx: RewriteContext, name: str) -> ir.Relation:
-    """
-    Extract into this Analysiscontext a new top level Logical that contains this body plus a
-    derive task into a new temporary relation, which is also registered with the ctx.
-    The exposed_vars determine the arguments of this temporary relation. The prefix
-    can be used to customize the name of the relation, which defaults to the task kind.
-
-    Return the temporary relation created for the extraction.
-    """
-    connection = create_connection_relation(task, exposed_vars, ctx, name)
-
-    # add derivation to the extracted body
-    body.add(f.derive(connection, exposed_vars))
-
-    # extract the body
-    ctx.top_level.append(clone_task(ir.Logical(task.engine, tuple(), tuple(body))))
-
-    return connection
-
-def create_connection_relation(task: ir.Task, exposed_vars: list[ir.Var], ctx: RewriteContext, name: str) -> ir.Relation:
-    """
-    Create a new relation with a name based off this task, with fields that represent
-    the types and names of these exposed vars, and register in the context.
-    """
-    connection = f.relation(name, [f.field(v.name, v.type) for v in exposed_vars])
-    ctx.relations.append(connection)
-
-    return connection
-
-
-def create_task_name(name_cache: NameCache, task: ir.Task, prefix: Optional[str]=None) -> str:
-    """ Helper to generate consistent names for tasks extracted from a logical. """
-    prefix = prefix if prefix else f"_{task.kind}"
-    return name_cache.get_name(task.id, prefix)
-
-
-CLONABLE = (ir.Var, ir.Default, ir.Task)
-T = TypeVar('T', bound=ir.Task)
-def clone_task(task: T) -> T:
-    """
-    Create a new task that is a clone of this task. This operation clones only sub-tasks
-    and variables, and preserves variable references.
-
-    This is useful when we are rewriting the metamodel and want to copy parts of a task to
-    some other place. It is important to clone to avoid having the same object present
-    multiple times in the metamodel.
-    """
-
-    # map from original object id to the rewritten object
-    cache = {}
-    def from_cache(original):
-        """ Lookup this object from the cache above, dealing with collections and with
-        objects that were not rewritten. """
-        if isinstance(original, tuple):
-            return tuple([from_cache(c) for c in original])
-        elif isinstance(original, FrozenOrderedSet):
-            return ordered_set(*[from_cache(c) for c in original])
-        elif isinstance(original, CLONABLE):
-            return cache.get(original.id, original)
-        else:
-            return original
-
-    # the last node that was processed and rewritten
-    prev_node = None
-    stack: list[ir.Node] = [task]
-    def to_stack(original):
-        """ Add this original node to the stack if it is clonable and was never processed;
-        return True iff the node was added to the stack. """
-        if isinstance(original, CLONABLE) and original.id not in cache:
-            stack.append(original)
-            return True
-        return False
-
-    while stack:
-        # peek the current node and get the initializable fields (i.e. ignore Node id)
-        curr = stack[-1]
-        curr_fields = list(filter(lambda f: f.init, fields(curr)))
-        stacked_children = False
-
-        # go over the fields adding to the stack the ones that we need to rewrite
-        for field in curr_fields:
-            field_value = getattr(curr, field.name)
-            if isinstance(field_value, (tuple, FrozenOrderedSet)):
-                # node field is a collection (tuple or set)
-                for s in field_value:
-                    if isinstance(s, tuple):
-                        # the value within the collection is a tuple (can happen for lookup args)
-                        for c in s:
-                            stacked_children = to_stack(c) or stacked_children
-                    else:
-                        # the value within the collection is a scalar
-                        stacked_children = to_stack(s) or stacked_children
-            else:
-                # node field is a scalar
-                stacked_children = to_stack(field_value) or stacked_children
-
-        # if no childrean were stacked, we rewrote all fields of curr, so we can pop it and rewrite it
-        if not stacked_children:
-            stack.pop()
-            if curr.id not in cache:
-                children = []
-                for f in curr_fields:
-                    children.append(from_cache(getattr(curr, f.name)))
-                # create a new prev_node with the cloned children
-                prev_node = curr.__class__(*children)
-                cache[curr.id] = prev_node
-
-    # the last node we processed is the rewritten original node
-    assert(isinstance(prev_node, type(task)))
-    return prev_node

relationalai/semantics/metamodel/rewrite/discharge_constraints.py

@@ -1,39 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from relationalai.semantics.metamodel import ir, compiler as c, visitor as v, builtins
-from relationalai.semantics.metamodel.util import rewrite_list
-
-
-class DischargeConstraints(c.Pass):
-    """
-    Since we should not generate code for `unique`, `exclusive`, `anyof` builtins all Require Check nodes marked with
-    `discharge` annotation when at least one of the builtins is presented in a check.
-    All Require/Check ir nodes marked with `discharge` annotation will be removed from the IR model in Flatten pass.
-    """
-
-    def rewrite(self, model: ir.Model, options: dict = {}) -> ir.Model:
-        return DischargeConstraintsVisitor().walk(model)
-
-
-@dataclass
-class DischargeConstraintsVisitor(v.Rewriter):
-    """
-    Visitor marks all nodes which should be removed from IR model with `discharge` annotation.
-    """
-
-    def handle_require(self, node: ir.Require, parent: ir.Node):
-        checks = rewrite_list(ir.Check, lambda n: self.walk(n, node), node.checks)
-        # discharge require if all the checks are discharged
-        if all(builtins.discharged_annotation in check.annotations for check in checks):
-            return node.reconstruct(node.engine, node.domain, checks, node.annotations | [builtins.discharged_annotation])
-        return node.reconstruct(node.engine, node.domain, checks, node.annotations)
-
-    def handle_check(self, node: ir.Check, parent: ir.Node):
-        check = self.walk(node.check, node)
-        assert isinstance(check, ir.Logical)
-        discharged_names = [builtins.unique.name, builtins.exclusive.name, builtins.anyof.name]
-        for item in check.body:
-            if isinstance(item, ir.Lookup) and item.relation.name in discharged_names:
-                return node.reconstruct(check, node.error, node.annotations | [builtins.discharged_annotation])
-        return node.reconstruct(check, node.error, node.annotations)

relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py

@@ -1,210 +0,0 @@
-from __future__ import annotations
-
-from relationalai.semantics.metamodel import ir
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel.visitor import Visitor, Rewriter
-from relationalai.semantics.metamodel.util import OrderedSet
-from relationalai.semantics.metamodel import helpers, factory as f
-from typing import Optional, Any
-
-"""
-Handle DNF decomposition of unions
-
-E.g., we go from
-
-Logical
-    Logical
-        Union
-            Logical
-                Foo(foo::Foo)
-                a(foo::Foo, a::Int128)
-                a::Int128 < 100::Int128
-                Union
-                    Logical
-                        b(foo::Foo, b::Int128)
-                        b::Int128 < 100::Int128
-                    Logical
-                        c(foo::Foo, c::Int128)
-                        c::Int128 > 0::Int128
-            Logical
-                Foo(foo::Foo)
-                id(foo::Foo, id::Int128)
-                id::Int128 < 100::Int128
-            Logical
-                Foo(foo::Foo)
-                x(foo::Foo, x::Int128)
-                x::Int128 < 100::Int128
-                Union
-                    Logical
-                        y(foo::Foo, y::Int128)
-                        y::Int128 < 100::Int128
-                    Logical
-                        z(foo::Foo, z::Int128)
-                        z::Int128 > 0::Int128
-        ...
-        -> output(...)
-
-to
-
-Logical
-    Logical
-        Logical
-            Foo(foo::Foo)
-            a(foo::Foo, a::Int128)
-            a::Int128 < 100::Int128
-            Logical
-                b(foo::Foo, b::Int128)
-                b::Int128 < 100::Int128
-        ...
-        -> output(...)
-    Logical
-        Logical
-            Foo(foo::Foo)
-            a(foo::Foo, a::Int128)
-            a::Int128 < 100::Int128
-            Logical
-                c(foo::Foo, c::Int128)
-                c::Int128 > 0::Int128
-        ...
-        -> output(...)
-    Logical
-        Logical
-            Foo(foo::Foo)
-            id(foo::Foo, id::Int128)
-            id::Int128 < 100::Int128
-        ...
-        -> output(...)
-    Logical
-        Logical
-            Foo(foo::Foo)
-            x(foo::Foo, x::Int128)
-            x::Int128 < 100::Int128
-            Logical
-                y(foo::Foo, y::Int128)
-                y::Int128 < 100::Int128
-        ...
-        -> output(...)
-    Logical
-        Logical
-            Foo(foo::Foo)
-            x(foo::Foo, x::Int128)
-            x::Int128 < 100::Int128
-            Logical
-                z(foo::Foo, z::Int128)
-                z::Int128 > 0::Int128
-        ...
-        -> output(...)
-"""
-class DNFUnionSplitter(Pass):
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        visitor = DNFExtractor()
-        model.accept(visitor)
-        return DNFRewriter(visitor).walk(model) if visitor.replaced_by else model
-
-class DNFExtractor(Visitor):
-    def __init__(self):
-        # The logical that contains the output.
-        # The assumption for the IR at this point is that there is only one output.
-        self.output_logical: Optional[ir.Logical] = None
-        self.active_negations: list[ir.Not] = []
-        # Nodes that have to split into multiple similar nodes, depending on the changes
-        # of sub-nodes.
-        self.should_split: OrderedSet[ir.Node] = OrderedSet()
-        # Track that a node has to be replaced by other nodes.
-        # If node X should be replaced by nodes Y and Z, then the parent should be replicated
-        # one time replacing X with Y and one time replacing X with Z.
-        self.replaced_by: dict[ir.Node, list[ir.Task]] = {}
-
-    def enter(self, node: ir.Node, parent: Optional[ir.Node]=None) -> Visitor:
-        if isinstance(node, ir.Logical):
-            if any(isinstance(x, ir.Output) for x in node.body):
-                assert not self.output_logical, "multiple outputs"
-                self.output_logical = node
-
-        elif isinstance(node, ir.Not):
-            self.active_negations.append(node)
-
-        return self
-
-    def leave(self, node: ir.Node, parent: Optional[ir.Node]=None) -> ir.Node:
-        if isinstance(node, ir.Logical) and node in self.should_split:
-            # The given logical, will be replaced by multiple logicals, each with a different
-            # group of tasks. We need to generate all the possible groups.
-            # A list of logical bodies (lists).
-            replacement_bodies: list[list[ir.Task]] = [[]]
-            for task in node.body:
-                if task in self.replaced_by:
-                    new_replacement_bodies: list[list[ir.Task]] = []
-                    replacement_tasks = self.replaced_by[task]
-                    for body in replacement_bodies:
-                        for new_task in replacement_tasks:
-                            # copy to mutate
-                            new_body = list(body)
-                            new_body.append(new_task.clone())
-                            new_replacement_bodies.append(new_body)
-                    replacement_bodies = new_replacement_bodies
-
-                else:
-                    for new_body in replacement_bodies:
-                        new_body.append(task.clone())
-
-            replacement_tasks: list[ir.Task] = []
-            for body in replacement_bodies:
-                new_task = f.logical(body, node.hoisted)
-                replacement_tasks.append(new_task)
-            self.replaced_by[node] = replacement_tasks
-
-            if node != self.output_logical:
-                self.should_split.add(parent)
-            elif node == self.output_logical:
-                assert parent and isinstance(parent, ir.Logical)
-                new_parent = f.logical(tuple(replacement_tasks), node.hoisted)
-                self.replaced_by[parent] = [new_parent]
-
-        elif isinstance(node, ir.Not) and self.active_negations[-1] == node:
-            self.active_negations.pop()
-
-        elif (isinstance(node, ir.Union) and
-              self.output_logical and
-              len(self.active_negations) % 2 == 0 and
-              len(node.tasks) > 1):
-            # We split the union when there is a branch with vars "X,Y" and another with "X,Z"
-            # If some branches have vars "X, Y, Z" and others have "X, Y" or "Y, Z" we don't split
-            should_split = False
-            all_vars = helpers.collect_vars(node.tasks[0])
-            for t in node.tasks[1:]:
-                vars = helpers.collect_vars(t)
-                curr_intersection = vars.get_set().intersection(all_vars.get_set())
-                should_split |= not (curr_intersection == vars.get_set() or curr_intersection == all_vars.get_set())
-                if should_split:
-                    replacements:list[ir.Task] = []
-                    for t in node.tasks:
-                        # If some branch should already be replaced, we flatten all the replacements here.
-                        if t in self.replaced_by:
-                            replacements.extend(self.replaced_by[t])
-                        else:
-                            replacements.append(t)
-                    self.replaced_by[node] = replacements
-                    self.should_split.add(parent)
-                    break
-                all_vars.update(vars)
-
-        if isinstance(node, ir.Logical) and node == self.output_logical:
-            self.output_logical = None
-
-        return node
-
-class DNFRewriter(Rewriter):
-    def __init__(self, visitor: DNFExtractor):
-        super().__init__()
-        self.visitor = visitor
-        self.outer_parent_logical: Optional[ir.Logical] = None
-
-    def handle_logical(self, node: ir.Logical, parent: ir.Node, ctx:Optional[Any]=None) -> ir.Logical:
-        if node in self.visitor.replaced_by:
-            new_tasks = self.visitor.replaced_by[node]
-            assert len(new_tasks) == 1
-            new_task = new_tasks[0]
-            assert isinstance(new_task, ir.Logical)
-            return new_task
-        return node

relationalai/semantics/metamodel/rewrite/extract_nested_logicals.py

@@ -1,78 +0,0 @@
-from __future__ import annotations
-
-from relationalai.semantics.metamodel import ir, factory as f, helpers
-from relationalai.semantics.metamodel.visitor import Rewriter, collect_by_type
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel.util import OrderedSet, ordered_set, NameCache
-from relationalai.semantics.metamodel import dependency
-
-class ExtractNestedLogicals(Pass):
-
-    def __init__(self):
-        super().__init__()
-
-    #--------------------------------------------------
-    # Public API
-    #--------------------------------------------------
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        # process the root node
-        extractor = LogicalExtractor(model)
-        root = extractor.walk(model.root, model)
-
-        # no extractions, just return the original model
-        if not extractor.ctx.top_level:
-            return model
-
-        # create a new model with the extracted top level + the new root
-        body = extractor.ctx.top_level + [root]
-        return ir.Model(
-                model.engines,
-                OrderedSet.from_iterable(model.relations).update(extractor.ctx.relations).frozen(),
-                model.types,
-                ir.Logical(model.root.engine, tuple(), tuple(body))
-            )
-
-class LogicalExtractor(Rewriter):
-    def __init__(self, model):
-        super().__init__()
-        self.ctx = helpers.RewriteContext()
-        self.info = dependency.analyze(model.root)
-        self.name_cache = NameCache()
-
-    def handle_logical(self, node: ir.Logical, parent: ir.Node):
-        # rewrite the children
-        logical = super().handle_logical(node, parent)
-
-        # logicals that hoist vars and all vars do not have a default value will be
-        # extracted, except for logicals that require special treatment of the exposed
-        # variables (which is currently done by flatten), such as when the parent is a Match
-        # or a Union, of if the logical has a Rank.
-        if not (
-            logical.hoisted and
-            not isinstance(parent, (ir.Match, ir.Union)) and
-            all(isinstance(v, ir.Var) for v in logical.hoisted) and
-            not any(isinstance(c, ir.Rank) for c in logical.body)
-            ):
-            return logical
-
-        # compute the vars to be exposed by the extracted logical; those are keys (what
-        # makes the values unique) + the values (the hoisted variables)
-        exposed_vars = ordered_set()
-
-        # if there are aggregations, make sure we don't expose the projected and input vars,
-        # but expose groupbys
-        for agg in collect_by_type(ir.Aggregate, logical):
-            exposed_vars.difference_update(agg.projection)
-            exposed_vars.difference_update(helpers.aggregate_inputs(agg))
-            exposed_vars.update(agg.group)
-
-        # add the values (hoisted)
-        exposed_vars.update(helpers.hoisted_vars(logical.hoisted))
-
-        body = ordered_set()
-        body.update(self.info.task_dependencies(node)) # notice info is based on the original node
-        body.update(logical.body)
-
-        name = helpers.create_task_name(self.name_cache, logical, "_nested_logical")
-        connection = helpers.extract(logical, body, exposed_vars.get_list(), self.ctx, name)
-        return f.lookup(connection, exposed_vars.get_list())