FEX-Emu--FEX/ThunkLibs/Generator/gen.cpp

#include "clang/AST/RecursiveASTVisitor.h"

#include <fstream>
#include <iostream>
#include <iomanip>
#include <string_view>

#include <openssl/sha.h>

#include "interface.h"

struct FunctionParams {
    std::vector<clang::QualType> param_types;
};

struct ThunkedCallback : FunctionParams {
    clang::QualType return_type;

    std::size_t callback_index;
    std::size_t user_arg_index;

    bool is_stub = false;
    bool is_variadic = false;
};

/**
 * Guest<->Host transition point.
 *
 * These are normally used to translate the public API of the guest to host
 * function calls (ThunkedAPIFunction), but a thunk library may also define
 * internal thunks that don't correspond to any function in the implemented
 * API.
 */
struct ThunkedFunction : FunctionParams {
    std::string function_name;
    clang::QualType return_type;

    // If true, param_types contains an extra size_t and the valist for marshalling through an internal function
    bool is_variadic = false;

    // If true, the unpacking function will call a custom fexfn_impl function
    // to be provided manually instead of calling the host library function
    // directly.
    // This is implied e.g. for thunks generated for variadic functions
    bool custom_host_impl = false;

    FunctionParams GetNonvariadicParams() const {
        if (!is_variadic) {
            return *this;
        }

        auto ret = param_types;
        ret.pop_back();
        ret.pop_back();
        return { std::move(ret) };
    }

    std::string GetOriginalFunctionName() const {
        const std::string suffix = "_internal";
        assert(function_name.length() > suffix.size());
        assert((std::string_view { &*function_name.end() - suffix.size(), suffix.size() } == suffix));
        return function_name.substr(0, function_name.size() - suffix.size());
    }

    // Maps parameter index to ThunkedCallback
    std::unordered_map<unsigned, ThunkedCallback> callbacks;

    clang::FunctionDecl* decl;
};

/**
 * Function that is part of the API of the thunked library.
 *
 * For each of these, there is:
 * - A publicly visible guest entrypoint (usually auto-generated but may be manually defined)
 * - A pointer to the native host library function loaded through dlsym (or a user-provided function specified via host_loader)
 * - A ThunkedFunction with the same function_name (possibly suffixed with _internal)
 */
struct ThunkedAPIFunction : FunctionParams {
    std::string function_name;

    clang::QualType return_type;

    // If true, no guest-side implementation of this function will be autogenerated
    bool custom_guest_impl;

    bool is_variadic;
};

static std::vector<ThunkedFunction> thunks;
static std::vector<ThunkedAPIFunction> thunked_api;

class ASTVisitor : public clang::RecursiveASTVisitor<ASTVisitor> {
    clang::ASTContext& context;

    enum class CallbackStrategy {
        Default,
        Stub,
    };

    struct Annotations {
        std::optional<clang::QualType> uniform_va_type;

        CallbackStrategy callback_strategy = CallbackStrategy::Default;
    };

    Annotations GetAnnotations(clang::CXXRecordDecl* decl) {
        Annotations ret;

        for (const auto& base : decl->bases()) {
            auto annotation = base.getType().getAsString();
            if (annotation == "fexgen::callback_stub") {
                ret.callback_strategy = CallbackStrategy::Stub;
            } else {
                throw Error(base.getSourceRange().getBegin(), "Unknown annotation");
            }
        }

        for (const auto& child_decl : decl->getPrimaryContext()->decls()) {
            if (auto field = llvm::dyn_cast_or_null<clang::FieldDecl>(child_decl)) {
                throw Error(field->getBeginLoc(), "Unknown field annotation");
            } else if (auto type_alias = llvm::dyn_cast_or_null<clang::TypedefNameDecl>(child_decl)) {
                auto name = type_alias->getNameAsString();
                if (name == "uniform_va_type") {
                    ret.uniform_va_type = type_alias->getUnderlyingType();
                } else {
                    throw Error(type_alias->getBeginLoc(), "Unknown type alias annotation");
                }
            }
        }

        return ret;
    }

    using ClangDiagnosticAsException = std::pair<clang::SourceLocation, unsigned>;

    template<std::size_t N>
    [[nodiscard]] ClangDiagnosticAsException Error(clang::SourceLocation loc, const char (&message)[N]) {
        auto id = context.getDiagnostics().getCustomDiagID(clang::DiagnosticsEngine::Error, message);
        return std::pair(loc, id);
    }

public:
    ASTVisitor(clang::ASTContext& context_) : context(context_) {
    }

    /**
     * Matches "template<> struct fex_gen_config<LibraryFunc> { ... }"
     */
    bool VisitClassTemplateSpecializationDecl(clang::ClassTemplateSpecializationDecl* decl) try {
        if (decl->getName() != "fex_gen_config") {
            return true;
        }

        if (decl->getSpecializationKind() == clang::TSK_ExplicitInstantiationDefinition) {
            throw Error(decl->getBeginLoc(), "fex_gen_config may not be partially specialized\n");
        }

        const auto& template_args = decl->getTemplateArgs();
        assert(template_args.size() == 1);

        auto emitted_function = llvm::dyn_cast<clang::FunctionDecl>(template_args[0].getAsDecl());
        assert(emitted_function && "Argument is not a function");
        auto return_type = emitted_function->getReturnType();

        const auto annotations = GetAnnotations(decl);
        if (return_type->isFunctionPointerType()) {
            throw Error(decl->getBeginLoc(),
                        "Function pointer return types require explicit annotation\n");
        }

        // TODO: Use the types as written in the signature instead?
        ThunkedFunction data;
        data.function_name = emitted_function->getName().str();
        data.return_type = return_type;
        data.is_variadic = emitted_function->isVariadic();

        data.decl = emitted_function;

        for (std::size_t param_idx = 0; param_idx < emitted_function->param_size(); ++param_idx) {
            auto* param = emitted_function->getParamDecl(param_idx);
            data.param_types.push_back(param->getType());

            if (param->getType()->isFunctionPointerType()) {
                auto funcptr = param->getFunctionType()->getAs<clang::FunctionProtoType>();
                ThunkedCallback callback;
                callback.return_type = funcptr->getReturnType();
                for (auto& cb_param : funcptr->getParamTypes()) {
                    callback.param_types.push_back(cb_param);
                }
                callback.is_stub = annotations.callback_strategy == CallbackStrategy::Stub;
                callback.is_variadic = funcptr->isVariadic();

                data.callbacks.emplace(param_idx, callback);
                // TODO: Support for more than one callback is untested
                assert(data.callbacks.size() == 1);
                if (funcptr->isVariadic() && !callback.is_stub) {
                    throw Error(decl->getBeginLoc(), "Variadic callbacks are not supported");
                }
            }
        }

        const bool has_nonstub_callbacks = std::any_of(data.callbacks.begin(), data.callbacks.end(),
                                                       [](auto& cb) { return !cb.second.is_stub; });
        thunked_api.push_back(ThunkedAPIFunction { (const FunctionParams&)data, data.function_name, data.return_type,
                                                    has_nonstub_callbacks || data.is_variadic,
                                                    data.is_variadic });

        if (data.is_variadic) {
            if (!annotations.uniform_va_type) {
                throw Error(decl->getBeginLoc(), "Variadic functions must be annotated with parameter type using uniform_va_type");
            }

            // Convert variadic argument list into a count + pointer pair
            data.param_types.push_back(context.getSizeType());
            data.param_types.push_back(context.getPointerType(*annotations.uniform_va_type));
        }

        if (has_nonstub_callbacks || data.is_variadic) {
            // This function is thunked through an "_internal" symbol since its signature
            // is different from the one in the native host/guest libraries.
            data.function_name = data.function_name + "_internal";
            data.custom_host_impl = true;
        }

        thunks.push_back(std::move(data));

        return true;
    } catch (ClangDiagnosticAsException& exception) {
        context.getDiagnostics().Report(exception.first, exception.second);
        return false;
    }
};

class ASTConsumer : public clang::ASTConsumer {
public:
    void HandleTranslationUnit(clang::ASTContext& context) override {
        ASTVisitor{context}.TraverseDecl(context.getTranslationUnitDecl());
    }
};

FrontendAction::FrontendAction(const std::string& libname_, const OutputFilenames& output_filenames_)
    : libname(libname_), output_filenames(output_filenames_) {
    thunks.clear();
    thunked_api.clear();
}

void FrontendAction::EndSourceFileAction() {
    static auto format_decl = [](clang::QualType type, const std::string_view& name) {
        if (type->isFunctionPointerType()) {
            auto signature = type.getAsString();
            const char needle[] = { '(', '*', ')' };
            auto it = std::search(signature.begin(), signature.end(), std::begin(needle), std::end(needle));
            if (it == signature.end()) {
                // It's *probably* a typedef, so this should be safe after all
                return signature + " " + std::string(name);
            } else {
                signature.insert(it + 2, name.begin(), name.end());
                return signature;
            }
        } else {
            return type.getAsString() + " " + std::string(name);
        }
    };

    auto format_struct_members = [](const FunctionParams& params, const char* indent) {
        std::string ret;
        for (std::size_t idx = 0; idx < params.param_types.size(); ++idx) {
            ret += indent + format_decl(params.param_types[idx].getUnqualifiedType(), "a_" + std::to_string(idx)) + ";\n";
        }
        return ret;
    };

    auto format_function_args = [](const FunctionParams& params) {
        std::string ret;
        for (std::size_t idx = 0; idx < params.param_types.size(); ++idx) {
            ret += "args->a_" + std::to_string(idx) + ", ";
        }
        // drop trailing ", "
        ret.resize(ret.size() > 2 ? ret.size() - 2 : 0);
        return ret;
    };

    auto format_function_params = [](const FunctionParams& params) {
        std::string ret;
        for (std::size_t idx = 0; idx < params.param_types.size(); ++idx) {
            auto& type = params.param_types[idx];
            ret += format_decl(type, "a_" + std::to_string(idx)) + ", ";
        }
        // drop trailing ", "
        ret.resize(ret.size() > 2 ? ret.size() - 2 : 0);
        return ret;
    };

    auto get_sha256 = [this](const std::string& function_name) {
        std::string sha256_message = libname + ":" + function_name;
        std::vector<unsigned char> sha256(SHA256_DIGEST_LENGTH);
        SHA256(reinterpret_cast<const unsigned char*>(sha256_message.data()),
               sha256_message.size(),
               sha256.data());
        return sha256;
    };

    if (!output_filenames.thunks.empty()) {
        std::ofstream file(output_filenames.thunks);

        file << "extern \"C\" {\n";
        for (auto& thunk : thunks) {
            const auto& function_name = thunk.function_name;
            auto sha256 = get_sha256(function_name);
            file << "MAKE_THUNK(" << libname << ", " << function_name << ", \"";
            bool first = true;
            for (auto c : sha256) {
                file << (first ? "" : ", ") << "0x" << std::hex << std::setw(2) << std::setfill('0') << +c;
                first = false;
            }
            file << "\")\n";
        }

        file << "}\n";
    }

    if (!output_filenames.function_packs_public.empty()) {
        std::ofstream file(output_filenames.function_packs_public);

        file << "extern \"C\" {\n";
        for (auto& data : thunked_api) {
            if (data.custom_guest_impl) {
                continue;
            }

            const auto& function_name = data.function_name;

            file << "__attribute__((alias(\"fexfn_pack_" << function_name << "\"))) auto " << function_name << "(";
            for (std::size_t idx = 0; idx < data.param_types.size(); ++idx) {
                auto& type = data.param_types[idx];
                file << (idx == 0 ? "" : ", ") << format_decl(type, "a_" + std::to_string(idx));
            }
            file << ") -> " << data.return_type.getAsString() << ";\n";
        }
        file << "}\n";
    }

    if (!output_filenames.function_packs.empty()) {
        std::ofstream file(output_filenames.function_packs);

        file << "extern \"C\" {\n";
        for (auto& data : thunks) {
            const auto& function_name = data.function_name;
            bool is_void = data.return_type->isVoidType();
            file << "static auto fexfn_pack_" << function_name << "(";
            for (std::size_t idx = 0; idx < data.param_types.size(); ++idx) {
                auto& type = data.param_types[idx];
                file << (idx == 0 ? "" : ", ") << format_decl(type, "a_" + std::to_string(idx));
            }
            // Using trailing return type as it makes handling function pointer returns much easier
            file << ") -> " << data.return_type.getAsString() << " {\n";
            file << "  struct {\n";
            for (std::size_t idx = 0; idx < data.param_types.size(); ++idx) {
                auto& type = data.param_types[idx];
                file << "    " << format_decl(type.getUnqualifiedType(), "a_" + std::to_string(idx)) << ";\n";
            }
            if (!is_void) {
                file << "    " << format_decl(data.return_type, "rv") << ";\n";
            } else if (data.param_types.size() == 0) {
                // Avoid "empty struct has size 0 in C, size 1 in C++" warning
                file << "    char force_nonempty;\n";
            }
            file << "  } args;\n";
            for (std::size_t idx = 0; idx < data.param_types.size(); ++idx) {
                file << "  args.a_" << idx << " = a_" << idx << ";\n";
            }
            file << "  fexthunks_" << libname << "_" << function_name << "(&args);\n";
            if (!is_void) {
                file << "  return args.rv;\n";
            }
            file << "}\n";
        }
        file << "}\n";
    }

    if (!output_filenames.function_unpacks.empty()) {
        std::ofstream file(output_filenames.function_unpacks);

        file << "extern \"C\" {\n";
        for (auto& thunk : thunks) {
            const auto& function_name = thunk.function_name;
            bool is_void = thunk.return_type->isVoidType();

            file << "struct fexfn_packed_args_" << libname << "_" << function_name << " {\n";
            file << format_struct_members(thunk, "  ");
            if (!is_void) {
                file << "  " << format_decl(thunk.return_type, "rv") << ";\n";
            } else if (thunk.param_types.size() == 0) {
                // Avoid "empty struct has size 0 in C, size 1 in C++" warning
                file << "    char force_nonempty;\n";
            }
            file << "};\n";

            for (auto& [cb_idx, cb] : thunk.callbacks) {
                if (cb.is_stub) {
                    bool is_first_cb = (cb_idx == thunk.callbacks.begin()->first);
                    const char* variadic_ellipsis = cb.is_variadic ? ", ..." : "";
                    auto cb_function_name = "fexfn_unpack_" + function_name + "CBFN" + (is_first_cb ? "" : std::to_string(cb_idx)) + "_stub";
                    file << "[[noreturn]] static " << cb.return_type.getAsString() << " "
                         << cb_function_name << "("
                         << format_function_params(cb) << variadic_ellipsis << ") {\n";
                    file << "  fprintf(stderr, \"SHOULD NOT HAVE BEEN CALLED\");\n"; // TODO: Reword
                    file << "  abort();\n";
                    file << "}\n";
                }
            }

            file << "static void fexfn_unpack_" << libname << "_" << function_name << "(fexfn_packed_args_" << libname << "_" << function_name << "* args) {\n";
            file << (is_void ? "  " : "  args->rv = ") << (thunk.custom_host_impl ? "fexfn_impl_" : "fexldr_ptr_") << libname << "_" << function_name << "(";
            {
                std::string ret;
                for (std::size_t idx = 0; idx < thunk.param_types.size(); ++idx) {
                    auto cb = thunk.callbacks.find(idx);
                    if (cb != thunk.callbacks.end() && cb->second.is_stub) {
                        bool is_first_cb = (cb->first == thunk.callbacks.begin()->first);
                        auto cb_function_name = "fexfn_unpack_" + function_name + "CBFN" + (is_first_cb ? "" : std::to_string(cb->first)) + "_stub";
                        ret += cb_function_name + ", ";
                    } else {
                        ret += "args->a_" + std::to_string(idx) + ", ";
                    }
                }
                // drop trailing ", "
                ret.resize(ret.size() > 2 ? ret.size() - 2 : 0);
                file << ret;
            }
            file << ");\n";
            file << "}\n";
        }

        file << "}\n";
    }

    if (!output_filenames.tab_function_unpacks.empty()) {
        std::ofstream file(output_filenames.tab_function_unpacks);

        for (auto& thunk : thunks) {
            const auto& function_name = thunk.function_name;
            auto sha256 = get_sha256(function_name);

            file << "{(uint8_t*)\"";
            for (auto c : sha256) {
                file << "\\x" << std::hex << std::setw(2) << std::setfill('0') << +c;
            }
            file << "\", &fexfn_type_erased_unpack<fexfn_unpack_" << libname << "_" << function_name << ">}, // " << libname << ":" << function_name << "\n";
        }
    }

    if (!output_filenames.ldr.empty()) {
        std::ofstream file(output_filenames.ldr);

        file << "static void* fexldr_ptr_" << libname << "_so;\n";
        file << "extern \"C\" bool fexldr_init_" << libname << "() {\n";
        file << "  fexldr_ptr_" << libname << "_so = dlopen(\"" << libname << ".so\", RTLD_LOCAL | RTLD_LAZY);\n";
        file << "  if (!fexldr_ptr_" << libname << "_so) { return false; }\n\n";
        for (auto& import : thunked_api) {
            file << "  (void*&)fexldr_ptr_" << libname << "_" << import.function_name << " = dlsym(fexldr_ptr_" << libname << "_so, \"" << import.function_name << "\");\n";
        }
        file << "  return true;\n";
        file << "}\n";
    }

    if (!output_filenames.ldr_ptrs.empty()) {
        std::ofstream file(output_filenames.ldr_ptrs);

        for (auto& import : thunked_api) {
            const auto& function_name = import.function_name;
            const char* variadic_ellipsis = import.is_variadic ? ", ..." : "";
            file << "using fexldr_type_" << libname << "_" << function_name << " = auto " << "(" << format_function_params(import) << variadic_ellipsis << ") -> " << import.return_type.getAsString() << ";\n";
            file << "static fexldr_type_" << libname << "_" << function_name << " *fexldr_ptr_" << libname << "_" << function_name << ";\n";
        }
    }

    if (!output_filenames.callback_structs.empty()) {
        std::ofstream file(output_filenames.callback_structs);

        for (auto& thunk : thunks) {
            for (const auto& [cb_idx, cb] : thunk.callbacks) {
                if (cb.is_stub) {
                    continue;
                }

                file << "struct " << thunk.GetOriginalFunctionName() << "CB_Args {\n";
                file << format_struct_members(cb, "  ");
                if (!cb.return_type->isVoidType()) {
                    file << "  " << format_decl(cb.return_type, "rv") << ";\n";
                }
                file << "};\n";
            }
        }
    }

    if (!output_filenames.callback_typedefs.empty()) {
        std::ofstream file(output_filenames.callback_typedefs);

        for (auto& thunk : thunks) {
            for (const auto& [cb_idx, cb] : thunk.callbacks) {
                if (cb.is_stub) {
                    continue;
                }

                bool is_first_cb = (cb_idx == thunk.callbacks.begin()->first);
                auto cb_function_name = thunk.GetOriginalFunctionName() + "CBFN" + (is_first_cb ? "" : std::to_string(cb_idx));
                file << "typedef " << cb.return_type.getAsString() << " "
                     << cb_function_name << "("
                     << format_function_params(cb) << ");\n";
            }
        }
    }

    if (!output_filenames.callback_unpacks.empty()) {
        std::ofstream file(output_filenames.callback_unpacks);

        for (auto& thunk : thunks) {
            for (const auto& [cb_idx, cb] : thunk.callbacks) {
                if (cb.is_stub) {
                    continue;
                }

                bool is_void = cb.return_type->isVoidType();
                bool is_first_cb = (cb_idx == thunk.callbacks.begin()->first);
                auto cb_function_name = thunk.function_name + "CB" + (is_first_cb ? "" : std::to_string(cb_idx));
                file << "static void fexfn_unpack_" << libname << "_" << cb_function_name << "(uintptr_t cb, void* argsv) {\n";
                file << "  typedef " << cb.return_type.getAsString() << " fn_t (" << format_function_params(cb) << ");\n";
                file << "  auto callback = reinterpret_cast<fn_t*>(cb);\n";
                file << "  struct arg_t {\n";
                file << format_struct_members(cb, "    ");
                if (!is_void) {
                    file << "    " << format_decl(cb.return_type, "rv") << ";\n";
                }
                file << "  };\n";
                file << "  auto args = (arg_t*)argsv;\n";
                file << (is_void ? "  " : "  args->rv = ") << "callback(" << format_function_args(cb) << ");\n";
                file << "}\n";
            }
        }
    }

    if (!output_filenames.callback_unpacks_header.empty()) {
        std::ofstream file(output_filenames.callback_unpacks_header);

        for (auto& thunk : thunks) {
            for (const auto& [cb_idx, cb] : thunk.callbacks) {
                if (cb.is_stub) {
                    continue;
                }

                bool is_first_cb = (cb_idx == thunk.callbacks.begin()->first);
                auto cb_function_name = thunk.GetOriginalFunctionName() + "CB" + (is_first_cb ? "" : std::to_string(cb_idx));
                file << "uintptr_t " << libname << "_" << cb_function_name << ";\n";
            }
        }
    }

    if (!output_filenames.callback_unpacks_header_init.empty()) {
        std::ofstream file(output_filenames.callback_unpacks_header_init);

        for (auto& thunk : thunks) {
            for (const auto& [cb_idx, cb] : thunk.callbacks) {
                if (cb.is_stub) {
                    continue;
                }

                bool is_first_cb = (cb_idx == thunk.callbacks.begin()->first);
                auto cb_function_name = thunk.function_name + "CB" + (is_first_cb ? "" : std::to_string(cb_idx));
                file << "(uintptr_t)&fexfn_unpack_" << libname << "_" << cb_function_name << ",\n";
            }
        }
    }
}

std::unique_ptr<clang::ASTConsumer> FrontendAction::CreateASTConsumer(clang::CompilerInstance&, clang::StringRef) {
    return std::make_unique<ASTConsumer>();
}