diff --git a/src/mcmc/execution/chain_runner.cpp b/src/mcmc/execution/chain_runner.cpp
index fef5ade8..fac671c6 100644
--- a/src/mcmc/execution/chain_runner.cpp
+++ b/src/mcmc/execution/chain_runner.cpp
@@ -6,6 +6,24 @@
 #include "mcmc/samplers/metropolis_sampler.h"
 
 
+namespace {
+
+// Store this sample's NUTS per-iteration diagnostics, if the chain is collecting
+// them and the step produced a NUTSDiagnostics object. Keeps the dynamic_cast
+// and the wide store_nuts_diagnostics call out of the main sampling loop.
+void store_nuts_diagnostics_if_present(ChainResult& chain_result, int sample_index,
+                                       const SamplerBase& sampler, const StepResult& result) {
+    if (!chain_result.has_nuts_diagnostics || !sampler.has_nuts_diagnostics()) return;
+    auto* diag = dynamic_cast<NUTSDiagnostics*>(result.diagnostics.get());
+    if (diag) {
+        chain_result.store_nuts_diagnostics(sample_index, diag->tree_depth, diag->divergent,
+                                            diag->non_reversible, diag->energy, diag->accept_prob);
+    }
+}
+
+}  // namespace
+
+
 SamplerSpec resolve_sampler_spec(const std::string& sampler_type) {
     if (sampler_type == "nuts") {
         return SamplerSpec{SamplerKind::NUTS, /*learn_sd=*/true, /*nuts_diag=*/true, /*am_diag=*/false};
@@ -88,12 +106,7 @@ void run_mcmc_chain(
         if (schedule.sampling(iter)) {
             int sample_index = iter - config.no_warmup;
 
-            if (chain_result.has_nuts_diagnostics && sampler->has_nuts_diagnostics()) {
-                auto* diag = dynamic_cast<NUTSDiagnostics*>(result.diagnostics.get());
-                if (diag) {
-                    chain_result.store_nuts_diagnostics(sample_index, diag->tree_depth, diag->divergent, diag->non_reversible, diag->energy, diag->accept_prob);
-                }
-            }
+            store_nuts_diagnostics_if_present(chain_result, sample_index, *sampler, result);
 
             if (chain_result.has_am_diagnostics) {
                 chain_result.store_am_diagnostics(sample_index, result.accept_prob);