from starkware.cairo.common.dict_access import DictAccess
from starkware.cairo.common.math import assert_lt_felt

// Verifies that dict_accesses lists valid chronological accesses (and updates)
// to a mutable dictionary and outputs a squashed dict with one DictAccess instance per key
// (value before and value after) which summarizes all the changes to that key.
//
// Example:
//   Input: {(key1, 0, 2), (key1, 2, 7), (key2, 4, 1), (key1, 7, 5), (key2, 1, 2)}
//   Output: {(key1, 0, 5), (key2, 4, 2)}
//
// Arguments:
// dict_accesses - a pointer to the beginning of an array of DictAccess instances. The format of
//   each entry is a triplet (key, prev_value, new_value).
// dict_accesses_end - a pointer to the end of said array.
// squashed_dict - a pointer to an output array, which will be filled with
//   DictAccess instances sorted by key with the first and last value for each key.
//
// Returns:
// squashed_dict - end pointer to squashed_dict.
//
// Implicit arguments:
// range_check_ptr - range check builtin pointer.
func squash_dict{range_check_ptr}(
    dict_accesses: DictAccess*, dict_accesses_end: DictAccess*, squashed_dict: DictAccess*
) -> (squashed_dict: DictAccess*) {
    alloc_locals;
    %{ vm_enter_scope() %}
    local ptr_diff = dict_accesses_end - dict_accesses;

    if (ptr_diff == 0) {
        // Access array is empty, nothing to check.
        %{ vm_exit_scope() %}
        return (squashed_dict=squashed_dict);
    }
    local first_key;
    local big_keys;
    tempvar n_accesses = ptr_diff / DictAccess.SIZE;
    %{
        dict_access_size = ids.DictAccess.SIZE
        address = ids.dict_accesses.address_
        assert ids.ptr_diff % dict_access_size == 0, \
            'Accesses array size must be divisible by DictAccess.SIZE'
        n_accesses = ids.n_accesses
        if '__squash_dict_max_size' in globals():
            assert n_accesses <= __squash_dict_max_size, \
                f'squash_dict() can only be used with n_accesses<={__squash_dict_max_size}. ' \
                f'Got: n_accesses={n_accesses}.'
        # A map from key to the list of indices accessing it.
        access_indices = {}
        for i in range(n_accesses):
            key = memory[address + dict_access_size * i]
            access_indices.setdefault(key, []).append(i)
        # Descending list of keys.
        keys = sorted(access_indices.keys(), reverse=True)
        # Are the keys used bigger than range_check bound.
        ids.big_keys = 1 if keys[0] >= range_check_builtin.bound else 0
        ids.first_key = key = keys.pop()
    %}

    // Call inner.
    if (big_keys != 0) {
        tempvar range_check_ptr = range_check_ptr;
    } else {
        assert first_key = [range_check_ptr];
        tempvar range_check_ptr = range_check_ptr + 1;
    }
    let (range_check_ptr, squashed_dict) = squash_dict_inner(
        range_check_ptr=range_check_ptr,
        dict_accesses=dict_accesses,
        dict_accesses_end_minus1=dict_accesses_end - 1,
        key=first_key,
        remaining_accesses=n_accesses,
        squashed_dict=squashed_dict,
        big_keys=big_keys,
    );
    %{ vm_exit_scope() %}
    return (squashed_dict=squashed_dict);
}

// Inner tail-recursive function for squash_dict.
//
// Arguments:
// range_check_ptr - range check builtin pointer.
// dict_accesses - a pointer to the beginning of an array of DictAccess instances.
// dict_accesses_end_minus1 - a pointer to the end of said array, minus 1.
// key - current DictAccess key to check.
// remaining_accesses - remaining number of accesses that need to be accounted for. Starts with
//   the total number of entries in dict_accesses array, and slowly decreases until it reaches 0.
// squashed_dict - a pointer to an output array, which will be filled with
// DictAccess instances sorted by key with the first and last value for each key.
//
// Hints:
// keys - a descending list of the keys for which we have accesses. Destroyed in the process.
// access_indices - A map from key to a descending list of indices in the dict_accesses array that
//   access this key. Destroyed in the process.
//
// Returns:
// range_check_ptr - updated range check builtin pointer.
// squashed_dict - end pointer to squashed_dict.
func squash_dict_inner(
    range_check_ptr,
    dict_accesses: DictAccess*,
    dict_accesses_end_minus1: felt*,
    key,
    remaining_accesses,
    squashed_dict: DictAccess*,
    big_keys,
) -> (range_check_ptr: felt, squashed_dict: DictAccess*) {
    alloc_locals;

    let dict_diff: DictAccess* = squashed_dict;

    // Loop to verify chronological accesses to the key.
    // These values are not needed from previous iteration.
    struct LoopTemps {
        index_delta_minus1: felt,
        index_delta: felt,
        ptr_delta: felt,
        should_continue: felt,
    }
    // These values are needed from previous iteration.
    struct LoopLocals {
        value: felt,
        access_ptr: DictAccess*,
        range_check_ptr: felt,
    }

    // Prepare first iteration.
    %{
        current_access_indices = sorted(access_indices[key])[::-1]
        current_access_index = current_access_indices.pop()
        memory[ids.range_check_ptr] = current_access_index
    %}
    // Check that first access_index >= 0.
    tempvar current_access_index = [range_check_ptr];
    tempvar ptr_delta = current_access_index * DictAccess.SIZE;

    let first_loop_locals = cast(ap, LoopLocals*);
    first_loop_locals.access_ptr = dict_accesses + ptr_delta, ap++;
    let first_access: DictAccess* = first_loop_locals.access_ptr;
    first_loop_locals.value = first_access.new_value, ap++;
    first_loop_locals.range_check_ptr = range_check_ptr + 1, ap++;

    // Verify first key.
    key = first_access.key;

    // Write key and first value to dict_diff.
    key = dict_diff.key;
    // Use a local variable, instead of a tempvar, to avoid increasing ap.
    local first_value = first_access.prev_value;
    assert first_value = dict_diff.prev_value;

    // Skip loop nondeterministically if necessary.
    // The verifier doesn't care if the loop is skipped or not. The only thing it checks is that
    // the function iterated over remaining_accesses accesses in total
    // with ascending keys and ascending indices for the same key.
    // This guarantees that all the entries were visited exactly once.
    local should_skip_loop;
    %{ ids.should_skip_loop = 0 if current_access_indices else 1 %}
    jmp skip_loop if should_skip_loop != 0;

    loop:
    // Define references to access the values from the previous iteration,
    // the temporary variables and the values for the current iteration.
    let prev_loop_locals = cast(ap - LoopLocals.SIZE, LoopLocals*);
    let loop_temps = cast(ap, LoopTemps*);
    let loop_locals = cast(ap + LoopTemps.SIZE, LoopLocals*);

    // Check access_index.
    %{
        new_access_index = current_access_indices.pop()
        ids.loop_temps.index_delta_minus1 = new_access_index - current_access_index - 1
        current_access_index = new_access_index
    %}
    // Check that new access_index > prev access_index.
    loop_temps.index_delta_minus1 = [prev_loop_locals.range_check_ptr], ap++;
    loop_temps.index_delta = loop_temps.index_delta_minus1 + 1, ap++;
    loop_temps.ptr_delta = loop_temps.index_delta * DictAccess.SIZE, ap++;
    loop_locals.access_ptr = prev_loop_locals.access_ptr + loop_temps.ptr_delta, ap++;

    // Check valid transition.
    let access: DictAccess* = loop_locals.access_ptr;
    prev_loop_locals.value = access.prev_value;
    loop_locals.value = access.new_value, ap++;

    // Verify key.
    key = access.key;

    // Next range_check_ptr.
    loop_locals.range_check_ptr = prev_loop_locals.range_check_ptr + 1, ap++;

    // The verifier doesn't care how many loop iterations are executed. See comment above.
    %{ ids.loop_temps.should_continue = 1 if current_access_indices else 0 %}
    jmp loop if loop_temps.should_continue != 0, ap++;

    skip_loop:
    let last_loop_locals = cast(ap - LoopLocals.SIZE, LoopLocals*);

    // Check if address is out of bounds.
    %{ assert len(current_access_indices) == 0 %}
    [ap] = dict_accesses_end_minus1 - cast(last_loop_locals.access_ptr, felt);
    [ap] = [last_loop_locals.range_check_ptr], ap++;
    // Calculating the number of used accesses from the number of range check usages for efficiency.
    tempvar n_used_accesses = last_loop_locals.range_check_ptr - range_check_ptr;
    %{ assert ids.n_used_accesses == len(access_indices[key]) %}

    // Write last value to dict_diff.
    last_loop_locals.value = dict_diff.new_value;

    let range_check_ptr = last_loop_locals.range_check_ptr + 1;
    tempvar remaining_accesses = remaining_accesses - n_used_accesses;

    // Exit recursion when done.
    if (remaining_accesses == 0) {
        %{ assert len(keys) == 0 %}
        return (range_check_ptr=range_check_ptr, squashed_dict=squashed_dict + DictAccess.SIZE);
    }

    // ap points to the address of the beginning of the next access in the list.
    let next_key = [ap];
    ap += 1;
    // Guess next_key and check that next_key > key.
    %{
        assert len(keys) > 0, 'No keys left but remaining_accesses > 0.'
        ids.next_key = key = keys.pop()
    %}

    if (big_keys != 0) {
        assert_lt_felt{range_check_ptr=range_check_ptr}(a=key, b=next_key);
        tempvar dict_accesses = dict_accesses;
        tempvar dict_accesses_end_minus1 = dict_accesses_end_minus1;
        tempvar next_key = next_key;
        tempvar remaining_accesses = remaining_accesses;
    } else {
        assert [range_check_ptr] = next_key - (key + 1);
        tempvar range_check_ptr = range_check_ptr + 1;
        tempvar dict_accesses = dict_accesses;
        tempvar dict_accesses_end_minus1 = dict_accesses_end_minus1;
        tempvar next_key = next_key;
        tempvar remaining_accesses = remaining_accesses;
    }

    return squash_dict_inner(
        range_check_ptr=range_check_ptr,
        dict_accesses=dict_accesses,
        dict_accesses_end_minus1=dict_accesses_end_minus1,
        key=next_key,
        remaining_accesses=remaining_accesses,
        squashed_dict=squashed_dict + DictAccess.SIZE,
        big_keys=big_keys,
    );
}
