3 files changed, 443 insertions, 0 deletions
diff --git a/include/klee/KDAlloc/allocator.h b/include/klee/KDAlloc/allocator.h
new file mode 100644
index 00000000..f59223ed
--- /dev/null
+++ b/include/klee/KDAlloc/allocator.h
@@ -0,0 +1,264 @@
+//===-- allocator.h ---------------------------------------------*- C++ -*-===//
+//
+//                     The KLEE Symbolic Virtual Machine
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef KDALLOC_ALLOCATOR_H
+#define KDALLOC_ALLOCATOR_H
+
+#include "define.h"
+#include "location_info.h"
+#include "mapping.h"
+#include "suballocators/loh.h"
+#include "suballocators/slot_allocator.h"
+#include "tagged_logger.h"
+
+#include "klee/ADT/Bits.h"
+#include "klee/ADT/Ref.h"
+
+#include <array>
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <ostream>
+#include <type_traits>
+
+namespace klee::kdalloc {
+/// Wraps a mapping that is shared with other allocators.
+class Allocator final : public TaggedLogger<Allocator> {
+public:
+  class Control final {
+    friend class Allocator;
+    friend class AllocatorFactory;
+
+    static constexpr const std::uint32_t unlimitedQuarantine =
+        static_cast<std::uint32_t>(-1);
+
+    /// @todo This should really be a data member `static constexpr const
+    /// std::array meta = { ... }`.
+    static inline const std::array<std::size_t, 8> &getMeta() noexcept {
+      static const std::array<std::size_t, 8> meta = {
+          1u,    // bool
+          4u,    // int
+          8u,    // pointer size
+          16u,   // double
+          32u,   // compound types #1
+          64u,   // compound types #2
+          256u,  // compound types #3
+          2048u, // reasonable buffers
+      };
+      return meta;
+    }
+
+    [[nodiscard]] static inline int
+    convertSizeToBinIndex(std::size_t const size) noexcept {
+      for (std::size_t i = 0; i < getMeta().size(); ++i) {
+        if (getMeta()[i] >= size) {
+          return i;
+        }
+      }
+      return getMeta().size();
+    }
+
+  public:
+    mutable class klee::ReferenceCounter _refCount;
+
+  private:
+    Mapping mapping;
+    std::array<suballocators::SlotAllocator::Control,
+               std::tuple_size<std::decay_t<decltype(getMeta())>>::value>
+        sizedBins;
+    suballocators::LargeObjectAllocator::Control largeObjectBin;
+
+  public:
+    Control() = delete;
+    Control(Control const &) = delete;
+    Control &operator=(Control const &) = delete;
+    Control(Control &&) = delete;
+    Control &operator=(Control &&) = delete;
+
+  private:
+    Control(Mapping &&mapping) : mapping(std::move(mapping)) {}
+  };
+
+  static constexpr const auto unlimitedQuarantine =
+      Control::unlimitedQuarantine;
+
+private:
+  klee::ref<Control> control;
+
+  std::array<suballocators::SlotAllocator,
+             std::tuple_size<std::decay_t<decltype(Control::getMeta())>>::value>
+      sizedBins;
+  suballocators::LargeObjectAllocator largeObjectBin;
+
+public:
+  std::ostream &logTag(std::ostream &out) const noexcept {
+    return out << "[alloc] ";
+  }
+
+  Allocator() = default;
+  Allocator(Allocator const &rhs) = default;
+  Allocator &operator=(Allocator const &rhs) = default;
+  Allocator(Allocator &&) = default;
+  Allocator &operator=(Allocator &&) = default;
+
+  Allocator(klee::ref<Control> control) : control(std::move(control)) {}
+
+  explicit operator bool() const noexcept { return !control.isNull(); }
+
+  auto const &getSizedBinInfo() const noexcept { return Control::getMeta(); }
+
+  [[nodiscard]] void *allocate(std::size_t size) {
+    assert(*this && "Invalid allocator");
+
+    auto const bin = Control::convertSizeToBinIndex(size);
+    traceLine("Allocating ", size, " bytes in bin ", bin);
+
+    void *result = nullptr;
+    if (bin < static_cast<int>(sizedBins.size())) {
+      result = sizedBins[bin].allocate(control->sizedBins[bin]);
+    } else {
+      result = largeObjectBin.allocate(control->largeObjectBin, size);
+    }
+    traceLine("Allocated ", result);
+    return result;
+  }
+
+  void free(void *ptr, std::size_t size) {
+    assert(*this && "Invalid allocator");
+    assert(ptr && "Freeing nullptrs is not supported"); // we are not ::free!
+
+    auto const bin = Control::convertSizeToBinIndex(size);
+    traceLine("Freeing ", ptr, " of size ", size, " in bin ", bin);
+
+    if (bin < static_cast<int>(sizedBins.size())) {
+      return sizedBins[bin].deallocate(control->sizedBins[bin], ptr);
+    } else {
+      return largeObjectBin.deallocate(control->largeObjectBin, ptr, size);
+    }
+  }
+
+  LocationInfo location_info(void const *const ptr,
+                             std::size_t const size) const noexcept {
+    assert(*this && "Invalid allocator");
+
+    if (!ptr || reinterpret_cast<std::uintptr_t>(ptr) < 4096) {
+      return LocationInfo::LI_NullPage;
+    }
+
+    // the following is technically UB if `ptr` does not actually point inside
+    // the mapping at all
+    for (std::size_t i = 0; i < Allocator::Control::getMeta().size(); ++i) {
+      if (control->sizedBins[i].mapping_begin() <= ptr &&
+          ptr < control->sizedBins[i].mapping_end()) {
+        if (reinterpret_cast<char const *>(ptr) + size <=
+            control->sizedBins[i].mapping_end()) {
+          return sizedBins[i].getLocationInfo(control->sizedBins[i], ptr, size);
+        } else {
+          return LocationInfo::LI_SpansSuballocators;
+        }
+      }
+    }
+    if (control->largeObjectBin.mapping_begin() <= ptr &&
+        ptr < control->largeObjectBin.mapping_end()) {
+      if (reinterpret_cast<char const *>(ptr) + size <=
+          control->largeObjectBin.mapping_end()) {
+        return largeObjectBin.getLocationInfo(control->largeObjectBin, ptr,
+                                              size);
+      } else {
+        // the loh is the suballocator using the largest addresses, therefore
+        // its mapping is not followed by that of another suballocator, but
+        // rather the end of the mapping
+        return LocationInfo::LI_NonNullOutsideMapping;
+      }
+    }
+
+    return LocationInfo::LI_NonNullOutsideMapping;
+  }
+};
+
+class AllocatorFactory {
+public:
+  static constexpr const auto unlimitedQuarantine =
+      Allocator::Control::unlimitedQuarantine;
+
+private:
+  klee::ref<Allocator::Control> control;
+
+public:
+  AllocatorFactory() = default;
+
+  AllocatorFactory(std::size_t const size, std::uint32_t const quarantineSize)
+      : AllocatorFactory(Mapping{0, size}, quarantineSize) {}
+
+  AllocatorFactory(std::uintptr_t const address, std::size_t const size,
+                   std::uint32_t const quarantineSize)
+      : AllocatorFactory(Mapping{address, size}, quarantineSize) {}
+
+  AllocatorFactory(Mapping &&mapping, std::uint32_t const quarantineSize) {
+    assert(mapping && "Invalid mapping");
+    assert(mapping.getSize() >
+               Allocator::Control::getMeta().size() * 4096 + 3 * 4096 &&
+           "Mapping is *far* to small");
+
+    control = new Allocator::Control(std::move(mapping));
+    auto const binSize =
+        static_cast<std::size_t>(1)
+        << (std::numeric_limits<std::size_t>::digits - 1 -
+            countLeadingZeroes(control->mapping.getSize() /
+                               (Allocator::Control::getMeta().size() + 1)));
+    char *const base = static_cast<char *>(control->mapping.getBaseAddress());
+    std::size_t totalSize = 0;
+    for (std::size_t i = 0; i < Allocator::Control::getMeta().size(); ++i) {
+      control->sizedBins[i].initialize(
+          base + totalSize, binSize, Allocator::Control::getMeta()[i],
+          quarantineSize == unlimitedQuarantine,
+          quarantineSize == unlimitedQuarantine ? 0 : quarantineSize);
+
+      totalSize += binSize;
+      assert(totalSize <= control->mapping.getSize() && "Mapping too small");
+    }
+
+    auto largeObjectBinSize = control->mapping.getSize() - totalSize;
+    assert(largeObjectBinSize > 0);
+    control->largeObjectBin.initialize(
+        base + totalSize, largeObjectBinSize,
+        quarantineSize == unlimitedQuarantine,
+        quarantineSize == unlimitedQuarantine ? 0 : quarantineSize);
+  }
+
+  explicit operator bool() const noexcept { return !control.isNull(); }
+
+  Mapping &getMapping() noexcept {
+    assert(!!*this && "Cannot get mapping of uninitialized factory.");
+    return control->mapping;
+  }
+
+  Mapping const &getMapping() const noexcept {
+    assert(!!*this && "Cannot get mapping of uninitialized factory.");
+    return control->mapping;
+  }
+
+  Allocator makeAllocator() const {
+    assert(!!*this &&
+           "Can only create an allocator from an initialized factory.");
+
+    return Allocator{control};
+  }
+
+  explicit operator Allocator() && {
+    assert(!!*this &&
+           "Can only create an allocator from an initialized factory.");
+
+    return Allocator{std::move(control)};
+  }
+};
+} // namespace klee::kdalloc
+
+#endif
diff --git a/include/klee/KDAlloc/kdalloc.h b/include/klee/KDAlloc/kdalloc.h
new file mode 100644
index 00000000..c04a4ebf
--- /dev/null
+++ b/include/klee/KDAlloc/kdalloc.h
@@ -0,0 +1,21 @@
+//===-- kdalloc.h -----------------------------------------------*- C++ -*-===//
+//
+//                     The KLEE Symbolic Virtual Machine
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef KDALLOC_KDALLOC_H
+#define KDALLOC_KDALLOC_H
+
+#include "allocator.h"
+#include "mapping.h"
+
+namespace klee::kdalloc {
+using StackAllocator = Allocator;
+using StackAllocatorFactory = AllocatorFactory;
+} // namespace klee::kdalloc
+
+#endif
diff --git a/include/klee/KDAlloc/mapping.h b/include/klee/KDAlloc/mapping.h
new file mode 100644
index 00000000..4afc9d86
--- /dev/null
+++ b/include/klee/KDAlloc/mapping.h
@@ -0,0 +1,158 @@
+//===-- mapping.h -----------------------------------------------*- C++ -*-===//
+//
+//                     The KLEE Symbolic Virtual Machine
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef KDALLOC_MAPPING_H
+#define KDALLOC_MAPPING_H
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <utility>
+
+#include <sys/mman.h>
+#include <unistd.h>
+
+#if defined(__linux__)
+#include <linux/version.h>
+#endif
+
+#if defined(__APPLE__)
+#include <mach/vm_map.h>
+#include <sys/types.h>
+#endif
+
+#include "klee/Support/ErrorHandling.h"
+
+namespace klee::kdalloc {
+class Mapping {
+  void *baseAddress = MAP_FAILED;
+  std::size_t size = 0;
+
+  bool try_map(std::uintptr_t baseAddress) noexcept {
+    assert(this->baseAddress == MAP_FAILED);
+
+    int flags = MAP_ANON | MAP_PRIVATE;
+#if defined(__linux__)
+    flags |= MAP_NORESERVE;
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 17, 0)
+    if (baseAddress != 0) {
+      flags |= MAP_FIXED_NOREPLACE;
+    }
+#endif
+#elif defined(__FreeBSD__)
+    if (baseAddress != 0) {
+      flags |= MAP_FIXED | MAP_EXCL;
+    }
+#endif
+
+    auto mappedAddress = ::mmap(reinterpret_cast<void *>(baseAddress), size,
+                                PROT_READ | PROT_WRITE, flags, -1, 0);
+    if (baseAddress != 0 &&
+        baseAddress != reinterpret_cast<std::uintptr_t>(mappedAddress)) {
+      [[maybe_unused]] int rc = ::munmap(mappedAddress, size);
+      assert(rc == 0 && "munmap failed");
+      this->baseAddress = MAP_FAILED;
+      return false;
+    }
+    if (mappedAddress == MAP_FAILED) {
+      this->baseAddress = MAP_FAILED;
+      return false;
+    }
+    this->baseAddress = mappedAddress;
+
+#if defined(__linux__)
+    {
+      [[maybe_unused]] int rc =
+          ::madvise(this->baseAddress, size,
+                    MADV_NOHUGEPAGE | MADV_DONTFORK | MADV_RANDOM);
+      assert(rc == 0 && "madvise failed");
+    }
+#elif defined(__FreeBSD__)
+    {
+      [[maybe_unused]] int rc =
+          ::minherit(this->baseAddress, size, INHERIT_NONE);
+      assert(rc == 0 && "minherit failed");
+    }
+#elif defined(__APPLE__)
+    {
+      [[maybe_unused]] int rc =
+          ::minherit(this->baseAddress, size, VM_INHERIT_NONE);
+      assert(rc == 0 && "minherit failed");
+    }
+#endif
+
+    return true;
+  }
+
+public:
+  Mapping() = default;
+
+  explicit Mapping(std::size_t size) noexcept : Mapping(0, size) {}
+
+  Mapping(std::uintptr_t baseAddress, std::size_t size) noexcept : size(size) {
+    try_map(baseAddress);
+    assert(*this && "failed to allocate mapping");
+    if (!*this) {
+      std::abort();
+    }
+  }
+
+  Mapping(Mapping const &) = delete;
+  Mapping &operator=(Mapping const &) = delete;
+
+  Mapping(Mapping &&other) noexcept
+      : baseAddress(other.baseAddress), size(other.size) {
+    other.baseAddress = MAP_FAILED;
+    other.size = 0;
+  }
+  Mapping &operator=(Mapping &&other) noexcept {
+    if (&other != this) {
+      using std::swap;
+      swap(other.baseAddress, baseAddress);
+      swap(other.size, size);
+    }
+    return *this;
+  }
+
+  [[nodiscard]] void *getBaseAddress() const noexcept {
+    assert(*this && "Invalid mapping");
+    return baseAddress;
+  }
+
+  [[nodiscard]] std::size_t getSize() const noexcept { return size; }
+
+  void clear() {
+    assert(*this && "Invalid mapping");
+
+#if defined(__linux__)
+    [[maybe_unused]] int rc = ::madvise(baseAddress, size, MADV_DONTNEED);
+    assert(rc == 0 && "madvise failed");
+#else
+    auto address = reinterpret_cast<std::uintptr_t>(baseAddress);
+    [[maybe_unused]] int rc = ::munmap(baseAddress, size);
+    assert(rc == 0 && "munmap failed");
+    baseAddress = MAP_FAILED;
+    auto success = try_map(address);
+    assert(success && "could not recreate the mapping");
+#endif
+  }
+
+  explicit operator bool() const noexcept { return baseAddress != MAP_FAILED; }
+
+  ~Mapping() {
+    if (*this) {
+      [[maybe_unused]] int rc = ::munmap(baseAddress, size);
+      assert(rc == 0 && "munmap failed");
+    }
+  }
+};
+} // namespace klee::kdalloc
+
+#endif