[stdlib] Use Indexer in normalize_index

Optimize normalize_index for UInt Indexer type.

Signed-off-by: Yinon Burgansky <yinonburgansky@gmail.com>

stdlib/src/collections/_index_normalization.mojo

@@ -13,16 +13,19 @@
 """The utilities provided in this module help normalize the access
 to data elements in arrays."""
 
+from sys.intrinsics import _type_is_eq
+
 
 @always_inline
 fn normalize_index[
-    ContainerType: Sized, //, container_name: StringLiteral
-](idx: Int, container: ContainerType) -> Int:
+    I: Indexer, ContainerType: Sized, //, container_name: StringLiteral
+](idx: I, container: ContainerType) -> UInt:
     """Normalize the given index value to a valid index value for the given container length.
 
     If the provided value is negative, the `index + container_length` is returned.
 
     Parameters:
+        I: A type that can be used as an index.
         ContainerType: The type of the container. Must have a `__len__` method.
         container_name: The name of the container. Used for the error message.
 
@@ -39,18 +42,33 @@ fn normalize_index[
         container_name,
         " that has 0 elements",
     )
-    debug_assert[assert_mode="safe", cpu_only=True](
-        -len(container) <= idx < len(container),
-        container_name,
-        " has length: ",
-        len(container),
-        " index out of bounds: ",
-        idx,
-        " should be between ",
-        -len(container),
-        " and ",
-        len(container) - 1,
-    )
-    if idx >= 0:
-        return idx
-    return idx + len(container)
+
+    @parameter
+    if _type_is_eq[I, UInt]():
+        var i = rebind[UInt](idx)
+        debug_assert[assert_mode="safe", cpu_only=True](
+            i < len(container),
+            container_name,
+            " index out of bounds: ",
+            i,
+            " should be less than ",
+            len(container),
+        )
+        return i
+    else:
+        var i = Int(idx)
+        debug_assert[assert_mode="safe", cpu_only=True](
+            -len(container) <= i < len(container),
+            container_name,
+            " has length: ",
+            len(container),
+            " index out of bounds: ",
+            i,
+            " should be between ",
+            -len(container),
+            " and ",
+            len(container) - 1,
+        )
+        if i >= 0:
+            return i
+        return i + len(container)

stdlib/src/collections/string/string.mojo

@@ -871,7 +871,7 @@ struct String(
             A new string containing the character at the specified position.
         """
         # TODO(#933): implement this for unicode when we support llvm intrinsic evaluation at compile time
-        var normalized_idx = normalize_index["String"](index(idx), self)
+        var normalized_idx = normalize_index["String"](idx, self)
         var buf = Self._buffer_type(capacity=1)
         buf.append(self._buffer[normalized_idx])
         buf.append(0)

stdlib/test/collections/test_index_normalization.mojo

@@ -21,12 +21,16 @@ def test_out_of_bounds_message():
     l = List[Int](1, 2)
     # CHECK: index out of bounds: 2
     _ = normalize_index["List"](2, l)
+    # CHECK: index out of bounds: 2
+    _ = normalize_index["List"](UInt(2), l)
     # CHECK: index out of bounds: -3
     _ = normalize_index["List"](-3, l)
 
     l2 = List[Int]()
     # CHECK: indexing into a List that has 0 elements
     _ = normalize_index["List"](2, l2)
+    # CHECK: indexing into a List that has 0 elements
+    _ = normalize_index["List"](UInt(2), l2)
 
 
 def test_normalize_index():
@@ -39,6 +43,8 @@ def test_normalize_index():
     assert_equal(normalize_index[""](1, container), 1)
     assert_equal(normalize_index[""](2, container), 2)
     assert_equal(normalize_index[""](3, container), 3)
+    assert_equal(normalize_index[""](UInt(0), container), 0)
+    assert_equal(normalize_index[""](UInt(3), container), 3)
 
 
 def main():