Curried Function example in Yul

CurryExample_yul.sol

object "ContractB" {
    code {
        datacopy(0, dataoffset("Runtime"), datasize("Runtime"))
        return(0, datasize("Runtime"))
    }
    object "Runtime" {
        code {
            let _calldata := 2048
            let _output_pointer := 0
            
            // This is where we keep our virtual functions
            // generated at runtime as partial function applications
            let _virtual_fns := 1024
            
            calldatacopy(_calldata, 0, calldatasize())
            
            let fn_sig := mslice(_calldata, 4)
            
            switch fn_sig
            case 0xffffffff {
                let internal_fn_sig := mslice(add(_calldata, 4), 4)
                let input_pointer := add(_calldata, 8)
                let input_size := sub(calldatasize(), 4)
                    
                let result_length := executeNative(
                    internal_fn_sig,
                    input_pointer,
                    input_size,
                    _output_pointer,
                    _virtual_fns
                )
                return (_output_pointer, result_length)
            }
            // other cases/function signatures
            default {
                mslicestore(_output_pointer, 0xeee1, 2)
                revert(_output_pointer, 2)
            }
        
            function executeNative(fsig, input_ptr, input_size, output_ptr, virtual_fns) -> result_length {
                switch fsig
                
                // a + b
                case 0xeeeeeeee {
                    let a := mload(input_ptr)
                    let b := mload(add(input_ptr, 32))
                    mstore(output_ptr, add(a, b))
                    result_length := 32
                }
                
                // a - 2
                case 0xdddddddd {
                    let a := mload(input_ptr)
                    mstore(output_ptr, sub(a, 2))
                    result_length := 32
                }
                
                case 0xcccccccc {
                    // e.g. 2 steps:
                    // 000000020000002800000020
                    // bbbbbbbbeeeeeeee000000000000000000000000000000000000000000000000000000000000004
                    // 00000000000000000000000000000000000000000000000000000000000000020
                    
                    // number of execution steps
                    let count := mslice(input_ptr, 4)
                    
                    // offsets/size in bytes for each step
                    let offsets_start := add(input_ptr, 4)
                    let input_inner := add(offsets_start, mul(count, 4))
                    
                    let temporary_ptr := 0x80
                    let existent_input_size := 0

                    for { let i := 0 } lt(i, count) { i := add(i, 1) } {
                        let step_length := mslice(add(offsets_start, mul(i, 4)), 4)
                        
                        // add current input after previous return value
                        mmultistore(add(temporary_ptr, existent_input_size), input_inner, step_length)
                       
                        result_length := executeInternal(temporary_ptr, add(existent_input_size, step_length), output_ptr, virtual_fns)
                        
                        // move termporary input after previous data
                        temporary_ptr := add(temporary_ptr, step_length)
                        
                        // store output as new input for the next step
                        mmultistore(temporary_ptr, output_ptr, result_length)
                        existent_input_size := result_length
                        
                        // move input pointer to the next step
                        input_inner := add(input_inner, step_length)
                    }
                }
                
                // curry1: fsig, partial application argument
                case 0xbbbbbbbb {
                    // first 32 bytes is the next free memory pointer
                    let fpointer := mload(virtual_fns)
                    if eq(fpointer, 0) {
                        fpointer := add(virtual_fns, 32)
                    }
                    
                    let internal_fsig := mslice(input_ptr, 4)
                    let arg := mload(add(input_ptr, 4))
                    
                    // virtual function marker
                    mslicestore(fpointer, 0xfefe, 2)
                    
                    // add input size (so we know how much to read)
                    mstore(add(fpointer, 2), input_size)
                    
                    // store the actual data - partial application argument
                    mmultistore(add(fpointer, 34), input_ptr, input_size)
                    
                    // update the free memory pointer for our curried functions references
                    mstore(virtual_fns, add(fpointer, 38))
                    
                    // return the virtual function pointer
                    mstore(output_ptr, fpointer)
                    result_length := 32
                }
                // other cases/function signatures
                default {
                    // revert with error code
                    mslicestore(output_ptr, 0xeee2, 2)
                    revert(output_ptr, 2)
                }
            }
            
            function executeInternal(input_ptr, input_size, output_ptr, virtual_fns) -> result_length {
                let fsig, offset := getfSig(input_ptr)
                
                switch offset
                case 4 {
                    result_length := executeNative(fsig, add(input_ptr, offset), sub(input_size, offset), output_ptr, virtual_fns)
                }
                case 32 {
                    result_length := executeCurriedFunction(fsig, add(input_ptr, offset), sub(input_size, offset), output_ptr, virtual_fns)
                }
                default {
                    // revert with error code
                    mslicestore(output_ptr, 0xeee3, 2)
                    revert(output_ptr, 2)
                }
            }
            
            function getfSig(input_ptr) -> fsig, offset {
                fsig := mslice(input_ptr, 4)
                offset := 4
                let fpointer := mload(input_ptr)
                
                if lt(fpointer, 10000000) {
                    // check if the curried function marker exists
                    if eq(mslice(fpointer, 2), 0xfefe) {
                        fsig := fpointer
                        offset := 32
                    }
                }
            }
            
            function executeCurriedFunction(fpointer, input_ptr, input_size, output_ptr, virtual_fns) -> result_length {
                // first 32 bytes are the input size
                let new_input_size := mload(add(fpointer, 2))
                        
                // exclude input size from input ptr 
                let new_input_ptr := add(fpointer, 34)
                
                // store the inputs for the curried function after the curried function arguments
                // effectively composing the input for the actual function that we need to run
                mmultistore(add(new_input_ptr, new_input_size), input_ptr, input_size)
                new_input_size := add(new_input_size, input_size)

                result_length := executeInternal(new_input_ptr, new_input_size, output_ptr, virtual_fns)
            }
            
            function mslice(position, length) -> result {
                result := div(mload(position), exp(2, sub(256, mul(length, 8))))
            }
            
            function mslicestore(_ptr, val, length) {
                let slot := 32
                mstore(_ptr, shl(mul(sub(slot, length), 8), val))
            }
            
            function mmultistore(_ptr_target, _ptr_source, sizeBytes) {
                let slot := 32
                let size := div(sizeBytes, slot)

                for { let i := 0 } lt(i, size)  { i := add(i, 1) } {
                    mstore(add(_ptr_target, mul(i, slot)), mload(add(_ptr_source, mul(i, slot))))
                }

                let current_length :=  mul(size, slot)
                let remaining := sub(sizeBytes, current_length)
                if gt(remaining, 0) {
                    mslicestore(
                        add(_ptr_target, current_length),
                        mslice(add(_ptr_source, current_length), remaining),
                        remaining
                    )
                }
            }
        }
    }
}