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

v0/relationalai/semantics/metamodel/helpers.py

@@ -0,0 +1,361 @@
+"""
+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 v0.relationalai.semantics.metamodel import ir, visitor, builtins, types, factory as f
+from v0.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

v0/relationalai/semantics/metamodel/rewrite/discharge_constraints.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from v0.relationalai.semantics.metamodel import ir, compiler as c, visitor as v, builtins
+from v0.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)

v0/relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+from v0.relationalai.semantics.metamodel import ir
+from v0.relationalai.semantics.metamodel.compiler import Pass
+from v0.relationalai.semantics.metamodel.visitor import Visitor, Rewriter
+from v0.relationalai.semantics.metamodel.util import OrderedSet
+from v0.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

v0/relationalai/semantics/metamodel/rewrite/extract_nested_logicals.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from v0.relationalai.semantics.metamodel import ir, factory as f, helpers
+from v0.relationalai.semantics.metamodel.visitor import Rewriter, collect_by_type
+from v0.relationalai.semantics.metamodel.compiler import Pass
+from v0.relationalai.semantics.metamodel.util import OrderedSet, ordered_set, NameCache
+from v0.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())