From f5eb0da701b31b286bebd4ba9abc5ba9d543701a Mon Sep 17 00:00:00 2001
From: Tyler Veness <calcmogul@gmail.com>
Date: Sat, 13 Apr 2024 11:41:28 -0700
Subject: [PATCH] Use function_ref where possible (#501)

---
 benchmarks/scalability/Util.hpp                    |  9 +++++----
 include/sleipnir/control/OCPSolver.hpp             | 10 +++++-----
 .../sleipnir/optimization/solver/InteriorPoint.hpp |  4 ++--
 src/optimization/solver/InteriorPoint.cpp          | 14 +++++++-------
 .../solver/util/FeasibilityRestoration.hpp         |  4 ++--
 5 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/benchmarks/scalability/Util.hpp b/benchmarks/scalability/Util.hpp
index 3df9bdb8..4f899cf8 100644
--- a/benchmarks/scalability/Util.hpp
+++ b/benchmarks/scalability/Util.hpp
@@ -5,7 +5,6 @@
 #include <chrono>
 #include <concepts>
 #include <fstream>
-#include <functional>
 #include <span>
 #include <string>
 #include <string_view>
@@ -13,6 +12,7 @@
 #include <casadi/casadi.hpp>
 #include <fmt/core.h>
 #include <sleipnir/optimization/OptimizationProblem.hpp>
+#include <sleipnir/util/FunctionRef.hpp>
 
 /**
  * Converts std::chrono::duration to a number of milliseconds rounded to three
@@ -38,8 +38,9 @@ constexpr double ToMilliseconds(
  *   solves it.
  */
 template <typename Problem>
-void RunBenchmark(std::ofstream& results, std::function<Problem()> setup,
-                  std::function<void(Problem&)> solve) {
+void RunBenchmark(std::ofstream& results,
+                  sleipnir::function_ref<Problem()> setup,
+                  sleipnir::function_ref<void(Problem&)> solve) {
   // Record setup time
   auto setupStartTime = std::chrono::system_clock::now();
   auto problem = setup();
@@ -85,7 +86,7 @@ template <typename Problem>
 int RunBenchmarksAndLog(
     std::string_view filename, bool diagnostics,
     std::chrono::duration<double> T, std::span<int> sampleSizesToTest,
-    std::function<Problem(std::chrono::duration<double>, int)> setup) {
+    sleipnir::function_ref<Problem(std::chrono::duration<double>, int)> setup) {
   std::ofstream results{std::string{filename}};
   if (!results.is_open()) {
     return 1;
diff --git a/include/sleipnir/control/OCPSolver.hpp b/include/sleipnir/control/OCPSolver.hpp
index 9eaab590..ef0243f4 100644
--- a/include/sleipnir/control/OCPSolver.hpp
+++ b/include/sleipnir/control/OCPSolver.hpp
@@ -5,13 +5,13 @@
 #include <stdint.h>
 
 #include <chrono>
-#include <functional>
 #include <utility>
 
 #include "sleipnir/autodiff/VariableMatrix.hpp"
 #include "sleipnir/optimization/OptimizationProblem.hpp"
 #include "sleipnir/util/Assert.hpp"
 #include "sleipnir/util/Concepts.hpp"
+#include "sleipnir/util/FunctionRef.hpp"
 #include "sleipnir/util/SymbolExports.hpp"
 
 namespace sleipnir {
@@ -25,8 +25,8 @@ namespace sleipnir {
  * - State transition: xₖ₊₁ = f(t, xₖ, uₖ, dt)
  */
 using DynamicsFunction =
-    std::function<VariableMatrix(const Variable&, const VariableMatrix&,
-                                 const VariableMatrix&, const Variable&)>;
+    function_ref<VariableMatrix(const Variable&, const VariableMatrix&,
+                                const VariableMatrix&, const Variable&)>;
 
 /**
  * Performs 4th order Runge-Kutta integration of dx/dt = f(t, x, u) for dt.
@@ -212,8 +212,8 @@ class SLEIPNIR_DLLEXPORT OCPSolver : public OptimizationProblem {
    *   vector, u is the input vector, and dt is the timestep duration.
    */
   void ForEachStep(
-      const std::function<void(const Variable&, const VariableMatrix&,
-                               const VariableMatrix&, const Variable&)>&
+      const function_ref<void(const Variable&, const VariableMatrix&,
+                              const VariableMatrix&, const Variable&)>
           callback) {
     Variable time = 0.0;
 
diff --git a/include/sleipnir/optimization/solver/InteriorPoint.hpp b/include/sleipnir/optimization/solver/InteriorPoint.hpp
index 72d29293..abc9f778 100644
--- a/include/sleipnir/optimization/solver/InteriorPoint.hpp
+++ b/include/sleipnir/optimization/solver/InteriorPoint.hpp
@@ -2,7 +2,6 @@
 
 #pragma once
 
-#include <functional>
 #include <span>
 
 #include <Eigen/Core>
@@ -11,6 +10,7 @@
 #include "sleipnir/optimization/SolverConfig.hpp"
 #include "sleipnir/optimization/SolverIterationInfo.hpp"
 #include "sleipnir/optimization/SolverStatus.hpp"
+#include "sleipnir/util/FunctionRef.hpp"
 #include "sleipnir/util/SymbolExports.hpp"
 
 namespace sleipnir {
@@ -48,7 +48,7 @@ SLEIPNIR_DLLEXPORT void InteriorPoint(
     std::span<Variable> decisionVariables,
     std::span<Variable> equalityConstraints,
     std::span<Variable> inequalityConstraints, Variable& f,
-    const std::function<bool(const SolverIterationInfo&)>& callback,
+    function_ref<bool(const SolverIterationInfo&)> callback,
     const SolverConfig& config, bool feasibilityRestoration, Eigen::VectorXd& x,
     Eigen::VectorXd& s, SolverStatus* status);
 
diff --git a/src/optimization/solver/InteriorPoint.cpp b/src/optimization/solver/InteriorPoint.cpp
index 122eb285..596988bf 100644
--- a/src/optimization/solver/InteriorPoint.cpp
+++ b/src/optimization/solver/InteriorPoint.cpp
@@ -34,13 +34,13 @@
 
 namespace sleipnir {
 
-void InteriorPoint(
-    std::span<Variable> decisionVariables,
-    std::span<Variable> equalityConstraints,
-    std::span<Variable> inequalityConstraints, Variable& f,
-    const std::function<bool(const SolverIterationInfo&)>& callback,
-    const SolverConfig& config, bool feasibilityRestoration, Eigen::VectorXd& x,
-    Eigen::VectorXd& s, SolverStatus* status) {
+void InteriorPoint(std::span<Variable> decisionVariables,
+                   std::span<Variable> equalityConstraints,
+                   std::span<Variable> inequalityConstraints, Variable& f,
+                   function_ref<bool(const SolverIterationInfo&)> callback,
+                   const SolverConfig& config, bool feasibilityRestoration,
+                   Eigen::VectorXd& x, Eigen::VectorXd& s,
+                   SolverStatus* status) {
   const auto solveStartTime = std::chrono::system_clock::now();
 
   // Map decision variables and constraints to VariableMatrices for Lagrangian
diff --git a/src/optimization/solver/util/FeasibilityRestoration.hpp b/src/optimization/solver/util/FeasibilityRestoration.hpp
index 16827c33..15f00881 100644
--- a/src/optimization/solver/util/FeasibilityRestoration.hpp
+++ b/src/optimization/solver/util/FeasibilityRestoration.hpp
@@ -4,7 +4,6 @@
 
 #include <algorithm>
 #include <cmath>
-#include <functional>
 #include <iterator>
 #include <span>
 #include <vector>
@@ -17,6 +16,7 @@
 #include "sleipnir/optimization/SolverIterationInfo.hpp"
 #include "sleipnir/optimization/SolverStatus.hpp"
 #include "sleipnir/optimization/solver/InteriorPoint.hpp"
+#include "sleipnir/util/FunctionRef.hpp"
 
 namespace sleipnir {
 
@@ -41,7 +41,7 @@ inline void FeasibilityRestoration(
     std::span<Variable> decisionVariables,
     std::span<Variable> equalityConstraints,
     std::span<Variable> inequalityConstraints, Variable& f, double μ,
-    const std::function<bool(const SolverIterationInfo&)>& callback,
+    function_ref<bool(const SolverIterationInfo&)> callback,
     const SolverConfig& config, Eigen::VectorXd& x, Eigen::VectorXd& s,
     SolverStatus* status) {
   // Feasibility restoration