compiletime-autodiff.hpp

#pragma once

#define F_INLINE __forceinline constexpr

template <int VarIndex>
struct Var
{
    template<int OtherVarIndex>
    static F_INLINE float eval(const Var<OtherVarIndex> & var, const float v)
    {
        if constexpr (OtherVarIndex == VarIndex) { return v; }
        else { static_assert(!"unsupported"); }
    }

    template<int OtherVarIndex>
    static F_INLINE float diffEval(const Var<OtherVarIndex> & var, const float v)
    {
        if constexpr (OtherVarIndex == VarIndex) { return 1.0f; }
        else { return 0.0f; }
    }
};

template <unsigned int WholeNumber, unsigned int Frac, unsigned int FracDenomTenth>
struct ConstantV2
{
    template<int OtherVarIndex>
    static F_INLINE float eval(const Var<OtherVarIndex> & var, const float v)
    {
        constexpr float fConstant = float(WholeNumber) + float(Frac) / float(FracDenomTenth);
        return fConstant;
    }

    template<int OtherVarIndex>
    static F_INLINE float diffEval(const Var<OtherVarIndex> & var, const float v)
    {
        return 0.0f;
    }
};

template<int W, int F, int FD>
constexpr auto ProcessFraction()
{
    return ConstantV2<W, F, FD>();
}

template<int W, int F, int FD, char D, char... Ds>
constexpr auto ProcessFraction()
{
    return ProcessFraction<W, F * 10 + D - '0', FD * 10, Ds...>();
}

template<int V>
constexpr auto ProcessWholeNumber()
{
    return ConstantV2<V, 0, 1>();
}

template<int V, char D, char... Ds>
constexpr auto ProcessWholeNumber()
{
    if constexpr (D == '.')
    {
        return ProcessFraction<V, 0, 1, Ds...>();
    }
    else
    {
        return ProcessWholeNumber<V * 10 + D - '0', Ds...>();
    }
}

template<char... Ds>
constexpr auto operator ""_f()
{
    return ProcessWholeNumber<0, Ds...>();
}

struct CTDiffOp
{
};

template <typename L, typename R>
struct OpAdd
{
    template<int OtherVarIndex>
    static F_INLINE float eval(const Var<OtherVarIndex> & var, const float v)
    {
        return L::eval(var, v) + R::eval(var, v);
    }

    template<int OtherVarIndex>
    static F_INLINE float diffEval(const Var<OtherVarIndex> & var, const float v)
    {
        return L::diffEval(var, v) + R::diffEval(var, v);
    }
};

template <typename L, typename R>
struct OpMul
{
    template<int OtherVarIndex>
    static F_INLINE float eval(const Var<OtherVarIndex> & var, const float v)
    {
        return L::eval(var, v) * R::eval(var, v);
    }

    template<int OtherVarIndex>
    static F_INLINE float diffEval(const Var<OtherVarIndex> & var, const float v)
    {
        float l = L::eval(var, v);
        float dl = L::diffEval(var, v);
        float r = R::eval(var, v);
        float dr = R::diffEval(var, v);
        return l * dr + r * dl;
    }
};

#define DECL_OPERATOR_OVERRIDE(NAME, OP)\
template <typename L, typename R>\
F_INLINE NAME<L, R> operator##OP##(const L & lhs, const R & rhs)\
{\
    return NAME<L, R>();\
}
DECL_OPERATOR_OVERRIDE(OpAdd, +)
DECL_OPERATOR_OVERRIDE(OpMul, *)
#undef DECL_OPERATOR_OVERRIDE