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

v0/relationalai/semantics/lqp/passes.py

@@ -1,680 +0,0 @@
-from v0.relationalai.semantics.metamodel.compiler import Pass
-from v0.relationalai.semantics.metamodel import ir, builtins as rel_builtins, factory as f, visitor
-from v0.relationalai.semantics.metamodel.typer import Checker, InferTypes, typer
-from v0.relationalai.semantics.metamodel import helpers, types
-from v0.relationalai.semantics.metamodel.util import FrozenOrderedSet
-
-from v0.relationalai.semantics.metamodel.rewrite import Flatten
-
-from ..metamodel.rewrite import DNFUnionSplitter, ExtractNestedLogicals, FormatOutputs
-from .rewrite import (
-    AnnotateConstraints, CDC, ExtractCommon, ExtractKeys, FunctionAnnotations, QuantifyVars,
-    Splinter, SplitMultiCheckRequires
-)
-from v0.relationalai.semantics.lqp.utils import output_names
-
-from typing import cast, List, Sequence, Tuple, Union, Optional, Iterable
-from collections import defaultdict
-import pandas as pd
-import hashlib
-
-def lqp_passes() -> list[Pass]:
-    return [
-        SplitMultiCheckRequires(),
-        FunctionAnnotations(),
-        AnnotateConstraints(),
-        Checker(),
-        CDC(), # specialize to physical relations before extracting nested and typing
-        ExtractNestedLogicals(), # before InferTypes to avoid extracting casts
-        InferTypes(),
-        DNFUnionSplitter(),
-        ExtractKeys(),
-        FormatOutputs(),
-        ExtractCommon(), # Extracts tasks that will become common after Flatten into their own definition
-        Flatten(),
-        Splinter(), # Splits multi-headed rules into multiple rules
-        QuantifyVars(), # Adds missing existentials
-        EliminateData(),  # Turns Data nodes into ordinary relations.
-        DeduplicateVars(),  # Deduplicates vars in Updates and Outputs.
-        PeriodMath(),  # Rewrite date period uses.
-        ConstantsToVars(),  # Turns constants in Updates and Outputs into vars.
-        UnifyDefinitions(),
-    ]
-
-# LQP does not support multiple definitions for the same relation. This pass unifies all
-# definitions for each relation into a single definition using a union.
-class UnifyDefinitions(Pass):
-    def __init__(self):
-        super().__init__()
-
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        # Maintain a cache of renamings for each relation. These need to be consistent
-        # across all definitions of the same relation.
-        self.renamed_relation_args: dict[Union[ir.Value, ir.Relation], list[ir.Var]] = {}
-
-        root = cast(ir.Logical, model.root)
-        new_tasks = self.get_combined_multidefs(root)
-        return ir.Model(
-            model.engines,
-            model.relations,
-            model.types,
-            f.logical(
-                tuple(new_tasks),
-                root.hoisted,
-                root.engine,
-            ),
-            model.annotations,
-        )
-
-    def _get_heads(self, logical: ir.Logical) -> list[Union[ir.Update, ir.Output]]:
-        derives = []
-        for task in logical.body:
-            if isinstance(task, ir.Update) and task.effect == ir.Effect.derive:
-                derives.append(task)
-            elif isinstance(task, ir.Output):
-                derives.append(task)
-        return derives
-
-    def _get_non_heads(self, logical: ir.Logical) -> list[ir.Task]:
-        non_derives = []
-        for task in logical.body:
-            if not(isinstance(task, ir.Update) and task.effect == ir.Effect.derive) and not isinstance(task, ir.Output):
-                non_derives.append(task)
-        return non_derives
-
-    def _get_head_identifier(self, head: Union[ir.Update, ir.Output]) -> Optional[ir.Value]:
-        if isinstance(head, ir.Update):
-            return head.relation
-        else:
-            assert isinstance(head, ir.Output)
-            if len(head.aliases) <= 2:
-                # For processing here, we need output to have at least the column markers
-                # `cols` and `col`, and also a key
-                return None
-
-            output_alias_names = helpers.output_alias_names(head.aliases)
-            output_vals = helpers.output_values(head.aliases)
-
-            # For normal outputs, the pattern is output[keys](cols, "col000" as 'col', ...)
-            if output_alias_names[0] == "cols" and output_alias_names[1] == "col":
-                return output_vals[1]
-
-            # For exports, the pattern is output[keys]("col000" as 'col', ...)
-            if rel_builtins.export_annotation in head.annotations:
-                if output_alias_names[0] == "col":
-                    return output_vals[0]
-
-        return None
-
-    def get_combined_multidefs(self, root: ir.Logical) -> list[ir.Logical]:
-        # Step 1: Group tasks by the relation they define.
-        relation_to_tasks: dict[Union[None, ir.Value, ir.Relation], list[ir.Logical]] = defaultdict(list)
-
-        for task in root.body:
-            task = cast(ir.Logical, task)
-            task_heads = self._get_heads(task)
-
-            # Some relations do not need to be grouped, e.g., if they don't contain a
-            # derive. Use `None` as a placeholder key for these cases.
-            if len(task_heads) != 1:
-                relation_to_tasks[None].append(task)
-                continue
-
-            head_id = self._get_head_identifier(task_heads[0])
-            relation_to_tasks[head_id].append(task)
-
-        # Step 2: For each relation, combine all of the body definitions into a union.
-        result_tasks = []
-        for relation, tasks in relation_to_tasks.items():
-            # If there's only one task for the relation, or if grouping is not needed, then
-            # just keep the original tasks.
-            if len(tasks) == 1 or relation is None:
-                result_tasks.extend(tasks)
-                continue
-
-            result_tasks.append(self._combine_tasks_into_union(tasks))
-        return result_tasks
-
-    def _get_variable_mapping(self, logical: ir.Logical) -> dict[ir.Value, ir.Var]:
-        heads = self._get_heads(logical)
-        assert len(heads) == 1, "should only have one head in a logical at this stage"
-        head = heads[0]
-
-        var_mapping = {}
-        head_id = self._get_head_identifier(head)
-
-        if isinstance(head, ir.Update):
-            args_for_renaming = head.args
-        else:
-            assert isinstance(head, ir.Output)
-            output_alias_names = helpers.output_alias_names(head.aliases)
-            if output_alias_names[0] == "cols" and output_alias_names[1] == "col":
-                assert len(head.aliases) > 2
-
-                # For outputs, we do not need to rename the `cols` and `col` markers or the
-                # keys.
-                output_values = helpers.output_values(head.aliases)[2:]
-
-            else:
-                assert rel_builtins.export_annotation in head.annotations and output_alias_names[0] == "col"
-                assert len(head.aliases) > 1
-
-                # For exports, we do not need to rename the `col` marker or the keys.
-                output_values = helpers.output_values(head.aliases)[1:]
-
-            args_for_renaming = []
-            for v in output_values:
-                if head.keys and isinstance(v, ir.Var) and v in head.keys:
-                    continue
-                args_for_renaming.append(v)
-
-        if head_id not in self.renamed_relation_args:
-            renamed_vars = []
-            for (i, arg) in enumerate(args_for_renaming):
-                typ = typer.to_type(arg)
-                assert arg not in var_mapping, "args of update should be unique"
-                if isinstance(arg, ir.Var):
-                    var_mapping[arg] = ir.Var(typ, arg.name)
-                else:
-                    var_mapping[arg] = ir.Var(typ, f"arg_{i}")
-
-                renamed_vars.append(var_mapping[arg])
-            self.renamed_relation_args[head_id] = renamed_vars
-        else:
-            for (arg, var) in zip(args_for_renaming, self.renamed_relation_args[head_id]):
-                var_mapping[arg] = var
-
-        return var_mapping
-
-    def _rename_variables(self, logical: ir.Logical) -> ir.Logical:
-        class RenameVisitor(visitor.Rewriter):
-            def __init__(self, var_mapping: dict[ir.Value, ir.Var]):
-                super().__init__()
-                self.var_mapping = var_mapping
-
-            def _get_mapped_value(self, val: ir.Value) -> ir.Value:
-                if isinstance(val, tuple):
-                    return tuple(self._get_mapped_value(t) for t in val)
-                return self.var_mapping.get(val, val)
-
-            def _get_mapped_values(self, vals: Iterable[ir.Value]) -> list[ir.Value]:
-                return [self._get_mapped_value(v) for v in vals]
-
-            def handle_var(self, node: ir.Var, parent: ir.Node) -> ir.Var:
-                return self.var_mapping.get(node, node)
-
-            # TODO: ideally, extend the rewriter class to allow rewriting PyValue to Var so
-            # we don't need to separately handle all cases containing them.
-            def handle_update(self, node: ir.Update, parent: ir.Node) -> ir.Update:
-                return ir.Update(
-                    node.engine,
-                    node.relation,
-                    tuple(self._get_mapped_values(node.args)),
-                    node.effect,
-                    node.annotations,
-                )
-
-            def handle_lookup(self, node: ir.Lookup, parent: ir.Node) -> ir.Lookup:
-                return ir.Lookup(
-                    node.engine,
-                    node.relation,
-                    tuple(self._get_mapped_values(node.args)),
-                    node.annotations,
-                )
-
-            def handle_output(self, node: ir.Output, parent: ir.Node) -> ir.Output:
-                new_aliases = FrozenOrderedSet(
-                    [(name, self._get_mapped_value(value)) for name, value in node.aliases]
-                )
-                if node.keys:
-                    new_keys = FrozenOrderedSet(
-                        [self.var_mapping.get(key, key) for key in node.keys]
-                    )
-                else:
-                    new_keys = node.keys
-
-                return ir.Output(
-                    node.engine,
-                    new_aliases,
-                    new_keys,
-                    node.annotations,
-                )
-
-            def handle_construct(self, node: ir.Construct, parent: ir.Node) -> ir.Construct:
-                new_values = tuple(self._get_mapped_values(node.values))
-                new_id_var = self.var_mapping.get(node.id_var, node.id_var)
-                return ir.Construct(
-                    node.engine,
-                    new_values,
-                    new_id_var,
-                    node.annotations,
-                )
-
-            def handle_aggregate(self, node: ir.Aggregate, parent: ir.Node) -> ir.Aggregate:
-                new_projection = tuple(self.var_mapping.get(arg, arg) for arg in node.projection)
-                new_group = tuple(self.var_mapping.get(arg, arg) for arg in node.group)
-                new_args = tuple(self._get_mapped_values(node.args))
-                return ir.Aggregate(
-                    node.engine,
-                    node.aggregation,
-                    new_projection,
-                    new_group,
-                    new_args,
-                    node.annotations,
-                )
-
-            def handle_rank(self, node: ir.Rank, parent: ir.Node) -> ir.Rank:
-                new_projection = tuple(self.var_mapping.get(arg, arg) for arg in node.projection)
-                new_group = tuple(self.var_mapping.get(arg, arg) for arg in node.group)
-                new_args = tuple(self.var_mapping.get(arg, arg) for arg in node.args)
-                new_result = self.var_mapping.get(node.result, node.result)
-
-                return ir.Rank(
-                    node.engine,
-                    new_projection,
-                    new_group,
-                    new_args,
-                    node.arg_is_ascending,
-                    new_result,
-                    node.limit,
-                    node.annotations,
-                )
-
-        var_mapping = self._get_variable_mapping(logical)
-
-        renamer = RenameVisitor(var_mapping)
-        result = renamer.walk(logical)
-
-        # Also need to append the equality for each renamed constant. E.g., if the mapping
-        # contains (50.0::FLOAT -> arg_2::FLOAT), we need to add
-        # `eq(arg_2::FLOAT, 50.0::FLOAT)` to the result.
-        value_eqs = []
-        for (old_var, new_var) in var_mapping.items():
-            if not isinstance(old_var, ir.Var):
-                value_eqs.append(f.lookup(rel_builtins.eq, [new_var, old_var]))
-
-        return ir.Logical(
-            result.engine,
-            result.hoisted,
-            tuple(value_eqs) + tuple(result.body),
-            result.annotations,
-        )
-
-    # This function is the main workhorse for this rewrite pass. It takes a list of tasks
-    # that define the same relation, and combines them into a single task that defines
-    # the relation using a union of all of the bodies.
-    def _combine_tasks_into_union(self, tasks: list[ir.Logical]) -> ir.Logical:
-        # Step 1: Rename the variables in all tasks so that they will match the final derive
-        # after reconstructing into a union
-        renamed_tasks = [self._rename_variables(task) for task in tasks]
-
-        # Step 2: Get the final derive
-        derives = self._get_heads(renamed_tasks[0])
-        assert len(derives) == 1, "should only have one derive in a logical at this stage"
-        # Also make sure that all the derives are the same. This should be the case because
-        # we renamed all the variables to be the same in step 1.
-        for task in renamed_tasks[1:]:
-            assert self._get_heads(task) == derives, "all derives should be the same"
-
-        derive = derives[0]
-
-        # Step 3: Remove the final `derive` from each task
-        renamed_task_bodies = [
-            f.logical(
-                tuple(self._get_non_heads(t)),  # Only keep non-head tasks
-                t.hoisted,
-                t.engine,
-            )
-            for t in renamed_tasks
-        ]
-
-        # Step 4: Construct a union of all the task bodies
-        union = f.union(
-            tuple(renamed_task_bodies),
-            [],
-            renamed_tasks[0].engine,
-        )
-
-        # Step 5: Add the final derive back
-        return f.logical(
-            (union, derive),
-            [],
-            renamed_tasks[0].engine,
-        )
-
-# Creates intermediary relations for all Data nodes and replaces said Data nodes
-# with a Lookup into these created relations. Reuse duplicate created relations.
-class EliminateData(Pass):
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        r = self.DataRewriter()
-        return r.walk(model)
-
-    # Does the actual work.
-    class DataRewriter(visitor.Rewriter):
-        new_relations: list[ir.Relation]
-        new_updates: list[ir.Logical]
-        # Counter for naming new relations.
-        # It must be that new_count == len new_updates == len new_relations.
-        new_count: int
-        # Cache for Data nodes to avoid creating duplicate intermediary relations
-        data_cache: dict[str, ir.Relation]
-
-        def __init__(self):
-            self.new_relations = []
-            self.new_updates = []
-            self.new_count = 0
-            self.data_cache = {}
-            super().__init__()
-
-        # Create a cache key for a Data node based on its structure and content
-        def _data_cache_key(self, node: ir.Data) -> str:
-            values = pd.util.hash_pandas_object(node.data).values
-            return hashlib.sha256(bytes(values)).hexdigest()
-
-        def _intermediary_relation(self, node: ir.Data) -> ir.Relation:
-            cache_key = self._data_cache_key(node)
-            if cache_key in self.data_cache:
-                return self.data_cache[cache_key]
-            self.new_count += 1
-            intermediary_name = f"formerly_Data_{self.new_count}"
-
-            intermediary_relation = f.relation(
-                intermediary_name,
-                [f.field(v.name, v.type) for v in node.vars]
-            )
-            self.new_relations.append(intermediary_relation)
-
-            intermediary_update = f.logical([
-                # For each row (union), equate values and their variable (logical).
-                f.union(
-                    [
-                        f.logical(
-                            [
-                                f.lookup(rel_builtins.eq, [f.literal(val), var])
-                                for (val, var) in zip(row, node.vars)
-                            ],
-                        )
-                        for row in node
-                    ],
-                    hoisted = node.vars,
-                ),
-                # And pop it back into the relation.
-                f.update(intermediary_relation, node.vars, ir.Effect.derive),
-            ])
-            self.new_updates.append(intermediary_update)
-
-            # Cache the result for reuse
-            self.data_cache[cache_key] = intermediary_relation
-
-            return intermediary_relation
-
-        # Create a new intermediary relation representing the Data (and pop it in
-        # new_updates/new_relations) and replace this Data with a Lookup of said
-        # intermediary.
-        def handle_data(self, node: ir.Data, parent: ir.Node) -> ir.Lookup:
-            intermediary_relation = self._intermediary_relation(node)
-            replacement_lookup = f.lookup(intermediary_relation, node.vars)
-
-            return replacement_lookup
-
-        # Walks the model for the handle_data work then updates the model with
-        # the new state.
-        def handle_model(self, model: ir.Model, parent: None):
-            walked_model = super().handle_model(model, parent)
-            assert len(self.new_relations) == len(self.new_updates) and self.new_count == len(self.new_relations)
-
-            # This is okay because its LQP.
-            assert isinstance(walked_model.root, ir.Logical)
-            root_logical = cast(ir.Logical, walked_model.root)
-
-            # We may need to add the new intermediaries from handle_data to the model.
-            if self.new_count  == 0:
-                return model
-            else:
-                return ir.Model(
-                    walked_model.engines,
-                    walked_model.relations | self.new_relations,
-                    walked_model.types,
-                    ir.Logical(
-                        root_logical.engine,
-                        root_logical.hoisted,
-                        root_logical.body + tuple(self.new_updates),
-                        root_logical.annotations,
-                    ),
-                    walked_model.annotations,
-                )
-
-# Deduplicate Vars in Updates and Outputs.
-class DeduplicateVars(Pass):
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        r = self.VarDeduplicator()
-        return r.walk(model)
-
-    # Return 1) a new list of Values with no duplicates (at the object level) and
-    # 2) equalities between any original Value and a deduplicated Value.
-    @staticmethod
-    def dedup_values(vals: Sequence[ir.Value]) -> Tuple[List[ir.Value], List[ir.Lookup]]:
-        # If a var is seen more than once, it is a duplicate and we will create
-        # a new Var and equate it with the seen one.
-        seen_vars = set()
-
-        new_vals = []
-        eqs = []
-
-        for i, val in enumerate(vals):
-            # Duplicates can only occur within Vars.
-            # TODO: we don't know for sure if these are the only relevant cases.
-            if isinstance(val, ir.Default) or isinstance(val, ir.Var):
-                var = val if isinstance(val, ir.Var) else val.var
-                if var in seen_vars:
-                    new_var = ir.Var(var.type, var.name + "_dup_" + str(i))
-                    new_val = new_var if isinstance(val, ir.Var) else ir.Default(new_var, val.value)
-                    new_vals.append(new_val)
-                    eqs.append(f.lookup(rel_builtins.eq, [new_var, var]))
-                else:
-                    seen_vars.add(var)
-                    new_vals.append(val)
-            else:
-                # No possibility of problematic duplication.
-                new_vals.append(val)
-
-        return new_vals, eqs
-
-    # Returns a reconstructed output with no duplicate variable objects
-    # (dedup_values) and now necessary equalities between any two previously
-    # duplicate variables.
-    @staticmethod
-    def dedup_output(output: ir.Output) -> List[Union[ir.Output, ir.Lookup]]:
-        vals = helpers.output_values(output.aliases)
-        deduped_vals, req_lookups = DeduplicateVars.dedup_values(vals)
-        # Need the names so we can recombine.
-        alias_names = output_names(output.aliases)
-        new_output = ir.Output(
-            output.engine,
-            FrozenOrderedSet(list(zip(alias_names, deduped_vals))),
-            output.keys,
-            output.annotations,
-        )
-        return req_lookups + [new_output]
-
-    # Returns a replacement update with no duplicate variable objects
-    # (dedup_values) and now necessary equalities between any two previously
-    # duplicate variables.
-    @staticmethod
-    def dedup_update(update: ir.Update) -> List[Union[ir.Update, ir.Lookup]]:
-        deduped_vals, req_lookups = DeduplicateVars.dedup_values(update.args)
-        new_update = ir.Update(
-            update.engine,
-            update.relation,
-            tuple(deduped_vals),
-            update.effect,
-            update.annotations,
-        )
-        return req_lookups + [new_update]
-
-    # Does the actual work.
-    class VarDeduplicator(visitor.Rewriter):
-        def __init__(self):
-            super().__init__()
-
-        # We implement handle_logical instead of handle_update/handle_output
-        # because in addition to modifying said update/output we require new
-        # lookups (equality between original and deduplicated variables).
-        def handle_logical(self, node: ir.Logical, parent: ir.Node):
-            # In order to recurse over subtasks.
-            node = super().handle_logical(node, parent)
-
-            new_body = []
-            for subtask in node.body:
-                if isinstance(subtask, ir.Output):
-                    new_body.extend(DeduplicateVars.dedup_output(subtask))
-                elif isinstance(subtask, ir.Update):
-                    new_body.extend(DeduplicateVars.dedup_update(subtask))
-                else:
-                    new_body.append(subtask)
-
-            return ir.Logical(
-                node.engine,
-                node.hoisted,
-                tuple(new_body),
-                node.annotations
-            )
-
-# Generate date arithmetic expressions, such as
-# `rel_primitive_date_add(:day, [date] delta, res_2)` by finding the period
-# expression for the delta and adding the period type to the date arithmetic expression.
-#
-# date_add and it's kin are generated by a period expression, e.g.,
-# `day(delta, res_1)`
-# followed by the date arithmetic expression using the period
-# `date_add([date] res_1 res_2)`
-class PeriodMath(Pass):
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        period_rewriter = self.PeriodRewriter()
-        model = period_rewriter.walk(model)
-        period_math_rewriter = self.PeriodMathRewriter(period_rewriter.period_vars)
-        model = period_math_rewriter.walk(model)
-        return model
-
-    # Find all period builtins. We need to make them safe for the emitter (either by
-    # translating to a cast, or removing) and store the variable and period type for use
-    # in the date/datetime add/subtract expressions.
-    class PeriodRewriter(visitor.Rewriter):
-        def __init__(self):
-            super().__init__()
-            self.period_vars: dict[ir.Var, str] = {}
-
-        def handle_lookup(self, node: ir.Lookup, parent: ir.Node) -> ir.Lookup:
-            if not rel_builtins.is_builtin(node.relation):
-                return node
-
-            if node.relation.name not in {
-                "year", "month", "week", "day", "hour", "minute", "second", "millisecond", "microsecond", "nanosecond"
-            }:
-                return node
-
-            assert len(node.args) == 2, "Expect 2 arguments for period builtins"
-            assert isinstance(node.args[1], ir.Var), "Expect result to be a variable"
-            period = node.relation.name
-            result_var = node.args[1]
-            self.period_vars[result_var] = period
-
-            # Ideally we could now remove the unused and unhandled period type construction
-            # but we may also need to cast the original variable to an Int64 for use by the
-            # date/datetime add/subtract expressions.
-            # TODO: Remove the node entirely where possible and update uses of the result
-            return f.lookup(rel_builtins.cast, [types.Int64, node.args[0], result_var])
-
-    # Update date/datetime add/subtract expressions with period information.
-    class PeriodMathRewriter(visitor.Rewriter):
-        def __init__(self, period_vars: dict[ir.Var, str]):
-            super().__init__()
-            self.period_vars: dict[ir.Var, str] = period_vars
-
-        def handle_lookup(self, node: ir.Lookup, parent: ir.Node) -> ir.Lookup:
-            if not rel_builtins.is_builtin(node.relation):
-                return node
-
-            if node.relation.name not in {
-                "date_add", "date_subtract", "datetime_add", "datetime_subtract"
-            }:
-                return node
-
-            if len(node.args) == 4:
-                # We've already visited this lookup
-                return node
-
-            assert isinstance(node.args[1], ir.Var), "Expect period to be a variable"
-            period_var = node.args[1]
-            assert period_var in self.period_vars, "datemath found, but no vars to insert"
-
-            period = self.period_vars[period_var]
-
-            new_args = [f.literal(period, types.Symbol)] + [arg for arg in node.args]
-
-            return f.lookup(node.relation, new_args)
-
-# Rewrite constants to vars in Updates. This results in a more normalized format where
-# updates contain only variables. This allows for easier rewrites in later passes.
-class ConstantsToVars(Pass):
-    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
-        r = self.ConstantToVarRewriter()
-        return r.walk(model)
-
-    # Return 1) a new list of Values with no duplicates (at the object level) and
-    # 2) equalities between any original Value and a deduplicated Value.
-    @staticmethod
-    def replace_constants_with_vars(vals: Sequence[ir.Value]) -> Tuple[List[ir.Value], List[ir.Lookup]]:
-        new_vals = []
-        eqs = []
-
-        for i, val in enumerate(vals):
-            if isinstance(val, ir.PyValue) or isinstance(val, ir.Literal):
-                # Replace constant with a new Var.
-                typ = typer.to_type(val)
-                assert isinstance(typ, ir.ScalarType), "can only replace scalar constants with vars"
-                new_var = ir.Var(typ, f"{typ.name.lower()}_{i}")
-                new_vals.append(new_var)
-                eqs.append(f.lookup(rel_builtins.eq, [new_var, val]))
-            else:
-                new_vals.append(val)
-
-        return new_vals, eqs
-
-    @staticmethod
-    def dedup_update(update: ir.Update) -> List[Union[ir.Update, ir.Lookup]]:
-        deduped_vals, req_lookups = ConstantsToVars.replace_constants_with_vars(update.args)
-        new_update = ir.Update(
-            update.engine,
-            update.relation,
-            tuple(deduped_vals),
-            update.effect,
-            update.annotations,
-        )
-        return req_lookups + [new_update]
-
-    # Does the actual work.
-    class ConstantToVarRewriter(visitor.Rewriter):
-        def __init__(self):
-            super().__init__()
-
-        # We implement handle_logical instead of handle_update because in
-        # addition to modifying said update we require new lookups (equality
-        # between original and deduplicated variables).
-        def handle_logical(self, node: ir.Logical, parent: ir.Node):
-            # In order to recurse over subtasks.
-            node = super().handle_logical(node, parent)
-
-            new_body = []
-            for subtask in node.body:
-                if isinstance(subtask, ir.Update):
-                    new_body.extend(ConstantsToVars.dedup_update(subtask))
-                else:
-                    new_body.append(subtask)
-
-            return ir.Logical(
-                node.engine,
-                node.hoisted,
-                tuple(new_body),
-                node.annotations
-            )