From af648c0330d460a220cb1b06e0712d7cf16b0613 Mon Sep 17 00:00:00 2001
From: Graham McIntire <graham@mcintire.me>
Date: Wed, 5 Aug 2026 10:05:41 -0500
Subject: [PATCH] fix: save full EXLA build cache (object files + libexla.so)
 in CI image

EXLA's cached_make copies libexla.so but Make still compiles .o files
when they're missing from deps/exla/cache/ (they don't survive
mix deps.get). Save the pre-built cache/ directory at /opt/exla-cache/
in the CI image and restore it after mix deps.get in the test step.
With both .o files and libexla.so present, Make skips all C++ compilation.
---
 .forgejo/workflows/build.yaml |  1 +
 Dockerfile.ci                 | 47 +++++++++++++++--------------------
 2 files changed, 21 insertions(+), 27 deletions(-)

diff --git a/.forgejo/workflows/build.yaml b/.forgejo/workflows/build.yaml
index 105cbe11..94e22aa0 100644
--- a/.forgejo/workflows/build.yaml
+++ b/.forgejo/workflows/build.yaml
@@ -141,6 +141,7 @@ jobs:
               sh -euc '\
               mkdir -p /app && cd /app && tar xf - && \
               mix deps.get && \
+              cp -a /opt/exla-cache/cache /app/deps/exla/cache && \
               mix ecto.create --quiet && \
               mix ecto.migrate --quiet && \
               mix test'
diff --git a/Dockerfile.ci b/Dockerfile.ci
index 424e8996..baee63e1 100644
--- a/Dockerfile.ci
+++ b/Dockerfile.ci
@@ -3,67 +3,60 @@
 # Pre-built CI test image. Extends the hexpm/elixir image with:
 #   * build-essential + git + curl + ca-certificates + cdo
 #   * hex + rebar
-#   * Precompiled EXLA NIF cached at ~/.cache/xla/
+#   * Precompiled EXLA NIF + object files at /opt/exla-cache/
 #
-# When mix test runs with this image, EXLA's cached_make finds the
-# precompiled libexla.so and skips the OOM-prone g++ C++ compilation
-# entirely. Rebuild this image whenever mix.exs changes elixir/exla/xla
-# versions or c_src/ files change (the cache key auto-invalidates).
+# EXLA's cached_make copies libexla.so to deps/exla/cache/, but Make
+# still tries to compile .o files when they're missing (they don't
+# survive mix deps.get). We save the full cache/ directory and restore
+# it in the test step before mix test runs.
 #
-# Built by .forgejo/workflows/build-ci-image.yaml.
+# Rebuild whenever mix.exs changes elixir/exla/xla versions or c_src/
+# files change (the cache key auto-invalidates).
+#
+# Built by .forgejo/workflows/build.yaml (build-ci-image job).
 
 ARG ELIXIR_IMAGE="docker.io/hexpm/elixir:1.20.1-erlang-29.0.2-debian-trixie-20260518-slim"
 
 FROM ${ELIXIR_IMAGE} AS builder
 
-# Install build tools and runtime test deps (cdo is needed by tests).
 RUN export DEBIAN_FRONTEND=noninteractive && \
     apt-get update -qq && \
     apt-get install -y -qq git curl ca-certificates build-essential cdo && \
     rm -rf /var/lib/apt/lists/*
 
-# Install hex + rebar
 RUN mix local.hex --force && mix local.rebar --force
 
-# Precompile EXLA to populate ~/.cache/xla/ with both the downloaded
-# XLA archive and the compiled libexla.so. We only need mix.exs +
-# mix.lock + vendor + minimal config — no full project checkout.
 WORKDIR /tmp/exla_warm
 
-# Copy only what's needed for deps.get + deps.compile exla
 COPY mix.exs mix.lock /tmp/exla_warm/
 COPY vendor /tmp/exla_warm/vendor
 
-# Minimal config — nx needs :default_backend to exist
 RUN mkdir -p /tmp/exla_warm/config && \
     printf 'import Config\nconfig :nx, :default_backend, EXLA.Backend\n' > /tmp/exla_warm/config/config.exs
 
-# Fetch and compile all deps. EXLA's cached_make compiler calls
-# Mix.Tasks.Compile.ElixirMake at build time, which requires elixir_make
-# to be compiled first. mix deps.compile without arguments handles
-# dependency ordering automatically. The cached_make step downloads the
-# XLA archive, compiles the C++ NIF, and caches libexla.so at
-# ~/.cache/xla/exla/{cache_key}/libexla.so.
+# Compile all deps. EXLA's cached_make downloads the XLA archive, compiles
+# the C++ NIF, and caches libexla.so at ~/.cache/xla/exla/{key}/.
 RUN mix deps.get && mix deps.compile
 
-# Clean up the temp project — the cache in ~/.cache/xla/ persists in
-# the image and is what matters.
-RUN rm -rf /tmp/exla_warm
+# Save the full build cache (libexla.so + .o files + xla_extension lib)
+# so Make finds all targets already built when restored in the test step.
+RUN mkdir -p /opt/exla-cache && \
+    cp -a deps/exla/cache /opt/exla-cache/cache
 
 FROM ${ELIXIR_IMAGE} AS final
 
-# Install runtime test deps
 RUN export DEBIAN_FRONTEND=noninteractive && \
     apt-get update -qq && \
     apt-get install -y -qq git curl ca-certificates build-essential cdo && \
     rm -rf /var/lib/apt/lists/*
 
-# Install hex + rebar
 RUN mix local.hex --force && mix local.rebar --force
 
-# Copy the precompiled EXLA cache from the builder stage. This is the
-# ~/.cache/xla/ directory containing the XLA archive download and the
-# compiled libexla.so keyed by exact Elixir/ERTS/exla/xla/c_src versions.
+# XLA archive download + cached libexla.so (cached_make fast path).
 COPY --from=builder /root/.cache/xla /root/.cache/xla
 
+# Pre-built EXLA object files + libexla.so + xla_extension/ lib.
+# Restored in the test step: cp -a /opt/exla-cache/cache /app/deps/exla/cache
+COPY --from=builder /opt/exla-cache /opt/exla-cache
+
 WORKDIR /app