Kakarot isolated Stack and Memory modules

kakarot_mem.cairo

// AUTHOR: Mario Iordanov
// Github: @marioiordanov
// Company: ShardLabs

%builtins output range_check

from starkware.cairo.common.serialize import serialize_word


from starkware.cairo.common.alloc import alloc

from starkware.cairo.common.dict import DictAccess, dict_read, dict_write
from starkware.cairo.common.default_dict import default_dict_new, default_dict_finalize
from starkware.cairo.common.uint256 import Uint256
from starkware.cairo.common.pow import pow
from starkware.cairo.common.math_cmp import is_le
from starkware.cairo.common.math import (assert_le, unsigned_div_rem)
from starkware.cairo.common.bool import FALSE
from starkware.cairo.common.registers import get_label_location

struct Stack{
    stack: DictAccess*,
    len: felt,
}

struct Memory{
    memory: DictAccess*,
    len: felt,
}

struct ExecutionContext{
    stack: Stack*,
    memory: Memory*,
}

func init_ctx() -> ExecutionContext*{
    alloc_locals;

    let stack = init_stack();
    let memory = init_memory();

    return new ExecutionContext(
        stack = stack,
        memory = memory
    );
}

func init_stack() -> Stack*{
    alloc_locals;
    let (stack:DictAccess*) = default_dict_new(0);

    return new Stack(
        stack = stack,
        len = 0
    );
}

func update_ctx_stack(
        self: ExecutionContext*, new_stack: Stack*
    ) -> ExecutionContext* {
        return new ExecutionContext(
            stack=new_stack,
            memory=self.memory,
            );
    }

func update_ctx_memory(
        self: ExecutionContext*, new_memory: Memory*
    ) -> ExecutionContext* {
        return new ExecutionContext(
            stack=self.stack,
            memory=new_memory,
            );
    }

func init_memory() -> Memory*{
    alloc_locals;
    let (memory:DictAccess*) = default_dict_new(0);

    return new Memory(
        memory = memory,
        len = 0
    );
}

func push{range_check_ptr}(self: Stack*, val: felt) -> Stack*{
    alloc_locals;

    let s = self.stack;
    let position = self.len;
    let (val_ptr) = alloc();
    [val_ptr] = val;
    let stack_element: Uint256 = bytes_i_to_uint256(val=val_ptr, i=1);

    dict_write{dict_ptr=s}(position, stack_element.high);
    dict_write{dict_ptr=s}(position + 1, stack_element.low);
    return new Stack(
        stack=s,
        len = self.len + 2
    );
}

func update_stack{range_check_ptr}(dict: DictAccess*, el: felt) -> Stack*{
    return new Stack(stack=dict, len=el);
}

func update_memory(dict: DictAccess*, el: felt) -> Memory*{
    return new Memory(memory=dict, len=el);
}

func pop{range_check_ptr}(self: Stack*) -> (new_stack: Stack*, val: felt){
    alloc_locals;
    let stack = self.stack;
    let position = self.len - 1;
    let (local popped: felt) = dict_read{dict_ptr=stack}(position);

    return (update_stack(stack, position), popped);
}

func stack_to_uint256{range_check_ptr}(
        word_dict: DictAccess*, stack_len: felt, n: felt, output: Uint256*
    ) -> (word_dict: DictAccess*) {
    if (n == 0) {
        return (word_dict=word_dict);
    }

    // Get Low and High of element at position N
    let (el_high) = dict_read{dict_ptr=word_dict}(stack_len - n);
    let (el_low) = dict_read{dict_ptr=word_dict}(stack_len - n + 1);

    // Save Uint256 value in array
    let n_index = n / 2 - 1;
    assert output[n_index] = Uint256(low=el_low, high=el_high);

    return stack_to_uint256(word_dict=word_dict, stack_len=stack_len, n=n - 2, output=output);
}

// @notice Pop N elements from the stack.
    // @param self - The pointer to the stack.
    // @param len - The len of elements to pop.
    // @return The new pointer to the stack.
    // @return elements the pointer to the first popped element.
    func pop_n{range_check_ptr}(self: Stack*, n: felt) -> (
        new_stack: Stack*, elements: Uint256*
    ) {
        alloc_locals;
        let word_dict = self.stack;
        let position_zero = self.len;

        // Check if there is underflow
        with_attr error_message("Kakarot: StackUnderflow") {
            assert_le(n * 2, position_zero);
        }

        let (new_elements: Uint256*) = alloc();

        // Generate an array of Uint256* to return
        let (word_dict) = stack_to_uint256(
            word_dict=word_dict, stack_len=position_zero, n=n * 2, output=new_elements
        );

        // Return Stack with updated Len
        let popped_len = 2 * n;
        return (
            new Stack(
            stack=word_dict,
            len=self.len - popped_len,
            ),
            new_elements,
        );
    }

func compute_half_uint256{range_check_ptr}(val: felt*, i: felt, res: felt) -> (res: felt) {
    if (i == 1) {
        return (res=res + [val]);
    }
    let (temp_pow) = pow(256, i - 1);
    let (res) = compute_half_uint256(val + 1, i - 1, res + [val] * temp_pow);
    return (res=res);
}

func bytes_i_to_uint256{range_check_ptr}(val: felt*, i: felt) -> Uint256 {
    alloc_locals;

    let is_sequence_32_bytes_or_less = is_le(i, 32);
    with_attr error_message("number must be shorter than 32 bytes") {
        assert is_sequence_32_bytes_or_less = 1;
    }

    let is_sequence_16_bytes_or_less = is_le(i, 16);

    // 1 - 16 bytes
    if (is_sequence_16_bytes_or_less != FALSE) {
        let (low) = compute_half_uint256(val=val, i=i, res=0);
        let res = Uint256(low=low, high=0);

        return res;
    }

    // 17 - 32 bytes
    let (low) = compute_half_uint256(val=val + i - 16, i=16, res=0);
    let (high) = compute_half_uint256(val=val, i=i - 16, res=0);
    let res = Uint256(low=low, high=high);

    return res;
}

func ceil_bytes_len_to_next_32_bytes_word{range_check_ptr}(bytes_len: felt) -> felt {
        let (q, _) = unsigned_div_rem(bytes_len + 31, 32);
        return q * 32;
    }

func pow256_rev(i: felt) -> felt {
        let (pow256_rev_address) = get_label_location(pow256_rev_table);
        return pow256_rev_address[i];

        pow256_rev_table:
        dw 340282366920938463463374607431768211456;
        dw 1329227995784915872903807060280344576;
        dw 5192296858534827628530496329220096;
        dw 20282409603651670423947251286016;
        dw 79228162514264337593543950336;
        dw 309485009821345068724781056;
        dw 1208925819614629174706176;
        dw 4722366482869645213696;
        dw 18446744073709551616;
        dw 72057594037927936;
        dw 281474976710656;
        dw 1099511627776;
        dw 4294967296;
        dw 16777216;
        dw 65536;
        dw 256;
        dw 1;
    }

func store{range_check_ptr}(
        self: Memory*, element: Uint256, offset: felt
    ) -> Memory* {
        let word_dict = self.memory;

        // Compute new bytes_len.
        let new_min_bytes_len = ceil_bytes_len_to_next_32_bytes_word(offset + 32);

        let fits = is_le(new_min_bytes_len, self.len);
        if (fits == FALSE) {
            tempvar new_bytes_len = new_min_bytes_len;
        } else {
            tempvar new_bytes_len = self.len;
        }

        // Check alignment of offset to 16B chunks.
        let (chunk_index, offset_in_chunk) = unsigned_div_rem(offset, 16);

        if (offset_in_chunk == 0) {
            // Offset is aligned. This is the simplest and most efficient case,
            // so we optimize for it. Note that no locals were allocated at all.
            dict_write{dict_ptr=word_dict}(chunk_index, element.high);
            dict_write{dict_ptr=word_dict}(chunk_index + 1, element.low);
            return (new Memory(
                memory=word_dict,
                len=new_bytes_len,
                ));
        }

        // Offset is misaligned.
        // |   W0   |   W1   |   w2   |
        //     |  EL_H  |  EL_L  |
        // ^---^
        //   |-- mask = 256 ** offset_in_chunk

        // Compute mask.
        tempvar mask = pow256_rev(offset_in_chunk);
        let mask_c = 2 ** 128 / mask;

        // Split the 2 input 16B chunks at offset_in_chunk.
        let (el_hh, el_hl) = unsigned_div_rem(element.high, mask_c);
        let (el_lh, el_ll) = unsigned_div_rem(element.low, mask_c);

        // Read the words at chunk_index, chunk_index + 2.
        let (w0) = dict_read{dict_ptr=word_dict}(chunk_index);
        let (w2) = dict_read{dict_ptr=word_dict}(chunk_index + 2);

        // Compute the new words.
        let (w0_h, _) = unsigned_div_rem(w0, mask);
        let (_, w2_l) = unsigned_div_rem(w2, mask);
        let new_w0 = w0_h * mask + el_hh;
        let new_w1 = el_hl * mask + el_lh;
        let new_w2 = el_ll * mask + w2_l;

        // Write new words.
        dict_write{dict_ptr=word_dict}(chunk_index, new_w0);
        dict_write{dict_ptr=word_dict}(chunk_index + 1, new_w1);
        dict_write{dict_ptr=word_dict}(chunk_index + 2, new_w2);
        return (new Memory(
            memory=word_dict,
            len=new_bytes_len,
            ));
    }

/// dummy store
func exec_mstore{range_check_ptr}(ctx: ExecutionContext*) -> ExecutionContext* {
        alloc_locals;

        let stack = ctx.stack;
        // Stack input:
        // 0 - offset: memory offset of the work we save.
        // 1 - value: value to store in memory.
        let (stack, popped) = pop_n(self=stack, n=2);
        let offset=popped[0];
        let value=popped[1];

        let memory:Memory* = store(self=ctx.memory, element=value, offset=offset.low);
        update_ctx_memory(self=ctx, new_memory=memory);
        update_ctx_stack(self=ctx, new_stack=stack);
        
        return ctx;
}

func exec_mload{
        range_check_ptr,
    }(ctx: model.ExecutionContext*) -> model.ExecutionContext* {
        alloc_locals;

        let stack = ctx.stack;

        // Stack input:
        // 0 - offset: memory offset of the word we read.
        let (stack, offset) = Stack.pop(stack);

        // Read word from memory at offset
        let (new_memory, value) = load(self=ctx.memory, offset=offset.low);

        let ctx = ExecutionContext.update_memory(ctx, new_memory);
        let ctx = ExecutionContext.update_stack(ctx, stack);

        return ctx;
}

func main{output_ptr: felt*, range_check_ptr}() {
    alloc_locals;
    let ctx: ExecutionContext* = init_ctx();

    let s: Stack* = ctx.stack;
    let s = push(self=s, val=0x80);
    let s = push(self=s, val=0x40);

    let ctx = update_ctx_stack(self = ctx, new_stack = s);
    
    let ctx = exec_mstore(ctx = ctx);

    return ();
}