Skip to content

Commit dfa3c63

Browse files
kaikaliikaikalii
committed
optimize all permutation in tuples
1 parent 36b679a commit dfa3c63

File tree

4 files changed

+115
-29
lines changed

4 files changed

+115
-29
lines changed

‎changelog.md

+1
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ This version is not yet released. If you are reading this on the website, then t
3939
- Optimizations
4040
- [`stencil ⧈`](https://uiua.org/docs/stencil) is now optimized when used with with monadic functions that are optimized for [`rows ≡`](https://uiua.org/docs/rows)
4141
- Add sortedness flags to arrays to allow short-circuiting some operations
42+
- Optimize [`tuples ⧅`](https://uiua.org/docs/tuples) with `⋅⋅1` and `⋅⧻`
4243
- LSP improvements
4344
- Completions now properly respect scoping
4445
- Unused private bindings are now dimmed

‎site/src/other.rs

+24-11
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,8 @@ pub fn Optimizations() -> impl IntoView {
251251
<tr><td><Prim prim=Box/></td></tr>
252252
<tr><td><Prim prim=First/></td></tr>
253253
<tr><td><Prims prims=[First, Reverse] show_names=true/></td></tr>
254+
<tr><td><Prims prims=[Sort]/>" / "<Prims prims=[Select, Rise, Dup]/>" / "<Prims prims=[Select, By, Rise]/></td></tr>
255+
<tr><td><Prims prims=[Reverse, Sort]/>" / "<Prims prims=[Select, Fall, Dup]/>" / "<Prims prims=[Select, By, Fall]/></td></tr>
254256
</table>
255257
<table class="bordered-table cell-centered-table">
256258
<tr><td><Prims prims=[Gap, Rand] show_names=true/></td></tr>
@@ -259,8 +261,6 @@ pub fn Optimizations() -> impl IntoView {
259261
<tr><td><Prim prim=Gap/><code>"constant"</code></td></tr>
260262
<tr><td><Prim prim=On/><code>"constant"</code></td></tr>
261263
<tr><td><Prim prim=By/><code>"constant"</code></td></tr>
262-
<tr><td><Prims prims=[Sort]/>" / "<Prims prims=[Select, Rise, Dup]/>" / "<Prims prims=[Select, By, Rise]/></td></tr>
263-
<tr><td><Prims prims=[Reverse, Sort]/>" / "<Prims prims=[Select, Fall, Dup]/>" / "<Prims prims=[Select, By, Fall]/></td></tr>
264264
<tr><td><Prims prims=[Un, Couple] show_names=true/></td></tr>
265265
<tr><td><Prims prims=[Un, Join] show_names=true/></td></tr>
266266
<tr><td><Prim prim=Rotate/></td></tr>
@@ -337,16 +337,29 @@ pub fn Optimizations() -> impl IntoView {
337337

338338
<Hd id="other-optimizations">"Other Optimizations"</Hd>
339339
<ul>
340-
<li><Prims prims=[Table, Fork]/><code>"F"</code><code>"G"</code>" is optimized to "<Prims prims=[Fork, Table]/><code>"F"</code><Prims prims=[Table]/><code>"G"</code>" for pure functions."</li>
341340
<li><Prim prim=Group/>" and "<Prim prim=Partition/>" are optimized to be fast with "<Prim prim=Len/>", "<Prim prim=First/>", "<Prim prim=Last/>"."</li>
342-
<li><Prim prim=Repeat/>" with a simple-enough function and a small constant count will compile to an unrolled loop."</li>
343-
<li>"The following splitting patterns are optimized for monadic function "<code>"F"</code>":"
344-
<table class="bordered-table">
345-
<tr><td><Prims prims=[Partition]/><code>"F"</code><Prims prims=[By, Ne]/></td></tr>
346-
<tr><td><Prims prims=[Partition]/><code>"F"</code><Prims prims=[Ne]/><code>"constant"</code><Prims prims=[Dup]/></td></tr>
347-
<tr><td><Prims prims=[Partition]/><code>"F"</code><Prims prims=[Not, By, Mask]/></td></tr>
348-
<tr><td><Prims prims=[Partition]/><code>"F"</code><Prims prims=[Not, Mask]/><code>"constant"</code><Prims prims=[Dup]/></td></tr>
349-
</table>
341+
<li>
342+
<Prim prim=Tuples/>" with the following functions:"
343+
<ul>
344+
<li><Prim prim=Lt/></li>
345+
<li><Prim prim=Le/></li>
346+
<li><Prim prim=Gt/></li>
347+
<li><Prim prim=Ge/></li>
348+
<li><Prim prim=Ne/></li>
349+
<li><Prim prim=Eq/></li>
350+
<li><Prim prim=Match/></li>
351+
<li><Prims prims=[Gap, Gap]/><code>"constant"</code></li>
352+
<li><Prims prims=[Gap, Len]/></li>
353+
</ul>
354+
</li>
355+
<li>
356+
"The following splitting patterns are optimized for monadic function "<code>"F"</code>":"
357+
<ul>
358+
<li><Prims prims=[Partition]/><code>"F"</code><Prims prims=[By, Ne]/></li>
359+
<li><Prims prims=[Partition]/><code>"F"</code><Prims prims=[Ne]/><code>"constant"</code><Prims prims=[Dup]/></li>
360+
<li><Prims prims=[Partition]/><code>"F"</code><Prims prims=[Not, By, Mask]/></li>
361+
<li><Prims prims=[Partition]/><code>"F"</code><Prims prims=[Not, Mask]/><code>"constant"</code><Prims prims=[Dup]/></li>
362+
</ul>
350363
</li>
351364
</ul>
352365
}

‎src/algorithm/tuples.rs

+84-18
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ use std::collections::{hash_map::Entry, HashMap};
33
use ecow::EcoVec;
44

55
use crate::{
6-
get_ops, grid_fmt::GridFmt, types::push_empty_rows_value, val_as_arr, Array, ArrayValue, Ops,
7-
Primitive, SigNode, Uiua, UiuaResult, Value,
6+
get_ops, grid_fmt::GridFmt, types::push_empty_rows_value, val_as_arr, Array, ArrayValue, Node,
7+
Ops, Primitive, SigNode, Uiua, UiuaResult, Value,
88
};
99

1010
use super::{monadic::range, table::table_impl, validate_size};
@@ -108,23 +108,44 @@ fn tuple2(f: SigNode, env: &mut Uiua) -> UiuaResult {
108108
};
109109

110110
'blk: {
111-
if let Some(prim) = f.node.as_primitive() {
112-
let res = match prim {
113-
Primitive::Lt => xs.choose(k, false, false, env)?,
114-
Primitive::Le if n >= k => xs.choose(k, false, true, env)?,
115-
Primitive::Gt => xs.choose(k, true, false, env)?,
116-
Primitive::Ge if n >= k => xs.choose(k, true, true, env)?,
117-
Primitive::Ne => xs.permute(k, env)?,
118-
Primitive::Eq | Primitive::Match if is_scalar => {
119-
let n = xs.as_nat(env, "Tuples of scalar must be a natural number")?;
120-
env.push(n);
121-
return Ok(());
111+
let res = match f.node.as_slice() {
112+
[Node::Prim(Primitive::Lt, _)] => xs.choose(k, false, false, env)?,
113+
[Node::Prim(Primitive::Le, _)] if n >= k => xs.choose(k, false, true, env)?,
114+
[Node::Prim(Primitive::Gt, _)] => xs.choose(k, true, false, env)?,
115+
[Node::Prim(Primitive::Ge, _)] if n >= k => xs.choose(k, true, true, env)?,
116+
[Node::Prim(Primitive::Ne, _)] => xs.permute(k, env)?,
117+
[Node::Prim(Primitive::Eq | Primitive::Match, _)] if is_scalar => {
118+
let n = xs.as_nat(env, "Tuples of scalar must be a natural number")?;
119+
env.push(n);
120+
return Ok(());
121+
}
122+
[Node::Prim(Primitive::Pop, _), Node::Prim(Primitive::Pop, _), Node::Push(val)] => {
123+
if let Ok(reps) = val.as_nat(env, None) {
124+
if k > 2 && reps > 1 {
125+
return Err(val
126+
.as_bool(env, "tuples of 3 or more must return a boolean")
127+
.unwrap_err());
128+
}
129+
if is_scalar {
130+
((n as f64).powi(k as i32) * reps as f64).into()
131+
} else {
132+
xs.permute_all(k, reps, env)?
133+
}
134+
} else {
135+
break 'blk;
122136
}
123-
_ => break 'blk,
124-
};
125-
env.push(res);
126-
return Ok(());
127-
}
137+
}
138+
[Node::Prim(Primitive::Pop, _), Node::Prim(Primitive::Len, _)] => {
139+
if is_scalar {
140+
((n as f64).powi(k as i32)).into()
141+
} else {
142+
xs.permute_all(k, 1, env)?
143+
}
144+
}
145+
_ => break 'blk,
146+
};
147+
env.push(res);
148+
return Ok(());
128149
}
129150
match k {
130151
0 if is_scalar => xs = 0.into(),
@@ -237,6 +258,7 @@ fn tuple2(f: SigNode, env: &mut Uiua) -> UiuaResult {
237258
count += 1;
238259
}
239260
// Increment curr
261+
env.respect_execution_limit()?;
240262
for i in (0..k).rev() {
241263
if curr[i] == row_count - 1 {
242264
curr[i] = 0;
@@ -294,6 +316,12 @@ impl Value {
294316
}
295317
val_as_arr!(self, |a| a.permute(k, env).map(Into::into))
296318
}
319+
fn permute_all(self, k: usize, reps: usize, env: &Uiua) -> UiuaResult<Self> {
320+
if let Ok(n) = self.as_nat(env, None) {
321+
return Ok((n as f64).powi(k as i32).into());
322+
}
323+
val_as_arr!(self, |a| a.permute_all(k, reps, env).map(Into::into))
324+
}
297325
}
298326

299327
fn combinations(n: usize, k: usize, same: bool) -> f64 {
@@ -486,4 +514,42 @@ impl<T: ArrayValue> Array<T> {
486514
}
487515
Ok(Array::new(shape, data))
488516
}
517+
fn permute_all(self, k: usize, reps: usize, env: &Uiua) -> UiuaResult<Self> {
518+
let n = self.row_count();
519+
let permutations = (n as f64).powi(k as i32) * reps as f64;
520+
if permutations.is_nan() {
521+
return Err(env.error("Combinatorial explosion"));
522+
}
523+
if permutations > usize::MAX as f64 {
524+
return Err(env.error(format!(
525+
"{} permutations would be too many",
526+
permutations.grid_string(false)
527+
)));
528+
}
529+
let mut shape = self.shape.clone();
530+
shape[0] = permutations.round() as usize;
531+
shape.insert(1, k);
532+
let elem_count = validate_size::<T>(shape.iter().copied(), env)?;
533+
let mut data = EcoVec::with_capacity(elem_count);
534+
let row_len = self.row_len();
535+
let mut curr = vec![0; k];
536+
'outer: loop {
537+
for _ in 0..reps {
538+
for &i in &curr {
539+
data.extend_from_slice(&self.data[i * row_len..][..row_len]);
540+
}
541+
}
542+
// Increment curr
543+
for i in (0..k).rev() {
544+
if curr[i] == n - 1 {
545+
curr[i] = 0;
546+
} else {
547+
curr[i] += 1;
548+
continue 'outer;
549+
}
550+
}
551+
break;
552+
}
553+
Ok(Array::new(shape, data))
554+
}
489555
}

‎tests/optimized.ua

+6
Original file line numberDiff line numberDiff line change
@@ -292,6 +292,12 @@ F ← ⬚10(/+◌1⊞(ׯ))
292292
⍤⤙≍ [0 3 2] △ ⧅(≤∘) 3 °△[0 2]
293293
⍤⤙≍ [1 3 2] △ ⧅(≤∘) 3 °△[1 2]
294294
⍤⤙≍ [4 3 2] △ ⧅(≤∘) 3 °△[2 2]
295+
⍤⤙≍ ⊃⧅(∘1◌◌)⧅⋅⋅1 4 3
296+
⍤⤙≍ ⊃⧅(∘1◌◌)⧅⋅⋅1 4 ⇡3
297+
⍤⤙≍ ⊃⧅(∘1◌◌)⧅⋅⋅1 4 [1_2 3_4]
298+
⍤⤙≍ ⊃⧅(∘2◌◌)⧅⋅⋅2 2 3
299+
⍤⤙≍ ⊃⧅(∘2◌◌)⧅⋅⋅2 2 ⇡3
300+
⍤⤙≍ ⊃⧅(∘2◌◌)⧅⋅⋅2 2 [1_2 3_4]
295301

296302
⍤⤙≍ ℂ5 0 ¯₄ 5
297303
⍤⤙≍ ℂ¯π 0 °¯₄ π

0 commit comments

Comments
 (0)