allow some shape mismatches in memberof and indexof

uiua-lang · Jan 9, 2025 · b961454 · b961454
1 parent 6d5a995
commit b961454
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 9 deletions.
diff --git a/src/algorithm/dyadic/search.rs b/src/algorithm/dyadic/search.rs
@@ -80,10 +80,9 @@ impl<T: ArrayValue> Array<T> {
             }
             Ordering::Less => {
                 if !of.shape.ends_with(&elems.shape) {
-                    return Err(env.error(format!(
-                        "Cannot look for array of shape {} in array of shape {}",
-                        elems.shape, of.shape
-                    )));
+                    let shape = Shape::from(&of.shape[..of.shape.len() - elems.shape.len() - 1]);
+                    let data = eco_vec![0; shape.elements()];
+                    return Ok(Array::new(shape, data));
                 }
                 if of.rank() - elems.rank() == 1 {
                     of.rows().any(|r| *elems == r).into()
@@ -187,11 +186,12 @@ impl<T: ArrayValue> Array<T> {
             }
             Ordering::Less => {
                 if !haystack.shape.ends_with(&needle.shape) {
-                    return Err(env.error(format!(
-                        "Cannot get index of array of shape {} in array of shape {}",
-                        needle.shape(),
-                        haystack.shape()
-                    )));
+                    let shape = Shape::from(
+                        &haystack.shape[..haystack.shape.len() - needle.shape.len() - 1],
+                    );
+                    let elem = haystack.shape.row_count() as f64;
+                    let data = eco_vec![elem; shape.elements()];
+                    return Ok(Array::new(shape, data));
                 }
                 if haystack.rank() - needle.rank() == 1 {
                     (haystack

diff --git a/tests/dyadic.ua b/tests/dyadic.ua
@@ -141,6 +141,8 @@
 ⍤⤙≍ ⟜(△∊5 °△)2_4
 ⍤⤙≍ ⟜(△∊5 °△)1_4
 ⍤⤙≍ ⟜(△∊5 °△)0_4
+⍤⤙≍ 0 ∊[¤1 ¤2] ¤5
+
 # Index of
 ⍤⤙≍ 1 ⊗ 5 [1 5 5]
 ⍤⤙≍ [1] ⊗ [5] [1 5 5]
@@ -157,6 +159,10 @@
 ⍤⤙≍ ⟜(△⊗⊙5 °△)2_4
 ⍤⤙≍ ⟜(△⊗⊙5 °△)1_4
 ⍤⤙≍ ⟜(△⊗⊙5 °△)0_4
+⍤⤙≍ 1 ⊗ ¤5 [¤1 ¤5]
+⍤⤙≍ 2 ⊗ ¤5 [[1] [2] [5] [3]]
+⍤⤙≍ 4 ⊗ ¤5 [[1] [2] [0] [3]]
+⍤⤙≍ 4 ⊗ ¤5 [1_0 2_0 5_0 3_0]
 
 # Find
 ⍤⤙≍ [0 1 0 0] ⌕ 2 [1 2 3 4]

diff --git a/tests/optimized.ua b/tests/optimized.ua
@@ -204,6 +204,7 @@ F ← ⬚10(/+◌1⊞(×¯))
 ⍤⤙≍ ⊃(∊∘♭₂⇡|∊♭₂⇡) [5 ¯5] [[2_¯π 3_∞ 0_¯5][NaN_0 5_¯5 ¯2_1]]
 ⍤⤙≍ ⊃(∊∘♭₂⇡|∊♭₂⇡) 3 °△2_3
 ⍤⤙≍ ⊃(∊∘♭₂⇡|∊♭₂⇡) 3_6 5
+⍤⤙≍ ⊃(∊∘♭₂⇡|∊♭₂⇡) 5_5 ¤5
 
 # Random Row
 ⍤⤙≍ ⊃(⊢°⍆|⊢∘°⍆) °△1_4_5_6