evaluator: Stop copying the whole heap on every formatter call

lib/include/pl/core/evaluator.hpp

@@ -393,6 +393,28 @@ namespace pl::core {
             return this->m_heap;
         }
 
+        /**
+         * @brief Starts recording heap changes so popHeapGuard() can undo them.
+         * Only cells that get written are copied, so the cost does not grow with
+         * the size of the heap.
+         */
+        void pushHeapGuard();
+
+        /**
+         * @brief Undoes all heap changes made since the matching pushHeapGuard().
+         * New cells are removed and changed cells get their old contents back.
+         */
+        void popHeapGuard();
+
+        /**
+         * @brief Gets a heap cell that is about to be changed. If a heap guard is
+         * active, the cell's current contents are saved first so the guard can
+         * restore them.
+         * @param address Address of the heap cell
+         * @return The heap cell
+         */
+        [[nodiscard]] std::vector<u8> &getHeapCellForWriting(u64 address);
+
         [[nodiscard]] std::map<u32, PatternLocalData> &getPatternLocalStorage() {
             return this->m_patternLocalStorage;
         }
@@ -606,6 +628,15 @@ namespace pl::core {
         std::map<u32, u32> m_heapReferenceCounts;
         std::map<u32, PatternLocalData> m_patternLocalStorage;
 
+        /**
+         * @brief Heap state saved by pushHeapGuard() and restored by popHeapGuard()
+         */
+        struct HeapGuard {
+            size_t heapSize;
+            std::map<u64, std::vector<u8>> originalCells;
+        };
+        std::vector<HeapGuard> m_heapGuards;
+
         std::map<std::string, std::set<ptrn::Pattern*>> m_attributedPatterns;
         std::vector<std::unique_ptr<Scope>> m_scopes;
         std::vector<std::shared_ptr<ptrn::Pattern>> m_patterns;

lib/include/pl/patterns/pattern.hpp

@@ -468,8 +468,8 @@ namespace pl::ptrn {
                 try {
                     const auto function = this->m_evaluator->findFunction(formatterFunctionName);
                     if (function.has_value()) {
-                        auto startHeap = this->m_evaluator->getHeap();
-                        ON_SCOPE_EXIT { this->m_evaluator->getHeap() = startHeap; };
+                        this->m_evaluator->pushHeapGuard();
+                        ON_SCOPE_EXIT { this->m_evaluator->popHeapGuard(); };
 
                         auto formatterResult = function->func(this->m_evaluator, { value });
                         if (formatterResult.has_value()) {
@@ -590,8 +590,8 @@ namespace pl::ptrn {
             auto evaluator = this->getEvaluator();
 
             if (auto transformFunc = evaluator->findFunction(this->getTransformFunction()); transformFunc.has_value()) {
-                auto startHeap = this->m_evaluator->getHeap();
-                ON_SCOPE_EXIT { this->m_evaluator->getHeap() = startHeap; };
+                this->m_evaluator->pushHeapGuard();
+                ON_SCOPE_EXIT { this->m_evaluator->popHeapGuard(); };
 
                 // Preserve pattern variable name
                 std::string patternName;
@@ -621,8 +621,8 @@ namespace pl::ptrn {
             else {
                 const auto function = this->m_evaluator->findFunction(formatterFunctionName);
                 if (function.has_value()) {
-                    auto startHeap = this->m_evaluator->getHeap();
-                    ON_SCOPE_EXIT { this->m_evaluator->getHeap() = startHeap; };
+                    this->m_evaluator->pushHeapGuard();
+                    ON_SCOPE_EXIT { this->m_evaluator->popHeapGuard(); };
 
                     // Preserve pattern variable name
                     std::string patternName;

lib/source/pl/core/evaluator.cpp

@@ -714,8 +714,8 @@ namespace pl::core {
 
             auto getStorage = [&, this]() -> auto& {
                 if (heapSection) {
-                    if (auto &heap = this->getHeap(); heap.size() > pattern->getHeapAddress())
-                        return heap[pattern->getHeapAddress()];
+                    if (this->getHeap().size() > pattern->getHeapAddress())
+                        return this->getHeapCellForWriting(pattern->getHeapAddress());
                     else
                         err::E0011.throwError(fmt::format("Tried accessing out of bounds heap cell {}. This is a bug.", pattern->getHeapAddress()));
                 } else if (patternLocalSection) {
@@ -886,6 +886,32 @@ namespace pl::core {
         this->m_scopes.pop_back();
     }
 
+    void Evaluator::pushHeapGuard() {
+        this->m_heapGuards.push_back({ this->m_heap.size(), {} });
+    }
+
+    void Evaluator::popHeapGuard() {
+        if (this->m_heapGuards.empty())
+            return;
+
+        auto guard = std::move(this->m_heapGuards.back());
+        this->m_heapGuards.pop_back();
+
+        this->m_heap.resize(guard.heapSize);
+        for (auto &[address, cell] : guard.originalCells)
+            this->m_heap[address] = std::move(cell);
+    }
+
+    std::vector<u8> &Evaluator::getHeapCellForWriting(u64 address) {
+        // Only the innermost guard needs a copy, since it restores the cell before outer guards can see it
+        if (!this->m_heapGuards.empty()) {
+            if (auto &guard = this->m_heapGuards.back(); address < guard.heapSize)
+                guard.originalCells.try_emplace(address, this->m_heap[address]);
+        }
+
+        return this->m_heap[address];
+    }
+
     std::vector<ast::ASTNode*> unpackCompoundStatements(const std::vector<std::shared_ptr<ast::ASTNode>> &nodes) {
         std::vector<ast::ASTNode*> result;
 
@@ -919,7 +945,8 @@ namespace pl::core {
             auto heapAddress = (address >> 32);
             auto storageAddress = address & 0xFFFF'FFFF;
             if (heapAddress < heap.size()) {
-                auto &storage = heap[heapAddress];
+                const bool modifiesCell = write || storageAddress + size > heap[heapAddress].size();
+                auto &storage = modifiesCell ? this->getHeapCellForWriting(heapAddress) : heap[heapAddress];
 
                 if (storageAddress + size > storage.size()) {
                     storage.resize(storageAddress + size);
@@ -1021,7 +1048,7 @@ namespace pl::core {
         if (id == ptrn::Pattern::MainSectionId)
             err::E0011.throwError("Cannot access main section.");
         else if (id == ptrn::Pattern::HeapSectionId)
-            return this->m_heap.back();
+            return this->getHeapCellForWriting(this->m_heap.size() - 1);
         else if (this->m_sections.contains(id))
             return this->m_sections[id].data;
         else if (id == ptrn::Pattern::InstantiationSectionId)

tests/CMakeLists.txt

@@ -40,6 +40,8 @@ set(AVAILABLE_TESTS
         Pragmas
         PragmasFail
         Format
+        FormatHeapGrowth
+        FormatHeapIsolation
         RValuesAssignmentInStruct
         TemplateParametersScope
         TypeNameOf
@@ -100,6 +102,7 @@ set(AVAILABLE_TESTS
 add_executable(pattern_language_tests
     source/main.cpp
     source/tests.cpp
+    source/test_pattern_format_heap.cpp
 )
 
 

tests/include/test_patterns/test_pattern_format_heap.hpp

@@ -0,0 +1,144 @@
+#pragma once
+
+#include "test_pattern.hpp"
+
+#include <pl/patterns/pattern.hpp>
+
+#include <chrono>
+#include <functional>
+
+namespace pl::test {
+
+    /**
+     * @brief Finds a top-level pattern by its variable name
+     * @param patterns Patterns to search
+     * @param name Variable name to look for
+     * @return The pattern, or nullptr if no pattern has that name
+     */
+    std::shared_ptr<ptrn::Pattern> findPattern(const std::vector<std::shared_ptr<ptrn::Pattern>> &patterns, const std::string &name);
+
+    /**
+     * @brief Calls a function on every pattern inside a pattern that has no
+     * entries of its own, such as the members of each struct in an array
+     * @param pattern Pattern to walk through
+     * @param callback Function to call on each of these patterns
+     */
+    void forEachLeaf(const std::shared_ptr<ptrn::Pattern> &pattern, const std::function<void(ptrn::Pattern&)> &callback);
+
+    class TestPatternFormatHeapGrowth : public TestPattern {
+    public:
+        TestPatternFormatHeapGrowth(core::Evaluator *evaluator) : TestPattern(evaluator, "FormatHeapGrowth") {
+        }
+        ~TestPatternFormatHeapGrowth() override = default;
+
+        [[nodiscard]] std::string getSourceCode() const override {
+            return R"(
+                str filler = "0123456789ABCDEF" * 0x100000;
+
+                fn format_entry(auto value) {
+                    return "x";
+                };
+
+                struct Entry {
+                    u8 value [[format("format_entry")]];
+                };
+
+                Entry entries[0x4000] @ 0x00;
+            )";
+        }
+
+        [[nodiscard]] bool runChecks(const std::vector<std::shared_ptr<ptrn::Pattern>> &patterns) const override {
+            auto entries = findPattern(patterns, "entries");
+            if (entries == nullptr)
+                return false;
+
+            // Each formatter call used to copy the whole heap, which made this take about 20 seconds
+            const auto start = std::chrono::steady_clock::now();
+
+            bool valid = true;
+            forEachLeaf(entries, [&](ptrn::Pattern &pattern) {
+                if (pattern.getFormattedValue() != "x")
+                    valid = false;
+            });
+
+            const auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - start);
+            if (elapsed > std::chrono::seconds(5)) {
+                fmt::print("Formatting took {} ms\n", elapsed.count());
+                return false;
+            }
+
+            return valid;
+        }
+    };
+
+    class TestPatternFormatHeapIsolation : public TestPattern {
+    public:
+        TestPatternFormatHeapIsolation(core::Evaluator *evaluator) : TestPattern(evaluator, "FormatHeapIsolation") {
+        }
+        ~TestPatternFormatHeapIsolation() override = default;
+
+        [[nodiscard]] std::string getSourceCode() const override {
+            return R"(
+                u32 counter = 0;
+                str label = "a";
+
+                fn bump_format(auto value) {
+                    counter = counter + 1;
+                    label = "b";
+                    return "bumped";
+                };
+
+                fn bump_transform(auto value) {
+                    counter = counter + 100;
+                    label = "c";
+                    return value;
+                };
+
+                fn show_state(auto value) {
+                    if (counter == 0 && label == "a")
+                        return "clean";
+                    else
+                        return "leaked";
+                };
+
+                struct Formatted {
+                    u8 value [[format("bump_format")]];
+                };
+
+                struct Transformed {
+                    u8 value [[transform("bump_transform")]];
+                };
+
+                Formatted formatted[4] @ 0x00;
+                Transformed transformed[4] @ 0x10;
+                u8 state @ 0x20 [[format("show_state")]];
+            )";
+        }
+
+        [[nodiscard]] bool runChecks(const std::vector<std::shared_ptr<ptrn::Pattern>> &patterns) const override {
+            auto formatted = findPattern(patterns, "formatted");
+            auto transformed = findPattern(patterns, "transformed");
+            auto state = findPattern(patterns, "state");
+            if (formatted == nullptr || transformed == nullptr || state == nullptr)
+                return false;
+
+            bool valid = true;
+            forEachLeaf(formatted, [&](ptrn::Pattern &pattern) {
+                if (pattern.getFormattedValue() != "bumped")
+                    valid = false;
+            });
+            forEachLeaf(transformed, [&](ptrn::Pattern &pattern) {
+                wolv::util::unused(pattern.getValue());
+            });
+
+            // Changes that formatters and transforms make to globals must not last after they return
+            if (auto stateValue = state->getFormattedValue(); stateValue != "clean") {
+                fmt::print("Formatter state leaked: {}\n", stateValue);
+                valid = false;
+            }
+
+            return valid;
+        }
+    };
+
+}

tests/source/test_pattern_format_heap.cpp

@@ -0,0 +1,24 @@
+#include "test_patterns/test_pattern_format_heap.hpp"
+
+namespace pl::test {
+
+    std::shared_ptr<ptrn::Pattern> findPattern(const std::vector<std::shared_ptr<ptrn::Pattern>> &patterns, const std::string &name) {
+        for (const auto &pattern : patterns) {
+            if (pattern->getVariableName() == name)
+                return pattern;
+        }
+
+        return nullptr;
+    }
+
+    void forEachLeaf(const std::shared_ptr<ptrn::Pattern> &pattern, const std::function<void(ptrn::Pattern&)> &callback) {
+        if (auto iterable = dynamic_cast<ptrn::IIterable*>(pattern.get()); iterable != nullptr) {
+            iterable->forEachEntry(0, iterable->getEntryCount(), [&](u64, const auto &entry) {
+                forEachLeaf(entry, callback);
+            });
+        } else {
+            callback(*pattern);
+        }
+    }
+
+}

tests/source/tests.cpp

@@ -28,6 +28,7 @@
 #include "test_patterns/test_pattern_pragmas_pr249.hpp"
 #include "test_patterns/test_pattern_pragmas_fail.hpp"
 #include "test_patterns/test_pattern_format.hpp"
+#include "test_patterns/test_pattern_format_heap.hpp"
 #include "test_patterns/test_pattern_rvalues_assignment_in_struct.hpp"
 #include "test_patterns/test_pattern_template_parameters_scope.hpp"
 #include "test_patterns/test_pattern_typenameof.hpp"
@@ -83,6 +84,8 @@ std::array Tests = {
     TEST(PragmasIssue249),
     TEST(PragmasFail),
     TEST(Format),
+    TEST(FormatHeapGrowth),
+    TEST(FormatHeapIsolation),
     TEST(RValuesAssignmentInStruct),
     TEST(TemplateParametersScope),
     TEST(TypeNameOf),