fix: support Optim v2 (NLSolversBase v8) + add CI smoketest

Author: MilesCranmerBot

.github/workflows/CI.yml

@@ -110,11 +110,53 @@ jobs:
           version: '1'
       - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
-      - name: Pin Optim v1 and run tests
+      - name: Pin Optim v1 + NLSolversBase v7
         run: |
-          julia --color=yes -e 'import Pkg; Pkg.activate("test"); Pkg.add(Pkg.PackageSpec(name="Optim", version="1")); Pkg.status(["Optim", "NLSolversBase"])'
-          julia --color=yes --threads=auto --check-bounds=yes --depwarn=yes -e 'import Pkg; Pkg.activate("."); Pkg.test()'
+          julia --color=yes -e 'import Pkg; Pkg.add("Coverage")'
+          julia --color=yes -e 'import Pkg; Pkg.activate("."); Pkg.add(Pkg.PackageSpec(name="Optim", version="1")); Pkg.add(Pkg.PackageSpec(name="NLSolversBase", version="7")); Pkg.status(["Optim", "NLSolversBase"])'
         shell: bash
+      - name: Run Optim tests (with coverage)
+        id: run-tests
+        run: |
+          SR_TEST=optim julia --color=yes --threads=auto --check-bounds=yes --depwarn=yes --code-coverage=user -e 'import Coverage; import Pkg; Pkg.activate("."); Pkg.test(coverage=true)'
+          julia --color=yes coverage.jl
+        shell: bash
+      - name: Upload coverage artifact
+        if: steps.run-tests.outcome == 'success'
+        uses: actions/upload-artifact@v6
+        with:
+          name: coverage-optim-v1-${{ runner.os }}-julia-1
+          path: lcov.info
+
+
+  optim_v2_smoketest:
+    name: Optim v2 (NLSolversBase v8) - ubuntu-latest
+    runs-on: ubuntu-latest
+    timeout-minutes: 60
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: '1'
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - name: Pin Optim v2 + NLSolversBase v8
+        run: |
+          julia --color=yes -e 'import Pkg; Pkg.add("Coverage")'
+          julia --color=yes -e 'import Pkg; Pkg.activate("."); Pkg.add(Pkg.PackageSpec(name="Optim", version="2")); Pkg.add(Pkg.PackageSpec(name="NLSolversBase", version="8")); Pkg.status(["Optim", "NLSolversBase"])'
+        shell: bash
+      - name: Run Optim tests (with coverage)
+        id: run-tests
+        run: |
+          SR_TEST=optim julia --color=yes --threads=auto --check-bounds=yes --depwarn=yes --code-coverage=user -e 'import Coverage; import Pkg; Pkg.activate("."); Pkg.test(coverage=true)'
+          julia --color=yes coverage.jl
+        shell: bash
+      - name: Upload coverage artifact
+        if: steps.run-tests.outcome == 'success'
+        uses: actions/upload-artifact@v6
+        with:
+          name: coverage-optim-v2-${{ runner.os }}-julia-1
+          path: lcov.info
 
 
   codecov:
@@ -123,6 +165,8 @@ jobs:
     needs:
       - test
       - additional_tests
+      - optim_v1_smoketest
+      - optim_v2_smoketest
     steps:
       # Codecov uploader expects a git checkout (commit metadata + repo root)
       - uses: actions/checkout@v4

ext/DynamicExpressionsOptimExt.jl

@@ -92,25 +92,33 @@ function wrap_func(
 end
 
 # `NLSolversBase.InplaceObjective` is an internal type whose field layout changed
-# between NLSolversBase versions.
+# between NLSolversBase versions (and therefore between Optim majors).
 #
-# - NLSolversBase v7 (Optim v1.x):  df, fdf, fgh, hv, fghv
-# - NLSolversBase v8 (Optim v2.x):  fdf, fgh, hvp, fghvp, fjvp
+# This extension supports:
+# - Optim v1.x (NLSolversBase v7.x): df, fdf, fgh, hv, fghv
+# - Optim v2.x (NLSolversBase v8.x): fdf, fgh, hvp, fghvp, fjvp
+#
+# We store the fields both as symbols (for runtime layout checks) and as `Val`s
+# (so the wrapper construction is type-stable and can compile-in the field set).
 const _INPLACEOBJECTIVE_SPEC_V8 = (
-    fields=(:fdf, :fgh, :hvp, :fghvp, :fjvp),
-    x_last=(:fdf, :fgh),
-    xv_tail=(:hvp, :fghvp, :fjvp),
+    field_syms = (:fdf, :fgh, :hvp, :fghvp, :fjvp),
+    fields = (Val(:fdf), Val(:fgh), Val(:hvp), Val(:fghvp), Val(:fjvp)),
+    x_last = (Val(:fdf), Val(:fgh)),
+    xv_tail = (Val(:hvp), Val(:fghvp), Val(:fjvp)),
 )
 const _INPLACEOBJECTIVE_SPEC_V7 = (
-    fields=(:df, :fdf, :fgh, :hv, :fghv), x_last=(:df, :fdf, :fgh), xv_tail=(:hv, :fghv)
+    field_syms = (:df, :fdf, :fgh, :hv, :fghv),
+    fields = (Val(:df), Val(:fdf), Val(:fgh), Val(:hv), Val(:fghv)),
+    x_last = (Val(:df), Val(:fdf), Val(:fgh)),
+    xv_tail = (Val(:hv), Val(:fghv)),
 )
 
 @inline function _wrap_inplaceobjective_field(
-    ::Val{field}, f::NLSolversBase.InplaceObjective, tree::N, refs, spec
+    v_field::Val{field}, f::NLSolversBase.InplaceObjective, tree::N, refs, spec
 ) where {field,N<:Union{AbstractExpressionNode,AbstractExpression}}
-    if field in spec.x_last
+    if v_field in spec.x_last
         return _wrap_objective_x_last(getfield(f, field), tree, refs)
-    elseif field in spec.xv_tail
+    elseif v_field in spec.xv_tail
         return _wrap_objective_xv_tail(getfield(f, field), tree, refs)
     else
         throw(
@@ -125,8 +133,8 @@ end
 @inline function _wrap_inplaceobjective(
     f::NLSolversBase.InplaceObjective, tree::N, refs, spec
 ) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
-    wrapped = map(spec.fields) do field
-        _wrap_inplaceobjective_field(Val(field), f, tree, refs, spec)
+    wrapped = map(spec.fields) do v_field
+        _wrap_inplaceobjective_field(v_field, f, tree, refs, spec)
     end
     return NLSolversBase.InplaceObjective(wrapped...)
 end
@@ -141,14 +149,15 @@ function wrap_func(
     # We use `@static` branching so that only the relevant layout for the *installed*
     # NLSolversBase version is compiled/instrumented.
     @static if fieldnames(NLSolversBase.InplaceObjective) ==
-        _INPLACEOBJECTIVE_SPEC_V8.fields
+        _INPLACEOBJECTIVE_SPEC_V8.field_syms
         # NLSolversBase v8 / Optim v2
         return _wrap_inplaceobjective(f, tree, refs, _INPLACEOBJECTIVE_SPEC_V8)
-    elseif fieldnames(NLSolversBase.InplaceObjective) == _INPLACEOBJECTIVE_SPEC_V7.fields
+    elseif fieldnames(NLSolversBase.InplaceObjective) == _INPLACEOBJECTIVE_SPEC_V7.field_syms
         # NLSolversBase v7 / Optim v1
         return _wrap_inplaceobjective(f, tree, refs, _INPLACEOBJECTIVE_SPEC_V7)
         # (Optim < 1 is no longer supported.)
     else
+        # LCOV_EXCL_START
         fields = fieldnames(NLSolversBase.InplaceObjective)
         throw(
             ArgumentError(
@@ -157,6 +166,7 @@ function wrap_func(
                 "Please open an issue at github.com/SymbolicML/DynamicExpressions.jl with your versions.",
             ),
         )
+        # LCOV_EXCL_END
     end
 end
 

test/runtests.jl

@@ -1,21 +1,32 @@
 using SafeTestsets
 using TestItemRunner
 
-# Check if SR_ENZYME_TEST is set in env
-test_name = split(get(ENV, "SR_TEST", "main"), ",")
+# Control which test groups run.
+#
+# Accepts a comma-separated list in SR_TEST (default: "main").
+#
+# - "main": full test suite (testitems)
+# - "optim": Optim-specific testitems only
+# - "narity": n-ary testitems only
+# - "jet": JET analysis
+# - "enzyme": Enzyme tests
 
-unknown_tests = filter(Base.Fix2(∉, ["enzyme", "jet", "main", "narity"]), test_name)
+test_names = split(get(ENV, "SR_TEST", "main"), ",")
+
+allowed = ["enzyme", "jet", "main", "narity", "optim"]
+unknown_tests = filter(Base.Fix2(∉, allowed), test_names)
 
 if !isempty(unknown_tests)
     error("Unknown test names: $unknown_tests")
 end
 
-if "enzyme" in test_name
+if "enzyme" in test_names
     @safetestset "Test enzyme derivatives" begin
         include("test_enzyme.jl")
     end
 end
-if "jet" in test_name
+
+if "jet" in test_names
     @safetestset "JET" begin
         using Preferences
         set_preferences!(
@@ -61,7 +72,25 @@ if "jet" in test_name
         end
     end
 end
-if "main" in test_name
-    include("unittest.jl")
+
+# Collect testitem definitions to load, then run them once.
+testitem_files = String[]
+
+if "main" in test_names
+    push!(testitem_files, "unittest.jl")
+else
+    if "optim" in test_names
+        push!(testitem_files, "test_optim.jl")
+    end
+    if "narity" in test_names
+        push!(testitem_files, "test_n_arity_nodes.jl")
+    end
+end
+
+for f in testitem_files
+    include(f)
+end
+
+if !isempty(testitem_files)
     @run_package_tests
 end

test/test_optim.jl

@@ -222,12 +222,14 @@ end
     dummy = (args...) -> nothing
     obj = Optim.NLSolversBase.InplaceObjective((dummy for _ in fields)...)
 
-    spec = if fields == ext._INPLACEOBJECTIVE_SPEC_V8.fields
+    spec = if fields == ext._INPLACEOBJECTIVE_SPEC_V8.field_syms
         ext._INPLACEOBJECTIVE_SPEC_V8
-    elseif fields == ext._INPLACEOBJECTIVE_SPEC_V7.fields
+    elseif fields == ext._INPLACEOBJECTIVE_SPEC_V7.field_syms
         ext._INPLACEOBJECTIVE_SPEC_V7
     else
-        ext._INPLACEOBJECTIVE_SPEC_OLD
+        # Should be unreachable for supported Optim majors, but pick one so the
+        # test still compiles if a future layout appears.
+        ext._INPLACEOBJECTIVE_SPEC_V8
     end
 
     @test_throws ArgumentError ext._wrap_inplaceobjective_field(