Skip to content

Commit

Permalink
Added fn to normalize index calculation
Browse files Browse the repository at this point in the history
Signed-off-by: Kaushal Phulgirkar <kaushalpp01@gmail.com>
  • Loading branch information
StandinKP committed Apr 30, 2024
1 parent ee52186 commit 6720cad
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 16 deletions.
26 changes: 11 additions & 15 deletions stdlib/src/collections/list.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -235,9 +235,7 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
"""
debug_assert(i <= self.size, "insert index out of range")

var normalized_idx = i
if i < 0:
normalized_idx = _max(0, len(self) + i)
var normalized_idx = self._normalize_index(i)

var earlier_idx = len(self)
var later_idx = len(self) - 1
Expand Down Expand Up @@ -307,9 +305,7 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
"""
debug_assert(-len(self) <= i < len(self), "pop index out of range")

var normalized_idx = i
if i < 0:
normalized_idx += len(self)
var normalized_idx = self._normalize_index(i)

var ret_val = move_from_pointee(self.data + normalized_idx)
for j in range(normalized_idx + 1, self.size):
Expand Down Expand Up @@ -505,9 +501,7 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
"""
debug_assert(-self.size <= i < self.size, "index must be within bounds")

var normalized_idx = i
if i < 0:
normalized_idx += len(self)
var normalized_idx = self._normalize_index(i)

destroy_pointee(self.data + normalized_idx)
initialize_pointee_move(self.data + normalized_idx, value^)
Expand Down Expand Up @@ -569,9 +563,7 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
"""
debug_assert(-self.size <= i < self.size, "index must be within bounds")

var normalized_idx = i
if i < 0:
normalized_idx += len(self)
var normalized_idx = self._normalize_index(i)

return (self.data + normalized_idx)[]

Expand All @@ -590,9 +582,7 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
Returns:
An immutable reference to the element at the given index.
"""
var normalized_idx = i
if i < 0:
normalized_idx += Reference(self)[].size
var normalized_idx = Reference(self)[]._normalize_index(i)

var offset_ptr = Reference(self)[].data + normalized_idx
return offset_ptr[]
Expand Down Expand Up @@ -666,3 +656,9 @@ struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
result += ", "
result += "]"
return result

@always_inline
fn _normalize_index(self, index: Int) -> Int:
if index < 0:
return _max(0, len(self) + index)
return index
2 changes: 1 addition & 1 deletion stdlib/test/collections/test_list.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -549,7 +549,7 @@ def test_2d_dynamic_list():

def test_list_explicit_copy():
var list = List[CopyCounter]()
list.append(CopyCounter()^)
list.append(CopyCounter())
var list_copy = List(list)
assert_equal(0, list.__get_ref(0)[].copy_count)
assert_equal(1, list_copy.__get_ref(0)[].copy_count)
Expand Down

0 comments on commit 6720cad

Please sign in to comment.