2/2: Add enumerate() to DSLX stdlib

xls/dslx/bytecode/builtins.cc

@@ -518,24 +518,6 @@ absl::Status RunBuiltinCtz(const Bytecode& bytecode, InterpreterStack& stack) {
   return absl::OkStatus();
 }
 
-absl::Status RunBuiltinEnumerate(const Bytecode& bytecode,
-                                 InterpreterStack& stack) {
-  XLS_RET_CHECK(!stack.empty());
-  XLS_ASSIGN_OR_RETURN(InterpValue input, stack.Pop());
-  XLS_ASSIGN_OR_RETURN(const std::vector<InterpValue>* values,
-                       input.GetValues());
-
-  std::vector<InterpValue> elements;
-  elements.reserve(values->size());
-  for (int32_t i = 0; i < values->size(); i++) {
-    elements.push_back(
-        InterpValue::MakeTuple({InterpValue::MakeU32(i), values->at(i)}));
-  }
-  XLS_ASSIGN_OR_RETURN(InterpValue result, InterpValue::MakeArray(elements));
-  stack.Push(result);
-  return absl::OkStatus();
-}
-
 absl::Status RunBuiltinOrReduce(const Bytecode& bytecode,
                                 InterpreterStack& stack) {
   VLOG(3) << "Executing builtin OrReduce.";

xls/dslx/bytecode/builtins.h

@@ -72,8 +72,6 @@ absl::Status RunBuiltinCeilLog2(const Bytecode& bytecode,
 absl::Status RunBuiltinClz(const Bytecode& bytecode, InterpreterStack& stack);
 absl::Status RunBuiltinCover(const Bytecode& bytecode, InterpreterStack& stack);
 absl::Status RunBuiltinCtz(const Bytecode& bytecode, InterpreterStack& stack);
-absl::Status RunBuiltinEnumerate(const Bytecode& bytecode,
-                                 InterpreterStack& stack);
 absl::Status RunBuiltinOrReduce(const Bytecode& bytecode,
                                 InterpreterStack& stack);
 absl::Status RunBuiltinRange(const Bytecode& bytecode, InterpreterStack& stack);

xls/dslx/bytecode/bytecode_interpreter.cc

@@ -1673,8 +1673,6 @@ absl::Status BytecodeInterpreter::RunBuiltinFn(const Bytecode& bytecode,
       return RunBuiltinCover(bytecode, stack_);
     case Builtin::kCtz:
       return RunBuiltinCtz(bytecode, stack_);
-    case Builtin::kEnumerate:
-      return RunBuiltinEnumerate(bytecode, stack_);
     case Builtin::kFail: {
       XLS_ASSIGN_OR_RETURN(InterpValue value, Pop());
       std::string message{value.ToString()};

xls/dslx/bytecode/bytecode_interpreter_test.cc

@@ -1940,29 +1940,6 @@ fn main(x: s10, y: s10) -> s10 {
   EXPECT_THAT(bits.ToInt64(), IsOkAndHolds(-507));
 }
 
-TEST_F(BytecodeInterpreterTest, BuiltinEnumerate) {
-  constexpr std::string_view kProgram = R"(
-fn main() -> (u32, u8)[4] {
-  let x = u8[4]:[5, 6, 7, 8];
-  enumerate(x)
-})";
-
-  XLS_ASSERT_OK_AND_ASSIGN(InterpValue actual, Interpret(kProgram, "main"));
-  XLS_ASSERT_OK_AND_ASSIGN(
-      InterpValue expected,
-      InterpValue::MakeArray({
-          InterpValue::MakeTuple(
-              {InterpValue::MakeUBits(32, 0), InterpValue::MakeUBits(8, 5)}),
-          InterpValue::MakeTuple(
-              {InterpValue::MakeUBits(32, 1), InterpValue::MakeUBits(8, 6)}),
-          InterpValue::MakeTuple(
-              {InterpValue::MakeUBits(32, 2), InterpValue::MakeUBits(8, 7)}),
-          InterpValue::MakeTuple(
-              {InterpValue::MakeUBits(32, 3), InterpValue::MakeUBits(8, 8)}),
-      }));
-  EXPECT_TRUE(expected.Eq(actual));
-}
-
 TEST_F(BytecodeInterpreterTest, BuiltinUMulp) {
   constexpr std::string_view kProgram = R"(
 fn main(x: u10, y: u10) -> u10 {

xls/dslx/dslx_builtins.h

@@ -46,7 +46,6 @@ namespace xls::dslx {
   X("decode", kDecode)                         \
   X("element_count", kElementCount)            \
   X("encode", kEncode)                         \
-  X("enumerate", kEnumerate)                   \
   X("fail!", kFail)                            \
   X("gate!", kGate)                            \
   X("map", kMap)                               \

xls/dslx/fmt/ast_fmt.cc

@@ -417,6 +417,12 @@ DocRef Fmt(const ChannelTypeAnnotation& n, Comments& comments,
   return ConcatNGroup(arena, pieces);
 }
 
+DocRef Fmt(const TypeVariableTypeAnnotation& n, Comments& comments,
+           DocArena& arena) {
+  std::vector<DocRef> pieces = {Fmt(*n.type_variable(), comments, arena)};
+  return ConcatNGroup(arena, pieces);
+}
+
 DocRef Fmt(const TypeAnnotation& n, Comments& comments, DocArena& arena) {
   if (auto* t = dynamic_cast<const BuiltinTypeAnnotation*>(&n)) {
     return Fmt(*t, comments, arena);
@@ -433,6 +439,12 @@ DocRef Fmt(const TypeAnnotation& n, Comments& comments, DocArena& arena) {
   if (auto* t = dynamic_cast<const ChannelTypeAnnotation*>(&n)) {
     return Fmt(*t, comments, arena);
   }
+  if (auto* t = dynamic_cast<const TypeVariableTypeAnnotation*>(&n)) {
+    return Fmt(*t, comments, arena);
+  }
+  if (dynamic_cast<const GenericTypeAnnotation*>(&n)) {
+    return arena.Make(Keyword::kType);
+  }
   if (dynamic_cast<const SelfTypeAnnotation*>(&n)) {
     return arena.Make(Keyword::kSelfType);
   }

xls/dslx/frontend/builtin_stubs.x

@@ -56,8 +56,6 @@ fn decode<T: type, N: u32>(x: uN[N]) -> T;
 
 fn element_count<T: type>() -> u32;
 
-fn enumerate<T: type, N: u32>(x: T[N]) -> (u32, T)[N];
-
 fn fail!<N: u32, T: type> (label: u8[N], fallback_value: T) -> T;
 
 fn gate!<T: type>(x: u1, y: T) -> T;

xls/dslx/frontend/builtins_metadata.cc

@@ -98,7 +98,6 @@ const absl::flat_hash_map<std::string, BuiltinsData>& GetParametricBuiltins() {
           {"signex", {}},
           {"array_slice", {}},
           {"update", {}},
-          {"enumerate", {}},
           {"widening_cast", {}},
           {"checked_cast", {}},
 

xls/dslx/frontend/parser_test.cc

@@ -4077,7 +4077,8 @@ TEST_F(ParserTest, ParseMapWithLambdaNoParamAnnotation) {
 
 TEST_F(ParserTest, ParseMapWithLambdaNoParamAnnotationMultipleParams) {
   RoundTrip(
-      R"(const ARR = map(enumerate(range(0, u16:5)), |i, j| { 2 * i * j });)");
+      R"(import std;
+const ARR = map(std::enumerate(range(0, u16:5)), |i, j| { 2 * i * j });)");
 }
 
 TEST_F(ParserTest, ParseLambdaWithNoBrackets) {

xls/dslx/ir_convert/ir_converter_test.cc

@@ -899,8 +899,9 @@ fn main(input: u8[2]) -> u8[2] {
 
 TEST_F(IrConverterTest, ArrayEnumerate) {
   constexpr std::string_view program = R"(
+import std;
 fn main(array: u8[4]) -> (u32, u8)[4] {
-  enumerate(array)
+  std::enumerate(array)
 }
 )";
   XLS_ASSERT_OK_AND_ASSIGN(std::string converted,

xls/dslx/ir_convert/testdata/ir_converter_test_ArrayEnumerate.ir

@@ -1,19 +1,25 @@
 package test_module
 
-file_number 0 "test_module.x"
+file_number 0 "xls/dslx/stdlib/std.x"
+file_number 1 "test_module.x"
 
-top fn __test_module__main(array: bits[8][4] id=1) -> (bits[32], bits[8])[4] {
-  literal.2: bits[32] = literal(value=0, id=2)
-  literal.5: bits[32] = literal(value=1, id=5)
-  literal.8: bits[32] = literal(value=2, id=8)
-  literal.11: bits[32] = literal(value=3, id=11)
-  array_index.3: bits[8] = array_index(array, indices=[literal.2], id=3)
-  array_index.6: bits[8] = array_index(array, indices=[literal.5], id=6)
-  array_index.9: bits[8] = array_index(array, indices=[literal.8], id=9)
-  array_index.12: bits[8] = array_index(array, indices=[literal.11], id=12)
-  tuple.4: (bits[32], bits[8]) = tuple(literal.2, array_index.3, id=4)
-  tuple.7: (bits[32], bits[8]) = tuple(literal.5, array_index.6, id=7)
-  tuple.10: (bits[32], bits[8]) = tuple(literal.8, array_index.9, id=10)
-  tuple.13: (bits[32], bits[8]) = tuple(literal.11, array_index.12, id=13)
-  ret array.14: (bits[32], bits[8])[4] = array(tuple.4, tuple.7, tuple.10, tuple.13, id=14)
+fn ____std__enumerate__4_u8_counted_for_0_body(i: bits[32] id=7, result: (bits[32], bits[8])[4] id=10, x: bits[8][4] id=11) -> (bits[32], bits[8])[4] {
+  literal.8: bits[32] = literal(value=0, id=8)
+  add.9: bits[32] = add(i, literal.8, id=9)
+  array_index.12: bits[8] = array_index(x, indices=[add.9], id=12)
+  tuple.13: (bits[32], bits[8]) = tuple(add.9, array_index.12, id=13)
+  ret array_update.14: (bits[32], bits[8])[4] = array_update(result, tuple.13, indices=[add.9], id=14)
+}
+
+fn __std__enumerate__4_u8(x: bits[8][4] id=1) -> (bits[32], bits[8])[4] {
+  literal.3: bits[32] = literal(value=0, id=3)
+  literal.4: bits[8] = literal(value=0, id=4)
+  tuple.5: (bits[32], bits[8]) = tuple(literal.3, literal.4, id=5)
+  array.6: (bits[32], bits[8])[4] = array(tuple.5, tuple.5, tuple.5, tuple.5, id=6)
+  N: bits[32] = literal(value=4, id=2)
+  ret counted_for.15: (bits[32], bits[8])[4] = counted_for(array.6, trip_count=4, stride=1, body=____std__enumerate__4_u8_counted_for_0_body, invariant_args=[x], id=15)
+}
+
+top fn __test_module__main(array: bits[8][4] id=16) -> (bits[32], bits[8])[4] {
+  ret invoke.17: (bits[32], bits[8])[4] = invoke(array, to_apply=__std__enumerate__4_u8, id=17)
 }

xls/dslx/stdlib/std.x

@@ -14,6 +14,43 @@
 
 // DSLX standard library routines.
 
+#![feature(generics)]
+
+pub fn enumerate<T: type, N: u32>(x: T[N]) -> (u32, T)[N] {
+    for (i, result) in 0..N {
+        update(result, i, (i, x[i]))
+    }([(u32:0, zero!<T>()), ...])
+}
+
+#[test]
+fn emumerate_test() {
+    let array = [1, 2, 4, 8];
+    let enumerated = enumerate(array);
+    assert_eq(enumerated[0], (0, 1));
+    assert_eq(enumerated[1], (1, 2));
+    assert_eq(enumerated[2], (2, 4));
+    assert_eq(enumerated[3], (3, 8));
+}
+
+#[test]
+fn enumerate_type_test() {
+    type RamData = uN[8];
+    const DATA = [RamData:1, 2, 4, 8];
+    let enumerated = enumerate(DATA);
+    assert_eq(enumerated[0], (0, 1));
+    assert_eq(enumerated[1], (1, 2));
+    assert_eq(enumerated[2], (2, 4));
+    assert_eq(enumerated[3], (3, 8));
+}
+
+#[test]
+fn enumerate_tuple_test() {
+    let x = [(true, 3), (false, 2)];
+    let enumerated = enumerate(x);
+    assert_eq(enumerated[0], (0, (true, 3)));
+    assert_eq(enumerated[1], (1, (false, 2)));
+}
+
 pub fn sizeof<S: bool, N: u32>(x: xN[S][N]) -> u32 { N }
 
 #[test]

xls/dslx/tests/BUILD

@@ -333,12 +333,6 @@ dslx_lang_test(
 
 dslx_lang_test(name = "for_over_range")
 
-dslx_lang_test(
-    name = "enumerate",
-    compare = "interpreter",
-    convert_to_ir = True,
-)
-
 dslx_lang_test(
     name = "character_conversion",
     # TODO: https://github.com/google/xls/issues/1526 - fails opportunistic_postcondition in autofmt

xls/dslx/type_system/typecheck_module_test.cc

@@ -2067,9 +2067,10 @@ fn f() -> u32[3] {
 
 TEST_F(TypecheckV2Test, EnumerateBuiltin) {
   XLS_EXPECT_OK(Typecheck(R"(
+import std;
 type MyTup = (u32, u2);
 fn f(x: u2[7]) -> MyTup[7] {
-  enumerate(x)
+  std::enumerate(x)
 }
 )"));
 }

xls/dslx/type_system_v2/typecheck_module_v2_builtin_test.cc

@@ -659,21 +659,6 @@ const_assert!(element_count<imported::T>() == 10);
   XLS_EXPECT_OK(TypecheckV2(kProgram, "main", &import_data));
 }
 
-TEST(TypecheckV2BuiltinTest, Enumerate) {
-  EXPECT_THAT(R"(const Y = enumerate<u16, u32:3>([u16:1, u16:2, u16:3]);)",
-              TypecheckSucceeds(HasNodeWithType("Y", "(uN[32], uN[16])[3]")));
-}
-
-TEST(TypecheckV2BuiltinTest, EnumerateImplicitSize) {
-  EXPECT_THAT(R"(const Y = enumerate<u16>([u16:1, u16:2, u16:3]);)",
-              TypecheckSucceeds(HasNodeWithType("Y", "(uN[32], uN[16])[3]")));
-}
-
-TEST(TypecheckV2BuiltinTest, EnumerateImplicitType) {
-  EXPECT_THAT(R"(const Y = enumerate([u16:1, u16:2, u16:3]);)",
-              TypecheckSucceeds(HasNodeWithType("Y", "(uN[32], uN[16])[3]")));
-}
-
 TEST(TypecheckV2BuiltinTest, Fail) {
   EXPECT_THAT(
       R"(

xls/examples/dslx_intro/prefix_scan_equality.x

@@ -14,10 +14,12 @@
 
 // Prefix scans an array of 8 32-bit values and produces a running count of
 // duplicate values in the run.
+import std;
+
 fn prefix_scan_eq(x: u32[8]) -> u3[8] {
   let (_, _, result) =
     for ((i, elem), (prior, count, result)): ((u32, u32), (u32, u3, u3[8]))
-          in enumerate(x) {
+          in std::enumerate(x) {
       let (to_place, new_count): (u3, u3) = match (i == u32:0, prior == elem) {
         // The first iteration always places 0 and propagates seen count of 1.
         (true, _) => (u3:0, u3:1),

xls/modules/rle/rle_dec.x

@@ -168,7 +168,7 @@ proc RunLengthDecoderTransactionTest {
     ];
     let tok = for ((counter, stimulus), tok):
         ((u32, (TestSymbol, TestCount)) , token)
-        in enumerate(TransactionTestStimuli) {
+        in std::enumerate(TransactionTestStimuli) {
       let last = counter == (array_size(TransactionTestStimuli) - u32:1);
       let data_in = TestDecInData{
         symbol: stimulus.0,
@@ -191,7 +191,7 @@ proc RunLengthDecoderTransactionTest {
     ];
     let tok = for ((counter, symbol), tok):
         ((u32, TestSymbol) , token)
-        in enumerate(TransationTestOutputs) {
+        in std::enumerate(TransationTestOutputs) {
       let last = counter == (array_size(TransationTestOutputs) - u32:1);
       let data_out = TestDecOutData{
         symbol: symbol,
@@ -243,7 +243,7 @@ proc RunLengthDecoderLastAfterLastTest {
     ];
     let tok = for ((counter, stimulus), tok):
         ((u32, TestDecInData) , token)
-        in enumerate(LastAfterLastTestStimuli) {
+        in std::enumerate(LastAfterLastTestStimuli) {
       let tok = send(tok, dec_input_s, stimulus);
       trace_fmt!("Sent {} stimuli, symbol: 0x{:x}, count:{}, last: {}",
           counter + u32:1, stimulus.symbol, stimulus.count, stimulus.last);
@@ -255,7 +255,7 @@ proc RunLengthDecoderLastAfterLastTest {
     ];
     let tok = for ((counter, output), tok):
         ((u32, TestDecOutData) , token)
-        in enumerate(LastAfterLastTestOutputs) {
+        in std::enumerate(LastAfterLastTestOutputs) {
       let (tok, dec_output) = recv(tok, dec_output_r);
       trace_fmt!(
           "Received {} transactions, symbol: 0x{:x}, last: {}",