implement subscripted on/by/with/off

changelog.md

@@ -6,6 +6,8 @@ Uiua is not yet stable.
 This version is not yet released. If you are reading this on the website, then these changes are live here.
 ### Language
 - **Breaking Change** - [`repeat ⍥`](https://uiua.org/docs/repeat) and [`do ⍢`](https://uiua.org/docs/do) with net-negative signatures now preserve lower stack values between iterations
+- Add subscripted [`on ⟜`](https://uiua.org/docs/on), [`by ⊸`](https://uiua.org/docs/by), [`with ⤙`](https://uiua.org/docs/with), and [`off ⤚`](https://uiua.org/docs/off)
+  - These preserve N arguments rather than just the first or last
 - Signature comments can now use a `$` rather than a `?` to automatically label arguments and outputs
 - Add sided subscripts for [`reach 𝄐`](https://uiua.org/docs/reach)
 - Add experimental subscripts to [`negate ¯`](https://uiua.org/docs/negate)

site/src/other.rs

@@ -517,6 +517,10 @@ pub fn Subscripts() -> impl IntoView {
         subscript(First, "First N values", "⊢₂ \"hello\""),
         subscript(Last, "Last N values", "⊣₂ \"hello\""),
         subscript(Rand, "Random integer", "# Experimental!\n⚂₁₀₀"),
+        subscript(On, "First N values", "{⟜₂[⊙⊙∘] 1 2 3}"),
+        subscript(By, "Last N values", "{⊸₂[⊙⊙∘] 1 2 3}"),
+        subscript(With, "Last N values", "{⤙₂[⊙⊙∘] 1 2 3}"),
+        subscript(Off, "First N values", "{⤚₂[⊙⊙∘] 1 2 3}"),
         subscript(Both, "Apply to N argument sets", "[∩₃+ 1 2 3 4 5 6]"),
         subscript(Each, "Apply to rank N subarrays", "∵₁□ °△2_3_4"),
         subscript(Rows, "Apply to subarrays N deep", "≡₂□ °△2_3_4"),

src/check.rs

@@ -585,6 +585,13 @@ impl VirtualEnv {
                 }
             },
             Node::ImplMod(prim, args, _) => match prim {
+                &OnSub(n) | &BySub(n) | &WithSub(n) | &OffSub(n) => {
+                    let [sn] = get_args_nodes(args)?;
+                    let args = sn.sig.args.max(n);
+                    self.handle_args_outputs(args, args);
+                    self.node(&sn.node)?;
+                    self.handle_args_outputs(0, n);
+                }
                 ReduceContent | ReduceDepth(_) => {
                     let [sig] = get_args(args)?;
                     let args = sig.args.saturating_sub(sig.outputs);

src/compile/mod.rs

@@ -1969,8 +1969,9 @@ code:
         let span = self.add_span(span);
         Node::Prim(prim, span)
     }
-    #[allow(clippy::match_single_binding)]
+    #[allow(clippy::match_single_binding, unused_parens)]
     fn subscript(&mut self, sub: Subscripted, span: CodeSpan) -> UiuaResult<Node> {
+        use Primitive::*;
         let Some(n) = self.subscript_n_or_side(sub.n) else {
             return self.word(sub.word);
         };
@@ -1987,15 +1988,10 @@ code:
                     _ => {
                         if !matches!(
                             prim,
-                            Primitive::Both
-                                | Primitive::Bracket
-                                | Primitive::Repeat
-                                | Primitive::Tuples
-                                | Primitive::Rows
-                                | Primitive::Each
-                                | Primitive::Inventory
-                                | Primitive::Stencil
-                                | Primitive::Reach
+                            (Both | Bracket)
+                                | (Reach | On | By | With | Off)
+                                | (Rows | Each | Inventory)
+                                | (Repeat | Tuples | Stencil)
                         ) {
                             self.add_error(
                                 m.modifier.span.clone().merge(n.span.clone()),
@@ -2043,10 +2039,8 @@ code:
                             self.primitive(prim, span),
                         ])
                     }
-                    Primitive::Deshape => {
-                        Node::ImplPrim(ImplPrimitive::DeshapeSub(n), self.add_span(span))
-                    }
-                    Primitive::Transpose => {
+                    Deshape => Node::ImplPrim(ImplPrimitive::DeshapeSub(n), self.add_span(span)),
+                    Transpose => {
                         self.subscript_experimental(prim, &span);
                         if n > 100 {
                             self.add_error(span.clone(), "Too many subscript repetitions");
@@ -2055,7 +2049,7 @@ code:
                             .map(|_| self.primitive(prim, span.clone()))
                             .collect()
                     }
-                    Primitive::Neg => {
+                    Neg => {
                         self.subscript_experimental(prim, &span);
                         if n == 0 {
                             self.add_error(span.clone(), "Cannot have the 0th root of unity");
@@ -2068,115 +2062,108 @@ code:
                             -4 => -crate::Complex::I,
                             _ => (crate::Complex::I * (std::f64::consts::TAU / n as f64)).exp(),
                         };
-                        Node::from_iter([
-                            Node::new_push(root_of_unity),
-                            self.primitive(Primitive::Mul, span),
-                        ])
+                        Node::from_iter([Node::new_push(root_of_unity), self.primitive(Mul, span)])
                     }
-                    Primitive::Sqrt => {
+                    Sqrt => {
                         if n == 0 {
                             self.add_error(span.clone(), "Cannot take 0th root");
                         }
                         Node::from_iter([
                             Node::new_push(1.0 / n.max(1) as f64),
-                            self.primitive(Primitive::Pow, span),
+                            self.primitive(Pow, span),
                         ])
                     }
-                    Primitive::Floor | Primitive::Ceil => {
+                    Floor | Ceil => {
                         self.subscript_experimental(prim, &span);
                         let mul = 10f64.powi(n);
                         Node::from_iter([
                             Node::new_push(mul),
-                            self.primitive(Primitive::Mul, span.clone()),
+                            self.primitive(Mul, span.clone()),
                             self.primitive(prim, span.clone()),
                             Node::new_push(mul),
-                            self.primitive(Primitive::Div, span),
+                            self.primitive(Div, span),
                         ])
                     }
-                    Primitive::Round => {
+                    Round => {
                         let mul = 10f64.powi(n);
                         Node::from_iter([
                             Node::new_push(mul),
-                            self.primitive(Primitive::Mul, span.clone()),
+                            self.primitive(Mul, span.clone()),
                             self.primitive(prim, span.clone()),
                             Node::new_push(mul),
-                            self.primitive(Primitive::Div, span),
+                            self.primitive(Div, span),
                         ])
                     }
-                    Primitive::Rand => {
+                    Rand => {
                         self.subscript_experimental(prim, &span);
                         Node::from_iter([
-                            self.primitive(Primitive::Rand, span.clone()),
+                            self.primitive(Rand, span.clone()),
                             Node::new_push(n),
-                            self.primitive(Primitive::Mul, span.clone()),
-                            self.primitive(Primitive::Floor, span),
+                            self.primitive(Mul, span.clone()),
+                            self.primitive(Floor, span),
                         ])
                     }
-                    Primitive::Utf8 => match n {
-                        8 => self.primitive(Primitive::Utf8, span),
+                    Utf8 => match n {
+                        8 => self.primitive(Utf8, span),
                         16 => Node::ImplPrim(ImplPrimitive::Utf16, self.add_span(span)),
                         _ => {
                             self.add_error(span.clone(), "Only UTF-8 and UTF-16 are supported");
-                            self.primitive(Primitive::Utf8, span)
+                            self.primitive(Utf8, span)
                         }
                     },
-                    Primitive::Couple => match n {
-                        1 => self.primitive(Primitive::Fix, span),
-                        2 => self.primitive(Primitive::Couple, span),
+                    Couple => match n {
+                        1 => self.primitive(Fix, span),
+                        2 => self.primitive(Couple, span),
                         n => Node::Array {
                             len: ArrayLen::Static(self.positive_subscript(
                                 n,
-                                Primitive::Couple,
+                                Couple,
                                 span.clone(),
                             )?),
                             inner: Node::empty().into(),
                             boxed: false,
-                            prim: Some(Primitive::Couple),
+                            prim: Some(Couple),
                             span: self.add_span(span),
                         },
                     },
-                    Primitive::Box => Node::Array {
-                        len: ArrayLen::Static(self.positive_subscript(
-                            n,
-                            Primitive::Box,
-                            span.clone(),
-                        )?),
+                    Box => Node::Array {
+                        len: ArrayLen::Static(self.positive_subscript(n, Box, span.clone())?),
                         inner: Node::empty().into(),
                         boxed: true,
-                        prim: Some(Primitive::Box),
+                        prim: Some(Box),
                         span: self.add_span(span),
                     },
-                    Primitive::Stack => Node::ImplPrim(
+                    Stack => Node::ImplPrim(
                         ImplPrimitive::StackN {
-                            n: self.positive_subscript(n, Primitive::Stack, span.clone())?,
+                            n: self.positive_subscript(n, Stack, span.clone())?,
                             inverse: false,
                         },
                         self.add_span(span),
                     ),
-                    Primitive::First | Primitive::Last => {
+                    First | Last => {
                         let n = self.positive_subscript(n, prim, span.clone())?;
                         let span = self.add_span(span);
                         match n {
-                            0 => Node::Prim(Primitive::Pop, span),
+                            0 => Node::Prim(Pop, span),
                             1 => Node::Prim(prim, span),
-                            n if prim == Primitive::First => Node::from_iter([
+                            n if prim == First => Node::from_iter([
                                 Node::new_push(n),
-                                Node::Prim(Primitive::Take, span),
+                                Node::Prim(Take, span),
                                 Node::Unpack {
                                     count: n,
                                     unbox: false,
-                                    prim: Some(Primitive::First),
+                                    prim: Some(First),
                                     span,
                                 },
                             ]),
                             n => Node::from_iter([
                                 Node::new_push(-(n as i32)),
-                                Node::Prim(Primitive::Take, span),
-                                Node::Prim(Primitive::Reverse, span),
+                                Node::Prim(Take, span),
+                                Node::Prim(Reverse, span),
                                 Node::Unpack {
                                     count: n,
                                     unbox: false,
-                                    prim: Some(Primitive::Last),
+                                    prim: Some(Last),
                                     span,
                                 },
                             ]),

src/compile/modifier.rs

@@ -444,13 +444,38 @@ impl Compiler {
             On => {
                 let (sn, _) = self.monadic_modifier_op(modified)?;
                 let span = self.add_span(modified.modifier.span.clone());
-                let prim = if sn.sig.args == 0 { Dip } else { On };
-                Node::Mod(prim, eco_vec![sn], span)
+                if let Some(sub) = subscript
+                    .and_then(|sub| self.subscript_n(sub, On))
+                    .filter(|n| n.value > 1)
+                {
+                    let n = self.positive_subscript(sub.value, On, sub.span)?;
+                    Node::ImplMod(ImplPrimitive::OnSub(n), eco_vec![sn], span)
+                } else {
+                    let prim = if sn.sig.args == 0 { Dip } else { On };
+                    Node::Mod(prim, eco_vec![sn], span)
+                }
             }
             By => {
                 let (mut sn, _) = self.monadic_modifier_op(modified)?;
                 let span = self.add_span(modified.modifier.span.clone());
-                if sn.sig.args == 0 {
+                if let Some(sub) = subscript
+                    .and_then(|sub| self.subscript_n(sub, By))
+                    .filter(|n| n.value > 1)
+                {
+                    let n = self.positive_subscript(sub.value, By, sub.span)?;
+                    if n == sn.sig.args {
+                        self.emit_diagnostic(
+                            format!(
+                                "Prefer {} over subscripted {} here",
+                                Below.format(),
+                                By.format()
+                            ),
+                            DiagnosticKind::Style,
+                            modified.modifier.span.clone(),
+                        )
+                    }
+                    Node::ImplMod(ImplPrimitive::BySub(n), eco_vec![sn], span)
+                } else if sn.sig.args == 0 {
                     sn.node.prepend(Node::Prim(Identity, span));
                     sn.node
                 } else {
@@ -593,6 +618,7 @@ impl Compiler {
             prim @ (With | Off) => {
                 let (mut sn, _) = self.monadic_modifier_op(modified)?;
                 let span = self.add_span(modified.modifier.span.clone());
+                let sig = sn.sig;
                 let (inner, before) = match sn.sig.args {
                     0 => (SigNode::new((2, 2), Node::Prim(Identity, span)), sn.node),
                     1 => {
@@ -602,7 +628,34 @@ impl Compiler {
                     }
                     _ => (sn, Node::empty()),
                 };
-                Node::from_iter([before, Node::Mod(prim, eco_vec![inner], span)])
+                Node::from_iter([
+                    before,
+                    if let Some(sub) = subscript
+                        .and_then(|sub| self.subscript_n(sub, prim))
+                        .filter(|n| n.value > 1)
+                    {
+                        let n = self.positive_subscript(sub.value, prim, sub.span)?;
+                        let prim = if prim == Off {
+                            if n == sig.args {
+                                self.emit_diagnostic(
+                                    format!(
+                                        "Prefer {} over subscripted {} here",
+                                        Below.format(),
+                                        Off.format()
+                                    ),
+                                    DiagnosticKind::Style,
+                                    modified.modifier.span.clone(),
+                                )
+                            }
+                            ImplPrimitive::OffSub(n)
+                        } else {
+                            ImplPrimitive::WithSub(n)
+                        };
+                        Node::ImplMod(prim, eco_vec![inner], span)
+                    } else {
+                        Node::Mod(prim, eco_vec![inner], span)
+                    },
+                ])
             }
             prim @ (Above | Below) => {
                 let (mut sn, _) = self.monadic_modifier_op(modified)?;

src/primitive/defs.rs

@@ -1848,10 +1848,10 @@ primitive!(
     /// ex: [⊙⟜⊙⋅⟜∘  1 2 3 4] # Easy to read with ⟜
     ///   : [⊙(.⊙⋅.) 1 2 3 4] # Hard to read with .
     ///   : [⊙.⊙⊙⋅.  1 2 3 4] # Hard to read with .
-    ///
     /// [on] can be used with a function pack. `on``(F|G)` becomes `on``F``on``G`.
     /// ex: [⟜(+1|×2|¯)] 5
-    ///
+    /// Subscripted [on] keeps the first N arguments on top of the stack.
+    /// ex: {⟜₂[⊙⊙∘] 1 2 3}
     /// [on] is equivalent to [fork][identity], but can often be easier to read.
     ([1], On, Stack, ("on", '⟜')),
     /// Duplicate a function's last argument before calling it
@@ -1870,6 +1870,8 @@ primitive!(
     /// ex: ⊂⊸↙ 2 [1 2 3 4 5]
     ///   : ⊜□⊸≠ @  "Hey there buddy"
     ///   : ⊕□⊸◿ 5 [2 9 5 21 10 17 3 35]
+    /// Subscripted [by] keeps the last N arguments below the outputs on the stack.
+    /// ex: {⊸₂[⊙⊙∘] 1 2 3}
     ([1], By, Stack, ("by", '⊸')),
     /// Call a function but keep its last argument on the top of the stack
     ///
@@ -1891,6 +1893,8 @@ primitive!(
     /// [with] with a noadic function will be coerced to `with``identity`.
     /// ex: [⤙1 2 3]
     /// If you do not want these behaviors, use [on] instead.
+    /// Subscripted [with] keeps the last N arguments above the outputs on the stack.
+    /// ex: {⤙₂[⊙⊙∘] 1 2 3}
     ([1], With, Stack, ("with", '⤙')),
     /// Call a function but keep its first argument under the outputs on the stack
     ///
@@ -1913,6 +1917,8 @@ primitive!(
     /// [off] with a noadic function will be coerced to `off``identity`.
     /// ex: [⤚1 2 3]
     /// If you do not want these behaviors, use [by] instead.
+    /// Subscripted [off] keeps the first N arguments below the outputs on the stack.
+    /// ex: {⤚₂[⊙⊙∘] 1 2 3}
     ([1], Off, Stack, ("off", '⤚')),
     /// Keep all arguments to a function above the outputs on the stack
     ///
@@ -3229,6 +3235,10 @@ macro_rules! impl_primitive {
             UndoRotate(usize),
             ReduceDepth(usize),
             StackN { n: usize, inverse: bool },
+            OnSub(usize),
+            BySub(usize),
+            WithSub(usize),
+            OffSub(usize),
         }
 
         impl ImplPrimitive {
@@ -3260,8 +3270,11 @@ macro_rules! impl_primitive {
             pub fn modifier_args(&self) -> Option<usize> {
                 match self {
                     $($(ImplPrimitive::$variant => Some($margs),)?)*
-                    ImplPrimitive::ReduceDepth(_) => Some(1),
-                    ImplPrimitive::EachSub(_) => Some(1),
+                    ImplPrimitive::ReduceDepth(_) | ImplPrimitive::EachSub(_) => Some(1),
+                    ImplPrimitive::OnSub(_)
+                    | ImplPrimitive::BySub(_)
+                    | ImplPrimitive::WithSub(_)
+                    | ImplPrimitive::OffSub(_) => Some(1),
                     _ => None
                 }
             }