Fix extent overflows (#502)

dsharlet · Nov 27, 2024 · d8a6cf5 · d8a6cf5
1 parent 9d4aced
commit d8a6cf5
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 11 deletions.
diff --git a/runtime/buffer.cc b/runtime/buffer.cc
@@ -27,7 +27,7 @@ index_t alloc_extent(const dim& dim) {
     // TODO: We can do better than this if the dim doesn't cross a fold boundary.
     return dim.fold_factor();
   } else {
-    return dim.max() >= dim.min() ? dim.extent() : 0;
+    return dim.extent();
   }
 }
 
@@ -263,8 +263,8 @@ void copy_impl(raw_buffer& src, raw_buffer& dst, const void* padding) {
     src.crop(0, dst_dim0.min(), dst_dim0.max());
 
     const index_t padded_size = dst_dim0.extent() * elem_size;
-    const index_t pad_before = (src_dim0.begin() - dst_dim0.begin()) * elem_size;
-    const index_t pad_after = (dst_dim0.end() - src_dim0.end()) * elem_size;
+    const index_t pad_before = src_dim0.begin() > dst_dim0.begin() ? (src_dim0.begin() - dst_dim0.begin()) * elem_size : 0;
+    const index_t pad_after = dst_dim0.end() > src_dim0.end() ? (dst_dim0.end() - src_dim0.end()) * elem_size : 0;
     const index_t size = padded_size - pad_before - pad_after;
     dst.slice(0);
     src.slice(0);

diff --git a/runtime/buffer.h b/runtime/buffer.h
@@ -19,8 +19,7 @@ namespace slinky {
 // alias to `long` under some compilers which can cause some not so fun
 // overloading issues with expr(), so let's use std::conditional
 // instead to make it an exact alias of either int32_t or int64_t.
-using index_t =
-    std::conditional<sizeof(void*) == 4, std::int32_t, std::int64_t>::type;
+using index_t = std::conditional<sizeof(void*) == 4, std::int32_t, std::int64_t>::type;
 
 // Helper to offset a pointer by a number of bytes.
 template <typename T>
@@ -65,10 +64,7 @@ class dim {
   index_t max() const { return max_; }
   index_t begin() const { return min_; }
   index_t end() const { return max_ + 1; }
-  index_t extent() const {
-    assert(!sub_overflows<index_t>(max_, min_) && !add_overflows<index_t>(max_ - min_, 1));
-    return (max_ - min_) + 1;
-  }
+  index_t extent() const { return end() > begin() ? end() - begin() : 0; }
   index_t stride() const { return stride_; }
   index_t fold_factor() const { return fold_factor_; }
   bool empty() const { return max_ < min_; }

diff --git a/runtime/print.cc b/runtime/print.cc
@@ -360,8 +360,7 @@ std::ostream& operator<<(std::ostream& os, const raw_buffer& buf) {
 }
 
 std::ostream& operator<<(std::ostream& os, const dim& d) {
-  const index_t extent = d.max() >= d.min() ? d.extent() : 0;
-  os << "{min=" << d.min() << ", max=" << d.max() << ", extent=" << extent << ", stride=" << d.stride();
+  os << "{min=" << d.min() << ", max=" << d.max() << ", extent=" << d.extent() << ", stride=" << d.stride();
   if (d.fold_factor() != dim::unfolded) {
     os << ", fold_factor=" << d.fold_factor();
   }