3 files changed, 64 insertions, 0 deletions

diff --git a/include/klee/KDAlloc/allocator.h b/include/klee/KDAlloc/allocator.h
index d7b1e0e2..2d8a07ac 100644
--- a/include/klee/KDAlloc/allocator.h
+++ b/include/klee/KDAlloc/allocator.h
@@ -61,6 +61,19 @@ public:
       return meta.size();
     }
 
+    [[nodiscard]] inline int
+    convertPtrToBinIndex(void const *const p) noexcept {
+      for (std::size_t i = 0; i < sizedBins.size(); ++i) {
+        if (p >= sizedBins[i].mapping_begin() &&
+            p < sizedBins[i].mapping_end()) {
+          return i;
+        }
+      }
+      assert(p >= largeObjectBin.mapping_begin() &&
+             p < largeObjectBin.mapping_end());
+      return meta.size();
+    }
+
   public:
     mutable class klee::ReferenceCounter _refCount;
 
@@ -134,6 +147,19 @@ public:
     }
   }
 
+  std::size_t getSize(void const *const ptr) const noexcept {
+    assert(!!ptr);
+
+    auto bin = control->convertPtrToBinIndex(ptr);
+    traceLine("Getting size for ", ptr, " in bin ", bin);
+
+    if (bin < static_cast<int>(sizedBins.size())) {
+      return Control::meta[bin];
+    } else {
+      return largeObjectBin.getSize(control->largeObjectBin, ptr);
+    }
+  }
+
   LocationInfo location_info(void const *const ptr,
                              std::size_t const size) const noexcept {
     assert(*this && "Invalid allocator");
diff --git a/include/klee/KDAlloc/suballocators/loh.h b/include/klee/KDAlloc/suballocators/loh.h
index 4b99942b..d52306fc 100644
--- a/include/klee/KDAlloc/suballocators/loh.h
+++ b/include/klee/KDAlloc/suballocators/loh.h
@@ -212,6 +212,16 @@ public:
     return out << "[LOH] ";
   }
 
+  std::size_t getSize(Control const &control,
+                      void const *const ptr) const noexcept {
+    assert(control.mapping_begin() <= ptr && ptr < control.mapping_end() &&
+           "This property should have been ensured by the caller");
+    assert(!!data &&
+           "Can only get size of an object if objects already exist...");
+
+    return data->regions.getSize(static_cast<char const *>(ptr));
+  }
+
   LocationInfo getLocationInfo(Control const &control, void const *const ptr,
                                std::size_t const size) const noexcept {
     assert(control.mapping_begin() <= ptr &&
diff --git a/include/klee/KDAlloc/suballocators/sized_regions.h b/include/klee/KDAlloc/suballocators/sized_regions.h
index 7adeea45..7236a619 100644
--- a/include/klee/KDAlloc/suballocators/sized_regions.h
+++ b/include/klee/KDAlloc/suballocators/sized_regions.h
@@ -159,6 +159,34 @@ public:
 
   [[nodiscard]] bool isEmpty() const noexcept { return !root; }
 
+  [[nodiscard]] std::size_t getSize(char const *const address) const noexcept {
+    assert(root && "Cannot get size from an empty treap");
+
+    Node const *currentNode = &*root;
+    Node const *closestPredecessor = nullptr;
+    Node const *closestSuccessor = nullptr;
+    while (currentNode) {
+      if (address < currentNode->getBaseAddress()) {
+        assert(!closestSuccessor || currentNode->getBaseAddress() <
+                                        closestSuccessor->getBaseAddress());
+        closestSuccessor = currentNode;
+        currentNode = &*currentNode->lhs;
+      } else {
+        assert(!closestPredecessor || currentNode->getBaseAddress() >
+                                          closestPredecessor->getBaseAddress());
+        closestPredecessor = currentNode;
+        currentNode = &*currentNode->rhs;
+      }
+    }
+
+    assert(closestPredecessor && closestSuccessor &&
+           "address must be in between two regions");
+
+    return closestSuccessor->getBaseAddress() -
+           (closestPredecessor->getBaseAddress() +
+            closestPredecessor->getSize());
+  }
+
   /// Computes the LocationInfo. This functionality really belongs to the
   /// `LargeObjectAllocator`, as it assumes that this treap contains free
   /// regions in between allocations. It also knows that there is a redzone at