-
Notifications
You must be signed in to change notification settings - Fork 146
Expand file tree
/
Copy pathmod.rs
More file actions
114 lines (94 loc) · 3.71 KB
/
mod.rs
File metadata and controls
114 lines (94 loc) · 3.71 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors
//! The optimizer applies metadata-only rewrite rules (`reduce` and `reduce_parent`) in a
//! fixpoint loop until no more transformations are possible.
//!
//! Optimization runs between execution steps, which is what enables cross-step optimizations:
//! after a child is decoded, new `reduce_parent` rules may match that were previously blocked.
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use crate::ArrayRef;
pub mod rules;
/// Extension trait for optimizing array trees using reduce/reduce_parent rules.
pub trait ArrayOptimizer {
/// Optimize the root array node only by running reduce and reduce_parent rules to fixpoint.
fn optimize(&self) -> VortexResult<ArrayRef>;
/// Optimize the entire array tree recursively (root and all descendants).
fn optimize_recursive(&self) -> VortexResult<ArrayRef>;
}
impl ArrayOptimizer for ArrayRef {
fn optimize(&self) -> VortexResult<ArrayRef> {
Ok(try_optimize(self)?.unwrap_or_else(|| self.clone()))
}
fn optimize_recursive(&self) -> VortexResult<ArrayRef> {
Ok(try_optimize_recursive(self)?.unwrap_or_else(|| self.clone()))
}
}
fn try_optimize(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
let mut current_array = array.clone();
let mut any_optimizations = false;
// Apply reduction rules to the current array until no more rules apply.
let mut loop_counter = 0;
'outer: loop {
if loop_counter > 100 {
vortex_bail!("Exceeded maximum optimization iterations (possible infinite loop)");
}
loop_counter += 1;
if let Some(new_array) = current_array.reduce()? {
current_array = new_array;
any_optimizations = true;
continue;
}
// Apply parent reduction rules to each slot in the context of the current array.
// Its important to take all slots here, as `current_array` can change inside the loop.
for (slot_idx, slot) in current_array.slots().iter().enumerate() {
let Some(child) = slot else { continue };
if let Some(new_array) = child.reduce_parent(¤t_array, slot_idx)? {
// If the parent was replaced, then we attempt to reduce it again.
current_array = new_array;
any_optimizations = true;
// Continue to the start of the outer loop
continue 'outer;
}
}
// No more optimizations can be applied
break;
}
if any_optimizations {
Ok(Some(current_array))
} else {
Ok(None)
}
}
fn try_optimize_recursive(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
let mut current_array = array.clone();
let mut any_optimizations = false;
if let Some(new_array) = try_optimize(¤t_array)? {
current_array = new_array;
any_optimizations = true;
}
let mut new_slots = Vec::with_capacity(current_array.slots().len());
let mut any_slot_optimized = false;
for slot in current_array.slots() {
match slot {
Some(child) => {
if let Some(new_child) = try_optimize_recursive(child)? {
new_slots.push(Some(new_child));
any_slot_optimized = true;
} else {
new_slots.push(Some(child.clone()));
}
}
None => new_slots.push(None),
}
}
if any_slot_optimized {
current_array = current_array.with_slots(new_slots)?;
any_optimizations = true;
}
if any_optimizations {
Ok(Some(current_array))
} else {
Ok(None)
}
}