builtin_runner.py

from abc import ABC, abstractmethod
from typing import Any, Callable, Dict, List, Optional, Tuple

from starkware.cairo.lang.builtins.instance_def import BuiltinInstanceDef
from starkware.cairo.lang.vm.relocatable import MaybeRelocatable, RelocatableValue
from starkware.cairo.lang.vm.utils import MemorySegmentRelocatableAddresses
from starkware.python.math_utils import div_ceil, next_power_of_2, safe_div


class InsufficientAllocatedCells(Exception):
    pass


class BuiltinRunner(ABC):
    @abstractmethod
    def initialize_segments(self, runner):
        """
        Adds memory segments for the builtin.
        """

    @abstractmethod
    def initial_stack(self) -> List[MaybeRelocatable]:
        """
        Returns the initial stack elements enforced by this builtin.
        """

    @abstractmethod
    def get_needed_number_allocated_zeros(self) -> int:
        """
        The number of allocated zeros needed.
        Needed when the empty/zero instances of the builtin use pointers.
        """

    @abstractmethod
    def set_address_allocated_zeros(self, addr: RelocatableValue):
        """
        The address of allocated zeros needed.
        Needed when the empty/zero instances of the builtin use pointers.
        """

    @abstractmethod
    def final_stack(self, runner, pointer: MaybeRelocatable) -> MaybeRelocatable:
        """
        Reads values from the end of the stack ([pointer - 1], [pointer - 2], ...), and returns
        the updated pointer (e.g., pointer - 2 if two values were read).
        This function may also do builtin specific validation of said values.
        """

    @abstractmethod
    def get_used_cells(self, runner) -> int:
        """
        Returns the number of used cells.
        """

    @abstractmethod
    def get_used_instances(self, runner) -> int:
        """
        Returns the number of used instances.
        """

    @abstractmethod
    def get_allocated_instances(self, runner) -> int:
        """
        Returns the number of instances in the trace (including unused instances).
        """

    @abstractmethod
    def get_allocated_memory_units(self, runner) -> int:
        """
        Returns the number of memory units used by the builtin.
        """

    @abstractmethod
    def get_used_cells_and_allocated_size(self, runner) -> Tuple[int, int]:
        """
        Returns the number of used cells and the allocated size, and raises
        InsufficientAllocatedCells if there are more used cells than allocated cells.
        """

    @abstractmethod
    def finalize_segments(self, runner):
        """
        Calls runner.segments.finalize for the memory segments added in initialize_segments.
        """

    @abstractmethod
    def get_memory_segment_addresses(self, runner) -> Dict[str, MemorySegmentRelocatableAddresses]:
        """
        Returns a dict from segment name to MemorySegmentRelocatableAddresses
        (begin_addr and stop_ptr of the corresponding segment).
        """

    @abstractmethod
    def run_security_checks(self, runner):
        """
        Runs some security checks to make sure a proof can be generated given the memory.
        """

    def relocate(self, relocate_value: Callable[[MaybeRelocatable], MaybeRelocatable]):
        """
        Relocates the internal values of the builtin using the given function relocate_value.
        """
        return

    def add_auto_deduction_rules(self, runner):
        """
        Adds auto-deduction rules for this builtin (if applicable).
        Auto deduction rules are applied when an unknown memory cell in the builtin segment is
        accessed.
        """
        return

    def add_validation_rules(self, runner):
        """
        Adds validation rules for this builtin (if applicable).
        Validation rules are applied once a builtin instance is written to memory.
        For more details, see 'add_validation_rule' in validated_memory_dict.py.
        """
        return

    def air_private_input(self, runner) -> Dict[str, Any]:
        """
        Returns information about the builtin that should be added to the AIR private input.
        """
        return {}

    def get_range_check_usage(self, runner) -> Optional[Tuple[int, int]]:
        """
        Returns (rc_min, rc_max), i.e., the minimal and maximal range-checked values, if the builtin
        used any range check cells. Otherwise, returns None.
        """
        return None

    def get_used_perm_range_check_units(self, runner) -> int:
        """
        Returns the number of range check units used by the builtin.
        """
        return 0

    def get_used_diluted_check_units(self, diluted_spacing: int, diluted_n_bits: int) -> int:
        """
        Returns the number of diluted check units used by the builtin.
        """
        return 0

    def get_additional_data(self) -> Any:
        """
        Returns additional data that was created in the builtin runner. This data can be loaded
        to another builtin runner of the same type using extend_additional_data().
        This data must be JSON-serializable.
        """
        return

    def extend_additional_data(
        self,
        data: Any,
        relocate_callback: Callable[[MaybeRelocatable], MaybeRelocatable],
        data_is_trusted: bool = True,
    ):
        """
        Adds the additional data created by another instance of the builtin runner.
        relocate_callback is a callback function used to relocate the addresses.
        If data_is_trusted is False, the function does not assume that instances
        that were processed by the other builtin runner were properly validated.
        """
        return

    def get_memory_accesses(self, runner):
        """
        Returns memory addresses that are used by the builtin itself and therefore should not count
        as memory holes.
        """
        return {}


class BuiltinVerifier(ABC):
    @abstractmethod
    def expected_stack(self, public_input) -> Tuple[List[int], List[int]]:
        """
        Returns a pair (initial_stack, final_stack).
        They contain the expected elements of the initial stack or final stack that are associated
        with this builtin.
        """


class SimpleBuiltinRunner(BuiltinRunner):
    """
    A base class for simple builtins that use a single segment.
    """

    def __init__(
        self,
        name: str,
        included: bool,
        ratio: Optional[int],
        cells_per_instance: int,
        n_input_cells: int,
        instances_per_component: int = 1,
        additional_memory_units_per_instance: int = 0,
    ):
        """
        Constructs a SimpleBuiltinRunner.
        cells_per_instance is the number of memory cells per invocation.
        n_input_cells is the number of the first memory cells in each invocation that form the
        input. The rest of the cells are considered output.
        instances_per_component is the number of invocations being handled in each call to the
        corresponding component. It must divide the total number of invocations.
        """
        self.name = name
        self.included = included
        self.ratio = ratio
        self.instances_per_component = instances_per_component
        self._base: Optional[RelocatableValue] = None
        self.stop_ptr: Optional[RelocatableValue] = None
        self.cells_per_instance = cells_per_instance
        self.n_input_cells = n_input_cells
        self.additional_memory_units_per_instance = additional_memory_units_per_instance

    def get_needed_number_allocated_zeros(self) -> int:
        """
        The number of allocated zeros needed.
        Needed when the empty/zero instances of the builtin use pointers.
        """
        return 0

    def set_address_allocated_zeros(self, addr: RelocatableValue):
        """
        The address of allocated zeros needed.
        Needed when the empty/zero instances of the builtin use pointers.
        """
        raise NotImplementedError(
            f"set_address_allocated_zeros should not be called when needed number of allocated"
            + f"zeros is 0."
        )

    def get_instance_def(self) -> Optional[BuiltinInstanceDef]:
        """
        Returns a BuiltinInstanceDef object that represents the builtin if such object exists.
        """
        return None

    def initialize_segments(self, runner):
        self._base = runner.segments.add()

    @property
    def base(self) -> RelocatableValue:
        assert self._base is not None, "Uninitialized self.base."
        return self._base

    def initial_stack(self) -> List[MaybeRelocatable]:
        return [self.base] if self.included else []

    def final_stack(self, runner, pointer):
        if self.included:
            self.stop_ptr = runner.vm_memory[pointer - 1]
            used = self.get_used_instances(runner=runner) * self.cells_per_instance
            assert self.stop_ptr == self.base + used, (
                f"Invalid stop pointer for {self.name}. "
                + f"Expected: {self.base + used}, found: {self.stop_ptr}"
            )
            return pointer - 1
        else:
            self.stop_ptr = self.base
            return pointer

    def get_used_cells(self, runner):
        used = runner.segments.get_segment_used_size(self.base.segment_index)
        return used

    def get_used_instances(self, runner):
        return div_ceil(self.get_used_cells(runner), self.cells_per_instance)

    def get_allocated_instances(self, runner):
        if self.ratio is None:
            # Dynamic layout has the exact number of instances it needs (up to a power of 2).
            instances = self.get_used_instances(runner)
            needed_components = div_ceil(instances, self.instances_per_component)
            components = next_power_of_2(needed_components) if needed_components > 0 else 0
            return self.instances_per_component * components
        assert isinstance(self.ratio, int), "ratio is not an int"
        if self.ratio == 0:
            # The builtin is not used.
            return 0
        min_steps = self.ratio * self.instances_per_component
        if runner.vm.current_step < min_steps:
            raise InsufficientAllocatedCells(
                f"Number of steps must be at least {min_steps} for the {self.name} builtin."
            )
        return safe_div(runner.vm.current_step, self.ratio)

    def get_allocated_memory_units(self, runner):
        return (
            self.cells_per_instance + self.additional_memory_units_per_instance
        ) * self.get_allocated_instances(runner)

    def get_used_cells_and_allocated_size(self, runner):
        used = self.get_used_cells(runner)
        size = self.cells_per_instance * self.get_allocated_instances(runner)
        if used > size:
            raise InsufficientAllocatedCells(
                f"The {self.name} builtin used {used} cells but the capacity is {size}."
            )
        return used, size

    def finalize_segments(self, runner):
        _, size = self.get_used_cells_and_allocated_size(runner)

        runner.segments.finalize(segment_index=self.base.segment_index, size=size)

    def get_memory_segment_addresses(self, runner):
        return {
            self.name: MemorySegmentRelocatableAddresses(
                begin_addr=self.base,
                stop_ptr=self.stop_ptr,
            )
        }

    def run_security_checks(self, runner):
        offsets = {
            addr.offset
            for addr in runner.vm_memory.keys()
            if isinstance(addr, RelocatableValue) and addr.segment_index == self.base.segment_index
        }
        n = (max(offsets) // self.cells_per_instance + 1) if len(offsets) > 0 else 0

        # Verify that n is not too large to make sure the expected_offsets set that is constructed
        # below is not too large.
        assert n <= len(offsets) // self.n_input_cells, f"Missing memory cells for {self.name}."

        # Check that the two inputs (x and y) of each instance are set.
        expected_offsets = {
            self.cells_per_instance * i + j for i in range(n) for j in range(self.n_input_cells)
        }
        if not expected_offsets <= offsets:
            missing_offsets = list(expected_offsets - offsets)
            dots = "..." if len(missing_offsets) > 20 else "."
            missing_offsets_str = ", ".join(map(str, missing_offsets[:20])) + dots
            raise AssertionError(f"Missing memory cells for {self.name}: {missing_offsets_str}")

        # Verify auto deduction rules for the unassigned output cells.
        # Assigned output cells are checked as part of the call to verify_auto_deductions().
        for i in range(n):
            for j in range(self.n_input_cells, self.cells_per_instance):
                addr = self.base + (self.cells_per_instance * i + j)
                if runner.vm.validated_memory.get(addr) is None:
                    runner.vm.verify_auto_deductions_for_addr(addr)

    def get_memory_accesses(self, runner):
        segment_size = runner.segments.get_segment_size(self.base.segment_index)
        return {self.base + x for x in range(segment_size)}


class SimpleBuiltinRunnerWithLowRatio(SimpleBuiltinRunner):
    def __init__(
        self,
        name: str,
        included: bool,
        ratio: Optional[int],
        cells_per_instance: int,
        n_input_cells: int,
        instances_per_component: int = 1,
        additional_memory_units_per_instance: int = 0,
        ratio_den: int = 1,
    ):
        super().__init__(
            name=name,
            included=included,
            ratio=ratio,
            cells_per_instance=cells_per_instance,
            n_input_cells=n_input_cells,
            instances_per_component=instances_per_component,
            additional_memory_units_per_instance=additional_memory_units_per_instance,
        )
        self.ratio_den = ratio_den

    def get_allocated_instances(self, runner):
        if self.ratio is None:
            # Dynamic layout has the exact number of instances it needs (up to a power of 2).
            instances = self.get_used_instances(runner)
            needed_components = div_ceil(instances, self.instances_per_component)
            components = next_power_of_2(needed_components) if needed_components > 0 else 0
            return self.instances_per_component * components
        assert isinstance(self.ratio, int), "ratio is not an int"
        if self.ratio == 0:
            # The builtin is not used.
            return 0
        min_steps = div_ceil(self.ratio * self.instances_per_component, self.ratio_den)
        if runner.vm.current_step < min_steps:
            raise InsufficientAllocatedCells(
                f"Number of steps must be at least {min_steps} for the {self.name} builtin."
            )
        return safe_div(runner.vm.current_step * self.ratio_den, self.ratio)