/*
 * Copyright (c) 2023 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "co2_context.h"

#include <errno.h>

void context_entry(void);

#if defined(__aarch64__)
asm(".global context_entry; .type context_entry, %function; context_entry:\n"
    "ldp x0, x1, [sp], #0x10\n"
    "br  x1\n"
    ".size context_entry, . - context_entry\n"
    ".global co2_save_context; .type co2_save_context, %function; co2_save_context:\n"
    "stp x19, x20, [x0,#0]\n"
    "stp x21, x22, [x0,#16]\n"
    "stp x23, x24, [x0,#32]\n"
    "stp x25, x26, [x0,#48]\n"
    "stp x27, x28, [x0,#64]\n"
    "stp x29, x30, [x0,#80]\n"
    "mov x2, sp\n"
    "str x2, [x0,#104]\n"
    "stp  d8,  d9, [x0,#112]\n"
    "stp d10, d11, [x0,#128]\n"
    "stp d12, d13, [x0,#144]\n"
    "stp d14, d15, [x0,#160]\n"
    "mov x0, #0\n"
    "ret\n"
    ".size co2_save_context, . - co2_save_context\n"
    ".global co2_restore_context; .type co2_restore_context, %function; co2_restore_context:\n"
    "ldp x19, x20, [x0,#0]\n"
    "ldp x21, x22, [x0,#16]\n"
    "ldp x23, x24, [x0,#32]\n"
    "ldp x25, x26, [x0,#48]\n"
    "ldp x27, x28, [x0,#64]\n"
    "ldp x29, x30, [x0,#80]\n"
    "ldr x2, [x0,#104]\n"
    "mov sp, x2\n"
    "ldp  d8,  d9, [x0,#112]\n"
    "ldp d10, d11, [x0,#128]\n"
    "ldp d12, d13, [x0,#144]\n"
    "ldp d14, d15, [x0,#160]\n"
    "mov x0, #1\n"
    "ret\n"
    ".size co2_restore_context, . - co2_restore_context\n");
#elif defined(__arm__)
asm(".global context_entry; .type context_entry, %function; context_entry:\n"
    "pop {r0, r1}\n"
    "bx  r1\n"
    ".size context_entry, . - context_entry\n"
    ".global co2_save_context; .type co2_save_context, %function; co2_save_context:\n"
    "mov ip, r0\n"
    "str sp, [ip], #4\n"
    "str lr, [ip], #4\n"
    "stmia ip!, {v1-v6, sl, fp}\n"
    ".fpu vfp\n"
    "vstmia ip!, {d8-d15}\n"
    "mov r0, #0\n"
    "bx lr\n"
    ".size co2_save_context, . - co2_save_context\n"
    ".global co2_restore_context; .type co2_restore_context, %function; co2_restore_context:\n"
    "mov ip, r0\n"
    "ldr a4, [ip], #4\n"
    "ldr r4, [ip], #4\n"
    "mov sp, a4\n"
    "mov lr, r4\n"
    "ldmia ip!, {v1-v6, sl, fp}\n"
    ".fpu vfp\n"
    "vldmia ip!, {d8-d15}\n"
    "mov r0, #1\n"
    "bx lr\n"
    ".size co2_restore_context, . - co2_restore_context\n");
#elif (defined(__x86_64__) && !defined _MSC_VER)
asm(".global context_entry; .type context_entry, %function; context_entry:\n"
    "pop %rdi\n"
    "jmp *(%rsp)\n"
    ".size context_entry, . - context_entry\n"
    ".global co2_save_context; .type co2_save_context, %function; co2_save_context:\n"
    "mov %rbx, (%rdi)\n"
    "mov %rbp, 8(%rdi)\n"
    "mov %r12, 16(%rdi)\n"
    "mov %r13, 24(%rdi)\n"
    "mov %r14, 32(%rdi)\n"
    "mov %r15, 40(%rdi)\n"
    "lea 8(%rsp), %rdx\n"
    "mov %rdx, 48(%rdi)\n"
    "mov (%rsp), %rdx\n"
    "mov %rdx, 56(%rdi)\n"
    "xor %rax, %rax\n"
    "ret\n"
    ".size co2_save_context, . - co2_save_context\n"
    ".global co2_restore_context; .type co2_restore_context, %function; co2_restore_context:\n"
    "xor %rax, %rax\n"
    "inc %rax\n"
    "mov (%rdi), %rbx\n"
    "mov 8(%rdi), %rbp\n"
    "mov 16(%rdi), %r12\n"
    "mov 24(%rdi), %r13\n"
    "mov 32(%rdi), %r14\n"
    "mov 40(%rdi), %r15\n"
    "mov 48(%rdi), %rdx\n"
    "mov %rdx, %rsp\n"
    "jmp *56(%rdi)\n"
    ".size co2_restore_context, . - co2_restore_context\n");
#endif

int co2_init_context(struct co2_context* ctx, void (*func)(void*), void* arg, void* stack, size_t stack_size)
{
    if (stack_size < 0x4 * sizeof(uintptr_t)) {
        return EINVAL;
    }

    uintptr_t stack_top = (uintptr_t)stack + stack_size - 0x2 * sizeof(uintptr_t);
    stack_top -= stack_top % (0x2 * sizeof(uintptr_t));
    uintptr_t* data = (uintptr_t*)stack_top;

    ctx->regs[REG_LR] = (uintptr_t)context_entry;
    ctx->regs[REG_SP] = stack_top;

    data[0] = (uintptr_t)arg;
    data[1] = (uintptr_t)func;

    return 0;
}