src/wasm/module_tier_ladder.cpp

#ifdef WITH_RIVE_SCRIPTING_WASM

#include "rive/wasm/module_tier_ladder.hpp"
#include "rive/wasm/aot_artifact.hpp"
#include "rive/wasm/wasm_lane.hpp"

#include <cstdlib>
#include <cstring>

namespace rive
{
TierSpecies bootTierSpecies()
{
  const char* tier = getenv("RIVE_WASM_AOT_SYNC");
  if (tier != nullptr && strcmp(tier, "o0") == 0)
  {
    return TierSpecies::o0;
  }
#ifdef RIVE_WASM_HW_BOUNDS
  return TierSpecies::hw;
#else
  return TierSpecies::o3;
#endif
}

int x86_64Level(
    uint32_t leaf1Ecx,
    uint32_t leaf7Ebx,
    uint32_t extEcx,
    uint64_t xcr0)
{
  // FMA, MOVBE, XSAVE, OSXSAVE, AVX and F16C, then BMI1, AVX2 and BMI2, then
  // LZCNT, then the XMM and YMM state.
  constexpr uint32_t kLeaf1 = (1u << 12) | (1u << 22) | (1u << 26) |
                              (1u << 27) | (1u << 28) | (1u << 29);
  constexpr uint32_t kLeaf7 = (1u << 3) | (1u << 5) | (1u << 8);
  constexpr uint32_t kExt = 1u << 5;
  constexpr uint64_t kXcr0 = 0x6;
  return (leaf1Ecx & kLeaf1) == kLeaf1 && (leaf7Ebx & kLeaf7) == kLeaf7 &&
                 (extEcx & kExt) == kExt && (xcr0 & kXcr0) == kXcr0
             ? 3
             : 2;
}
} // namespace rive

#ifdef RIVE_WASM_TIER_LADDER

#include "wasm_export.h"

#include <algorithm>
#include <atomic>
#include <cerrno>
#include <chrono>
#include <cstdio>
#include <filesystem>
#include <fstream>
#include <iterator>
#include <cstdlib>
#include <cstring>
#include <thread>
#include <sys/stat.h>
#ifdef _WIN32
#include <windows.h>
#else
#include <fcntl.h>
#include <signal.h>
#include <spawn.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <unistd.h>
#if !defined(RIVE_ANDROID) && !defined(RIVE_NX)
#include <dlfcn.h>
#endif
#if defined(__linux__) && !defined(__ANDROID__)
#include <link.h>
#endif

extern char** environ;
#endif
#if defined(__x86_64__) || defined(_M_X64)
#if defined(_MSC_VER) && !defined(__clang__)
#include <intrin.h>
#else
#include <cpuid.h>
#endif
#endif

namespace rive
{

constexpr const char* kCancelledReason = "the compile was cancelled";

static uint64_t processId()
{
#ifdef _WIN32
  return GetCurrentProcessId();
#else
  return (uint64_t)getpid();
#endif
}

static void makeDir(const std::string& path)
{
#ifdef RIVE_NX
  // Our NX libc++ has no std::filesystem directory ops.
  mkdir(path.c_str(), 0755);
#elif defined(_WIN32)
  std::error_code ec;
  std::filesystem::create_directories(path, ec);
#else
  // Private from the start, so no other user can plant artifacts.
  if (mkdir(path.c_str(), 0700) != 0 && errno == ENOENT)
  {
    size_t slash = path.find_last_of('/');
    if (slash != std::string::npos && slash != 0)
    {
      makeDir(path.substr(0, slash));
      mkdir(path.c_str(), 0700);
    }
  }
#endif
}

static bool fileExists(const std::string& path)
{
  struct stat st;
  return stat(path.c_str(), &st) == 0;
}

// Artifacts run as native code, so a cache another user can write is never
// read. One that does not exist yet gets created private.
static bool trustedCacheDir(const std::string& path)
{
#if defined(_WIN32) || defined(RIVE_NX)
  return true;
#else
  struct stat st;
  return stat(path.c_str(), &st) != 0 ||
         (S_ISDIR(st.st_mode) && st.st_uid == geteuid() &&
          (st.st_mode & (S_IWGRP | S_IWOTH)) == 0);
#endif
}

#ifndef RIVE_NX
// The temp dir is per user on Apple and Windows but shared on Linux, where
// the user's own cache dir holds the artifacts instead.
static std::string defaultCacheDir()
{
#if defined(_WIN32) || defined(__APPLE__)
  return (std::filesystem::temp_directory_path() / "rive_aot_cache").string();
#else
  const char* xdg = getenv("XDG_CACHE_HOME");
  if (xdg != nullptr && xdg[0] == '/')
  {
    return std::string(xdg) + "/rive/aot_cache";
  }
  const char* home = getenv("HOME");
  return home != nullptr && home[0] == '/'
             ? std::string(home) + "/.cache/rive/aot_cache"
             : std::string();
#endif
}
#endif

// Windows refuses to rename over an existing file, and another compile of
// the same artifact landing first is still a success.
static bool publishFile(const std::string& tmpPath, const std::string& path)
{
  if (rename(tmpPath.c_str(), path.c_str()) == 0)
  {
    return true;
  }
  std::remove(tmpPath.c_str());
  return fileExists(path);
}

// The x86-64 level this machine's host artifacts compile for, 0 off x86.
static int hostX86_64Level()
{
#if defined(__x86_64__) || defined(_M_X64)
  static const int level = [] {
    // CI holds a v3 host at the baseline, where relaxed madd calls fmaf.
    const char* cap = getenv("RIVE_AOT_X86_LEVEL");
    if (cap != nullptr && strcmp(cap, "2") == 0)
    {
      return 2;
    }
    uint32_t leaf1Ecx = 0;
    uint32_t leaf7Ebx = 0;
    uint32_t extEcx = 0;
    uint64_t xcr0 = 0;
#if defined(_MSC_VER) && !defined(__clang__)
    int regs[4];
    __cpuid(regs, 0);
    int maxLeaf = regs[0];
    __cpuid(regs, 1);
    leaf1Ecx = (uint32_t)regs[2];
    if (maxLeaf >= 7)
    {
      __cpuidex(regs, 7, 0);
      leaf7Ebx = (uint32_t)regs[1];
    }
    __cpuid(regs, 0x80000000);
    if ((uint32_t)regs[0] >= 0x80000001)
    {
      __cpuid(regs, 0x80000001);
      extEcx = (uint32_t)regs[2];
    }
    if ((leaf1Ecx & (1u << 27)) != 0)
    {
      xcr0 = _xgetbv(0);
    }
#else
    unsigned a, b, c, d;
    if (__get_cpuid(1, &a, &b, &c, &d))
    {
      leaf1Ecx = c;
    }
    if (__get_cpuid_count(7, 0, &a, &b, &c, &d))
    {
      leaf7Ebx = b;
    }
    if (__get_cpuid(0x80000001, &a, &b, &c, &d))
    {
      extEcx = c;
    }
    // xgetbv faults unless the OS turned on xsave.
    if ((leaf1Ecx & (1u << 27)) != 0)
    {
      uint32_t lo;
      uint32_t hi;
      __asm__ volatile("xgetbv" : "=a"(lo), "=d"(hi) : "c"(0));
      xcr0 = ((uint64_t)hi << 32) | lo;
    }
#endif
    return x86_64Level(leaf1Ecx, leaf7Ebx, extEcx, xcr0);
  }();
  return level;
#else
  return 0;
#endif
}

std::string wamrcHwGuardFlag()
{
  return "--rive-hw-guard=" + std::to_string(WASM_RIVE_HW_GUARD_SIZE);
}

std::string loaderBinaryPath()
{
  std::string path;
#if defined(_WIN32)
  HMODULE module = nullptr;
  char name[MAX_PATH];
  if (GetModuleHandleExA(
          GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
              GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
          (LPCSTR)&loaderBinaryPath,
          &module))
  {
    DWORD length = GetModuleFileNameA(module, name, MAX_PATH);
    if (length > 0 && length < MAX_PATH)
    {
      path.assign(name, length);
    }
  }
#elif defined(__linux__) && !defined(__ANDROID__)
  Dl_info info;
  link_map* map = nullptr;
  if (dladdr1(
          (const void*)&loaderBinaryPath,
          &info,
          (void**)&map,
          RTLD_DL_LINKMAP) != 0 &&
      map != nullptr)
  {
    // glibc names the main program by argv[0], bare on a PATH launch.
    path = map->l_name[0] == '\0' ? "/proc/self/exe" : map->l_name;
  }
#elif !defined(RIVE_ANDROID) && !defined(RIVE_NX)
  Dl_info info;
  if (dladdr((const void*)&loaderBinaryPath, &info) != 0 &&
      info.dli_fname != nullptr)
  {
    path = info.dli_fname;
  }
#endif
  return path;
}

// Every relink rewrites the binary holding the loader, so its stat tells
// builds apart without reading it.
static const std::string& loaderBuild()
{
  static const std::string build = [] {
    std::string path = loaderBinaryPath();
    struct stat st;
    uint64_t key = 0;
    if (!path.empty() && stat(path.c_str(), &st) == 0)
    {
      const uint64_t fields[] = {
          (uint64_t)st.st_ino,
          (uint64_t)st.st_size,
          (uint64_t)st.st_mtime};
      key = wasmModuleKey(
          Span<const uint8_t>((const uint8_t*)fields, sizeof(fields)));
    }
    char text[17];
    snprintf(text, sizeof(text), "%016llx", (unsigned long long)key);
    return std::string(text);
  }();
  return build;
}

// A loader fix changes no compiler key, so each build keeps its own marker.
static std::string rejectedMarker(const std::string& artifactPath)
{
  return artifactPath + "." + loaderBuild() + ".rejected";
}

#ifdef _WIN32
// CreateProcess takes one command line; quote each argument so the child's
// CRT splits it back into the same argv.
static std::string quoteArg(const std::string& arg)
{
  if (!arg.empty() && arg.find_first_of(" \t\"") == std::string::npos)
  {
    return arg;
  }
  std::string out = "\"";
  size_t slashes = 0;
  for (char c : arg)
  {
    if (c == '\\')
    {
      slashes++;
      continue;
    }
    out.append(c == '"' ? slashes * 2 + 1 : slashes, '\\');
    slashes = 0;
    out += c;
  }
  out.append(slashes * 2, '\\');
  return out + "\"";
}

static HANDLE inheritable(HANDLE handle)
{
  HANDLE copy = nullptr;
  if (handle == nullptr || handle == INVALID_HANDLE_VALUE ||
      !DuplicateHandle(
          GetCurrentProcess(),
          handle,
          GetCurrentProcess(),
          &copy,
          0,
          TRUE,
          DUPLICATE_SAME_ACCESS))
  {
    return nullptr;
  }
  return copy;
}

// Starts args below interactive priority with no console window. stdout goes
// to out, or nowhere; stderr is the host's unless quiet. Only these handles
// are inherited, so a concurrent child never holds another's pipe open.
static intptr_t startChild(
    const std::vector<std::string>& args,
    std::string* failure,
    HANDLE out = nullptr,
    bool quiet = false)
{
  std::string line;
  for (const std::string& arg : args)
  {
    line += (line.empty() ? "" : " ") + quoteArg(arg);
  }
  SECURITY_ATTRIBUTES sa{sizeof(sa), nullptr, TRUE};
  HANDLE nul = CreateFileA(
      "NUL",
      GENERIC_READ | GENERIC_WRITE,
      FILE_SHARE_READ | FILE_SHARE_WRITE,
      &sa,
      OPEN_EXISTING,
      0,
      nullptr);
  HANDLE err = quiet ? nullptr : inheritable(GetStdHandle(STD_ERROR_HANDLE));
  STARTUPINFOEXA si{};
  si.StartupInfo.cb = sizeof(si);
  si.StartupInfo.dwFlags = STARTF_USESTDHANDLES;
  si.StartupInfo.hStdInput = nul;
  si.StartupInfo.hStdOutput = out != nullptr ? out : nul;
  si.StartupInfo.hStdError = err != nullptr ? err : nul;
  HANDLE handles[3] = {
      nul,
      si.StartupInfo.hStdOutput,
      si.StartupInfo.hStdError};
  std::sort(std::begin(handles), std::end(handles));
  DWORD handleCount = (DWORD)(std::unique(
                                  std::begin(handles),
                                  std::end(handles)) -
                              handles);
  SIZE_T attrSize = 0;
  InitializeProcThreadAttributeList(nullptr, 1, 0, &attrSize);
  std::vector<uint8_t> attrs(attrSize);
  si.lpAttributeList = (LPPROC_THREAD_ATTRIBUTE_LIST)attrs.data();
  PROCESS_INFORMATION pi{};
  bool listReady = InitializeProcThreadAttributeList(
      si.lpAttributeList,
      1,
      0,
      &attrSize);
  bool ok = listReady &&
            UpdateProcThreadAttribute(
                si.lpAttributeList,
                0,
                PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
                handles,
                handleCount * sizeof(HANDLE),
                nullptr,
                nullptr) &&
            CreateProcessA(
                nullptr,
                line.data(),
                nullptr,
                nullptr,
                TRUE,
                CREATE_NO_WINDOW | BELOW_NORMAL_PRIORITY_CLASS |
                    EXTENDED_STARTUPINFO_PRESENT,
                nullptr,
                nullptr,
                &si.StartupInfo,
                &pi);
  DWORD error = ok ? 0 : GetLastError();
  if (listReady)
  {
    DeleteProcThreadAttributeList(si.lpAttributeList);
  }
  CloseHandle(nul);
  if (err != nullptr)
  {
    CloseHandle(err);
  }
  if (!ok)
  {
    *failure = "could not run " + args[0] + ", windows error " +
               std::to_string(error);
    return -1;
  }
  CloseHandle(pi.hThread);
  return (intptr_t)pi.hProcess;
}

static void killChild(intptr_t child) { TerminateProcess((HANDLE)child, 1); }

// The exit code, or -1 when the child did not exit on its own.
static int waitChild(intptr_t child)
{
  DWORD code = 1;
  if (WaitForSingleObject((HANDLE)child, INFINITE) != WAIT_OBJECT_0 ||
      !GetExitCodeProcess((HANDLE)child, &code))
  {
    return -1;
  }
  return (int)code;
}

static void releaseChild(intptr_t child) { CloseHandle((HANDLE)child); }

// Runs wamrc with flag and --version: true when it exits cleanly, with the
// first line of its output when asked.
static bool probeWamrc(
    const std::string& wamrc,
    const char* flag,
    std::string* firstLine)
{
  std::vector<std::string> args = {wamrc};
  if (flag != nullptr)
  {
    args.push_back(flag);
  }
  args.push_back("--version");
  SECURITY_ATTRIBUTES sa{sizeof(sa), nullptr, TRUE};
  HANDLE readEnd = nullptr;
  HANDLE writeEnd = nullptr;
  if (!CreatePipe(&readEnd, &writeEnd, &sa, 0))
  {
    return false;
  }
  SetHandleInformation(readEnd, HANDLE_FLAG_INHERIT, 0);
  std::string failure;
  intptr_t child = startChild(args, &failure, writeEnd, true);
  CloseHandle(writeEnd);
  std::string text;
  char buffer[256];
  DWORD read = 0;
  while (ReadFile(readEnd, buffer, sizeof(buffer), &read, nullptr) && read > 0)
  {
    text.append(buffer, read);
  }
  CloseHandle(readEnd);
  if (child == -1)
  {
    return false;
  }
  int code = waitChild(child);
  releaseChild(child);
  if (firstLine != nullptr)
  {
    *firstLine = text.substr(0, text.find('\n'));
  }
  return code == 0;
}
#elif defined(RIVE_ANDROID) || defined(RIVE_NX)
// No wamrc on device, and neither platform can spawn one.
static intptr_t startChild(
    const std::vector<std::string>&,
    std::string* failure)
{
  *failure = "wamrc cannot run on this platform";
  return -1;
}

static void killChild(intptr_t) {}

static int waitChild(intptr_t) { return -1; }

static void releaseChild(intptr_t) {}
#else
// Starts args below interactive priority; a compiler's stdout is noise at
// editor runtime, stderr stays for diagnosing failed compiles.
static intptr_t startChild(
    const std::vector<std::string>& args,
    std::string* failure)
{
  std::vector<char*> argv;
  for (const std::string& arg : args)
  {
    argv.push_back(const_cast<char*>(arg.c_str()));
  }
  argv.push_back(nullptr);
  pid_t pid = -1;
  posix_spawn_file_actions_t actions;
  posix_spawn_file_actions_init(&actions);
  posix_spawn_file_actions_addopen(
      &actions,
      STDOUT_FILENO,
      "/dev/null",
      O_WRONLY,
      0);
  int rc = posix_spawn(&pid, argv[0], &actions, nullptr, argv.data(), environ);
  posix_spawn_file_actions_destroy(&actions);
  if (rc != 0)
  {
    *failure = "could not run " + args[0] + ", " + strerror(rc);
    return -1;
  }
  // The editor's frame loop must never contend with LLVM.
  setpriority(PRIO_PROCESS, pid, 10);
  return pid;
}

static void killChild(intptr_t child) { kill((pid_t)child, SIGKILL); }

// The exit code, or -1 when the child did not exit on its own.
static int waitChild(intptr_t child)
{
  int status = 0;
  while (waitpid((pid_t)child, &status, 0) < 0 && errno == EINTR)
  {
  }
  return WIFEXITED(status) ? WEXITSTATUS(status) : -1;
}

static void releaseChild(intptr_t) {}
#endif

#if !defined(_WIN32) && !defined(RIVE_NX)
// Installed apps put wamrc under paths with spaces.
static std::string shellQuote(const std::string& text)
{
  std::string quoted = "'";
  for (char c : text)
  {
    quoted += c == '\'' ? std::string("'\\''") : std::string(1, c);
  }
  return quoted + "'";
}

// Runs wamrc with flag and --version: true when it exits cleanly, with the
// first line of its output when asked.
static bool probeWamrc(
    const std::string& wamrc,
    const char* flag,
    std::string* firstLine)
{
  std::string cmd = shellQuote(wamrc) +
                    (flag != nullptr ? " " + std::string(flag) : "") +
                    " --version" +
                    (firstLine != nullptr ? " 2>/dev/null"
                                          : " >/dev/null 2>&1");
  FILE* pipe = popen(cmd.c_str(), "r");
  if (pipe == nullptr)
  {
    return false;
  }
  char line[128] = {0};
  if (firstLine != nullptr && fgets(line, sizeof(line), pipe) != nullptr)
  {
    *firstLine = line;
  }
  return pclose(pipe) == 0;
}
#endif

ModuleTierLadder& ModuleTierLadder::instance()
{
  static ModuleTierLadder* ladder = new ModuleTierLadder();
  return *ladder;
}

void ModuleTierLadder::configure(
    const std::string& wamrcPath,
    const std::string& cacheDir,
    bool compileOnBoot)
{
  // Before a rebuild can replace the loader's binary on disk.
  loaderBuild();
  std::unique_lock<std::mutex> lock(m_mutex);
  m_wamrcPath = wamrcPath;
  m_compileOnBoot = compileOnBoot;
  m_versionProbed = false;
  if (m_wamrcPath.empty())
  {
    if (const char* env = getenv("RIVE_WAMRC"))
    {
      m_wamrcPath = env;
    }
  }
  m_cacheDir = cacheDir;
  m_failures.clear();
  m_compiled.clear();
  m_pruned = false;
  if (m_cacheDir.empty() && !m_wamrcPath.empty())
  {
    const char* dirEnv = getenv("RIVE_AOT_CACHE_DIR");
#ifdef RIVE_NX
    // No temp_directory_path in this libc++, and the platform cannot map
    // executable pages anyway, so the ladder only runs when pointed at a
    // directory explicitly.
    m_cacheDir = dirEnv != nullptr ? dirEnv : "";
#else
    m_cacheDir = dirEnv != nullptr ? dirEnv : defaultCacheDir();
#endif
  }
}

void ModuleTierLadder::unconfigure()
{
  std::unique_lock<std::mutex> lock(m_mutex);
  m_wamrcPath.clear();
  m_cacheDir.clear();
  m_failures.clear();
  m_compiled.clear();
  m_versionProbed = false;
  m_compileOnBoot = false;
}

bool ModuleTierLadder::enabled()
{
  std::unique_lock<std::mutex> lock(m_mutex);
  return enabledLocked();
}

bool ModuleTierLadder::enabledLocked()
{
  if (m_wamrcPath.empty())
  {
    m_disabledReason =
        "no wamrc, set RIVE_WAMRC or pass --aot with our wamrc on PATH";
    return false;
  }
  if (m_cacheDir.empty())
  {
    m_disabledReason = "no aot cache, set RIVE_AOT_CACHE_DIR";
    return false;
  }
  if (!trustedCacheDir(m_cacheDir))
  {
    m_disabledReason = "the aot cache " + m_cacheDir +
                       " is not private to this user, run chmod go-w on it";
    return false;
  }
  wamrcVersion();
  m_disabledReason = m_wamrcUsable ? "" : m_wamrcPath + " " + m_wamrcProblem;
  return m_wamrcUsable;
}

std::string ModuleTierLadder::disabledReason()
{
  if (enabled())
  {
    return std::string();
  }
  std::unique_lock<std::mutex> lock(m_mutex);
  return m_disabledReason;
}

std::string ModuleTierLadder::untrustedCacheDir()
{
  std::unique_lock<std::mutex> lock(m_mutex);
  return m_cacheDir.empty() || trustedCacheDir(m_cacheDir) ? std::string()
                                                           : m_cacheDir;
}

std::string ModuleTierLadder::compilerPath()
{
  std::unique_lock<std::mutex> lock(m_mutex);
  return m_wamrcPath;
}

const std::string& ModuleTierLadder::wamrcVersion()
{
  // Callers hold m_mutex. The version keys the cache so a compiler swap
  // reads as misses, never as stale native code.
  if (!m_versionProbed)
  {
    m_versionProbed = true;
    m_wamrcVersion = "unknown";
#ifndef RIVE_NX
    std::string v;
    probeWamrc(m_wamrcPath, nullptr, &v);
    while (!v.empty() &&
           (v.back() == '\n' || v.back() == '\r' || v.back() == ' '))
    {
      v.pop_back();
    }
    for (char& c : v)
    {
      if (c == ' ' || c == '/')
      {
        c = '-';
      }
    }
    if (!v.empty())
    {
      m_wamrcVersion = v;
    }
    // Every build of our patched wamrc prints the same version.
    uint64_t binaryKey = 0;
    fileContentKey(m_wamrcPath, binaryKey);
    char binary[20];
    snprintf(binary, sizeof(binary), "-%016llx", (unsigned long long)binaryKey);
    m_wamrcVersion += binary;
    auto accepts = [&](const char* flag) {
      return probeWamrc(m_wamrcPath, flag, nullptr);
    };
    // Only our patched wamrc takes --rive-interrupt, and only a current one
    // the guard flag; say so once rather than fail each compile quietly.
    m_wamrcProblem.clear();
    if (!accepts("--rive-interrupt"))
    {
      m_wamrcProblem = "is not our patched wamrc; build wamrc from our WAMR "
                       "patches (runtime/scripting/wamr_patches)";
    }
    else if (!accepts(wamrcHwGuardFlag().c_str()))
    {
      m_wamrcProblem = "predates our current WAMR patches; rebuild wamrc "
                       "from the current patches "
                       "(runtime/scripting/wamr_patches)";
    }
    m_wamrcUsable = m_wamrcProblem.empty();
    m_wamrcTunes = accepts("--tune-cpu=generic");
    if (!m_wamrcUsable)
    {
      fprintf(
          stderr,
          "wasm aot: %s %s. AOT is off.\n",
          m_wamrcPath.c_str(),
          m_wamrcProblem.c_str());
    }
#endif
  }
  return m_wamrcVersion;
}

// wamrc flags naming a device target, null for one we cannot compile for.
static const char* const* crossTargetArgs(const std::string& target)
{
  // Android reserves x18 for its shadow call stack.
  static const char* const androidAarch64[] = {
      "--target=aarch64v8",
      "--target-abi=gnu",
      "--cpu=generic",
      "--cpu-features=+reserve-x18",
      nullptr};
  // wamrc compiles SIMD only for SSE4.1, which the NDK x86_64 ABI
  // guarantees.
  static const char* const androidX64[] =
      {"--target=x86_64", "--target-abi=gnu", "--cpu=x86-64-v2", nullptr};
  if (target == kAotTargetAndroidAarch64)
  {
    return androidAarch64;
  }
  if (target == kAotTargetAndroidX64)
  {
    return androidX64;
  }
  return nullptr;
}

static bool asksHostBudget(const std::string& wasmPath)
{
  std::ifstream in(wasmPath, std::ios::binary);
  std::vector<uint8_t> bytes(
      (std::istreambuf_iterator<char>(in)),
      std::istreambuf_iterator<char>());
  return exportsFunction(bytes, "__riveHostBudget");
}

std::string ModuleTierLadder::keyedCacheDir()
{
  // The revision folds our wamrc flag choices into the cache key; bump
  // it whenever species flags change or stale artifacts get served on a
  // cache hit.
  const std::string& version = wamrcVersion();
  return m_cacheDir + "/" + version + "-r18-g" +
         std::to_string(WASM_RIVE_HW_GUARD_SIZE) + (m_wamrcTunes ? "t" : "");
}

// Device artifacts sit apart from this machine's, which share their names.
std::string ModuleTierLadder::artifactDir(const std::string& target)
{
  std::string dir = keyedCacheDir();
  if (!target.empty())
  {
    return dir + "/" + target;
  }
  // A host sharing the cache finds only code its cpu runs.
  return hostX86_64Level() >= 3 ? dir + "/x86-64-v3" : dir;
}

void ModuleTierLadder::createArtifactDir(const std::string& target)
{
  makeDir(m_cacheDir);
  makeDir(keyedCacheDir());
  makeDir(artifactDir(target));
}

std::string ModuleTierLadder::artifactName(
    uint64_t moduleKey,
    TierSpecies species)
{
  if (species == TierSpecies::o0)
  {
    return aotArtifactName(moduleKey, ".o0");
  }
  if (species == TierSpecies::hw)
  {
    return aotArtifactName(moduleKey, ".hw");
  }
  return aotArtifactName(moduleKey);
}

std::string ModuleTierLadder::artifactPath(
    uint64_t moduleKey,
    TierSpecies species,
    const std::string& target)
{
  std::unique_lock<std::mutex> lock(m_mutex);
  // Targets come from link peers, so only known ones reach a path.
  if (m_wamrcPath.empty() || m_cacheDir.empty() ||
      !trustedCacheDir(m_cacheDir) ||
      (!target.empty() && crossTargetArgs(target) == nullptr))
  {
    return std::string();
  }
  std::string path = artifactDir(target) + "/" +
                     artifactName(moduleKey, species);
  return fileExists(path) ? path : std::string();
}

#ifndef RIVE_NX
// Another process may still be using an older compiler's key dir this long
// after its last compile.
constexpr auto kStaleKeyDirAge = std::chrono::hours(24 * 7);
constexpr uintmax_t kKeyDirBytes = 512ull << 20;

// The names every keyedCacheDir has had: a wamrc version line, then a
// revision.
static bool isKeyDirName(const std::string& name)
{
  size_t revision = name.rfind("-r");
  if ((name.rfind("wamrc-", 0) != 0 && name.rfind("unknown-", 0) != 0) ||
      revision == std::string::npos)
  {
    return false;
  }
  std::string digits = name.substr(revision + 2);
  if (!digits.empty() && digits.back() == 't')
  {
    digits.pop_back();
  }
  return !digits.empty() &&
         digits.find_first_not_of("0123456789") == std::string::npos;
}

// Every module and every wamrc update leave artifacts behind, so stale key
// dirs go, then this one's oldest compiles until it fits.
static void pruneCache(
    const std::string& cacheDir,
    const std::string& keyedDir,
    bool staleKeyDirs)
{
  namespace fs = std::filesystem;
  std::error_code ec;
  auto staleBefore = fs::file_time_type::clock::now() - kStaleKeyDirAge;
  if (staleKeyDirs)
  {
    std::string current = fs::path(keyedDir).filename().string();
    std::vector<fs::path> stale;
    for (fs::directory_iterator it(cacheDir, ec), end; !ec && it != end;
         it.increment(ec))
    {
      std::string name = it->path().filename().string();
      std::error_code statError;
      if (name != current && isKeyDirName(name) &&
          it->is_directory(statError) &&
          it->last_write_time(statError) < staleBefore && !statError)
      {
        stale.push_back(it->path());
      }
    }
    for (const fs::path& dir : stale)
    {
      fs::remove_all(dir, ec);
    }
  }
  struct Artifact
  {
    fs::file_time_type written;
    uintmax_t size;
    fs::path path;
  };
  std::vector<Artifact> artifacts;
  std::vector<fs::path> staleMarkers;
  uintmax_t total = 0;
  ec.clear();
  for (fs::recursive_directory_iterator it(keyedDir, ec), end; !ec && it != end;
       it.increment(ec))
  {
    std::string extension = it->path().extension().string();
    std::error_code statError;
    // Each build leaves its own markers, so old ones expire like key dirs.
    if (extension == ".rejected" &&
        it->last_write_time(statError) < staleBefore && !statError)
    {
      staleMarkers.push_back(it->path());
    }
    else if (
        (extension == ".aot" || extension == ".wasm") &&
        it->is_regular_file(statError))
    {
      Artifact artifact{
          it->last_write_time(statError),
          it->file_size(statError),
          it->path()};
      if (!statError)
      {
        total += artifact.size;
        artifacts.push_back(std::move(artifact));
      }
    }
  }
  for (const fs::path& marker : staleMarkers)
  {
    fs::remove(marker, ec);
  }
  if (total <= kKeyDirBytes)
  {
    return;
  }
  std::sort(
      artifacts.begin(),
      artifacts.end(),
      [](const Artifact& a, const Artifact& b) {
        return a.written < b.written;
      });
  for (const Artifact& artifact : artifacts)
  {
    if (total <= kKeyDirBytes)
    {
      break;
    }
    if (fs::remove(artifact.path, ec))
    {
      total -= artifact.size;
    }
  }
}
#endif

// The cache is shared across processes and a name must never be visible
// half-written: write a process-unique temp, then atomically rename.
static bool writeFileAtomic(const std::string& path, Span<const uint8_t> bytes)
{
  std::string tmpPath = path + "." + std::to_string(processId()) + ".tmp";
  FILE* f = fopen(tmpPath.c_str(), "wb");
  if (f == nullptr)
  {
    return false;
  }
  size_t written = fwrite(bytes.data(), 1, bytes.size(), f);
  if (fclose(f) != 0 || written != bytes.size())
  {
    std::remove(tmpPath.c_str());
    return false;
  }
  return publishFile(tmpPath, path);
}

TierCompile ModuleTierLadder::compileSync(
    uint64_t moduleKey,
    Span<const uint8_t> moduleBytes,
    TierSpecies species,
    const std::string& target,
    const std::atomic<bool>* cancel)
{
  TierCompile result;
  Job job{moduleKey, species};
  job.target = target;
  job.cancel = cancel;
  std::string cacheDir;
  std::string keyedDir;
  bool staleKeyDirs = false;
  {
    std::unique_lock<std::mutex> lock(m_mutex);
    if (lookupLocked(job, result))
    {
      return result;
    }
    createArtifactDir(target);
    // Another user may have made the dir first.
    if (!enabledLocked())
    {
      result.reason = m_disabledReason;
      return result;
    }
    cacheDir = m_cacheDir;
    keyedDir = keyedCacheDir();
    job.wasmPath = keyedDir + "/" + wasmModuleName(moduleKey);
    if (!fileExists(job.wasmPath) &&
        !writeFileAtomic(job.wasmPath, moduleBytes))
    {
      result.reason = "could not write the module into the aot cache " +
                      m_cacheDir;
      m_failures[{moduleKey, species, target}] = result.reason;
      return result;
    }
    if (cancel != nullptr)
    {
      m_cancelableJobs.push_back(&job);
    }
  }
  auto compileStart = std::chrono::steady_clock::now();
  result.reason = runWamrc(job);
  {
    std::unique_lock<std::mutex> lock(m_mutex);
    if (cancel != nullptr)
    {
      m_cancelableJobs.erase(
          std::find(m_cancelableJobs.begin(), m_cancelableJobs.end(), &job));
    }
    if (!result.reason.empty() && !job.cancelled)
    {
      m_failures[{moduleKey, species, target}] = result.reason;
    }
    if (result.reason.empty() && target.empty())
    {
      m_compiled.insert({moduleKey, species});
    }
    staleKeyDirs = result.reason.empty() && !m_pruned;
    m_pruned = m_pruned || staleKeyDirs;
  }
  if (result.reason.empty())
  {
    result.path = job.finalPath;
    result.compiled = true;
    if (target.empty())
    {
      auto compileMs = std::chrono::duration_cast<std::chrono::milliseconds>(
                           std::chrono::steady_clock::now() - compileStart)
                           .count();
      fprintf(
          stderr,
          "wasm aot: compiled %s in %lld ms\n",
          result.path.c_str(),
          (long long)compileMs);
    }
#ifndef RIVE_NX
    pruneCache(cacheDir, keyedDir, staleKeyDirs);
#endif
  }
  return result;
}

std::vector<TierCompile> ModuleTierLadder::compileAll(
    const std::vector<Span<const uint8_t>>& modules,
    TierSpecies species,
    const std::atomic<bool>* cancel)
{
  std::vector<uint64_t> keys;
  for (Span<const uint8_t> module : modules)
  {
    keys.push_back(wasmModuleKey(module));
  }
  std::vector<TierCompile> results(modules.size());
  std::vector<size_t> misses;
  {
    std::unique_lock<std::mutex> lock(m_mutex);
    for (size_t i = 0; i < modules.size(); i++)
    {
      Job job{keys[i], species};
      if (!lookupLocked(job, results[i]))
      {
        misses.push_back(i);
      }
    }
  }
  auto compile = [&](size_t i) {
    if (cancel != nullptr && cancel->load())
    {
      results[i].reason = kCancelledReason;
      return;
    }
    results[i] = compileSync(
        keys[i],
        modules[i],
        species,
        std::string(),
        cancel);
  };
  std::vector<std::thread> threads;
  for (size_t i = 1; i < misses.size(); i++)
  {
    threads.emplace_back(compile, misses[i]);
  }
  if (!misses.empty())
  {
    compile(misses.front());
  }
  for (std::thread& thread : threads)
  {
    thread.join();
  }
  return results;
}

std::vector<TierCompile> ModuleTierLadder::compileWithLanes(
    Span<const uint8_t> moduleBytes,
    TierSpecies species,
    const std::atomic<bool>* cancel)
{
  std::vector<std::vector<uint8_t>> lanes = laneModules(moduleBytes);
  std::vector<Span<const uint8_t>> modules = {moduleBytes};
  for (const std::vector<uint8_t>& lane : lanes)
  {
    modules.emplace_back(lane.data(), lane.size());
  }
  return compileAll(modules, species, cancel);
}

bool ModuleTierLadder::lookupLocked(Job& job, TierCompile& result)
{
  if (!enabledLocked())
  {
    result.reason = m_disabledReason;
    return true;
  }
  if (!job.target.empty() && crossTargetArgs(job.target) == nullptr)
  {
    result.reason = "wamrc cannot compile for " + job.target;
    return true;
  }
  auto failure = m_failures.find({job.moduleKey, job.species, job.target});
  if (failure != m_failures.end())
  {
    result.reason = failure->second;
    return true;
  }
  job.finalPath = artifactDir(job.target) + "/" +
                  artifactName(job.moduleKey, job.species);
  if (fileExists(job.finalPath))
  {
    result.path = job.finalPath;
    return true;
  }
  std::ifstream rejected(rejectedMarker(job.finalPath));
  if (rejected && std::getline(rejected, result.reason))
  {
    m_failures[{job.moduleKey, job.species, job.target}] = result.reason;
    return true;
  }
  return false;
}

TierCompile ModuleTierLadder::bootCompile(
    uint64_t moduleKey,
    Span<const uint8_t> moduleBytes,
    TierSpecies species)
{
  {
    std::unique_lock<std::mutex> lock(m_mutex);
    if (!m_compileOnBoot)
    {
      TierCompile result;
      Job job{moduleKey, species};
      if (!lookupLocked(job, result))
      {
        result.reason = "no aot artifact cached for this module yet";
      }
      return result;
    }
  }
  // The boot loads its lanes next, and compiles of one module never merge,
  // so they all finish before it returns.
  return compileWithLanes(moduleBytes, species).front();
}

bool ModuleTierLadder::failed(
    uint64_t moduleKey,
    TierSpecies species,
    const std::string& target)
{
  std::unique_lock<std::mutex> lock(m_mutex);
  return m_failures.count({moduleKey, species, target}) != 0;
}

void ModuleTierLadder::rejectArtifact(
    uint64_t moduleKey,
    TierSpecies species,
    const std::string& reason)
{
  std::unique_lock<std::mutex> lock(m_mutex);
  m_failures[{moduleKey, species, std::string()}] = reason;
  if (m_cacheDir.empty())
  {
    return;
  }
  std::string path = artifactDir(std::string()) + "/" +
                     artifactName(moduleKey, species);
  std::remove(path.c_str());
  // A recompile would make the same artifact again, so later sessions skip
  // it. One an older session left, maybe torn, compiles afresh.
  if (m_compiled.count({moduleKey, species}) != 0)
  {
    writeFileAtomic(
        rejectedMarker(path),
        Span<const uint8_t>((const uint8_t*)reason.data(), reason.size()));
  }
}

void ModuleTierLadder::cancelCompiles()
{
  std::unique_lock<std::mutex> lock(m_mutex);
  for (Job* job : m_cancelableJobs)
  {
    if (job->cancel == nullptr || !job->cancel->load())
    {
      continue;
    }
    job->cancelled = true;
    if (job->child != -1)
    {
      killChild(job->child);
    }
  }
}

std::string ModuleTierLadder::runWamrc(Job& job)
{
  // Callers dropped m_mutex; only job fields and immutable config are
  // touched off-lock, except child which cancelCompiles reads under
  // lock.
  const std::string& finalPath = job.finalPath;
  std::string wamrc;
  std::string cacheDir;
  [[maybe_unused]] bool tunes = false;
  {
    std::unique_lock<std::mutex> lock(m_mutex);
    wamrc = m_wamrcPath;
    cacheDir = m_cacheDir;
    tunes = m_wamrcTunes;
  }
  // Unique temp per compile: the same artifact can be raced by two boots
  // on different threads, or by two processes sharing the cache.
  static std::atomic<uint64_t> tmpSerial{0};
  std::string tmpPath = finalPath + "." + std::to_string(processId()) + "-" +
                        std::to_string(++tmpSerial) + ".tmp";

  std::vector<std::string> args = {wamrc};
  if (job.species == TierSpecies::hw)
  {
    // Guard-page bounds; native stack checks stay sw (the runtime
    // builds with WASM_DISABLE_STACK_HW_BOUND_CHECK).
    args.push_back("--bounds-checks=0");
    args.push_back("--stack-bounds-checks=1");
    args.push_back(wamrcHwGuardFlag());
  }
  else
  {
    // sw bounds run on every build; wamrc's default hw-bounds output
    // segfaults on runtimes without the guard-page trap handler.
    args.push_back("--bounds-checks=1");
  }
  if (job.species == TierSpecies::o0)
  {
    args.push_back("--opt-level=0");
  }
  else
  {
    args.push_back("--opt-level=3");
  }
  if (!job.target.empty())
  {
    const char* const* targetArgs = crossTargetArgs(job.target);
    if (targetArgs == nullptr)
    {
      return "wamrc cannot compile for " + job.target;
    }
    for (; *targetArgs != nullptr; targetArgs++)
    {
      args.push_back(*targetArgs);
    }
  }
#if (defined(__APPLE__) || defined(_WIN32)) &&                                 \
    (defined(__aarch64__) || defined(_M_ARM64))
  else
  {
    // The platform clobbers x18; reserve it here so safety does not
    // depend on the wamrc we were handed. Features need an explicit cpu.
    args.push_back("--cpu=generic");
#if defined(__APPLE__)
    // Every Apple silicon Mac has single instruction atomics, which make
    // parallel calls' barriers cheaper than load and store exclusive loops.
    args.push_back("--cpu-features=+reserve-x18,+lse");
#else
    args.push_back("--cpu-features=+reserve-x18");
#endif
#if defined(__APPLE__)
    // Generic scheduling turns float compare chains into mispredicted
    // branches (up to 1.5x slower); tuning leaves the ISA baseline alone.
    if (tunes)
    {
      args.push_back("--tune-cpu=apple-m1");
    }
#endif
  }
#elif defined(__x86_64__) || defined(_M_X64)
  else
  {
    // Left alone wamrc targets this machine's exact cpu, which the cache
    // does not key. Without FMA, relaxed madd calls fmaf.
    args.push_back("--target=x86_64");
    args.push_back(
        hostX86_64Level() >= 3 ? "--cpu=x86-64-v3" : "--cpu=x86-64-v2");
#if defined(__APPLE__)
    // Code maps above 4GB here, out of reach of 32 bit relocations.
    args.push_back("--size-level=0");
#elif defined(_WIN32)
    // wamrc's own default for this host, which naming a target drops.
    args.push_back("--size-level=1");
#endif
  }
#endif
  // Loop back-edges and calls check the watchdog only in modules that ask
  // for it; the module bytes, and so its cache key, decide.
  if (asksHostBudget(job.wasmPath))
  {
    args.push_back("--rive-interrupt");
  }
  args.push_back("-o");
  args.push_back(tmpPath);
  args.push_back(job.wasmPath);
  std::string failure;
  intptr_t child = startChild(args, &failure);
  if (child == -1)
  {
    fprintf(stderr, "[wasm] %s\n", failure.c_str());
    return failure;
  }
  {
    std::unique_lock<std::mutex> lock(m_mutex);
    job.child = child;
    // A cancel that landed before the spawn had nothing to kill.
    if (job.cancel != nullptr && job.cancel->load())
    {
      job.cancelled = true;
    }
    if (job.cancelled)
    {
      killChild(child);
    }
  }
  int exitCode = waitChild(child);
  {
    std::unique_lock<std::mutex> lock(m_mutex);
    job.child = -1;
  }
  releaseChild(child);
  if (job.cancelled || exitCode != 0)
  {
    std::remove(tmpPath.c_str());
    if (job.cancelled)
    {
      return kCancelledReason;
    }
    // wamrc prints usage errors to stdout, which is discarded; an
    // unlogged exit here is otherwise a silently dead lane.
    fprintf(
        stderr,
        "[wasm] wamrc failed (exit %d) on %s\n",
        exitCode,
        job.wasmPath.c_str());
    return "wamrc could not compile the module";
  }
  // Atomic arrival: a partial artifact can never carry the final name.
  if (!publishFile(tmpPath, finalPath))
  {
    return "could not write the artifact into the aot cache " + cacheDir;
  }
  return std::string();
}

} // namespace rive

#endif // RIVE_WASM_TIER_LADDER
#endif // WITH_RIVE_SCRIPTING_WASM