Code generation¶

def render_c_module(
  fun: Function | ConcreteFunction,
  *,
  header_name: str | None = None,
  source_name: str | None = None,
  lang: str = "c",
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
  frame_budget: int = 65536,
) -> CModule:
  """Render ``fun`` into its header / ``.c`` pair from a single lowering. The kernel is always C.
  ``lang`` picks the header a caller includes (``f.h`` or ``f.hpp``) and ``casadi`` adds the
  CasADi 3.8 compatible symbols to both."""
  fun = as_concrete(fun)
  _check_lang(lang)
  ctx, body = _render_observed(
    fun,
    casadi=casadi,
    recipe=BuildRecipe(cpu=cpu, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal, frame_budget=frame_budget),
    source_name=source_name,
  )
  symbol = c_ident(fun.name)
  return CModule(
    fun=ctx.fun,
    header_name=f"{symbol}.{'hpp' if lang == 'cpp' else 'h'}" if header_name is None else header_name,
    source_name=f"{symbol}.c" if source_name is None else source_name,
    body=body,
    program=ctx.prog,
    workspace_size=ctx.workspace_size,
    backends=ctx.backends,
    lang=lang,
    casadi=casadi,
    recipe=ctx.recipe,
    _libraries=ctx.libraries,
  )

def render_c_source(
  fun: Function | ConcreteFunction,
  *,
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
  frame_budget: int = 65536,
) -> str:
  """Render a standalone pointer-ABI C implementation of ``fun`` and its callees. ``casadi`` adds
  the CasADi 3.8 compatible symbols.

  A ``LoweringError`` for a function that cannot be lowered propagates. There is no fallback.
  """
  fun = as_concrete(fun)
  recipe = BuildRecipe(cpu=cpu, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal, frame_budget=frame_budget)
  return _render_observed(fun, casadi=casadi, recipe=recipe)[1]

def render_c_api_header(
  fun: Function | ConcreteFunction,
  *,
  lang: str = "c",
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
  frame_budget: int = 65536,
) -> str:
  """Render the public header for ``fun``: the ABI declarations, ``SZ_*`` constants, the typed
  buffers of the chosen ``lang``, the
  sparse-output tables, and with ``casadi`` the CasADi query prototypes."""
  fun = as_concrete(fun)
  _check_lang(lang)
  if casadi:
    fun = casadi_layout(fun)
  ctx = _lower(
    fun, recipe=BuildRecipe(cpu=cpu, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal, frame_budget=frame_budget)
  )
  return ctx.recipe.comment(f"{c_ident(fun.name)}.c") + _render_header(
    ctx.fun, ctx.backends, ctx.workspace_size, _frame_bytes(ctx.prog), lang=lang, casadi=casadi
  )

def write_module(
  fun: Function | ConcreteFunction | Solver,
  out_dir: Path,
  *,
  lang: str = "c",
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
  frame_budget: int = 65536,
) -> CModule:
  """Write ``fun``'s header / ``.c`` into ``out_dir`` and return the module. A ``Solver`` renders its ``function``."""
  if isinstance(fun, Solver):
    fun = fun.function
  module = render_c_module(
    fun, lang=lang, casadi=casadi, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal, cpu=cpu, frame_budget=frame_budget
  )
  out_dir.mkdir(parents=True, exist_ok=True)
  (out_dir / module.header_name).write_text(module.header)
  (out_dir / module.source_name).write_text(module.source)
  if module._libraries:
    manifest = {
      "source": module.source_name,
      "compile_flags": ["-O2", *module.recipe.cpu_flags, "-fno-math-errno"],
      "link": ["m", *(["mvec"] if module.recipe.vector_libm == "glibc" else [])],
      "libraries": export_libraries(module._libraries, out_dir),
    }
    (out_dir / f"{c_ident(module.fun.name)}.build.json").write_text(json.dumps(manifest, indent=2) + "\n")
  return module

def workspace_size(
  fun: Function | ConcreteFunction,
  *,
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
  frame_budget: int = 65536,
) -> int:
  """Entries of scratch ``fun`` needs in ``w[]``, the value its header's ``SZ_W`` quotes: doubles
  when every leaf is float64, floats when every leaf is float32, and bytes otherwise.
  ``CModule.workspace_size`` is the same number without a second lowering, so prefer it when the
  module is already in hand."""
  fun = as_concrete(fun)
  if casadi:
    fun = casadi_layout(fun)
  return _lower(
    fun, recipe=BuildRecipe(cpu=cpu, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal, frame_budget=frame_budget)
  ).workspace_size

def c_api_signature(fun: Function | ConcreteFunction, symbol: str | None = None, *, solver_options: bool = False) -> str:
  """The entry signature of ``fun``, named ``symbol`` or after ``fun``. ``arg``, ``res`` and ``w``
  point to ``double`` when every leaf is float64, to ``float`` when every leaf is float32, and to
  ``void`` otherwise, with ``SZ_W`` then counting bytes. ``mem`` follows CasADi 3.8: an ``int``
  memory handle a stateful function would index a pool with. Every Scaly function is stateless and
  ignores it."""
  fun = as_concrete(fun)
  t = ENTRY_C_TYPES[workspace_dtype(fun)]
  options = ", const scaly_solver_option* const* solver_options" if solver_options else ""
  return f"int {symbol or c_ident(fun.name)}(const {t}** arg, {t}** res, int* iw, {t}* w, int mem{options})"

def load_library(path: Path, *, isolated: bool, library_dirs: tuple[str, ...] = ()) -> ctypes.CDLL:
  """Load a shared library, isolating incompatible Linux dependency builds from one another.

  External kernels reuse a namespace when shared names identify the same path and content,
  including transitive dependencies. Loaded handles and
  namespaces live until process exit, so old compiled handles remain callable after a rebuild.
  Environment changes after the last load remain subject to the limitation tracked in #159.
  """
  if isolated and sys.platform == "win32":
    with ExitStack() as stack:
      for directory in library_dirs:
        stack.enter_context(os.add_dll_directory(str(Path(directory).resolve())))
      return ctypes.CDLL(str(path))
  if not isolated or sys.platform != "linux":
    return ctypes.CDLL(str(path))
  libc = ctypes.CDLL(None)
  dlmopen = libc.dlmopen
  dlmopen.argtypes = [ctypes.c_long, ctypes.c_char_p, ctypes.c_int]
  dlmopen.restype = ctypes.c_void_p
  dlerror = libc.dlerror
  dlerror.argtypes = []
  dlerror.restype = ctypes.c_char_p

  def open_library(namespace_id: int, name: str) -> ctypes.CDLL:
    dlerror()
    handle = dlmopen(namespace_id, os.fsencode(name), os.RTLD_NOW | os.RTLD_LOCAL)
    if not handle:
      error = dlerror()
      reason = error.decode(errors="replace") if error else "unknown loader error"
      advice = (
        ". Use a fresh process with fewer incompatible dependency versions." if "no more namespaces" in reason or "static TLS" in reason else ""
      )
      raise OSError(f"dlmopen failed for {name}: {reason}{advice}")
    # Python 3.14 reopens a path even when CDLL receives a handle, so bind it directly.
    lib = ctypes.CDLL(None)
    lib._handle = handle
    lib._name = name
    return lib

  with _LINKER_NAMESPACE_LOCK:
    dependencies = _namespace_dependencies(path)
    previous = next(
      (
        candidate
        for candidate in _LINKER_NAMESPACES
        if all(name not in candidate.dependencies or candidate.dependencies[name] == identity for name, identity in dependencies.items())
      ),
      None,
    )
    if previous is None:
      # A libc-only anchor leaves other symbols in each kernel's local dependency scope.
      # DEEPBIND can split libc's allocator calls from a kernel's free; the extern buffer
      # contract does not prevent that internal mismatch. Keep ordinary LOCAL lookup.
      anchor = open_library(_LM_ID_NEWLM, _linux_loader()[1])
      dlinfo = libc.dlinfo
      dlinfo.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p]
      dlinfo.restype = ctypes.c_int
      resolved = ctypes.c_long()
      if dlinfo(anchor._handle, _RTLD_DI_LMID, ctypes.byref(resolved)) != 0:
        raise OSError(f"dlinfo failed for {path}")
      previous = _LinkerNamespace(resolved.value, [anchor], {}, ctypes.c_void_p.in_dll(anchor, "environ"))
      _LINKER_NAMESPACES.append(previous)
    # A namespace's libc retains the process environment array from its first load. Refresh it
    # before another initializer calls getenv, since Python may have replaced that array.
    previous.environ.value = ctypes.c_void_p.in_dll(libc, "environ").value
    lib = open_library(previous.id, str(path))
    previous.libraries.append(lib)
    previous.dependencies.update(dependencies)
  return lib

def invalidate_cache(fun: ConcreteFunction) -> None:
  """Make ``fun``'s next build in this process use another cache path.

  Existing source files and libraries stay on disk for processes still building or loading them.
  Loaded handles remain callable, including on Windows, where loaded DLLs cannot be deleted or
  replaced. After ``recompile()`` in one process, a fresh process still reuses the original shared
  library instead of rebuilding it. Scaly writes each library under a temporary name and renames
  it into place only when complete, so no process finds a partial library.

  Scaly does not prune old libraries after process exit. To reclaim the space, delete the cache
  manually while no process is using it. Safe to call before the first build.
  """
  with _artifact_lock:
    _cache_epochs[fun.name] = _cache_epochs.get(fun.name, 0) + 1

def find_c_compiler() -> Compiler | None:
  """Find the C compiler the JIT uses: ``SCALY_CC``, else ``zig cc`` from the ``zig-bin`` package, else ``CC``, else ``cc`` on ``PATH``.

  Returns the compiler command and which of those settings supplied it, or ``None`` if nothing is found.
  """
  if override := os.environ.get("SCALY_CC"):
    command = _command(override)
    return Compiler(command, "SCALY_CC") if command is not None else None
  if (zig := _zig_bin_executable()) is not None:
    return Compiler((zig, "cc"), "zig-bin")
  if override := os.environ.get("CC"):
    command = _command(override)
    return Compiler(command, "CC") if command is not None else None
  found = shutil.which("cc")
  return Compiler((found,), "PATH") if found is not None else None

def cache_root() -> Path:
  """Return ``SCALY_CACHE_DIR``, else ``$XDG_CACHE_HOME/scaly/jit``, else ``~/.cache/scaly/jit``.

  Windows returns an absolute extended path so deeply nested caches work beyond the legacy path limit.
  """
  override = env_path("SCALY_CACHE_DIR")
  xdg = env_path("XDG_CACHE_HOME")
  base = xdg if xdg is not None else Path.home() / ".cache"
  root = override if override is not None else base / "scaly" / "jit"
  if sys.platform == "win32":
    # Repeated benchmark runs nest the cache deeply enough to exceed legacy MAX_PATH.
    path = str(root.resolve())
    if not path.startswith("\\\\?\\"):
      path = "\\\\?\\UNC\\" + path[2:] if path.startswith("\\\\") else "\\\\?\\" + path
    return Path(path)
  return root