From ba53b52d55eb9408261c290263461e0a898a2595 Mon Sep 17 00:00:00 2001
From: Alex Light <allight@google.com>
Date: Fri, 1 Nov 2024 14:33:49 -0700
Subject: [PATCH] Remove chained update-index when the same element is indexed
 and both are known to be in bounds.

PiperOrigin-RevId: 692293570
---
 xls/passes/array_simplification_pass.cc      | 22 ++++++++++++++++++++
 xls/passes/array_simplification_pass_test.cc | 19 +++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/xls/passes/array_simplification_pass.cc b/xls/passes/array_simplification_pass.cc
index d9fe13d837..a5e1048b48 100644
--- a/xls/passes/array_simplification_pass.cc
+++ b/xls/passes/array_simplification_pass.cc
@@ -227,6 +227,28 @@ absl::StatusOr<SimplifyResult> SimplifyArrayIndex(
     return SimplifyResult::Changed({array_index->array()});
   }
 
+  // An array-index of an array-update of the same index can be better
+  // thought of as a mux between the last element of the array (in the case the
+  // update & index are both out of bounds) and the updated value. In the case
+  // where the update/index location is in bounds we can just replace with the
+  // update value.
+  //
+  // Only checking single-dimensional array-indexes because its significantly
+  // simpler and is the common case.
+  //
+  if (array_index->array()->Is<ArrayUpdate>()) {
+    ArrayUpdate* update = array_index->array()->As<ArrayUpdate>();
+    if (IndicesAreDefinitelyEqual(array_index->indices(), update->indices(),
+                                  query_engine) &&
+        // Check both since if either has the known-in-bounds mark we can go
+        // ahead.
+        (IndicesAreDefinitelyInBounds(array_index, query_engine) ||
+         IndicesAreDefinitelyInBounds(update, query_engine))) {
+      XLS_RETURN_IF_ERROR(array_index->ReplaceUsesWith(update->update_value()));
+      return SimplifyResult::Changed({});
+    }
+  }
+
   // An array index which indexes into a 1-element array can have its index
   // replaced with a literal 0, since the clamping behavior makes all indexing
   // equivalent:
diff --git a/xls/passes/array_simplification_pass_test.cc b/xls/passes/array_simplification_pass_test.cc
index ae7b28d92f..7723ba8e4e 100644
--- a/xls/passes/array_simplification_pass_test.cc
+++ b/xls/passes/array_simplification_pass_test.cc
@@ -1602,5 +1602,24 @@ TEST_F(ArraySimplificationPassTest, MultidimensionalUnitArrayUpdate) {
           /*default_value=*/m::Param("array")));
 }
 
+TEST_F(ArraySimplificationPassTest, ArrayUpdateIndexInBounds) {
+  auto p = CreatePackage();
+  FunctionBuilder fb(TestName(), p.get());
+  BValue arr = fb.Param("array", p->GetArrayType(10, p->GetBitsType(32)));
+  BValue idx = fb.Param("upd_idx", p->GetBitsType(32));
+  // NB This sort of selector requires context-sensitive range analysis to see.
+  BValue idx_bound = fb.Select(fb.ULt(idx, fb.Literal(UBits(10, 32))), {idx},
+                               fb.Literal(UBits(0, 32)));
+  BValue val = fb.Param("val", p->GetBitsType(32));
+  BValue update_arr =
+      fb.ArrayUpdate(arr, val, {idx_bound}, /*known_in_bounds=*/true);
+  fb.ArrayIndex(update_arr, {idx_bound}, /*known_in_bounds=*/true);
+  XLS_ASSERT_OK_AND_ASSIGN(Function * f, fb.Build());
+  ScopedVerifyEquivalence sve(f);
+  ScopedRecordIr sri(p.get());
+  ASSERT_THAT(Run(f), IsOkAndHolds(true));
+  EXPECT_THAT(f->return_value(), m::Param("val"));
+}
+
 }  // namespace
 }  // namespace xls