Fix kompile -O3 (#1009)

There exists a bug in `kompile -O3` which occurs if a specialized step
function has a residual of collection sort (i.e., its rhs has a function
that returns a collection) and garbage collection is triggered
immediately after this rule applies. In this case, the collection in
question has been allocated with koreAllocAlwaysGC, and the memory it is
using gets freed. As a result, subsequent allocations can write to that
memory and corrupt the term pointed to by that collection pointer.

The fix is to detect when this occurs just prior to invoking the garbage
collector, and copying the collection pointer into a block with a block
header, which then gets passed as the updated garbage collection root to
the gc. GC keeps that memory alive, and then the post-gc address seen by
the function that invoked the collector corresponds to the migrated
pointer on the main kore heap.

We implement this with the following steps:

1. Create a family of functions `store_<cat>_for_gc` and
`load_<cat>_for_gc` which handle the process of storing and loading the
collection pointer to the heap when the GC is invoked.
2. Create new reserved symbols used to store these values for all
collection sort categories present in the definition.
3. Ensure that the garbage collector knows that a root is a block
containing a collection rather than the collection pointer itself.

---------

Co-authored-by: Bruce Collie <[email protected]>

Author: Dwight Guth

include/kllvm/ast/AST.h

@@ -222,6 +222,7 @@ class kore_symbol {
   /* A unique integer representing the layout of the symbol in memory.
      See create_term.cpp for more information about the layout of K terms. */
   uint16_t layout_{};
+  bool instantiated_ = false;
 
 public:
   static ptr<kore_symbol> create(std::string const &name) {
@@ -1017,6 +1018,8 @@ void read_multimap(
 
 sptr<kore_pattern> strip_raw_term(sptr<kore_pattern> const &term);
 
+std::string get_raw_symbol_name(sort_category);
+
 namespace detail {
 
 template <typename T>

include/runtime/collect.h

@@ -36,7 +36,15 @@ void migrate_rangemap(void *m);
 void migrate_set(void *s);
 void migrate_collection_node(void **node_ptr);
 void set_kore_memory_functions_for_gmp(void);
-void kore_collect(void **, uint8_t, layoutitem *);
+void kore_collect(void **, uint8_t, layoutitem *, bool *);
+bool store_map_for_gc(void **, map *);
+bool store_set_for_gc(void **, set *);
+bool store_list_for_gc(void **, list *);
+bool store_rangemap_for_gc(void **, rangemap *);
+map *load_map_for_gc(void **, bool);
+set *load_set_for_gc(void **, bool);
+list *load_list_for_gc(void **, bool);
+rangemap *load_rangemap_for_gc(void **, bool);
 }
 
 #ifdef GC_DBG

include/runtime/collections/RBTree.h

@@ -130,6 +130,17 @@ class RBTree {
     root_ = t.root_;
   }
 
+  RBTree(RBTree const &other) = default;
+
+  RBTree(RBTree &&other) { other.root_.swap(root_); }
+
+  RBTree &operator=(RBTree const &other) = default;
+
+  RBTree &operator=(RBTree &&other) {
+    other.root_.swap(root_);
+    return *this;
+  }
+
   // Return true if this tree is empty.
   [[nodiscard]] bool empty() const { return root_->is_leaf(); }
 

include/runtime/collections/rangemap.h

@@ -328,6 +328,14 @@ class RangeMap {
     treemap_ = m.treemap_;
   }
 
+  RangeMap(RangeMap const &other) = default;
+
+  RangeMap(RangeMap &&other) = default;
+
+  RangeMap &operator=(RangeMap const &other) = default;
+
+  RangeMap &operator=(RangeMap &&other) = default;
+
   // Getter for the rb-tree underlying this rangemap.
   [[nodiscard]] rb_tree::RBTree<Range<T>, V> treemap() const {
     return treemap_;

lib/ast/AST.cpp

@@ -319,6 +319,9 @@ bool kore_symbol::is_builtin() const {
 }
 
 void kore_symbol::instantiate_symbol(kore_symbol_declaration *decl) {
+  if (instantiated_) {
+    return;
+  }
   std::vector<sptr<kore_sort>> instantiated;
   int i = 0;
   kore_sort::substitution vars;
@@ -332,6 +335,7 @@ void kore_symbol::instantiate_symbol(kore_symbol_declaration *decl) {
   sort_ = return_sort->substitute(vars);
 
   arguments_ = instantiated;
+  instantiated_ = true;
 }
 
 std::string kore_variable::get_name() const {

lib/ast/definition.cpp

@@ -70,15 +70,37 @@ void kore_definition::add_module(sptr<kore_module> module) {
   modules_.push_back(std::move(module));
 }
 
+std::string get_raw_symbol_name(sort_category cat) {
+  return "rawCollection_" + std::to_string((int)cat);
+}
+
 void kore_definition::insert_reserved_symbols() {
   auto mod = kore_module::create("K-RAW-TERM");
   auto decl = kore_symbol_declaration::create("rawTerm", true);
-  auto sort = kore_composite_sort::create("SortKItem");
+  auto kitem = kore_composite_sort::create("SortKItem");
 
-  decl->get_symbol()->add_sort(sort);
-  decl->get_symbol()->add_argument(sort);
+  decl->get_symbol()->add_sort(kitem);
+  decl->get_symbol()->add_argument(kitem);
   mod->add_declaration(std::move(decl));
 
+  for (auto const &cat : hooked_sorts_) {
+    switch (cat.first.cat) {
+    case sort_category::Map:
+    case sort_category::List:
+    case sort_category::Set:
+    case sort_category::RangeMap: {
+      auto decl = kore_symbol_declaration::create(
+          get_raw_symbol_name(cat.first.cat), true);
+      auto sort = cat.second;
+      decl->get_symbol()->add_sort(kitem);
+      decl->get_symbol()->add_argument(sort);
+      mod->add_declaration(std::move(decl));
+      break;
+    }
+    default: break;
+    }
+  }
+
   add_module(std::move(mod));
 }
 
@@ -121,24 +143,11 @@ SymbolMap kore_definition::get_overloads() const {
 
 // NOLINTNEXTLINE(*-function-cognitive-complexity)
 void kore_definition::preprocess() {
-  insert_reserved_symbols();
-
   for (auto *axiom : axioms_) {
     axiom->pattern_ = axiom->pattern_->expand_aliases(this);
   }
   auto symbols = std::map<std::string, std::vector<kore_symbol *>>{};
   unsigned next_ordinal = 0;
-  for (auto const &decl : symbol_declarations_) {
-    if (decl.second->attributes().contains(
-            attribute_set::key::FreshGenerator)) {
-      auto sort = decl.second->get_symbol()->get_sort();
-      if (sort->is_concrete()) {
-        fresh_functions_[dynamic_cast<kore_composite_sort *>(sort.get())
-                             ->get_name()]
-            = decl.second->get_symbol();
-      }
-    }
-  }
   for (auto iter = axioms_.begin(); iter != axioms_.end();) {
     auto *axiom = *iter;
     axiom->ordinal_ = next_ordinal;
@@ -163,12 +172,68 @@ void kore_definition::preprocess() {
       }
     }
   }
+
   for (auto const &entry : symbols) {
     for (auto *symbol : entry.second) {
       auto *decl = symbol_declarations_.at(symbol->get_name());
       symbol->instantiate_symbol(decl);
     }
   }
+
+  for (auto const &entry : symbols) {
+    for (auto *symbol : entry.second) {
+      for (auto const &sort : symbol->get_arguments()) {
+        if (sort->is_concrete()) {
+          hooked_sorts_[dynamic_cast<kore_composite_sort *>(sort.get())
+                            ->get_category(this)]
+              = std::dynamic_pointer_cast<kore_composite_sort>(sort);
+        }
+      }
+      if (symbol->get_sort()->is_concrete()) {
+        hooked_sorts_[dynamic_cast<kore_composite_sort *>(
+                          symbol->get_sort().get())
+                          ->get_category(this)]
+            = std::dynamic_pointer_cast<kore_composite_sort>(
+                symbol->get_sort());
+      }
+    }
+  }
+
+  insert_reserved_symbols();
+
+  for (auto &module : modules_) {
+    auto const &declarations = module->get_declarations();
+    for (auto const &declaration : declarations) {
+      auto *decl = dynamic_cast<kore_symbol_declaration *>(declaration.get());
+      if (decl == nullptr) {
+        continue;
+      }
+      if (decl->is_hooked() && decl->get_object_sort_variables().empty()) {
+        kore_symbol *symbol = decl->get_symbol();
+        symbols.emplace(symbol->get_name(), std::vector<kore_symbol *>{symbol});
+      }
+    }
+  }
+
+  for (auto const &entry : symbols) {
+    for (auto *symbol : entry.second) {
+      auto *decl = symbol_declarations_.at(symbol->get_name());
+      symbol->instantiate_symbol(decl);
+    }
+  }
+
+  for (auto const &decl : symbol_declarations_) {
+    if (decl.second->attributes().contains(
+            attribute_set::key::FreshGenerator)) {
+      auto sort = decl.second->get_symbol()->get_sort();
+      if (sort->is_concrete()) {
+        fresh_functions_[dynamic_cast<kore_composite_sort *>(sort.get())
+                             ->get_name()]
+            = decl.second->get_symbol();
+      }
+    }
+  }
+
   uint32_t next_symbol = 0;
   uint16_t next_layout = 1;
   auto instantiations
@@ -202,20 +267,6 @@ void kore_definition::preprocess() {
   for (auto const &entry : symbols) {
     auto range = variables.at(entry.first);
     for (auto *symbol : entry.second) {
-      for (auto const &sort : symbol->get_arguments()) {
-        if (sort->is_concrete()) {
-          hooked_sorts_[dynamic_cast<kore_composite_sort *>(sort.get())
-                            ->get_category(this)]
-              = std::dynamic_pointer_cast<kore_composite_sort>(sort);
-        }
-      }
-      if (symbol->get_sort()->is_concrete()) {
-        hooked_sorts_[dynamic_cast<kore_composite_sort *>(
-                          symbol->get_sort().get())
-                          ->get_category(this)]
-            = std::dynamic_pointer_cast<kore_composite_sort>(
-                symbol->get_sort());
-      }
       if (!symbol->is_concrete()) {
         if (symbol->is_polymorphic()) {
           symbol->first_tag_ = range.first;