Skip to content

Commit 38c8e2f

Browse files
kaikaliikaikalii
committed
factor out some shared pervasion code
1 parent 29c9555 commit 38c8e2f

File tree

1 file changed

+83
-59
lines changed

1 file changed

+83
-59
lines changed

‎src/algorithm/pervade.rs

Lines changed: 83 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -92,44 +92,59 @@ where
9292
let _a_depth = a_depth.min(a.rank());
9393
let _b_depth = b_depth.min(b.rank());
9494

95-
let new_rank = a.rank().max(b.rank());
96-
let mut new_shape = Shape::with_capacity(new_rank);
97-
let a_fill = env.scalar_fill::<A>();
98-
let b_fill = env.scalar_fill::<B>();
99-
for i in 0..new_rank {
100-
let c = match (a.shape.get(i).copied(), b.shape.get(i).copied()) {
101-
(None, None) => unreachable!(),
102-
(Some(a), None) => a,
103-
(None, Some(b)) => b,
104-
(Some(ad), Some(bd)) => {
105-
if ad == bd || ad == 1 || bd == 1 {
106-
pervade_dim(ad, bd)
107-
} else if ad < bd {
108-
match &a_fill {
109-
Ok(_) => pervade_dim(ad, bd),
110-
Err(e) => {
111-
return Err(env.error(format!(
112-
"Shapes {} and {} are not compatible{e}",
113-
a.shape, b.shape
114-
)))
95+
fn derive_new_shape(
96+
ash: &Shape,
97+
bsh: &Shape,
98+
a_fill_err: Option<&'static str>,
99+
b_fill_err: Option<&'static str>,
100+
env: &Uiua,
101+
) -> UiuaResult<Shape> {
102+
let new_rank = ash.len().max(bsh.len());
103+
let mut new_shape = Shape::with_capacity(new_rank);
104+
for i in 0..new_rank {
105+
let c = match (ash.get(i).copied(), bsh.get(i).copied()) {
106+
(None, None) => unreachable!(),
107+
(Some(a), None) => a,
108+
(None, Some(b)) => b,
109+
(Some(ad), Some(bd)) => {
110+
if ad == bd || ad == 1 || bd == 1 {
111+
pervade_dim(ad, bd)
112+
} else if ad < bd {
113+
match a_fill_err {
114+
None => pervade_dim(ad, bd),
115+
Some(e) => {
116+
return Err(env.error(format!(
117+
"Shapes {ash} and {bsh} are not compatible{e}"
118+
)))
119+
}
115120
}
116-
}
117-
} else {
118-
match &b_fill {
119-
Ok(_) => pervade_dim(ad, bd),
120-
Err(e) => {
121-
return Err(env.error(format!(
122-
"Shapes {} and {} are not compatible{e}",
123-
a.shape, b.shape
124-
)))
121+
} else {
122+
match b_fill_err {
123+
None => pervade_dim(ad, bd),
124+
Some(e) => {
125+
return Err(env.error(format!(
126+
"Shapes {ash} and {bsh} are not compatible{e}"
127+
)))
128+
}
125129
}
126130
}
127131
}
128-
}
129-
};
130-
new_shape.push(c);
132+
};
133+
new_shape.push(c);
134+
}
135+
Ok(new_shape)
131136
}
132137

138+
let a_fill = env.scalar_fill::<A>();
139+
let b_fill = env.scalar_fill::<B>();
140+
let new_shape = derive_new_shape(
141+
&a.shape,
142+
&b.shape,
143+
a_fill.as_ref().err().copied(),
144+
b_fill.as_ref().err().copied(),
145+
env,
146+
)?;
147+
133148
// dbg!(&a.shape, &b.shape, &new_shape);
134149

135150
let mut new_data = eco_vec![C::default(); new_shape.elements()];
@@ -273,38 +288,47 @@ where
273288
let _a_depth = a_depth.min(a.rank());
274289
let _b_depth = b_depth.min(b.rank());
275290

276-
let new_rank = a.rank().max(b.rank());
277-
let mut new_shape = Shape::with_capacity(new_rank);
278-
let fill = env.scalar_fill::<T>();
279-
let mut requires_fill = false;
280-
for i in 0..new_rank {
281-
let c = match (a.shape.get(i).copied(), b.shape.get(i).copied()) {
282-
(None, None) => unreachable!(),
283-
(Some(a), None) => a,
284-
(None, Some(b)) => b,
285-
(Some(ad), Some(bd)) => {
286-
if ad == bd || ad == 1 || bd == 1 {
287-
requires_fill |= ad != bd && fill.is_ok();
288-
pervade_dim(ad, bd)
289-
} else {
290-
match &fill {
291-
Ok(_) => {
292-
requires_fill = true;
293-
pervade_dim(ad, bd)
294-
}
295-
Err(e) => {
296-
return Err(env.error(format!(
297-
"Shapes {} and {} are not compatible{e}",
298-
a.shape, b.shape
299-
)))
291+
fn derive_new_shape(
292+
ash: &Shape,
293+
bsh: &Shape,
294+
fill_err: Option<&str>,
295+
env: &Uiua,
296+
) -> UiuaResult<(Shape, bool)> {
297+
let new_rank = ash.len().max(bsh.len());
298+
let mut new_shape = Shape::with_capacity(new_rank);
299+
let mut requires_fill = false;
300+
for i in 0..new_rank {
301+
let c = match (ash.get(i).copied(), bsh.get(i).copied()) {
302+
(None, None) => unreachable!(),
303+
(Some(a), None) => a,
304+
(None, Some(b)) => b,
305+
(Some(ad), Some(bd)) => {
306+
if ad == bd || ad == 1 || bd == 1 {
307+
requires_fill |= ad != bd && fill_err.is_none();
308+
pervade_dim(ad, bd)
309+
} else {
310+
match &fill_err {
311+
None => {
312+
requires_fill = true;
313+
pervade_dim(ad, bd)
314+
}
315+
Some(e) => {
316+
return Err(env.error(format!(
317+
"Shapes {ash} and {bsh} are not compatible{e}"
318+
)))
319+
}
300320
}
301321
}
302322
}
303-
}
304-
};
305-
new_shape.push(c);
323+
};
324+
new_shape.push(c);
325+
}
326+
Ok((new_shape, requires_fill))
306327
}
307328

329+
let fill = env.scalar_fill::<T>();
330+
let (new_shape, requires_fill) =
331+
derive_new_shape(a.shape(), b.shape(), fill.as_ref().err().copied(), env)?;
308332
let fill = if requires_fill { fill.ok() } else { None };
309333

310334
fn reuse_no_fill<T: ArrayValue + Copy>(

0 commit comments

Comments
 (0)