Remove mallocs for our lambda passing

This commit works by adding a SmallFunction class that behaves
like std::function, but allocates its data inline rather than through
a separate allocation.  It probably could have also worked by taking
a custom allocator and using the new RegionAllocator.

It adds a bit more restrictions than std::function does (the types
caught by the closure have to be more "trivial" than std::function
supports), so some of this change is marking some types as trivial,
or copying data into a trivial format from things that aren't (ex SmallVector).

Author: kmod

src/asm_writing/rewriter.cpp

@@ -816,12 +816,30 @@ RewriterVar* Rewriter::call(bool has_side_effects, void* func_addr, const Rewrit
         type = ActionType::MUTATION;
     else
         type = ActionType::NORMAL;
-    addAction([=]() { this->_call(result, has_side_effects, func_addr, args, args_xmm); }, uses, type);
+
+    // It's not nice to pass llvm::SmallVectors through a closure, especially with our SmallFunction
+    // optimization, so just regionAlloc them and copy the data in:
+    RewriterVar** _args = (RewriterVar**)this->regionAlloc(sizeof(RewriterVar*) * args.size());
+    memcpy(_args, args.begin(), sizeof(RewriterVar*) * args.size());
+    RewriterVar** _args_xmm = (RewriterVar**)this->regionAlloc(sizeof(RewriterVar*) * args_xmm.size());
+    memcpy(_args_xmm, args_xmm.begin(), sizeof(RewriterVar*) * args_xmm.size());
+
+    int args_size = args.size();
+    assert(args_xmm.size() <= 0x7fff);
+    // Hack: pack this into a short to make sure it fits in the closure
+    short xmm_args_size = args_xmm.size();
+
+    // Hack: explicitly order the closure arguments so they pad nicer
+    addAction([args_size, xmm_args_size, has_side_effects, this, result, func_addr, _args, _args_xmm]() {
+        this->_call(result, has_side_effects, func_addr, llvm::ArrayRef<RewriterVar*>(_args, args_size),
+                    llvm::ArrayRef<RewriterVar*>(_args_xmm, xmm_args_size));
+    }, uses, type);
+
     return result;
 }
 
-void Rewriter::_setupCall(bool has_side_effects, const RewriterVar::SmallVector& args,
-                          const RewriterVar::SmallVector& args_xmm) {
+void Rewriter::_setupCall(bool has_side_effects, llvm::ArrayRef<RewriterVar*> args,
+                          llvm::ArrayRef<RewriterVar*> args_xmm) {
     if (has_side_effects)
         assert(done_guarding);
 
@@ -966,8 +984,8 @@ void Rewriter::_setupCall(bool has_side_effects, const RewriterVar::SmallVector&
 #endif
 }
 
-void Rewriter::_call(RewriterVar* result, bool has_side_effects, void* func_addr, const RewriterVar::SmallVector& args,
-                     const RewriterVar::SmallVector& args_xmm) {
+void Rewriter::_call(RewriterVar* result, bool has_side_effects, void* func_addr, llvm::ArrayRef<RewriterVar*> args,
+                     llvm::ArrayRef<RewriterVar*> args_xmm) {
     assembler->comment("_call");
 
     // RewriterVarUsage scratch = createNewVar(Location::any());

src/asm_writing/rewriter.h

@@ -69,13 +69,9 @@ struct Location {
         int32_t _data;
     };
 
-    constexpr Location() : type(Uninitialized), _data(-1) {}
-    constexpr Location(const Location& r) : type(r.type), _data(r._data) {}
-    Location operator=(const Location& r) {
-        type = r.type;
-        _data = r._data;
-        return *this;
-    }
+    constexpr Location() noexcept : type(Uninitialized), _data(-1) {}
+    constexpr Location(const Location& r) = default;
+    Location& operator=(const Location& r) = default;
 
     constexpr Location(LocationType type, int32_t data) : type(type), _data(data) {}
 
@@ -283,11 +279,50 @@ class RewriterVar {
     friend class JitFragmentWriter;
 };
 
+// A utility class that is similar to std::function, but stores any closure data inline rather
+// than in a separate allocation.  It's similar to SmallVector, but will just fail to compile if
+// you try to put more bytes in than you allocated.
+// Currently, it only works for functions with the signature "void()"
+template <int N = 24> class SmallFunction {
+private:
+    void (*func)(void*);
+    char data[N];
+
+    template <typename Functor> struct Caller {
+        static void call(void* data) { (*(Functor*)data)(); }
+    };
+
+public:
+    template <typename Functor> SmallFunction(Functor&& f) noexcept {
+        static_assert(std::has_trivial_copy_constructor<typename std::remove_reference<Functor>::type>::value,
+                      "SmallFunction currently only works with simple types");
+        static_assert(std::is_trivially_destructible<typename std::remove_reference<Functor>::type>::value,
+                      "SmallFunction currently only works with simple types");
+        static_assert(sizeof(Functor) <= sizeof(data), "Please increase N");
+        new (data) typename std::remove_reference<Functor>::type(std::forward<Functor>(f));
+        func = Caller<Functor>::call;
+    }
+
+    SmallFunction() = default;
+    SmallFunction(const SmallFunction<N>& rhs) = default;
+    SmallFunction(SmallFunction<N>&& rhs) = default;
+    SmallFunction& operator=(const SmallFunction<N>& rhs) = default;
+    SmallFunction& operator=(SmallFunction<N>&& rhs) = default;
+
+    void operator()() { func(data); }
+};
+
 class RewriterAction {
 public:
-    std::function<void()> action;
+    SmallFunction<48> action;
+
+    template <typename F> RewriterAction(F&& action) : action(std::forward<F>(action)) {}
 
-    RewriterAction(std::function<void()> f) : action(std::move(f)) {}
+    RewriterAction() = default;
+    RewriterAction(const RewriterAction& rhs) = default;
+    RewriterAction(RewriterAction&& rhs) = default;
+    RewriterAction& operator=(const RewriterAction& rhs) = default;
+    RewriterAction& operator=(RewriterAction&& rhs) = default;
 };
 
 enum class ActionType { NORMAL, GUARD, MUTATION };
@@ -449,10 +484,9 @@ class Rewriter : public ICSlotRewrite::CommitHook {
 
     void _trap();
     void _loadConst(RewriterVar* result, int64_t val);
-    void _setupCall(bool has_side_effects, const RewriterVar::SmallVector& args,
-                    const RewriterVar::SmallVector& args_xmm);
-    void _call(RewriterVar* result, bool has_side_effects, void* func_addr, const RewriterVar::SmallVector& args,
-               const RewriterVar::SmallVector& args_xmm);
+    void _setupCall(bool has_side_effects, llvm::ArrayRef<RewriterVar*> args, llvm::ArrayRef<RewriterVar*> args_xmm);
+    void _call(RewriterVar* result, bool has_side_effects, void* func_addr, llvm::ArrayRef<RewriterVar*> args,
+               llvm::ArrayRef<RewriterVar*> args_xmm);
     void _add(RewriterVar* result, RewriterVar* a, int64_t b, Location dest);
     int _allocate(RewriterVar* result, int n);
     void _allocateAndCopy(RewriterVar* result, RewriterVar* array, int n);

src/codegen/baseline_jit.cpp

@@ -668,8 +668,15 @@ RewriterVar* JitFragmentWriter::emitPPCall(void* func_addr, llvm::ArrayRef<Rewri
     RewriterVar::SmallVector args_vec(args.begin(), args.end());
 #if ENABLE_BASELINEJIT_ICS
     RewriterVar* result = createNewVar();
-    addAction([=]() { this->_emitPPCall(result, func_addr, args_vec, num_slots, slot_size); }, args,
-              ActionType::NORMAL);
+
+    int args_size = args.size();
+    RewriterVar** _args = (RewriterVar**)regionAlloc(sizeof(RewriterVar*) * args_size);
+    memcpy(_args, args.begin(), sizeof(RewriterVar*) * args_size);
+
+    addAction([=]() {
+        this->_emitPPCall(result, func_addr, llvm::ArrayRef<RewriterVar*>(_args, args_size), num_slots, slot_size);
+    }, args, ActionType::NORMAL);
+
     if (type_recorder) {
         RewriterVar* type_recorder_var = imm(type_recorder);
         RewriterVar* obj_cls_var = result->getAttr(offsetof(Box, cls));
@@ -813,7 +820,7 @@ void JitFragmentWriter::_emitOSRPoint(RewriterVar* result, RewriterVar* node_var
     assertConsistent();
 }
 
-void JitFragmentWriter::_emitPPCall(RewriterVar* result, void* func_addr, const RewriterVar::SmallVector& args,
+void JitFragmentWriter::_emitPPCall(RewriterVar* result, void* func_addr, llvm::ArrayRef<RewriterVar*> args,
                                     int num_slots, int slot_size) {
     assembler::Register r = allocReg(assembler::R11);
 

src/codegen/baseline_jit.h

@@ -294,7 +294,7 @@ class JitFragmentWriter : public Rewriter {
     void _emitGetLocal(RewriterVar* val_var, const char* name);
     void _emitJump(CFGBlock* b, RewriterVar* block_next, int& size_of_exit_to_interp);
     void _emitOSRPoint(RewriterVar* result, RewriterVar* node_var);
-    void _emitPPCall(RewriterVar* result, void* func_addr, const RewriterVar::SmallVector& args, int num_slots,
+    void _emitPPCall(RewriterVar* result, void* func_addr, llvm::ArrayRef<RewriterVar*> args, int num_slots,
                      int slot_size);
     void _emitRecordType(RewriterVar* type_recorder_var, RewriterVar* obj_cls_var);
     void _emitReturn(RewriterVar* v);