feat: under support for regex (#809)

Author: magistau

changelog.md

@@ -48,6 +48,7 @@ This version is not yet released. If you are reading this on the website, then t
   - [`under ⍜`](https://uiua.org/docs/under)[`&cd`](https://uiua.org/docs/&cd) will return to the original directory afterward
 - Add experimental [`reciprocal ⨪`](https://uiua.org/docs/reciprocal) function, which computes the multiplicative inverse AKA reciprocal of a number
 - Make subscripted [`stack ?`](https://uiua.org/docs/stack) merge adjacent non-subscripted [`stack ?`](https://uiua.org/docs/stack) chains
+- Implement [`under ⍜`](https://uiua.org/docs/under)[`regex`](https://uiua.org/docs/regex) for replacing using regex (called `gsub` in some other languages)
 ### Interpreter
 - Speed up the implementation of [`or ∨`](https://uiua.org/docs/or)
 - The fomatter no longer truncates trailing decimal `0`s from number literals

parser/src/defs.rs

@@ -2610,12 +2610,15 @@ primitive!(
     /// ex: regex "hi" "dog"
     ///   : △.
     /// ex: regex "[a-z]+" "hello world"
-    /// Escaped regex characters must be double-escaped.
+    /// If the pattern contains escaped characters such as `\w`, either these must be double escaped or the whole pattern must be represented with a raw string.
     /// ex: regex "\\d+" "123"
-    /// ex: P ← $"(\\d{_})"
+    /// ex: P ← $$ (\d{_})
     ///   : regex $"_-_-_"P3P3P4 "123-456-7890"
     /// Regex patterns with optional captures can be used with [fill].
     /// ex: ⬚""regex "a(b)?" "a ab"
+    /// [under] can be used to run arbitrary regex-based substitutions.
+    /// ex: Lorem
+    ///   : ⍜regex≡(□⊂⋅⊙⇌°□₃) $ (\w)(\w+)
     ///
     /// Uiua uses the [Rust regex crate](https://docs.rs/regex/latest/regex/) internally.
     (2, Regex, Algorithm, "regex"),

site/src/primitive.rs

@@ -353,6 +353,7 @@ fn all_unders() -> impl IntoView {
             { inverse_row([Repeat], Required, "Inner function must be invertible", "⍜⍥(×2). 5 1") }
             { inverse_row([Repeat], No, "Inner function must be invertible", "°⍥°(↧1000×2) 5") }
             { inverse_row([Switch], No, "", "⍜(⨬⊢⊣|×10) 1 [1 2 3 4]") }
+            { inverse_row([Regex], Optional, "Substitution", r#"⍜regex≡(□⊂⋅⊙⇌°□₃) ⊙"Hello, world!" $ (\w)(\w+)"#) }
             { inverse_row([Now], No, "Times execution", "⍜now(&sl 0.005)") }
             { inverse_row([Sys(FOpen)], Optional, view!("Calls "<Prim prim=Sys(Close)/>" on handle"), None) }
             { inverse_row([Sys(FCreate)], Optional, view!("Calls "<Prim prim=Sys(Close)/>" on handle"), None) }

src/compile/invert/mod.rs

@@ -91,6 +91,8 @@ as_node!(A, B, C);
 as_node!(A, B, C, D);
 as_node!(A, B, C, D, E);
 as_node!(A, B, C, D, E, F);
+as_node!(A, B, C, D, E, F, G);
+as_node!(A, B, C, D, E, F, G, H);
 as_node!(A, B, C, D, E, F, G, H, I);
 
 trait SpanFromNodes: Sized + fmt::Debug + Sync {

src/compile/invert/under.rs

@@ -163,6 +163,19 @@ static UNDER_PATTERNS: &[&dyn UnderPattern] = &[
     &Stash(2, Pick, UndoPick),
     &Stash(2, Select, UndoSelect),
     &Stash(2, AntiOrient, UndoAntiOrient),
+    &(
+        Regex,
+        (
+            Over,
+            Flip,
+            DoRegex,
+            PushUnd(1),
+            CopyUnd(1),
+            Flip,
+            PushUnd(1),
+        ),
+        (PopUnd(3), UndoRegex),
+    ),
     // Map control
     &MaybeVal((Get, (CopyUnd(2), Get), (PopUnd(1), Flip, PopUnd(1), Insert))),
     &Stash(2, Remove, UndoRemove),

src/impl_prim.rs

@@ -186,6 +186,8 @@ impl_primitive!(
     (2, UndoWhere),
     (2, AntiOrient),
     (3, UndoAntiOrient),
+    (2(2), DoRegex),
+    (4, UndoRegex),
     (3(2), UndoJoin),
     (1(1)[1], UndoPartition1),
     (3, UndoPartition2),
@@ -374,6 +376,8 @@ impl fmt::Display for ImplPrimitive {
             UndoWhere => write!(f, "{Under}{Where}"),
             AntiOrient => write!(f, "{Anti}{Orient}"),
             UndoAntiOrient => write!(f, "{Under}{Orient}"),
+            DoRegex => write!(f, "{Regex}"),
+            UndoRegex => write!(f, "{Under}{Regex}"),
             UndoInsert => write!(f, "{Under}{Insert}"),
             UndoRemove => write!(f, "{Under}{Remove}"),
             UndoPartition1 | UndoPartition2 => write!(f, "{Under}{Partition}"),

src/run_prim.rs

@@ -3,6 +3,7 @@
 //! For the meat of the actual array algorithms, see [`crate::algorithm`].
 
 use ecow::EcoVec;
+use rayon::iter::Either;
 use regex::Regex;
 use smallvec::SmallVec;
 
@@ -288,7 +289,12 @@ pub fn run_prim_func(prim: &Primitive, env: &mut Uiua) -> UiuaResult {
             env.push(vals);
         }
         Primitive::Stack => stack(env, false)?,
-        Primitive::Regex => regex(env)?,
+        Primitive::Regex => {
+            regex(env)?;
+            // NOTE: if you want to expose the match locations, n.t. they are given in bytes rather than codepoints
+            // that is, you should also convert between byte and codepoint positions in `regex` and `undo_regex`
+            env.pop("regex locations")?;
+        }
         Primitive::Hsv => env.monadic_env(Value::rgb_to_hsv)?,
         Primitive::Json => env.monadic_ref_env(Value::to_json_string)?,
         Primitive::Binary => env.monadic_ref_env(Value::to_binary)?,
@@ -834,6 +840,8 @@ impl ImplPrimitive {
                 let from = env.pop(3)?;
                 env.push(from.undo_anti_orient(indices, into, env)?);
             }
+            ImplPrimitive::DoRegex => regex(env)?,
+            ImplPrimitive::UndoRegex => undo_regex(env)?,
             ImplPrimitive::UndoRerank => {
                 let rank = env.pop(1)?;
                 let shape = Shape::from(
@@ -1565,11 +1573,12 @@ fn regex(env: &mut Uiua) -> UiuaResult {
 
         let mut matches: Value =
             Array::<Boxed>::new([0, regex.captures_len()].as_slice(), []).into();
+        let mut locations: Value = Array::<f64>::new([0].as_slice(), []).into();
 
         for caps in regex.captures_iter(&target) {
-            let row: EcoVec<Boxed> = caps
+            let row: EcoVec<_> = caps
                 .iter()
-                .flat_map(|m| {
+                .filter_map(|m| {
                     m.map(|m| Boxed(Value::from(m.as_str()))).or_else(|| {
                         env.value_fill()
                             .map(|fv| fv.value.clone())
@@ -1578,13 +1587,138 @@ fn regex(env: &mut Uiua) -> UiuaResult {
                 })
                 .collect();
             matches.append(row.into(), false, env)?;
+            locations.append(
+                (caps
+                    .get(0)
+                    .expect("existence of 0 group is guaranteed")
+                    .start() as f64)
+                    .into(),
+                false,
+                env,
+            )?;
         }
 
         env.push(matches);
+        env.push(locations);
         Ok(())
     })
 }
 
+fn undo_regex(env: &mut Uiua) -> UiuaResult {
+    use std::iter::{once, repeat, zip, Repeat};
+    let locations = env
+        .pop(1)?
+        .as_nats(env, "Capture locations should be natural numbers")?;
+    let captures = env.pop(2)?;
+    let haystack = env.pop(3)?.as_string(env, "Haystack should be a string")?;
+    let mut repls = env.pop(4)?;
+    debug_assert_eq!(locations.len(), captures.row_count());
+    let captures: Vec<_> = captures
+        .into_rows()
+        .map(|x| {
+            x.first(env)
+                .expect("Capture group rows should have at least one element")
+                .as_string_opt()
+                .expect("Capture group content should be a string")
+                .len()
+        })
+        .collect();
+    enum OneOrMany<T> {
+        One(T),
+        Many(Vec<T>),
+    }
+    impl<T: Clone> IntoIterator for OneOrMany<T> {
+        type Item = T;
+        type IntoIter = Either<Repeat<T>, std::vec::IntoIter<T>>;
+
+        fn into_iter(self) -> Self::IntoIter {
+            match self {
+                Self::One(x) => Either::Left(repeat(x)),
+                Self::Many(items) => Either::Right(items.into_iter()),
+            }
+        }
+    }
+    use OneOrMany::{Many, One};
+    impl<'a, T> IntoIterator for &'a OneOrMany<T> {
+        type Item = &'a T;
+        type IntoIter = Either<Repeat<&'a T>, std::slice::Iter<'a, T>>;
+
+        fn into_iter(self) -> Self::IntoIter {
+            match self {
+                One(x) => Either::Left(repeat(x)),
+                Many(items) => Either::Right(items.iter()),
+            }
+        }
+    }
+    let repls: OneOrMany<String> = loop {
+        let rank = repls.rank();
+        repls =
+            match (repls, rank) {
+                (Value::Char(arr), 0 | 1) => break One(arr.data.into_iter().collect()),
+                (Value::Char(arr), 2) => {
+                    break Many(arr.row_slices().map(|x| x.iter().collect()).collect())
+                }
+                (Value::Char(arr), 3) => {
+                    break Many(
+                        arr.into_rows()
+                            .map(|x| x.first(env).map(|x| x.data.into_iter().collect()))
+                            .collect::<UiuaResult<Vec<_>>>()?,
+                    )
+                }
+                (Value::Box(bx), 0) => bx
+                    .into_unboxed()
+                    .expect("Scalar box array should be unboxable"),
+                (Value::Box(arr), 1) => {
+                    break Many(
+                        arr.data
+                            .into_iter()
+                            .map(|x| x.0)
+                            .map(|x| x.as_string(env, "Expected boxed replacements to be strings"))
+                            .collect::<UiuaResult<_>>()?,
+                    )
+                }
+                (Value::Box(arr), 2) => {
+                    break Many(
+                        arr.into_rows()
+                            .map(|x| {
+                                x.first(env).and_then(|x| {
+                                    x.into_unboxed()
+                                        .expect("rank 0 box array should be unboxable")
+                                        .as_string(env, "Expected boxed replacements to be strings")
+                                })
+                            })
+                            .collect::<UiuaResult<_>>()?,
+                    )
+                }
+                _ => return Err(env.error(
+                    "Expected replacements to be a string, array of strings or 2d array of strings",
+                )),
+            }
+    };
+    if let Many(ref x) = repls {
+        if locations.len() != x.len() {
+            return Err(env.error("Expected to have the same amount of captures and replacements"));
+        }
+    }
+    let repls = (&repls)
+        .into_iter()
+        .take(locations.len())
+        .map(String::as_str)
+        .chain(once(""));
+    let ends = zip(&locations, captures).map(|(&loc, len)| loc + len);
+    let chunks = zip(
+        once(0).chain(ends),
+        locations.iter().copied().chain(once(haystack.len())),
+    )
+    .map(|(s, e)| &haystack[s..e]);
+    env.push(
+        zip(chunks, repls)
+            .flat_map(|(chunk, repl)| [chunk, repl])
+            .collect::<String>(),
+    );
+    Ok(())
+}
+
 thread_local! {
     pub(crate) static RNG: RefCell<SmallRng> = RefCell::new(SmallRng::from_entropy());
 }

tests/under.ua

@@ -482,3 +482,25 @@ F ← ⌅(+|⌅(-|⟜-|×10+))
 ⍤⤙≍ [¯1 ∞ ¯2] [⍜𝄐⊟¯ 1 ∞ 2]
 ⍤⤙≍ [4 ∞ 5] [⍜(⊟⊙◌)+ 1 ∞ 2 3]
 ⍤⤙≍ [¯1 ∞ ¯2] [⍜(⊟⊙◌)¯ 1 ∞ 2]
+
+# Regex
+⍤⤙≍ Lorem ⍜(regex"\\w+")∘ Lorem
+H ← "Hello,   world!!"
+⍤⤙≍ "world,   Hello!!" ⍜regex⇌ ⊙H$ \w+
+⍤⤙≍ "Holle,   wdlro!!" ⍜regex≡(□⊂⋅⊙⇌°□₃) ⊙H$ (\w)(\w+)
+⍤⤙≍ 0 ⍤⤙≍ H ⍜regex⟜⧻ ⊙H$ \d+
+⍤⤙≍ "@,   w!!" ⬚@@⍜regex≡⊣ ⊙H$ \b([a-z])?\w*\s?
+⍤⤙≍ "HHello,   wworld!!" ⍜regex≡(▽2°□⊣) ⊙H$ \b(\w)
+⍤⤙≍ ".H.e.l.l.o.,. . . .w.o.r.l.d.!.!." ⍜regex⋅@. "" H
+# different replacement forms:
+⊙H $ (\w)(\w*)
+⍤⤙≍ "olleH,   dlrow!!" ◡⍜regex≡⍜⊢⇌ # 2d, boxed strings
+# TODO: remove ∘ when the misoptimization is fixed
+⍤⤙≍ "olleH,   dlrow!!" ◡⍜regex≡(∘⇌⊢) # 1d boxed strings
+⍤⤙≍ "bip,   bip!!" ◡⍜regex⋅(□"bip") # boxed string
+⍤⤙≍ "*,   *!!" ◡⍜regex⋅(□@*) # boxed char
+⍤⤙≍ "olleH,   dlrow!!" ◡⍜regex≡(⇌°□⊢) # 2d strings
+⍤⤙≍ "olleH,   dlrow!!" ◡⍜regex≡(⊟⊸⇌°□⊢) # 1d strings
+⍤⤙≍ "bip,   bip!!" ◡⍜regex⋅"bip" # string
+⍤⤙≍ "*,   *!!" ◡⍜regex⋅@* # char
+⋅◌