make deshape subscript 0 give the first scalar

uiua-lang · Nov 30, 2024 · d8bb8c9 · d8bb8c9
1 parent 0a7a223
commit d8bb8c9
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 34 deletions.
diff --git a/src/algorithm/monadic.rs b/src/algorithm/monadic.rs
@@ -52,52 +52,70 @@ impl Value {
         self.shape_mut().push(deshaped);
     }
     pub(crate) fn deshape_sub(&mut self, irank: i32, extend: bool, env: &Uiua) -> UiuaResult {
-        if irank == 0 || irank > 0 && irank as usize == self.rank() {
+        if irank > 0 && irank as usize == self.rank() {
             return Ok(());
         }
         self.take_map_keys();
         let shape = self.shape_mut();
         let rank = irank.unsigned_abs() as usize;
-        if irank > 0 {
-            // Positive rank
-            match rank.cmp(&shape.len()) {
-                Ordering::Equal => {}
-                Ordering::Less => {
-                    let mid = shape.len() + 1 - rank;
-                    let new_first_dim: usize = shape[..mid].iter().product();
-                    *shape = once(new_first_dim)
-                        .chain(shape[mid..].iter().copied())
-                        .collect();
+        match irank.cmp(&0) {
+            Ordering::Equal => {
+                // First scalar
+                if shape.contains(&0) {
+                    if let Some(fill) = env.value_fill() {
+                        *self = fill.clone();
+                    } else {
+                        return Err(env.error(format!(
+                            "Cannot get first scalar of an empty array (shape {shape})"
+                        )));
+                    }
+                } else {
+                    *shape = [].into();
+                    val_as_arr!(self, |arr| arr.data.truncate(1));
                 }
-                Ordering::Greater => {
-                    if extend {
-                        for _ in 0..rank - shape.len() {
+            }
+            Ordering::Greater => {
+                // Positive rank
+                match rank.cmp(&shape.len()) {
+                    Ordering::Equal => {}
+                    Ordering::Less => {
+                        let mid = shape.len() + 1 - rank;
+                        let new_first_dim: usize = shape[..mid].iter().product();
+                        *shape = once(new_first_dim)
+                            .chain(shape[mid..].iter().copied())
+                            .collect();
+                    }
+                    Ordering::Greater => {
+                        if extend {
+                            for _ in 0..rank - shape.len() {
+                                shape.insert(0, 1);
+                            }
+                        } else {
                             shape.insert(0, 1);
                         }
-                    } else {
-                        shape.insert(0, 1);
                     }
                 }
             }
-        } else {
-            // Negative rank
-            if rank + 1 > shape.len() {
-                return if extend {
-                    Err(env.error(format!(
-                        "Negative {} has magnitude {}, but the \
+            Ordering::Less => {
+                // Negative rank
+                if rank + 1 > shape.len() {
+                    return if extend {
+                        Err(env.error(format!(
+                            "Negative {} has magnitude {}, but the \
                         rank-{} array cannot be reduced that much",
-                        Primitive::Deshape.format(),
-                        rank,
-                        shape.len()
-                    )))
-                } else {
-                    Ok(())
-                };
+                            Primitive::Deshape.format(),
+                            rank,
+                            shape.len()
+                        )))
+                    } else {
+                        Ok(())
+                    };
+                }
+                let new_first_dim: usize = shape[..=rank].iter().product();
+                *shape = once(new_first_dim)
+                    .chain(shape[rank + 1..].iter().copied())
+                    .collect();
             }
-            let new_first_dim: usize = shape[..=rank].iter().product();
-            *shape = once(new_first_dim)
-                .chain(shape[rank + 1..].iter().copied())
-                .collect();
         }
         self.validate_shape();
         Ok(())

diff --git a/src/primitive/defs.rs b/src/primitive/defs.rs
@@ -1010,9 +1010,11 @@ primitive!(
     /// ex: △ ♭₋₁ °△2_3_4_5
     ///   : △ ♭₋₂ °△2_3_4_5
     ///   : △ ♭₋₃ °△2_3_4_5
-    /// If the subscript rank is greater than the rank of the array, length-1 axes are added to the front for the shape
+    /// If the subscript rank is greater than the rank of the array, length-1 axes are added to the front for the shape.
     /// ex: ♭₂ [1 2 3]
     ///   : ♭₃ [1 2 3]
+    /// A subscript of `0` gives the first scalar in the array.
+    /// ex: ♭₀ [4_2_6 0_3_7]
     ///
     /// It looks like `♭` because it *flat*tens the array.
     ///