add orient inverses

src/algorithm/dyadic/mod.rs

@@ -1728,7 +1728,11 @@ impl Value {
  return Err(env.error("Orient indices must be unique"));
  }
  }
- let mut orientation: Vec<usize> = (0..target.rank()).collect();
+ target.orient_impl(&undices);
+ Ok(())
+ }
+ fn orient_impl(&mut self, undices: &[usize]) {
+ let mut orientation: Vec<usize> = (0..self.rank()).collect();
  let mut depth_rotations: Vec<(usize, i32)> = Vec::new();
  for (i, &u) in undices.iter().enumerate() {
  let j = orientation.iter().position(|&o| o == u).unwrap();
@@ -1743,8 +1747,43 @@ impl Value {
  depth_rotations.push((i, -1));
  }
  for (depth, amnt) in depth_rotations {
- target.transpose_depth(depth, amnt);
+ self.transpose_depth(depth, amnt);
+ }
+ }
+ pub(crate) fn undo_orient(&self, target: &mut Self, env: &Uiua) -> UiuaResult {
+ let indices = self.as_ints(env, "Unorient indices must be integers")?;
+ let mut undices = Vec::with_capacity(indices.len());
+ for i in indices {
+ let u = i.unsigned_abs();
+ if u >= target.rank() {
+ return Err(env.error(format!(
+ "Cannot unorient axis {i} in array of rank {}",
+ target.rank()
+ )));
+ }
+ if i >= 0 {
+ undices.push(u);
+ } else {
+ undices.push(target.rank() - u);
+ }
+ }
+ if undices.len() > target.rank() {
+ return Err(env.error(format!(
+ "Cannot unorient array of rank {} with {} indices",
+ target.rank(),
+ undices.len()
+ )));
+ }
+ for i in 0..target.rank() {
+ if !undices.contains(&i) {
+ undices.push(i);
+ }
+ }
+ let mut inverted_undices = undices.clone();
+ for (i, u) in undices.into_iter().enumerate() {
+ inverted_undices[u] = i;
  }
+ target.orient_impl(&inverted_undices);
  Ok(())
  }
 }

src/algorithm/invert.rs

@@ -225,6 +225,7 @@ static ON_INVERT_PATTERNS: &[&dyn InvertPattern] = {
  &pat!((Flip, Log), (Flip, 1, Flip, Div, Pow)),
  &([Min], [Min]),
  &([Max], [Max]),
+ &([Orient], [UndoOrient]),
  &pat!(
  Join,
  (
@@ -535,6 +536,12 @@ pub(crate) fn under_instrs(
  (Dup, Shape, PushToUnder(1), Where),
  (PopUnder(1), UndoWhere)
  ),
+ // Orient
+ &maybe_val!(pat!(
+ Orient,
+ (CopyToUnder(1), Orient),
+ (PopUnder(1), UndoOrient)
+ )),
  // System stuff
  &pat!(Now, (Now, PushToUnder(1)), (PopUnder(1), Now, Flip, Sub)),
  &maybe_val!(store1copy!(Sys(SysOp::FOpen), Sys(SysOp::Close))),

src/primitive/defs.rs

@@ -2747,6 +2747,7 @@ impl_primitive!(
  (3, UndoRerank),
  (2, UndoReshape),
  (2, UndoWhere),
+ (2, UndoOrient),
  (3(2), UndoJoin),
  (1[1], UndoPartition1),
  (3, UndpPartition2),

src/primitive/mod.rs

@@ -184,6 +184,7 @@ impl fmt::Display for ImplPrimitive {
  UndoSelect => write!(f, "{Under}{Select}"),
  UndoPick => write!(f, "{Under}{Pick}"),
  UndoWhere => write!(f, "{Under}{Where}"),
+ UndoOrient => write!(f, "{Under}{Orient}"),
  UndoInsert => write!(f, "{Under}{Insert}"),
  UndoRemove => write!(f, "{Under}{Remove}"),
  UndoPartition1 | UndpPartition2 => write!(f, "{Under}{Partition}"),
@@ -920,6 +921,12 @@ impl ImplPrimitive {
  let mask = indices.undo_where(&shape, env)?;
  env.push(mask);
  }
+ ImplPrimitive::UndoOrient => {
+ let indices = env.pop(1)?;
+ let mut array = env.pop(2)?;
+ indices.undo_orient(&mut array, env)?;
+ env.push(array);
+ }
  ImplPrimitive::UndoRerank => {
  let rank = env.pop(1)?;
  let shape = env.pop(2)?;