[External] [stdlib] Add Indexer trait (#40103)

[External] [stdlib] Add Indexer trait

First half of the split off from modularml#2384.
Partially addresses modularml#2337

Introduces the `Indexer` trait and applies to `Int`, `IntLiteral`,
`Bool`, and integral scalar `SIMD` types.

---------

Co-authored-by: bgreni <42788181+bgreni@users.noreply.github.com>
Closes modularml#2685
MODULAR_ORIG_COMMIT_REV_ID: 224c700e9ff9b28abf9e7d72ff3b3271d7e9de09

docs/changelog.md

@@ -238,16 +238,39 @@ what we publish.
   ```
 
 - `List` can now be converted to a `String` with a simplified syntax:
-  
+
   ```mojo
   var my_list = List[Int](2, 3)
-  print(my_list.__str__())  # prints [2, 3] 
+  print(my_list.__str__())  # prints [2, 3]
   ```
 
   Note that `List` doesn't conform to the `Stringable` trait yet so you cannot
   use `str(my_list)` yet.
     ([PR #2673](https://github.com/modularml/mojo/pull/2673) by [@gabrieldemarmiesse](https://github.com/gabrieldemarmiesse))
 
+- Added the `Indexer` trait to denote types that implement the `__index__()`
+  method which allow these types to be accepted in common `__getitem__` and
+  `__setitem__` implementations, as well as allow a new builtin `index` function
+  to be called on them. For example:
+
+  ```mojo
+  @value
+  struct AlwaysZero(Indexer):
+      fn __index__(self) -> Int:
+          return 0
+
+  struct MyList:
+      var data: List[Int]
+
+      fn __init__(inout self):
+          self.data = List[Int](1, 2, 3, 4)
+
+      fn __getitem__[T: Indexer](self, idx: T) -> T:
+          return self.data[index(idx)]
+
+  print(MyList()[AlwaysZero()])  # prints `1`
+  ```
+
 ### 🦋 Changed
 
 - The `let` keyword has been completely removed from the language. We previously

stdlib/src/builtin/bool.mojo

@@ -59,7 +59,12 @@ trait Boolable:
 @value
 @register_passable("trivial")
 struct Bool(
-    Stringable, CollectionElement, Boolable, EqualityComparable, Intable
+    Stringable,
+    CollectionElement,
+    Boolable,
+    EqualityComparable,
+    Intable,
+    Indexer,
 ):
     """The primitive Bool scalar value used in Mojo."""
 
@@ -151,6 +156,15 @@ struct Bool(
         """
         return __mlir_op.`pop.select`[_type=Int](self.value, Int(1), Int(0))
 
+    @always_inline("nodebug")
+    fn __index__(self) -> Int:
+        """Convert this Bool to an integer for indexing purposes.
+
+        Returns:
+            1 if the Bool is True, 0 otherwise.
+        """
+        return self.__int__()
+
     @always_inline("nodebug")
     fn __eq__(self, rhs: Bool) -> Bool:
         """Compare this Bool to RHS.

stdlib/src/builtin/int.mojo

@@ -27,6 +27,52 @@ from utils._visualizers import lldb_formatter_wrapping_type
 from utils._format import Formattable, Formatter
 from utils.inlined_string import _ArrayMem
 
+# ===----------------------------------------------------------------------=== #
+#  Indexer
+# ===----------------------------------------------------------------------=== #
+
+
+trait Indexer:
+    """This trait denotes a type that can be used to index a container that
+    handles integral index values.
+
+    This solves the issue of being able to index data structures such as `List`
+    with the various integral types without being too broad and allowing types
+    that are coercible to `Int` (e.g. floating point values that have `__int__`
+    method). In contrast to `Intable`, types conforming to `Indexer` must be
+    convertible to `Int` in a lossless way.
+    """
+
+    fn __index__(self) -> Int:
+        """Return the index value.
+
+        Returns:
+            The index value of the object.
+        """
+        ...
+
+
+# ===----------------------------------------------------------------------=== #
+#  index
+# ===----------------------------------------------------------------------=== #
+
+
+@always_inline("nodebug")
+fn index[T: Indexer](idx: T) -> Int:
+    """Returns the value of `__index__` for the given value.
+
+    Parameters:
+        T: A type conforming to the `Indexer` trait.
+
+    Args:
+        idx: The value.
+
+    Returns:
+        An `Int` respresenting the index value.
+    """
+    return idx.__index__()
+
+
 # ===----------------------------------------------------------------------=== #
 #  Intable
 # ===----------------------------------------------------------------------=== #
@@ -205,6 +251,7 @@ struct Int(
     Roundable,
     Stringable,
     Truncable,
+    Indexer,
 ):
     """This type represents an integer value."""
 

stdlib/src/builtin/int_literal.mojo

@@ -29,6 +29,7 @@ struct IntLiteral(
     Roundable,
     Stringable,
     Truncable,
+    Indexer,
 ):
     """This type represents a static integer literal value with
     infinite precision.  They can't be materialized at runtime and

stdlib/src/builtin/simd.mojo

@@ -139,6 +139,7 @@ struct SIMD[type: DType, size: Int = simdwidthof[type]()](
     Sized,
     Stringable,
     Truncable,
+    Indexer,
 ):
     """Represents a small vector that is backed by a hardware vector element.
 
@@ -179,6 +180,22 @@ struct SIMD[type: DType, size: Int = simdwidthof[type]()](
         _simd_construction_checks[type, size]()
         self = _unchecked_zero[type, size]()
 
+    @always_inline("nodebug")
+    fn __index__(self) -> Int:
+        """Returns the value as an int if it is an integral value.
+
+        Constraints:
+            Must be a scalar integral value.
+
+        Returns:
+            The value as an integer.
+        """
+        constrained[
+            type.is_integral() or type.is_bool(),
+            "expected integral or bool type",
+        ]()
+        return self.__int__()
+
     @always_inline("nodebug")
     fn __init__(inout self, value: SIMD[DType.float64, 1]):
         """Initializes the SIMD vector with a float.

stdlib/test/builtin/test_bool.mojo

@@ -100,10 +100,16 @@ def test_neg():
     assert_equal(0, -False)
 
 
+def test_indexer():
+    assert_equal(1, Bool.__index__(True))
+    assert_equal(0, Bool.__index__(False))
+
+
 def main():
     test_bool_cast_to_int()
     test_bool_none()
     test_convert_from_boolable()
     test_bool_to_string()
     test_bitwise()
     test_neg()
+    test_indexer()

stdlib/test/builtin/test_int.mojo

@@ -143,6 +143,11 @@ def test_int_representation():
     assert_equal(repr(Int(-100)), "-100")
 
 
+def test_indexer():
+    assert_equal(5, Int(5).__index__())
+    assert_equal(987, Int(987).__index__())
+
+
 def main():
     test_constructors()
     test_properties()
@@ -160,3 +165,4 @@ def main():
     test_abs()
     test_string_conversion()
     test_int_representation()
+    test_indexer()

stdlib/test/builtin/test_int_literal.mojo

@@ -78,6 +78,11 @@ def test_abs():
     assert_equal(abs(0), 0)
 
 
+def test_indexer():
+    assert_equal(1, IntLiteral.__index__(1))
+    assert_equal(88, IntLiteral.__index__(88))
+
+
 def main():
     test_int()
     test_ceil()
@@ -88,3 +93,4 @@ def main():
     test_mod()
     test_bit_width()
     test_abs()
+    test_indexer()

stdlib/test/builtin/test_simd.mojo

@@ -1009,6 +1009,13 @@ def test_min_max_clamp():
     assert_equal(i.clamp(-7, 4), I(-7, -5, 4, 4))
 
 
+def test_indexer():
+    assert_equal(5, Int8(5).__index__())
+    assert_equal(56, UInt32(56).__index__())
+    assert_equal(1, Scalar[DType.bool](True).__index__())
+    assert_equal(0, Scalar[DType.bool](False).__index__())
+
+
 def main():
     test_cast()
     test_simd_variadic()
@@ -1044,3 +1051,4 @@ def main():
     test_mul_with_overflow()
     test_abs()
     test_min_max_clamp()
+    test_indexer()