Skip to content

Commit

Permalink
[mojo-stdlib] Add index bounds checking to List struct methods
Browse files Browse the repository at this point in the history
Signed-off-by: Kaushal Phulgirkar <kaushalpp01@gmail.com>
  • Loading branch information
StandinKP committed May 7, 2024
1 parent 57a9113 commit 4c83594
Showing 1 changed file with 12 additions and 8 deletions.
20 changes: 12 additions & 8 deletions stdlib/src/collections/list.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,8 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
value: The value to insert.
"""
debug_assert(i <= self.size, "insert index out of range")

if i < 0 or i > len(self):
abort("Index out of bounds")
var normalized_idx = self._normalize_index(i)

var earlier_idx = len(self)
Expand Down Expand Up @@ -289,7 +290,8 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
The popped value.
"""
debug_assert(-len(self) <= i < len(self), "pop index out of range")

if i > len(self) - 1:
abort("Index out of bounds")
var normalized_idx = self._normalize_index(i)

var ret_val = move_from_pointee(self.data + normalized_idx)
Expand Down Expand Up @@ -485,7 +487,8 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
value: The value to assign.
"""
debug_assert(-self.size <= i < self.size, "index must be within bounds")

if i < 0 or i >= len(self):
abort("Index out of bounds")
var normalized_idx = self._normalize_index(i)

destroy_pointee(self.data + normalized_idx)
Expand Down Expand Up @@ -547,7 +550,8 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
A copy of the element at the given index.
"""
debug_assert(-self.size <= i < self.size, "index must be within bounds")

if i < 0 or i >= len(self):
abort("Index out of bounds")
var normalized_idx = self._normalize_index(i)

return (self.data + normalized_idx)[]
Expand All @@ -564,7 +568,9 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
Returns:
An immutable reference to the element at the given index.
"""
var normalized_idx = Reference(self)[]._normalize_index(i)
if i < 0 or i >= len(self[]):
abort("Index out of bounds")
var normalized_idx = Reference(self)._normalize_index(i)

return (self[].data + normalized_idx)[]

Expand Down Expand Up @@ -664,6 +670,4 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):

@always_inline
fn _normalize_index(self, index: Int) -> Int:
if index < 0:
return _max(0, len(self) + index)
return index
return index if index >= 0 else index + len(self)

0 comments on commit 4c83594

Please sign in to comment.