perf(deque): remove redundant modulo and bounds checks

Author: peter-jerry-ye

deque/deque.mbt

@@ -15,16 +15,73 @@
 ///|
 fn[T] set_null(buffer : UninitializedArray[T], index : Int) = "%fixedarray.set_null"
 
+///|
+/// The deque representation invariant guarantees these indices are in bounds.
+/// Keep these private so callers cannot bypass bounds checks.
+#inline
+fn[T] unsafe_buffer_get(buffer : UninitializedArray[T], index : Int) -> T = "%fixedarray.unsafe_get"
+
+///|
+#inline
+fn[T] unsafe_buffer_set(
+  buffer : UninitializedArray[T],
+  index : Int,
+  value : T,
+) -> Unit = "%fixedarray.unsafe_set"
+
 ///|
 fn[A] new_deque(capacity : Int) -> Deque[A] {
   { buf: UninitializedArray::make(capacity), len: 0, head: 0 }
 }
 
+///|
+/// Wraps an in-buffer offset without division. Comparing the offset with the
+/// contiguous space first also avoids overflowing `head + offset`.
+#inline
+fn wrap_index(
+  head : Int,
+  offset : Int,
+  capacity : Int,
+) -> Int where {
+  proof_require: wrap_index_pre(head, offset, capacity),
+  proof_ensure: result => wrap_index_post(head, offset, capacity, result),
+  proof_ensure: result => wrap_index_mod_equiv(head, offset, capacity, result),
+} {
+  let contiguous = capacity - head
+  if offset < contiguous {
+    proof_assert (head + offset) % capacity == head + offset
+    head + offset
+  } else {
+    proof_assert (head + offset) % capacity == offset - contiguous
+    offset - contiguous
+  }
+}
+
+///|
+/// Moves an in-buffer index back by one without division.
+#inline
+fn decrement_index(
+  index : Int,
+  capacity : Int,
+) -> Int where {
+  proof_require: decrement_index_pre(index, capacity),
+  proof_ensure: result => decrement_index_post(index, capacity, result),
+  proof_ensure: result => decrement_index_mod_equiv(index, capacity, result),
+} {
+  if index == 0 {
+    proof_assert (index + capacity - 1) % capacity == capacity - 1
+    capacity - 1
+  } else {
+    proof_assert (index + capacity - 1) % capacity == index - 1
+    index - 1
+  }
+}
+
 ///|
 /// Computes the tail index (index of last element) on demand.
 /// Only valid when len > 0.
 fn[A] Deque::tail_index(self : Deque[A]) -> Int {
-  (self.head + self.len - 1) % self.buf.length()
+  wrap_index(self.head, self.len - 1, self.buf.length())
 }
 
 ///|
@@ -286,15 +343,15 @@ pub fn[A] Deque::blit_to(
     }
     // Copy in reverse order
     for i in len>..0 {
-      let dst_idx = (dst.head + dst_offset + i) % dst.buf.length()
-      let src_idx = (self.head + src_offset + i) % self.buf.length()
-      dst.buf[dst_idx] = self.buf[src_idx]
+      let dst_idx = wrap_index(dst.head, dst_offset + i, dst.buf.length())
+      let src_idx = wrap_index(self.head, src_offset + i, self.buf.length())
+      unsafe_buffer_set(dst.buf, dst_idx, unsafe_buffer_get(self.buf, src_idx))
     }
   } else {
     for i in 0..<len {
-      let dst_idx = (dst.head + dst_offset + i) % dst.buf.length()
-      let src_idx = (self.head + src_offset + i) % self.buf.length()
-      dst.buf[dst_idx] = self.buf[src_idx]
+      let dst_idx = wrap_index(dst.head, dst_offset + i, dst.buf.length())
+      let src_idx = wrap_index(self.head, src_offset + i, self.buf.length())
+      unsafe_buffer_set(dst.buf, dst_idx, unsafe_buffer_get(self.buf, src_idx))
       if dst_offset + i >= dst_len {
         dst.len += 1
       }
@@ -359,9 +416,13 @@ pub fn[A] Deque::append(self : Deque[A], other : Deque[A]) -> Unit {
   let cap = self.buf.length()
   // Use captured state to read from other, avoiding aliasing issues
   for i in 0..<other_len {
-    let read_idx = (other_head + i) % other_buf_len
-    let write_idx = (self.head + self.len) % cap
-    self.buf[write_idx] = other_buf[read_idx]
+    let read_idx = wrap_index(other_head, i, other_buf_len)
+    let write_idx = wrap_index(self.head, self.len, cap)
+    unsafe_buffer_set(
+      self.buf,
+      write_idx,
+      unsafe_buffer_get(other_buf, read_idx),
+    )
     self.len += 1
   }
 }
@@ -422,22 +483,22 @@ pub fn[A] Deque::insert(self : Deque[A], index : Int, value : A) -> Unit {
   let cap = self.buf.length()
   if index < self.len / 2 {
     // Shift front elements left
-    let new_head = (self.head - 1 + cap) % cap
+    let new_head = decrement_index(self.head, cap)
     for i in 0..<index {
-      let to = (new_head + i) % cap
-      let from = (self.head + i) % cap
-      self.buf[to] = self.buf[from]
+      let to = wrap_index(new_head, i, cap)
+      let from = wrap_index(self.head, i, cap)
+      unsafe_buffer_set(self.buf, to, unsafe_buffer_get(self.buf, from))
     }
     self.head = new_head
   } else {
     // Shift back elements right
     for i = self.len; i > index; i = i - 1 {
-      let from = (self.head + i - 1) % cap
-      let to = (self.head + i) % cap
-      self.buf[to] = self.buf[from]
+      let from = wrap_index(self.head, i - 1, cap)
+      let to = wrap_index(self.head, i, cap)
+      unsafe_buffer_set(self.buf, to, unsafe_buffer_get(self.buf, from))
     }
   }
-  self.buf[(self.head + index) % cap] = value
+  unsafe_buffer_set(self.buf, wrap_index(self.head, index, cap), value)
   self.len += 1
 }
 
@@ -498,21 +559,21 @@ pub fn[A] Deque::remove(self : Deque[A], index : Int) -> A {
   let cap = self.buf.length()
   if index < self.len / 2 {
     // Shift front elements right
-    let new_head = (self.head + 1) % cap
+    let new_head = wrap_index(self.head, 1, cap)
     for i in index>..0 {
-      let to = (self.head + i + 1) % cap
-      let from = (self.head + i) % cap
-      self.buf[to] = self.buf[from]
+      let to = wrap_index(self.head, i + 1, cap)
+      let from = wrap_index(self.head, i, cap)
+      unsafe_buffer_set(self.buf, to, unsafe_buffer_get(self.buf, from))
     }
     set_null(self.buf, self.head)
     self.head = new_head
   } else {
     // Shift back elements left
-    let tail_idx = (self.head + self.len - 1) % cap
+    let tail_idx = wrap_index(self.head, self.len - 1, cap)
     for i in (index + 1)..<self.len {
-      let to = (self.head + i - 1) % cap
-      let from = (self.head + i) % cap
-      self.buf[to] = self.buf[from]
+      let to = wrap_index(self.head, i - 1, cap)
+      let from = wrap_index(self.head, i, cap)
+      unsafe_buffer_set(self.buf, to, unsafe_buffer_get(self.buf, from))
     }
     set_null(self.buf, tail_idx)
   }
@@ -596,7 +657,7 @@ pub fn[A] Deque::front(self : Deque[A]) -> A? {
   if self.len == 0 {
     None
   } else {
-    Some(self.buf[self.head])
+    Some(unsafe_buffer_get(self.buf, self.head))
   }
 }
 
@@ -614,7 +675,7 @@ pub fn[A] Deque::back(self : Deque[A]) -> A? {
   if self.len == 0 {
     None
   } else {
-    Some(self.buf[self.tail_index()])
+    Some(unsafe_buffer_get(self.buf, self.tail_index()))
   }
 }
 
@@ -636,8 +697,8 @@ pub fn[A] Deque::push_front(self : Deque[A], value : A) -> Unit {
     self.realloc()
   }
   let cap = self.buf.length()
-  self.head = (self.head - 1 + cap) % cap
-  self.buf[self.head] = value
+  self.head = decrement_index(self.head, cap)
+  unsafe_buffer_set(self.buf, self.head, value)
   self.len += 1
 }
 
@@ -659,8 +720,8 @@ pub fn[A] Deque::push_back(self : Deque[A], value : A) -> Unit {
     self.realloc()
   }
   let cap = self.buf.length()
-  let write_idx = (self.head + self.len) % cap
-  self.buf[write_idx] = value
+  let write_idx = wrap_index(self.head, self.len, cap)
+  unsafe_buffer_set(self.buf, write_idx, value)
   self.len += 1
 }
 
@@ -682,7 +743,7 @@ pub fn[A] Deque::unsafe_pop_front(self : Deque[A]) -> Unit {
   guard self.len > 0 else { abort("The deque is empty!") }
   set_null(self.buf, self.head)
   let cap = self.buf.length()
-  self.head = (self.head + 1) % cap
+  self.head = wrap_index(self.head, 1, cap)
   self.len -= 1
 }
 
@@ -771,10 +832,10 @@ pub fn[A] Deque::unsafe_pop_back(self : Deque[A]) -> Unit {
 /// ```
 pub fn[A] Deque::pop_front(self : Deque[A]) -> A? {
   guard self.len > 0 else { return None }
-  let value = self.buf[self.head]
+  let value = unsafe_buffer_get(self.buf, self.head)
   set_null(self.buf, self.head)
   let cap = self.buf.length()
-  self.head = (self.head + 1) % cap
+  self.head = wrap_index(self.head, 1, cap)
   self.len -= 1
   Some(value)
 }
@@ -792,7 +853,7 @@ pub fn[A] Deque::pop_front(self : Deque[A]) -> A? {
 pub fn[A] Deque::pop_back(self : Deque[A]) -> A? {
   guard self.len > 0 else { return None }
   let tail_idx = self.tail_index()
-  let value = self.buf[tail_idx]
+  let value = unsafe_buffer_get(self.buf, tail_idx)
   set_null(self.buf, tail_idx)
   self.len -= 1
   Some(value)
@@ -815,11 +876,8 @@ pub fn[A] Deque::at(self : Deque[A], index : Int) -> A {
   if index < 0 || index >= self.len {
     index_out_of_bounds(self.len, index)
   }
-  if self.head + index < self.buf.length() {
-    self.buf[self.head + index]
-  } else {
-    self.buf[self.head + index - self.buf.length()]
-  }
+  let physical_index = wrap_index(self.head, index, self.buf.length())
+  unsafe_buffer_get(self.buf, physical_index)
 }
 
 ///|
@@ -840,11 +898,8 @@ pub fn[A] Deque::set(self : Deque[A], index : Int, value : A) -> Unit {
   if index < 0 || index >= self.len {
     index_out_of_bounds(self.len, index)
   }
-  if self.head + index < self.buf.length() {
-    self.buf[self.head + index] = value
-  } else {
-    self.buf[self.head + index - self.buf.length()] = value
-  }
+  let physical_index = wrap_index(self.head, index, self.buf.length())
+  unsafe_buffer_set(self.buf, physical_index, value)
 }
 
 ///|
@@ -2312,8 +2367,8 @@ pub fn[A] Deque::binary_search_by(
 /// Safe element access with bounds checking
 pub fn[A] Deque::get(self : Deque[A], index : Int) -> A? {
   if index >= 0 && index < self.len {
-    let physical_index = (self.head + index) % self.buf.length()
-    Some(self.buf[physical_index])
+    let physical_index = wrap_index(self.head, index, self.buf.length())
+    Some(unsafe_buffer_get(self.buf, physical_index))
   } else {
     None
   }
@@ -2435,10 +2490,12 @@ pub fn[A] Deque::rev_in_place(self : Deque[A]) -> Unit {
   guard self.len > 0 else { return }
   let cap = self.buf.length()
   for _ in 0..<(self.len / 2); left = self.head, right = self.tail_index() {
-    let temp = self.buf[left]
-    self.buf[left] = self.buf[right]
-    self.buf[right] = temp
-    continue (left + 1) % cap, (right - 1 + cap) % cap
+    let temp = unsafe_buffer_get(self.buf, left)
+    unsafe_buffer_set(self.buf, left, unsafe_buffer_get(self.buf, right))
+    unsafe_buffer_set(self.buf, right, temp)
+    let left = wrap_index(left, 1, cap)
+    let right = decrement_index(right, cap)
+    continue left, right
   }
 }
 
@@ -2476,8 +2533,8 @@ pub fn[A] Deque::rev(self : Deque[A]) -> Deque[A] {
   let new_buf = UninitializedArray::make(len)
   // Copy elements in reverse order
   for i in 0..<len {
-    let src_idx = (self.head + len - i - 1) % self.buf.length()
-    new_buf[i] = self.buf[src_idx]
+    let src_idx = wrap_index(self.head, len - i - 1, self.buf.length())
+    unsafe_buffer_set(new_buf, i, unsafe_buffer_get(self.buf, src_idx))
   }
   // Create new deque with reversed elements
   { buf: new_buf, len, head: 0 }

deque/deque_test.mbt

@@ -2798,3 +2798,40 @@ test "Deque::append/self_alias" {
   dq.append(dq)
   debug_inspect(dq.to_array(), content="[1, 2, 3, 1, 2, 3]")
 }
+
+///|
+test "deque operations match array model across wrap boundaries" {
+  for initial_capacity in [0, 1, 3, 5, 15] {
+    let deque = @deque.Deque([], capacity=initial_capacity)
+    let model : Array[Int] = []
+    for step in 0..<1300 {
+      let cycle = step / 13
+      match step % 13 {
+        0 | 1 | 4 => {
+          deque.push_back(step)
+          model.push(step)
+        }
+        2 | 5 => {
+          deque.push_front(step)
+          model.insert(0, step)
+        }
+        3 => {
+          let index = (cycle + initial_capacity) % (model.length() + 1)
+          deque.insert(index, step)
+          model.insert(index, step)
+        }
+        6 | 10 => @test.assert_eq(deque.pop_front(), Some(model.remove(0)))
+        7 | 12 => @test.assert_eq(deque.pop_back(), model.pop())
+        8 | 11 => {
+          let index = (cycle + initial_capacity) % model.length()
+          @test.assert_eq(deque.remove(index), model.remove(index))
+        }
+        _ => {
+          deque.rev_in_place()
+          model.rev_in_place()
+        }
+      }
+      @test.assert_eq(deque.to_array(), model)
+    }
+  }
+}