/*
 * Wizer interface for Wasm module to be initialized.
 *
 * This header provides several macros that allow a Wasm module written in C/C++
 * to declare an initializer function, and ensure that global constructors (in
 * C++'s case) are run at initialization time rather than on startup of the
 * pre-initialized module.
 */
#ifndef _WIZER_H_
#define _WIZER_H_

#ifdef __cplusplus
#define __WIZER_EXTERN_C extern "C"
#else
#define __WIZER_EXTERN_C extern
#endif

#ifdef __clang_major__
// wasi-sdk-16 was the first wasi-sdk version that shipped with a version of
// wasi-libc that did not include __original_main. However, wasi-sdk-15 shipped
// with clang-14.0.0. To correctly identify the boundary where __original_main
// no longer exists, we check for either clang-15+ or specifically clang-14.0.4.
//
// wasi-sdk-17 ships with clang-15.0.6
// wasi-sdk-16 ships with clang-14.0.4
#if __clang_major__ >= 15 || (__clang_major__ == 14 && __clang_patchlevel__ == 4)
#define WIZER_MAIN_VOID __main_void
#else
#define WIZER_MAIN_VOID __original_main
#endif
#endif

// We default to assuming that the compiler is new enough to provide
// __main_void.
#ifndef WIZER_MAIN_VOID
#define WIZER_MAIN_VOID __main_void
#endif

/*
 * This macro inserts the exported functions necessary to allow Wizer to
 * pre-initialize a Wasm module.
 *
 * To use, simply invoke the macro in exactly one compilation unit (C/C++ file)
 * that is compiled into the Wasm module to be pre-initialized:
 *
 *     static void my_init_function() { ... }
 *
 *     WIZER_INIT(my_init_function);
 *
 * (The macro refers to the provided init function, so it must have been defined
 * or must have a forward declaration at the point the macro is used.)
 *
 * The resulting module should be processed by Wizer as follows:
 *
 *     $ wizer -r _start=wizer.resume -o out.wasm in.wasm
 *
 * The result of this will be the following behavior:
 *
 * - If the `in.wasm` (the direct compilation output of a program including this
 *   macro invocation) is run directly according to the WASI ABI (i.e., by
 *   invoking `_start`), then nothing changes: global constructors are run,
 *   `main()` is invoked, then global destructors are run. The initialization
 *   function is *not* run in this case.
 *
 * - During pre-initialization (i.e., during this `wizer` invocation), global
 *   constructors will run, and then the provided initialization function will
 *   run. The module's memory and global-variable state is then snapshotted and
 *   saved into `out.wasm`.
 *
 *   All other Wizer restrictions apply (see Wizer documentation for details):
 *   for example, WASI hostcalls may be blocked, depending on options, and
 *   invoking any other imported function will result in an immediate trap
 *   and failure of the Wizer run.
 *
 * - If the resulting `out.wasm` is then run using the WASI ABI, the program's
 *   global constructors are *not* re-run. Instead, execution starts directly at
 *   `main()`, using the heap and global-variable state left by the global
 *   constructor and initialization function execution during the Wizer
 *   invocation.
 *
 * If no initialization function is needed (i.e., only C++ global constructors
 * should be run), use `WIZER_DEFAULT_INIT()` instead.
 */
#define WIZER_INIT(init_func)                                                  \
    __WIZER_EXTERN_C void __wasm_call_ctors();                                 \
    __WIZER_EXTERN_C void __wasm_call_dtors();                                 \
    __WIZER_EXTERN_C void __wasi_proc_exit(int);                               \
    __WIZER_EXTERN_C int WIZER_MAIN_VOID();                                    \
    /* This function's export name `wizer.initialize` is specially          */ \
    /* recognized by Wizer. It is the direct entry point for pre-init.      */ \
    __attribute__((export_name("wizer.initialize"))) void                      \
    __wizer_initialize() {                                                     \
        /* `__wasm_call_ctors()` is generated by `wasm-ld` and invokes all  */ \
        /* of the global constructors. It is safe (and in fact necessary)   */ \
        /* to manually invoke it here because `wizer.initialize` is the     */ \
        /* direct entry point, and no libc startup (crt1.o or equivalent)   */ \
        /* is executed before this code does. */                               \
        __wasm_call_ctors();                                                   \
        /* We now invoke the provided init function before returning.       */ \
        init_func();                                                           \
    }                                                                          \
    /* This function replaces `_start` (the WASI-specified entry point) in  */ \
    /* the pre-initialized Wasm module.                                     */ \
    __attribute__((export_name("wizer.resume"))) void __wizer_resume() {       \
        /* `__main_void()` is defined by the WASI SDK toolchain due to      */ \
        /* special semantics in C/C++ for the `main()` function, i.e., ito  */ \
        /* can either take argc/argv or not. It collects arguments using    */ \
        /* the appropriate WASI calls and then invokes the user program's   */ \
        /* `main()`. This may change in the future; when it does, we will   */ \
        /* coordinate with the WASI-SDK toolchain to implement this entry   */ \
        /* point in an alternate way. */                                       \
        int r = WIZER_MAIN_VOID();                                             \
        /* Because we are replacing `_start()`, we need to manually invoke  */ \
        /* destructors as well.                                             */ \
        __wasm_call_dtors();                                                   \
        /* If main returned non-zero code, call `__wasi_proc_exit`. */         \
        if (r != 0) {                                                          \
          __wasi_proc_exit(r);                                                 \
        }                                                                      \
    }

/*
 * This macro is like `WIZER_INIT()`, but takes no initialization function.
 * Instead, the pre-initialization phase only executes C++ global constructors
 * before snapshotting the module state.
 *
 * See documentation for `WIZER_INIT()` for more details and usage instructions.
 */
#define WIZER_DEFAULT_INIT()      \
    static void __empty_init() {} \
    WIZER_INIT(__empty_init)

#endif  // _WIZER_H_