Poor man's JIT using nested lambdas

While answering this code review question, I came up with a way to convert an equation given at runtime to a std::function<double(double)> that would evaluate that equation at a given point. Assuming the equation was already split into tokens and converted to postfix notation, this is done by evaluating those tokens like you would do in a normal implementation of a RPN calculator by maintaining a stack of intermediate results, but instead of just storing the values, the stack stores lambda expressions that can calculate those values. At the end of the evaluation, the stack should thus contain a single lambda expression that implements the desired equation. Here is a simplified version of the parser than only supports a few operations:

#include <cmath>
#include <functional>
#include <iostream>
#include <numbers>
#include <stack>
#include <string>
#include <string_view>

std::function<double(double)> build_function(const std::vector<std::string_view> &rpn_tokens) {
std::stack<std::function<double(double)>> subexpressions;

for (const auto &token: rpn_tokens) {
if (token.empty())
throw std::runtime_error("empty token");

if (token == "x") { // Variable
subexpressions.push([](double x){
return x;
});
} else if (isdigit(token)) { // Literal number
double value = std::stof(token.data());

subexpressions.push([=](double /* ignored */){
return value;
});
} else if (token == "sin") { // Example unary operator
if (subexpressions.size() < 1) {
throw std::runtime_error("invalid expression");
}

auto operand = subexpressions.top();
subexpressions.pop();

subexpressions.push([=](double x){
return std::sin(operand(x));
});
} else if (token == "+") { // Example binary operator
if (subexpressions.size() < 2) {
throw std::runtime_error("invalid expression");
}

auto right_operand = subexpressions.top();
subexpressions.pop();
auto left_operand = subexpressions.top();
subexpressions.pop();

subexpressions.push([=](double x){
return left_operand(x) + right_operand(x);
});
} else {
throw std::runtime_error("invalid token");
}
}

if (subexpressions.size() != 1) {
throw std::runtime_error("invalid expression");
}

return subexpressions.top();
}

int main() {
auto function = build_function({"1", "x", "2", "+", "sin", "+"}); // 1 + sin(2 + x)

for (double x = 0; x < 2 * std::numbers::pi; x += 0.1) {
std::cout << x << ' ' << function(x) << '\n';
}
}

I call this a poor man's JIT because while it looks like we get a function object that is seemingly created at runtime, it basically is a bunch of precompiled functions (one for each lambda body in the above code) strung together by the lambda captures. So there is quite a lot more overhead when calling the above function() than if one would write the following:

auto function = [](double x){
return 1 + sin(2 + x);
}

But it should still be much faster than parsing the bunch of tokens that make up the equation each time you want to evaluate it.

Some questions:

• Is there a better way to do this?
• Is this technique already implemented in some library?
• The argument to resulting function has to be passed down to most of the lambdas (the exception being the one returning literal numbers). Is there a better way to do this?
• How should one handle building a function that takes multiple arguments? Can we make build_function a template somehow that can return a std::function with a variable number of arguments?

Performance measurements on an AMD Ryzen 3900X, code compiled with GCC 10.2.1 with -O2, after a warm-up of 60 million invocations, averaged over another 60 million invocations of function():

• A simple lambda: 10 ns per evaluation
• The above code: 21 ns per evaluation
• Deduplicator's original answer: 77 ns per evaluation
• Adding stack.reserve(): 39 ns per evaluation
• Making stack static: 26 ns per evaluation
• Deduplicator's version using separate bytecode and data stacks: 19 ns per evaluation
• Hardcoding using a fixed small stack: 17 ns per evaluation
• You have reinvented Forth, I believe! :) Apr 3 '21 at 15:40
• @Edward You mean threaded code? Apr 3 '21 at 17:47
• There are two basic aspects to Forth. One is to have a series of RPN-style instructions that operate on a stack. The second is to have threaded code. There are, of course, other things in Forth, but these are the most fundamental and unique. Apr 3 '21 at 18:01
• @G.Sliepen This is probably obvious, but if you're going to try to benchmark this, make sure your tokens are specified as command-line arguments, rather than as string literals in main(). Otherwise the compiler can just substitute the hell out of everything.
– Will
Apr 3 '21 at 23:49
• @CarstenS Sure, an answer that proposes an alternative way to get the same result in a better way (whether that is better looking code, better performance, less memory usage, and so on) is perfectly fine for this site. Apr 6 '21 at 13:56

Is there a better way to do this?

Well, you have dynamic dispatch and yet another scattered island of memory for every single token. That is massively inefficient.

A simple way to more efficiency is writing your own mini-VM for executing expressions. You will still only have to parse once, but now all the memory is in a single compact chunk which will be linearly used, and the compiler can see all the code.

Switch-statements with a single compact range of valid values are simplicity itself, and you avoid the argument shuffling, register saving, and all the other overhead of dynamically calling arbitrary functions.

The next step for efficiency would be compiling to native code instead.

Also, putting all the arguments in an array and then using an instruction with operand or consecutive simple instructions to get them should be a simple modification.

I also added extensive X-Macros to reduce repetition and avoid defining all applicable transformations all over the place.

#include <iostream>
#include <numbers>

#include <charconv>
#include <cmath>
#include <exception>
#include <memory>
#include <span>
#include <string_view>
#include <string>
#include <vector>

#define XX() \
/* token, tag, args, results, code */ \
X("x", x, 0, 1, *stack = x) \
X("sin", sin, 1, 1, *stack = std::sin(*stack)) \
X("+", plus, 2, 1, *stack += stack[-1])

enum class instruction : char {
end,
push,
#define X(a, b, c, d, e) b,
XX()
#undef X
};

auto build_function(std::span<const std::string_view> rpn_tokens) {
std::vector<instruction> code;
std::vector<double> data;
double temp;
unsigned count = 0, max_count = 0;
auto process = [&](instruction id, unsigned args, unsigned results) {
if (count < args)
throw std::runtime_error("invalid expression: underflow");
code.push_back(id);
count += results - args;
max_count = std::max(max_count, count);
};
for (auto token : rpn_tokens) {
#define X(a, b, c, d, e) \
if (token == a) \
process(instruction::b, c, d); \
else
XX()
#undef X
//      if (auto [p, ec] = std::from_chars(token.begin(), token.end(), temp); !ec && p = token.end()) {
//          process(instruction::push, 0, 1);
//          data.push_back(temp);
//      } else
//          throw std::runtime_error("invalid token");
try {
temp = std::stod(std::string(token));
data.push_back(temp);
process(instruction::push, 0, 1);
} catch(...) {
std::throw_with_nested(std::runtime_error("invalid token"));
}
}
process(instruction::end, 1, 0);
if (count)
throw std::runtime_error("invalid expression: overflow");
code.shrink_to_fit();
data.shrink_to_fit();
return [code, data, max_count](double x) {
constexpr auto small_stack = 128;

auto core = [](auto code, auto data, auto stack, auto x){
for (;;) {
switch(*code++) {
case instruction::end:
return *stack;
case instruction::push:
*--stack = *data++;
break;

#define X(a, b, c, d, e) \
case instruction::b: \
stack += c - d; \
e; \
break;
XX()
#undef X
default:
throw std::runtime_error("unexpected");
}
}
};
auto small = [&]{
double stack[small_stack];
return core(code.cbegin(), data.cbegin(), std::end(stack), x);
};
auto big = [&]{
auto stack = std::make_unique<double[]>(max_count);
return core(code.cbegin(), data.cbegin(), &stack + max_count, x);
};
return max_count <= small_stack ? small() : big();
};
}
#undef XX

int main() {
std::string_view rpn[] = {"1", "x", "2", "+", "sin", "+"}; // 1 + sin(2 + x)
auto function = build_function(rpn);

for (double x = 0; x < 2 * std::numbers::pi; x += 0.1) {
std::cout << x << ' ' << function(x) << '\n';
}
}
• I benchmarked your code against mine, and I was surprised my code was about 3.7 times faster! Of course, those numbers might be meaningless, as this code is just a very simplified expression parser, I should check it with more functions added and trying to evaluate more complex equations. Apr 3 '21 at 22:21
• Did you benchmark parsing, invoking, or both? Anyway, seems there is no substitute for measuring. I think a big part of the problem is the dynamic allocation while executing. Must experiment. Apr 3 '21 at 23:27
• It seems the stack is indeed a big part of the problem. With a static stack it goes down to 26 ns, very close to my solution. I guess a stack should not be necessary; the order and number of the stack operations is always the same for each invokation of the returned function, my solution kind of "hardcodes" that, but I don't see how to do that with yours while keeping bin as compact as it is... Apr 4 '21 at 8:04
• Threaded code can be surprisingly fast and surprisingly dense at the same time. There are many tales of embedded systems programmers who discovered that a FORTH interpreter plus a FORTH program are both faster and more compact than a hand-optimized C program for the same problem. And once you have to solve more than just one single problem, the density advantage only increases because you only need one copy of the FORTH interpreter. Apr 4 '21 at 9:39
• @PasserBy X macros are kind of neat. If you know them, and don't get too overboard. But yes, the preprocessor deserves to be treated as its own completely different language, which has to be understood. And it will only really start to pay off after adding a few more commands. Apr 4 '21 at 15:38

I will not try to investigate all different kinds of alternatives, but just try to illuminate what your implementation does.

Your code is definitely fine for a quick implementation of something. I have seen code similar to this in production.

What you are doing is essentially to create a tree representation of the expression but in such a way that it only allows to evaluate it and not to inspect the tree in any other way. If you create a std::function from a lambda, it needs to contain the lambda (i.e. its bound variables), and in general this will have to be in memory allocated on the heap. A good std::function implementation will likely have some small object optimization, so that the lambda can be stored directly in the std::function if it is small enough, but you know that you definitely exceed this size of your lambda has captured a std::function by value, because a std::function cannot possibly fit directly into a std::function. So what you end up with is not dissimilar to what you would get with

class Expr {
public:
virtual double evaluate(double x_val) const = 0;

virtual ~Expr() = default;
};
using ExprPtr = std::unique_ptr<Expr>;

and then for example

class PlusExpr final : public Expr {
const ExprPtr<Expr> lhs;
const ExprPtr<Expr> rhs;
public:
PlusExpr(ExprPtr lhs, ExprPtr rhs)
:lhs(std::move(lhs)), rhs(std::move(rhs)) {}
double evaluate(double x_val) const override
{
return lhs->evaluate(x_val) + rhs->evaluate(x_val);
}
};

Note that I had to make some decisions here. The Expr class cannot be copied (that would have been extra work), so it needs to be used together with pointers. On the other hand, this yields some clarity. Do we really want to copy these expression trees? With the above setup I would have to use something like

std::vector<ExprPtr> subexpressions;
...
} else if (token == "+") { // Example binary operator
if (subexpressions.size() < 2) {
throw std::runtime_error("invalid expression");
}

auto right_operand = std::move(subexpressions.back());
subexpressions.pop_back();
auto left_operand = std::move(subexpressions.back());
subexpressions.pop_back();

subexpressions.emplace_back(
new PlusExpr(std::move(left_operand), std::move(right_operand)));
} else ...

Now all of this is a bit more cumbersome, so for a quick solution std::functions are definitely nice. This points us to something else however. Your code copied the std::functions of the subexpressions at that point (when you captured them by value) and the cost was somehow hidden, because std::functions look so harmless. Indeed if you try your code on input like

0 1 + 1 + 1 + 1 + ... 1 +

you should notice that the cost of build_function is quadratic in the size of the input. You can of course fix that by also moving the subexpressions into the lambda instead of copying them.

The cost of evaluation an expression should be very similar for both approaches. There are subtle differences, but I could not tell which will come out in front.

Another advantage of the class based approach is that it can be extended to support more functionality than just evaluation. You could add a function to print the expression, or one could have functions modifying the expression, say for having an optimization pass that folded constants like another answer mentioned.

In summary, your approach has its place, but you should be aware that you are really constructing an expression tree in disguise and beware of costs hidden by the convenience of std::functions.

• I was aware that my implementation basically builds a tree of expressions at runtime, but thanks for this excellent explaination, as well as pointing out the problem with the quadratic cost of build_function() due to copying std::functions. Apr 5 '21 at 12:43
• @G.Sliepen, I could probably have been briefer then ;) Well, it may be useful to someone else this way. Apr 5 '21 at 13:02

Technically isdigit(token):

• Needs #include <cctype>.
• Should be std::isdigit(static_cast<unsigned char>(token)).

I guess passing a std::map<std::string_view, double> const& gives us flexibility with arguments.

You can do better. But this is a deep rabbit hole.

template<auto v>
struct int_constant {
constexpr double operator()(double)const{return v;}
};
struct literal {
double v = 0.;
constexpr double operator()(double)const{return v;}
};
struct variable {
constexpr double operator()(double arg)const{return arg;}
};
template<class Op, class T>
struct monoop {
Op op;
T t;
constexpr double operator()(double arg)const{ return op(t(arg)); }
};
template<class Op, class T>
monoop(Op, T)->monoop<Op,T>;
constexpr auto negate( auto t ) { return monoop{ std::negate<>{}, t }; }

template<class Op, class Lhs, class Rhs>
struct binop {
Op op;
Lhs lhs;
Rhs rhs;
constexpr double operator()(double arg)const{ return op(lhs(arg),rhs(arg)); }
};
template<class Op, class Lhs, class Rhs>
binop(Op,Lhs,Rhs)->binop<Op,Lhs,Rhs>;

auto sum( auto Lhs, auto Rhs ){
return binop{ std::plus<>{}, Lhs, Rhs };
}
auto product( auto Lhs, auto Rhs ){
return binop{ std::multiplies<>{}, Lhs, Rhs };
}

I can then write

auto program = product( sum( literal{3.0}, variable{} ), int_constant<2>{} );
auto squared = product( program, program );
auto again = product( squared, squared );
auto recurse = monoop{ program, program };
auto recurse2 = monoop{ recurse, recurse };

which are done at compile time. Cute and all, but not quite practical, because you want to read in a script.

But you can see that I can now build more complex expressions out of primitives. What more, I can do it at compile time, automatically, with some work.

We have this:

using subprogram = std::function<double(double)>;

now, everything above is a subprogram. We can add compilers:

using binop_reducer = std::function<subprogram(subprogram, subprogram)>;
using monoop_reducer = std::function<subprogram(subprogram)>;

these take subprograms, and compiles them into other subprograms. You did it manually with a bunch of switch cases.

With the idea of binary and unary compilers, you can build a map from the token "sin" to a sin compiler. Then your parser looks for an operation, be it "x" or "sin" or "+" or "-", determines what part of the grammar it matches, finds the appropriate compiler, and calls it, producing the result.

This gets you back to roughly where you started, using a more complex system.

To go a step further, you have to build a parse tree. Next, write a reducer -- it takes pieces of the parse tree, and reduces it to smaller ones.

struct binary_node {
e_binary_operation op;
parse_tree left;
parse_tree right;
};

The first pass version of the binary node compiler might compile the left, compile the right, look up the compiler for op, and connect them.

The next step would be to write non-trivial reduction steps, like one for a+b*c, or a+-b. A compiler that sees + as the binary op might know how to look in the right; if it sees a * instead of compiling it seperate, it compiles all of a b and c, then implements a+b*c.

This is manual work and it sucks. But you can write non-type erased compilers (!), make a compile time map from an enum class describing the binary and unary operation to the non-type erased compiler, then write metaprogramming code that composes them together.

It can then describe the parse tree that it was built out of, store that at run time how to recognize it in a reduction map, and the parse tree reducer can check each node in the parse tree for more complex reductions, and apply them.

How far down this rabbit hole do you want to go?

For a simple version of this, we'll look at RPN.

We name our tokens with std::string, because compiling can be slow, that is ok.

std::tuple<
std::map< std::string, nary_reducer<0> >,
std::map< std::string, nary_reducer<1> >,
std::map< std::string, nary_reducer<2> >
> reducers;

(an nary_reducer<n> takes n subprograms produces a subprogram).

for (auto token : tokens) {
// todo:handle constants here
if (std::get<0>( reducers ).count(token)) {
// the variable 'x' say
active.push( std::get<0>( reducers )[token]() );
} else if (std::get<1>( reducers ).count(token)) {
// sin of something
auto arg = active.pop(); // throw on failure, or whatever
active.push( std::get<0>( reducers )[token]( arg ));
} else if (std::get<2>( reducers ).count(token)) {

auto arg1 = active.pop(); // throw on failure, or whatever
auto arg2 = active.pop(); // throw on failure, or whatever
active.push( std::get<1>( reducers )[token]( arg1, arg2 ));
}
}

maybe throw in some special code to handle variables and constants.

Now we imaging keeping a buffer of 2 tokens at a time. Then building a map that maps std::tuple<std::string, std::string> to 0, 1, 2, or 3 ary reducers.

We then build every one of those pairs of reducers from our primitives as the start of this answer.

Then we go to every set of 3.

The shorter patterns still exit; if we don't spot the 3 pattern, we try the 2 pattern on the first 2, then the individual pattern on the first, then we add another token and try the 3 pattern. Etc.

Now, instead of using a std::function on every step, we only call one every 2 or 3 steps. This happens at the cost of exponentially (literally) larger program code, so the trade off eventually fails to generate speed.

Compiling also gets more expensive.

More intelligent reductions can be done as well; like knowing that 1+2 is the constant 3.

• I see where you are going. I do think that any optimization like constant folding should be done in a step before build_function(). And instead of trying to match arbitrary n-ary reducers at build_function() time, maybe it would be better to just define useful n-ary operations, and specify them explicitly in the RPN input. For example, 1 2 x FMA as input would push 1, 2 and the value of x to the stack, and then the token FMA will call a ternary operator that will then evaluate 1 + 2 * x. Apr 5 '21 at 13:08
• @g.st Sure, but that requires hand writing both the reducer and the function in the script. My system permits generating an exponentially large number of reducers from linear code, and lets the script writer write code and not be forced to kniw exponentially many optimized compound functions. What more, you can take fixed code and tweak compiler settings. As an example, imagine using this for a ROM interpreter (nary bytecode). You can even use the above to do real JIT (ship compiler and code, sanitize script, feed it to constexpr version, load resulting dll). The rabbit hole doesn;t stop.
– Yakk
Apr 5 '21 at 15:22
• Doesn’t C++20 do away with trivial deduction guides like binop(Op,Lhs,Rhs)->binop<Op,Lhs,Rhs>? Apr 6 '21 at 1:15