[xls] RangeQueryEngine: Add fast path for ArrayIndex

If all array indices are precise, i.e. the IntervalSet has only one value,
then we don't need to scan all array indices and can simply look up our
result.

PiperOrigin-RevId: 808137825

Author: James Molloy

xls/passes/range_query_engine.cc

@@ -674,6 +674,26 @@ absl::Status RangeQueryVisitor::HandleArrayIndex(ArrayIndex* array_index) {
 
   IntervalSetTree result = EmptyIntervalSetTree(array_index->GetType());
 
+  // Fast path: if all of the index intervals are precise, we can just perform
+  // a lookup.
+  if (absl::c_all_of(index_intervals,
+                     [](const IntervalSet& is) { return is.IsPrecise(); })) {
+    std::vector<int64_t> indexes;
+    indexes.reserve(index_intervals.size());
+    for (int64_t i = 0; i < index_intervals.size(); ++i) {
+      const IntervalSet& intervals = index_intervals[i];
+      XLS_ASSIGN_OR_RETURN(uint64_t index,
+                           intervals.GetPreciseValue()->ToUint64());
+      int64_t bound = dimension[i] - 1;
+      // Clamp to boundary if out of bounds.
+      indexes.push_back(std::clamp<uint64_t>(index, 0, bound));
+    }
+    result.Set({}, array_interval_set_tree->Get(indexes));
+
+    SetIntervalSetTree(array_index, std::move(result));
+    return absl::OkStatus();
+  }
+
   // Returns true if the given interval set covers the given index value for an
   // array of dimension `dim`. Includes handling of OOB conditions.
   auto intervals_cover_index = [&](const IntervalSet& intervals, int64_t index,