From 0393637eea8034f5a84ca00ecfc1b96f1c319c4c Mon Sep 17 00:00:00 2001
From: Chris Leary <cdleary@openai.com>
Date: Mon, 11 Nov 2024 20:13:32 -0800
Subject: [PATCH 1/3] [DSLX:stdlib] Add `std::distinct`.

---
 docs_src/dslx_std.md  | 11 +++++++-
 xls/dslx/stdlib/std.x | 60 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/docs_src/dslx_std.md b/docs_src/dslx_std.md
index 9e57c3100d..3e5121d974 100644
--- a/docs_src/dslx_std.md
+++ b/docs_src/dslx_std.md
@@ -46,7 +46,7 @@ fn test_array_size() {
 }
 ```
 
-### `widening_cast` and `checked_cast`
+### `widening_cast`, `checked_cast`
 
 `widening_cast` and `checked_cast` cast bits-type values to bits-type values
 with additional checks compared to casting with `as`.
@@ -1188,6 +1188,15 @@ index is used in a match expression (which will eagerly evaluate all of its
 arms), to prevent it from creating an error at simulation time if the value is
 ultimately discarded from the unselected match arm.
 
+#### `std::distinct`
+
+```dslx-snippet
+pub fn distinct<COUNT: u32, N: u32, S: bool>(items: xN[S][N][COUNT], valid: bool[COUNT]) -> bool
+```
+
+Returns true iff there are multiple equivalent elements in `items` that are
+also `valid`. That is, `valid` acts as an item mask.
+
 ## `import acm_random`
 
 Port of
diff --git a/xls/dslx/stdlib/std.x b/xls/dslx/stdlib/std.x
index e5fa6e309e..2e765f9323 100644
--- a/xls/dslx/stdlib/std.x
+++ b/xls/dslx/stdlib/std.x
@@ -1268,3 +1268,63 @@ fn clzt_test() {
 
 #[quickcheck]
 fn prop_clzt_same_as_clz(x: u64) -> bool { clz(x) == clzt(x) as u64 }
+
+/// Returns whether all the `items` are distinct (i.e. there are no duplicate
+/// items) after the `valid` mask is applied.
+pub fn distinct<COUNT: u32, N: u32, S: bool>(items: xN[S][N][COUNT], valid: bool[COUNT]) -> bool {
+    const INIT_ALL_DISTINCT = true;
+    for (i, all_distinct) in range(u32:0, COUNT) {
+        for (j, all_distinct) in range(u32:0, COUNT) {
+            if i != j && valid[i] && valid[j] && items[i] == items[j] {
+                false
+            } else {
+                all_distinct
+            }
+        }(all_distinct)
+    }(INIT_ALL_DISTINCT)
+}
+
+#[test]
+fn test_simple_nondistinct() { assert_eq(distinct(u2[2]:[1, 1], bool[2]:[true, true]), false) }
+
+#[test]
+fn test_distinct_unsigned() {
+    let items = u8[4]:[1, 2, 3, 2];
+    let valid = bool[4]:[true, true, true, true];
+    assert_eq(distinct(items, valid), false);
+}
+
+#[test]
+fn test_distinct_signed() {
+    let items = s8[3]:[-1, 0, 1];
+    let valid = bool[3]:[true, true, true];
+    assert_eq(distinct(items, valid), true);
+}
+
+#[test]
+fn test_distinct_with_invalid() {
+    let items = u8[4]:[1, 2, 3, 1];
+    let valid = bool[4]:[true, true, true, false];
+    assert_eq(distinct(items, valid), true);
+}
+
+#[quickcheck]
+fn quickcheck_forced_duplicate(xs: u4[4], to_dupe: u2) -> bool {
+    const ALL_VALID = bool[4]:[true, ...];
+    let forced_dupe = update(xs, (to_dupe as u32 + u32:1) % u32:4, xs[to_dupe]);
+    distinct(forced_dupe, ALL_VALID) == false
+}
+
+#[quickcheck]
+fn quickcheck_distinct_all_valid_items_same(value: u4, valid: bool[4]) -> bool {
+    let items = u4[4]:[value, ...];  // All items are the same.
+    let num_valid = popcount(valid as u4) as u32;
+
+    if num_valid <= u32:1 {
+        // With 0 or 1 valid items, they are trivially distinct.
+        distinct(items, valid) == true
+    } else {
+        // Since all valid items are the same, 'distinct' should return false.
+        distinct(items, valid) == false
+    }
+}

From f5a3ec6e442daa34ea2315cb4edd61dc30efba6d Mon Sep 17 00:00:00 2001
From: Chris Leary <cdleary@openai.com>
Date: Tue, 12 Nov 2024 07:43:18 -0800
Subject: [PATCH 2/3] Fix docs.

---
 docs_src/dslx_std.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs_src/dslx_std.md b/docs_src/dslx_std.md
index 3e5121d974..e0223dad51 100644
--- a/docs_src/dslx_std.md
+++ b/docs_src/dslx_std.md
@@ -1194,8 +1194,8 @@ ultimately discarded from the unselected match arm.
 pub fn distinct<COUNT: u32, N: u32, S: bool>(items: xN[S][N][COUNT], valid: bool[COUNT]) -> bool
 ```
 
-Returns true iff there are multiple equivalent elements in `items` that are
-also `valid`. That is, `valid` acts as an item mask.
+Returns whether all the `items` are distinct (i.e. there are no duplicate
+items) after the `valid` mask is applied.
 
 ## `import acm_random`
 

From afe585bdd4472f7e1c30f35b9f69902c9d8e674b Mon Sep 17 00:00:00 2001
From: Chris Leary <cdleary@openai.com>
Date: Tue, 12 Nov 2024 08:12:23 -0800
Subject: [PATCH 3/3] Duplicate a unit test into docs.

---
 docs_src/dslx_std.md | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/docs_src/dslx_std.md b/docs_src/dslx_std.md
index e0223dad51..6705387f59 100644
--- a/docs_src/dslx_std.md
+++ b/docs_src/dslx_std.md
@@ -1197,6 +1197,17 @@ pub fn distinct<COUNT: u32, N: u32, S: bool>(items: xN[S][N][COUNT], valid: bool
 Returns whether all the `items` are distinct (i.e. there are no duplicate
 items) after the `valid` mask is applied.
 
+```dslx
+import std;
+
+#[test]
+fn test_distinct_with_invalid() {
+    let items = u8[4]:[1, 2, 3, 1];
+    let valid = bool[4]:[true, true, true, false];
+    assert_eq(std::distinct(items, valid), true);
+}
+```
+
 ## `import acm_random`
 
 Port of