patricia.cairo

from starkware.cairo.common.alloc import alloc
from starkware.cairo.common.cairo_builtins import HashBuiltin
from starkware.cairo.common.dict_access import DictAccess
from starkware.cairo.common.hash import hash2
from starkware.cairo.common.math import (
    assert_in_range,
    assert_le,
    assert_lt_felt,
    assert_nn,
    assert_nn_le,
    assert_not_zero,
)
from starkware.cairo.common.patricia_utils import (
    MAX_LENGTH,
    NodeEdge,
    ParticiaGlobals,
    PatriciaUpdateConstants,
)

// ADDITIONAL_IMPORTS_MACRO()

// Given an edge node hash, opens the hash using the preimage hint, and returns a NodeEdge object.
func open_edge{hash_ptr: HashBuiltin*, range_check_ptr}(globals: ParticiaGlobals*, node: felt) -> (
    edge: NodeEdge*
) {
    alloc_locals;
    local edge: NodeEdge*;

    %{
        ids.edge = segments.add()
        ids.edge.length, ids.edge.path, ids.edge.bottom = preimage[ids.node]
        ids.hash_ptr.result = ids.node - ids.edge.length
        if __patricia_skip_validation_runner is not None:
            # Skip validation of the preimage dict to speed up the VM. When this flag is set,
            # mistakes in the preimage dict will be discovered only in the prover.
            __patricia_skip_validation_runner.verified_addresses.add(
                ids.hash_ptr + ids.HashBuiltin.result)
    %}
    // Validity checks.
    assert_in_range(edge.length, 1, MAX_LENGTH + 1);
    assert_lt_felt(edge.path, globals.pow2[edge.length]);
    // Note: we do not explicitly verify that bottom_hash is binary or leaf here, since it will be
    // verified later in the algorithm if necessary.
    assert hash_ptr.x = edge.bottom;
    assert hash_ptr.y = edge.path;
    // PREPARE_ADDITIONAL_HASH_INPUTS_MACRO(hash_ptr)
    assert node = hash_ptr.result + edge.length;
    let hash_ptr = hash_ptr + HashBuiltin.SIZE;
    return (edge=edge);
}

// Traversal:
// See patricia_utils.py for details about the node representation.
// The patricia updates algorithm is composed of 2 traversals: on the previous and the new trees.
// Each traversal outputs:
// 1. update_ptr - Sorted list of leaf indices and values of the traversal.
// 2. siblings - An encoding of the list of untraversed siblings of the nodes we traversed:
//      For every node in the underlying merkle tree (in-order), if we only traverse one of its
//      children, the other child will be encoded in this list.
//      The encoding is as follows:
//      * A value of 0 means a 0-sibling - empty subtree.
//      * A value of n, for 2 <= n <= MAX_LENGTH means n non leaf 0-siblings.
//      * Otherwise, the value is the sibling value.
//          It may be a hash (which we assume to be > MAX_LENGTH) for an inner node or a value
//          (any felt) for a leaf sibling.
// If the tree traversals output the same leaf indices, and the same siblings, it is guaranteed
// that all the non traversed leaves have the same value.
// All traversal functions do at least one leaf update in the given subtree.
//
// Args:
// * height - The height of the subtree rooted at this node.
// * path - The path from the root to this node.
//
// Hint args:
// * node - The current traversal node update tree. See build_update_tree() for more info.
// * descent_map - A map from for all the nodes on which we should descend. See get_descents()
//     in patricia_utils.py.
// * common_args - Contains itself and other relevant hint variables in the context of a traversal.

// Traverses an empty subtree.
func traverse_empty{update_ptr: DictAccess*, range_check_ptr, siblings: felt*}(
    globals: ParticiaGlobals*, height: felt, path: felt
) {
    if (height == 0) {
        assert update_ptr.key = path;
        let value = [cast(update_ptr, felt*) + globals.access_offset];
        assert value = 0;
        let update_ptr = update_ptr + DictAccess.SIZE;
        return ();
    }

    // Decide non deterministically if we should traverse both children or just one.
    // This is more efficient that trying to determine it from update_ptr.
    %{
        from starkware.python.merkle_tree import decode_node
        left_child, right_child, case = decode_node(node)
        memory[ap] = 1 if case != 'both' else 0
    %}
    jmp skip_both if [ap] != 0, ap++;

    // Traverse both children.
    %{ vm_enter_scope(dict(node=left_child, **common_args)) %}
    traverse_empty(globals=globals, height=height - 1, path=path * 2);
    %{ vm_exit_scope() %}

    %{ vm_enter_scope(dict(node=right_child, **common_args)) %}
    traverse_empty(globals=globals, height=height - 1, path=path * 2 + 1);
    %{ vm_exit_scope() %}
    return ();

    skip_both:
    %{
        descend = descent_map.get((ids.height, ids.path))
        memory[ap] = 0 if descend is None else 1
    %}
    jmp skip_single if [ap] != 0, ap++;

    // Single.
    let child_bit = [ap];
    %{
        ids.child_bit = 0 if case == 'left' else 1
        new_node = left_child if case == 'left' else right_child
        vm_enter_scope(dict(node=new_node, **common_args))
    %}
    child_bit = child_bit * child_bit, ap++;
    assert [siblings] = 0;
    let siblings = siblings + 1;
    traverse_empty(globals=globals, height=height - 1, path=path * 2 + child_bit);
    %{ vm_exit_scope() %}
    return ();

    skip_single:
    // Descend.
    tempvar word;
    %{ memory[ids.siblings], ids.word = descend %}
    tempvar length = [siblings];
    let siblings = siblings + 1;

    assert_in_range(length, 2, height + 1);
    tempvar length_pow2 = globals.pow2[length];
    assert_lt_felt(word, length_pow2);
    %{
        new_node = node
        for i in range(ids.length - 1, -1, -1):
            new_node = new_node[(ids.word >> i) & 1]
        vm_enter_scope(dict(node=new_node, **common_args))
    %}
    traverse_empty(globals=globals, height=height - length, path=path * length_pow2 + word);
    %{ vm_exit_scope() %}
    return ();
}

// Traverses a subtree rooted at the given NodeEdge.
func traverse_edge{
    hash_ptr: HashBuiltin*, range_check_ptr, update_ptr: DictAccess*, siblings: felt*
}(globals: ParticiaGlobals*, height: felt, path: felt, edge: NodeEdge) {
    if (edge.length == 0) {
        return traverse_binary_or_leaf(globals=globals, height=height, path=path, node=edge.bottom);
    }

    alloc_locals;

    %{
        descend = descent_map.get((ids.height, ids.path))
        memory[ap] = 0 if descend is None else 1
    %}
    jmp descend if [ap] != 0, ap++;

    // Extract one bit from the edge. edge.length is guaranteed to be >= 1.
    local new_length = edge.length - 1;
    local bound = globals.pow2[new_length];
    local bit;
    %{ ids.bit = (ids.edge.path >> ids.new_length) & 1 %}
    bit * bit = bit;
    local new_path = edge.path - bit * bound;
    assert_lt_felt(new_path, bound);
    local new_edge: NodeEdge = NodeEdge(length=new_length, path=new_path, bottom=edge.bottom);

    // Decide case.
    %{
        from starkware.python.merkle_tree import decode_node
        left_child, right_child, case = decode_node(node)
        memory[ap] = int(case != 'both')
    %}
    jmp skip_both if [ap] != 0, ap++;

    // Traverse both children.
    if (bit == 0) {
        %{ vm_enter_scope(dict(node=left_child, **common_args)) %}
        traverse_edge(globals=globals, height=height - 1, path=path * 2, edge=new_edge);
        %{ vm_exit_scope() %}
        local hash_ptr: HashBuiltin* = hash_ptr;

        %{ vm_enter_scope(dict(node=right_child, **common_args)) %}
        traverse_empty(globals=globals, height=height - 1, path=path * 2 + 1);
        %{ vm_exit_scope() %}
        return ();
    } else {
        %{ vm_enter_scope(dict(node=left_child, **common_args)) %}
        traverse_empty(globals=globals, height=height - 1, path=path * 2);
        %{ vm_exit_scope() %}
        ap += 0;

        %{ vm_enter_scope(dict(node=right_child, **common_args)) %}
        traverse_edge(globals=globals, height=height - 1, path=path * 2 + 1, edge=new_edge);
        %{ vm_exit_scope() %}
        return ();
    }

    skip_both:
    %{ memory[ap] = int(case == 'right') ^ ids.bit %}
    jmp skip_non_empty_child if [ap] != 0, ap++;

    // Traverse the non-empty child.
    assert [siblings] = 0;
    let siblings = siblings + 1;
    %{
        new_node = left_child if ids.bit == 0 else right_child
        vm_enter_scope(dict(node=new_node, **common_args))
    %}
    traverse_edge(globals=globals, height=height - 1, path=path * 2 + bit, edge=new_edge);
    %{ vm_exit_scope() %}
    return ();

    skip_non_empty_child:
    // Traverse the empty child.
    local range_check_ptr = range_check_ptr;
    local new_hash_ptr: HashBuiltin*;

    // Reserve a spot for the sibling. It is more efficient to compute it after the recursive
    // traversal.
    let current_sibling = siblings;
    let siblings = siblings + 1;

    // Traverse empty side.
    %{
        new_node = left_child if ids.bit == 1 else right_child
        vm_enter_scope(dict(node=new_node, **common_args))
    %}
    traverse_empty(globals=globals, height=height - 1, path=path * 2 + 1 - bit);
    %{ vm_exit_scope() %}

    if (edge.length == 1) {
        // In this case, the sibling is the bottom of our edge.
        // Make sure edge.bottom is binary or leaf.
        assert [current_sibling] = edge.bottom;
        if (height != 1) {
            // This check should only be done on the new tree.
            if (globals.access_offset == 2) {
                // PREPARE_ADDITIONAL_HASH_INPUTS_MACRO(hash_ptr)
                hash_ptr.result = edge.bottom;
                %{
                    ids.hash_ptr.x, ids.hash_ptr.y = preimage[ids.edge.bottom]
                    if __patricia_skip_validation_runner:
                        # Skip validation of the preimage dict to speed up the VM. When this flag is
                        # set, mistakes in the preimage dict will be discovered only in the prover.
                        __patricia_skip_validation_runner.verified_addresses.add(
                            ids.hash_ptr + ids.HashBuiltin.result)
                %}
                let hash_ptr = hash_ptr + HashBuiltin.SIZE;
                return ();
            } else {
                return ();
            }
        }
        return ();
    }

    // This is the case where we split an edge.
    let (hash) = hash2(edge.bottom, new_edge.path);
    assert [current_sibling] = hash + new_edge.length;
    return ();

    // Descend.
    descend:
    local length;
    local word;
    %{ ids.length, ids.word = descend %}
    assert [siblings] = length;
    let siblings = siblings + 1;

    // Check that the descend is valid.
    assert_in_range(length, 2, edge.length + 1);
    tempvar length_pow2 = globals.pow2[length];
    tempvar new_length = edge.length - length;
    tempvar new_length_pow2 = globals.pow2[new_length];
    tempvar new_path = edge.path - word * new_length_pow2;
    assert_lt_felt(word, length_pow2);
    assert_lt_felt(new_path, new_length_pow2);

    let new_edge: NodeEdge = NodeEdge(length=new_length, path=new_path, bottom=edge.bottom);
    %{
        new_node = node
        for i in range(ids.length - 1, -1, -1):
            new_node = new_node[(ids.word >> i) & 1]
        vm_enter_scope(dict(node=new_node, **common_args))
    %}
    traverse_edge(
        globals=globals, height=height - length, path=path * length_pow2 + word, edge=new_edge
    );
    %{ vm_exit_scope() %}

    return ();
}

// Traverses a subtree rooted at the given binary or leaf node with given hash/value.
func traverse_binary_or_leaf{
    hash_ptr: HashBuiltin*, range_check_ptr, update_ptr: DictAccess*, siblings: felt*
}(globals: ParticiaGlobals*, height: felt, path: felt, node: felt) {
    if (height == 0) {
        // Leaf.
        assert update_ptr.key = path;
        tempvar value = [cast(update_ptr, felt*) + globals.access_offset];
        assert value = node;
        assert_not_zero(value);
        let update_ptr = update_ptr + DictAccess.SIZE;
        return ();
    }

    alloc_locals;

    // Binary.
    let current_hash = hash_ptr;
    let hash_ptr = hash_ptr + HashBuiltin.SIZE;
    // PREPARE_ADDITIONAL_HASH_INPUTS_MACRO(current_hash)
    assert current_hash.result = node;

    %{
        from starkware.python.merkle_tree import decode_node
        left_child, right_child, case = decode_node(node)
        left_hash, right_hash = preimage[ids.node]

        # Fill non deterministic hashes.
        hash_ptr = ids.current_hash.address_
        memory[hash_ptr + ids.HashBuiltin.x] = left_hash
        memory[hash_ptr + ids.HashBuiltin.y] = right_hash

        if __patricia_skip_validation_runner:
            # Skip validation of the preimage dict to speed up the VM. When this flag is set,
            # mistakes in the preimage dict will be discovered only in the prover.
            __patricia_skip_validation_runner.verified_addresses.add(
                hash_ptr + ids.HashBuiltin.result)

        memory[ap] = int(case != 'both')
    %}
    jmp skip_both if [ap] != 0, ap++;

    // Traverse both children.
    %{ vm_enter_scope(dict(node=left_child, **common_args)) %}
    traverse_non_empty(globals=globals, height=height - 1, path=path * 2, node=current_hash.x);
    %{ vm_exit_scope() %}
    tempvar left_child = path * 2;
    %{ vm_enter_scope(dict(node=right_child, **common_args)) %}
    traverse_non_empty(
        globals=globals, height=height - 1, path=left_child + 1, node=current_hash.y
    );
    %{ vm_exit_scope() %}
    return ();

    skip_both:
    %{ memory[ap] = int(case != 'left') %}
    jmp skip_left if [ap] != 0, ap++;

    // Left.
    tempvar sib = current_hash.y;
    assert_not_zero(sib);
    assert [siblings] = sib;
    let siblings = siblings + 1;
    %{ vm_enter_scope(dict(node=left_child, **common_args)) %}
    traverse_non_empty(globals=globals, height=height - 1, path=path * 2, node=current_hash.x);
    %{ vm_exit_scope() %}
    return ();

    skip_left:
    %{ assert case == 'right' %}
    // Right.
    tempvar sib = current_hash.x;
    assert_not_zero(sib);
    assert [siblings] = sib;
    let siblings = siblings + 1;
    %{ vm_enter_scope(dict(node=right_child, **common_args)) %}
    traverse_non_empty(globals=globals, height=height - 1, path=path * 2 + 1, node=current_hash.y);
    %{ vm_exit_scope() %}
    return ();
}

// Traverses some of the leaves in the subtree rooted at node.
func traverse_node{
    hash_ptr: HashBuiltin*, range_check_ptr, update_ptr: DictAccess*, siblings: felt*
}(globals: ParticiaGlobals*, height: felt, path: felt, node: felt) {
    if (node == 0) {
        // Empty:
        traverse_empty(globals=globals, height=height, path=path);
        return ();
    }

    return traverse_non_empty(globals=globals, height=height, path=path, node=node);
}

// Same as traverse_node, but disallows empty nodes.
func traverse_non_empty{
    hash_ptr: HashBuiltin*, range_check_ptr, update_ptr: DictAccess*, siblings: felt*
}(globals: ParticiaGlobals*, height: felt, path: felt, node: felt) {
    %{ memory[ap] = 1 if ids.height == 0 or len(preimage[ids.node]) == 2 else 0 %}
    jmp binary if [ap] != 0, ap++;
    // Edge.
    let (edge) = open_edge(globals=globals, node=node);
    traverse_edge(globals=globals, height=height, path=path, edge=[edge]);
    return ();

    // Binary.
    binary:
    return traverse_binary_or_leaf(globals=globals, height=height, path=path, node=node);
}

// Computes a power of 2 array. In other words, writes the sequence 1, 2, 4, 8, ... to the given
// pointer (with the given length).
func compute_pow2_array(pow2_ptr: felt*, cur: felt, n: felt) {
    if (n == 0) {
        return ();
    }
    assert [pow2_ptr] = cur;
    return compute_pow2_array(pow2_ptr=pow2_ptr + 1, cur=cur * 2, n=n - 1);
}

// Performs an efficient update of multiple leaves in a Patricia Merkle tree.
//
// Arguments:
// update_ptr - a list of DictAccess instances sorted by key (e.g., the result of squash_dict).
// height - the height of the merkle tree.
// prev_root - the value of the root before the update.
// new_root - the value of the root after the update.
//
// Hint arguments:
// preimage - a dictionary from the hash value of a Patricia node to either
//   1. a pair of children values, for binary nodes.
//   2. a triplet of (length, path, bottom), for edge nodes.
//
// Implicit arguments:
// hash_ptr - hash builtin pointer.
//
// Assumptions: The keys in the update_ptr list are unique and sorted.
// Guarantees: All the keys in the update_ptr list are < 2**height.
func patricia_update{hash_ptr: HashBuiltin*, range_check_ptr}(
    update_ptr: DictAccess*, n_updates: felt, height: felt, prev_root: felt, new_root: felt
) {
    let (patricia_update_constants: PatriciaUpdateConstants*) = patricia_update_constants_new();
    patricia_update_using_update_constants(
        patricia_update_constants=patricia_update_constants,
        update_ptr=update_ptr,
        n_updates=n_updates,
        height=height,
        prev_root=prev_root,
        new_root=new_root,
    );

    return ();
}

func patricia_update_constants_new() -> (patricia_update_constants: PatriciaUpdateConstants*) {
    // Compute power-of-2 array for patricia updates.
    alloc_locals;
    let (local globals_pow2: felt*) = alloc();
    compute_pow2_array(pow2_ptr=globals_pow2, cur=1, n=MAX_LENGTH + 1);
    return (patricia_update_constants=new PatriciaUpdateConstants(globals_pow2=globals_pow2));
}

func patricia_update_using_update_constants{hash_ptr: HashBuiltin*, range_check_ptr}(
    patricia_update_constants: PatriciaUpdateConstants*,
    update_ptr: DictAccess*,
    n_updates: felt,
    height: felt,
    prev_root: felt,
    new_root: felt,
) {
    if (n_updates == 0) {
        prev_root = new_root;
        return ();
    }

    %{
        from starkware.cairo.common.patricia_utils import canonic, patricia_guess_descents
        from starkware.python.merkle_tree import build_update_tree

        # Build modifications list.
        modifications = []
        DictAccess_key = ids.DictAccess.key
        DictAccess_new_value = ids.DictAccess.new_value
        DictAccess_SIZE = ids.DictAccess.SIZE
        for i in range(ids.n_updates):
            curr_update_ptr = ids.update_ptr.address_ + i * DictAccess_SIZE
            modifications.append((
                memory[curr_update_ptr + DictAccess_key],
                memory[curr_update_ptr + DictAccess_new_value]))

        node = build_update_tree(ids.height, modifications)
        descent_map = patricia_guess_descents(
            ids.height, node, preimage, ids.prev_root, ids.new_root)
        del modifications
        __patricia_skip_validation_runner = globals().get(
            '__patricia_skip_validation_runner')

        common_args = dict(
            preimage=preimage, descent_map=descent_map,
            __patricia_skip_validation_runner=__patricia_skip_validation_runner)
        common_args['common_args'] = common_args
    %}
    alloc_locals;
    local update_end: DictAccess* = update_ptr + n_updates * DictAccess.SIZE;

    // Traverse prev tree.
    let (local siblings) = alloc();
    let original_update_ptr = update_ptr;
    let original_siblings = siblings;
    let (local globals_prev: ParticiaGlobals*) = alloc();
    assert [globals_prev] = ParticiaGlobals(
        pow2=patricia_update_constants.globals_pow2, access_offset=DictAccess.prev_value
    );

    assert_le(height, MAX_LENGTH);
    %{ vm_enter_scope(dict(node=node, **common_args)) %}
    with update_ptr, siblings {
        traverse_node(globals=globals_prev, height=height, path=0, node=prev_root);
    }
    %{ vm_exit_scope() %}
    assert update_ptr = update_end;
    local siblings_end: felt* = siblings;

    // Traverse new tree.
    let update_ptr = original_update_ptr;
    let siblings = original_siblings;
    let (local globals_new: ParticiaGlobals*) = alloc();
    assert [globals_new] = ParticiaGlobals(
        pow2=patricia_update_constants.globals_pow2, access_offset=DictAccess.new_value
    );

    %{ vm_enter_scope(dict(node=node, **common_args)) %}
    with update_ptr, siblings {
        traverse_node(globals=globals_new, height=height, path=0, node=new_root);
    }
    %{ vm_exit_scope() %}
    assert update_ptr = update_end;
    assert siblings = siblings_end;
    return ();
}