relationalai 0.13.5__py3-none-any.whl → 1.0.0a2__py3-none-any.whl

relationalai/semantics/lqp/rewrite/flatten_script.py

@@ -1,301 +0,0 @@
-from __future__ import annotations
-
-from relationalai.semantics.metamodel import ir, factory as f, helpers
-from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
-from relationalai.semantics.metamodel.rewrite.flatten import Flatten, negate, extend_body
-from relationalai.semantics.lqp.algorithms import is_script, mk_assign
-
-class FlattenScript(Flatten):
-    """
-    Flattens Match nodes inside @script Sequence blocks. This pass extends Flatten
-    of standard Logicals, to reuse a number of utilities, but DOES NOT flatten Match
-    nodes outside of scripts (which are handled by the Flatten pass).
-
-    Unlike the regular Flatten pass which extracts to top-level, FlattenScript
-    maintains order by inserting intermediate relations right before they're used
-    within the Sequence. This is necessary because order matters in a script Sequence.
-
-    Additionally, if the original Logical is a Loopy instruction, the flattened Logicals are
-    made Loopy too (`@assign`).
-
-    Example:
-
-    === BEFORE ===
-    Logical
-        Sequence @script @algorithm
-            Logical
-                dom(n)
-                Match ⇑[k]
-                    Logical ⇑[k]
-                        value(n, k)
-                    Logical ⇑[k]
-                        k = 0
-                → derive result(n, k) @assign @global
-                filter(n)
-
-    === AFTER ===
-    Logical
-        Sequence @script @algorithm
-            Logical
-                dom(n)
-                filter(n)
-                Logical ⇑[v]
-                    value(n, v)
-                → derive _match_1(n, v) @assign
-            Logical
-                dom(n)
-                filter(n)
-                Logical ⇑[v]
-                    v = 0
-                Not
-                    _match_1(n, _)
-                → derive _match_2(n, v) @assign
-            Logical
-                dom(n)
-                filter(n)
-                Union ⇑[v]
-                    _match_1(n, v)
-                    _match_2(n, v)
-                → derive result(n, v) @assign @global
-    """
-
-    class Context(Flatten.Context):
-        """Extended context with script tracking."""
-        def __init__(self, model: ir.Model, options: dict):
-            super().__init__(model, options)
-            self.in_script: bool = False
-
-    def rewrite(self, model: ir.Model, options: dict = {}) -> ir.Model:
-        """Traverse the model and flatten Match nodes inside script Sequences."""
-        ctx = FlattenScript.Context(model, options)
-        result = self.handle(model.root, ctx)
-
-        if result.replacement is None:
-            return model
-
-        # Convert relations list to FrozenOrderedSet (adding any new intermediate relations)
-        new_relations = OrderedSet.from_iterable(model.relations).update(ctx.rewrite_ctx.relations).frozen()
-
-        return ir.Model(
-            model.engines,
-            new_relations,
-            model.types,
-            result.replacement
-        )
-
-    def handle(self, task: ir.Task, ctx: Flatten.Context) -> Flatten.HandleResult:
-        """Override handle to add Loop support."""
-        if isinstance(task, ir.Loop):
-            return self.handle_loop(task, ctx)
-        return super().handle(task, ctx)
-
-    def handle_loop(self, task: ir.Loop, ctx: Flatten.Context) -> Flatten.HandleResult:
-        """Recursively handle the body of the loop."""
-        result = self.handle(task.body, ctx)
-
-        assert(result.replacement)
-
-        # If body unchanged, return original loop
-        if result.replacement is task.body:
-            return Flatten.HandleResult(task)
-
-        # Return new loop with handled body
-        return Flatten.HandleResult(ir.Loop(
-            task.engine,
-            task.hoisted,
-            task.iter,
-            result.replacement,
-            task.concurrency,
-            task.annotations
-        ))
-
-    def handle_logical(self, task: ir.Logical, ctx: Context):  # type: ignore[override]
-        """
-        Handle Logical nodes.
-
-        Outside scripts: simple traversal to find nested scripts.
-        Inside scripts: prevent extraction to top-level and keep everything in sequence.
-
-        Note: Type checker complains about parameter type narrowing, but this is safe
-        because we only create FlattenScript.Context in our own rewrite() method.
-        """
-        # Recursively process children
-        body: OrderedSet[ir.Task] = ordered_set()
-        for child in task.body:
-            result = self.handle(child, ctx)
-            if result.replacement is not None:
-                if ctx.in_script and isinstance(result.replacement, ir.Logical) and not result.replacement.hoisted:
-                    # Inside script: inline simple logicals without hoisting
-                    body.update(result.replacement.body)
-                else:
-                    body.add(result.replacement)
-
-        if not body:
-            return Flatten.HandleResult(None)
-
-        return Flatten.HandleResult(
-            ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations)
-        )
-
-    def flatten_match_in_logical(self, logical: ir.Logical, match: ir.Match, match_idx: int, ctx: Context) -> list[ir.Logical]:
-        """
-        Flatten a Match inside a Logical within a script Sequence.
-        Returns a list of Logicals to be inserted in sequence.
-        """
-        if not match.tasks:
-            return [logical]
-
-        # Split the logical into: tasks_before, match, tasks_after
-        tasks_before = list(logical.body[:match_idx])
-        tasks_after = list(logical.body[match_idx + 1:])
-
-        # Separate tasks_after into filters (non-Update tasks) and updates (Update tasks)
-        # Filters are constraints that should be included in all branches
-        filters = [task for task in tasks_after if not isinstance(task, ir.Update)]
-        updates = [task for task in tasks_after if isinstance(task, ir.Update)]
-
-        # Compute exposed variables
-        exposed_vars = self.compute_exposed_vars(match, match.tasks, ctx)
-
-        # Use dependency analysis for branch bodies (like flatten.py does)
-        branch_dependencies = ctx.info.task_dependencies(match)
-
-        # Collect all dependencies for the final Logical (tasks before + filters after)
-        final_dependencies = tasks_before + filters
-
-        # Negation length for wildcards
-        outputs = ctx.info.task_outputs(match)
-        negation_len = len(outputs) if outputs else 0
-
-        # Result: list of Logicals to insert
-        result_logicals = []
-        references = []
-        negated_reference = None
-
-        # Process each branch
-        for branch in match.tasks:
-            # Create connection relation for this branch
-            name = helpers.create_task_name(self.name_cache, branch, "_match")
-            relation = helpers.create_connection_relation(branch, exposed_vars, ctx.rewrite_ctx, name)
-
-            # Handle the branch (recursively process nested structures)
-            result = self.handle(branch, ctx)
-            branch_content = result.replacement if result.replacement else branch
-
-            # Update dependency tracking if branch was transformed
-            if result.replacement:
-                ctx.info.replaced(branch, result.replacement)
-
-            # Build logical for this branch
-            branch_body: OrderedSet[ir.Task] = ordered_set()
-
-            # Add dependencies (using dependency analysis, not all tasks)
-            branch_body.update(branch_dependencies)
-
-            # Add branch content using extend_body helper
-            extend_body(branch_body, branch_content)
-
-            # Add negation of previous branches (after branch content)
-            if negated_reference:
-                branch_body.add(negated_reference)
-
-            # Add derive to connection relation
-            branch_update = f.derive(relation, exposed_vars)
-            branch_body.add(branch_update)
-
-            # Create the Logical for this branch
-            branch_logical = mk_assign(ir.Logical(match.engine, tuple(), tuple(branch_body)))
-
-            result_logicals.append(branch_logical)
-
-            # Update references for final union
-            reference = f.lookup(relation, exposed_vars)
-            negated_reference = negate(reference, negation_len)
-            references.append(reference)
-
-        # Create final Logical with Union and remaining tasks
-        final_body: OrderedSet[ir.Task] = ordered_set()
-        final_body.update(final_dependencies)
-
-        # Add union of all branches
-        union = f.union(references, match.hoisted)
-        final_body.add(union)
-
-        # Add updates that came after the match (filters are already in dependencies)
-        final_body.update(updates)
-
-        # Create final logical preserving the original annotations
-        final_logical = ir.Logical(logical.engine, logical.hoisted, tuple(final_body), logical.annotations)
-
-        result_logicals.append(final_logical)
-
-        return result_logicals
-
-    def handle_sequence(self, task: ir.Sequence, ctx: Context):  # type: ignore[override]
-        """
-        Handle a Sequence.
-
-        If it's a script: set context flag and flatten Match nodes within.
-        If not a script: simple traversal (Flatten pass already processed it).
-
-        Note: Type checker complains about parameter type narrowing, but this is safe
-        because we only create FlattenScript.Context in our own rewrite() method.
-        """
-        if not is_script(task):
-            # Not a script sequence - already processed by Flatten, just return as-is
-            return Flatten.HandleResult(task)
-
-        # This is a script - mark context and process with flattening
-        old_in_script = ctx.in_script
-        ctx.in_script = True
-
-        # Process the sequence tasks
-        new_tasks: list[ir.Task] = []
-        for child in task.tasks:
-            # Check if this child is a Logical with Match that needs flattening
-            if isinstance(child, ir.Logical):
-                new_tasks.extend(self.try_flatten_logical(child, ctx))
-                continue
-
-            # No flattening needed, process normally
-            result = self.handle(child, ctx)
-            if result.replacement is not None:
-                new_tasks.append(result.replacement)
-
-        # Restore context
-        ctx.in_script = old_in_script
-
-        return Flatten.HandleResult(
-            ir.Sequence(task.engine, task.hoisted, tuple(new_tasks), task.annotations)
-        )
-
-    def try_flatten_logical(self, logical: ir.Logical, ctx: Context) -> list[ir.Logical]:
-        """
-        Flatten all Matches in a Logical.
-        Iteratively flattens until no more Matches remain in any of the resulting Logicals.
-        """
-        worklist = [logical]
-        result = []
-
-        while worklist:
-            current = worklist.pop()
-
-            # Find first Match in current logical
-            match = None
-            match_idx = -1
-            for i, child in enumerate(current.body):
-                if isinstance(child, ir.Match):
-                    match = child
-                    match_idx = i
-                    break
-
-            if match is None:
-                # No Match found - this logical is done
-                result.append(current)
-            else:
-                # Flatten and add results back to worklist for further processing.
-                # Reverse so that pop() returns them in the original order.
-                flattened = self.flatten_match_in_logical(current, match, match_idx, ctx)
-                worklist.extend(reversed(flattened))
-
-        return result

relationalai/semantics/lqp/rewrite/function_annotations.py

@@ -1,114 +0,0 @@
-from __future__ import annotations
-
-from typing import Optional
-from relationalai.semantics.metamodel import builtins
-from relationalai.semantics.metamodel.ir import (
-    Node, Model, Require, Logical, Relation, Annotation, Update
-)
-from relationalai.semantics.metamodel.compiler import Pass
-from relationalai.semantics.metamodel.visitor import Rewriter, Visitor
-from relationalai.semantics.lqp.rewrite.functional_dependencies import (
-    is_valid_unique_constraint, normalized_fd
-)
-
-# In the future iterations of PyRel metamodel, `Require` nodes will have a single `Check`
-# (and no `errors`). Currently, however, the unique constraints may result in multiple
-# `Check` nodes and for simplicity we split them in to separate `Require` nodes. This step
-# will be removed in the future.
-#
-# Note that unique constraints always have an empty `domain` so apply the splitting only
-# to such `Require` nodes.
-class SplitMultiCheckRequires(Pass):
-    """
-    Pass splits unique Require nodes that have empty domain but multiple checks into multiple
-    Require nodes with single check each.
-    """
-
-    def rewrite(self, model: Model, options: dict = {}) -> Model:
-        return SplitMultiCheckRequiresRewriter().walk(model)
-
-
-class SplitMultiCheckRequiresRewriter(Rewriter):
-    """
-    Splits unique Require nodes that have empty domain but multiple checks into multiple
-    Require nodes with single check each.
-    """
-    def handle_require(self, node: Require, parent: Node):
-
-        if isinstance(node.domain, Logical) and not node.domain.body and len(node.checks) > 1:
-            require_nodes = []
-            for check in node.checks:
-                single_check = self.walk(check, node)
-                require_nodes.append(
-                    node.reconstruct(node.engine, node.domain, (single_check,), node.annotations)
-                )
-            return require_nodes
-
-        return node
-
-
-class FunctionAnnotations(Pass):
-    """
-    Pass marks all appropriate relations with `function` annotation. Collects functional
-    dependencies from unique Require nodes and uses this information to identify functional
-    relations.
-    """
-
-    def rewrite(self, model: Model, options: dict = {}) -> Model:
-        collect_fds = CollectFDsVisitor()
-        collect_fds.visit_model(model, None)
-        annotated_model = FunctionalAnnotationsRewriter(collect_fds.functional_relations).walk(model)
-        return annotated_model
-
-
-class CollectFDsVisitor(Visitor):
-    """
-    Visitor collects all unique constraints.
-    """
-
-    # Currently, only information about k-functional fd is collected.
-    def __init__(self):
-        super().__init__()
-        self.functional_relations:dict[Relation, int] = {}
-
-    def visit_require(self, node: Require, parent: Optional[Node]):
-        if is_valid_unique_constraint(node):
-            fd = normalized_fd(node)
-            assert fd is not None
-            if fd.is_structural:
-                relation = fd.structural_relation
-                k = fd.structural_rank
-                current_k = self.functional_relations.get(relation, 0)
-                self.functional_relations[relation] = max(current_k, k)
-
-
-class FunctionalAnnotationsRewriter(Rewriter):
-    """
-    This visitor marks functional_relations with `@function(:checked [, k])` annotation.
-    """
-
-    def __init__(self, functional_relations: dict[Relation, int]):
-        super().__init__()
-        self.functional_relations = functional_relations
-
-    def get_functional_annotation(self, rel: Relation) -> Optional[Annotation]:
-        k = self.functional_relations.get(rel, None)
-        if k is None:
-            return None
-        if k == 1:
-            return builtins.function_checked_annotation
-        return builtins.function_ranked_checked_annotation(k)
-
-    def handle_relation(self, node: Relation, parent: Node):
-        function_annotation = self.get_functional_annotation(node)
-        if function_annotation:
-            return node.reconstruct(node.name, node.fields, node.requires,
-                                    node.annotations | [function_annotation], node.overloads)
-        return node.reconstruct(node.name, node.fields, node.requires, node.annotations, node.overloads)
-
-    def handle_update(self, node: Update, parent: Node):
-        function_annotation = self.get_functional_annotation(node.relation)
-        if function_annotation:
-            return node.reconstruct(node.engine, node.relation, node.args, node.effect,
-                                    node.annotations | [function_annotation])
-        return node.reconstruct(node.engine, node.relation, node.args, node.effect, node.annotations)