1use sva_formula::closed_form::{Part, map_children};
4use sva_formula::spectral_sum::atom::SpectralAtom;
5use sva_formula::{Body, ClosedForm, Lane, Line, SpectralSum, Var, normalize_closed_form};
6
7use super::truncate::Audible;
8use super::{Horizon, atoms, lines, point, truncate};
9use crate::error::CollapseError;
10use crate::label::Rule;
11use crate::profile::Profile;
12
13pub fn direct_is_cheaper(lines: usize, n: usize, len: usize) -> bool {
16 (lines as f64) * (len as f64) <= (n as f64) * (n as f64).log2()
17}
18
19pub struct LinePlan {
20 pub placed: Vec<Vec<Line>>,
21 pub summed: Vec<Vec<Line>>,
22 pub dropped: Vec<Vec<Line>>,
23 pub bins: usize,
24 pub tail_db: Option<f64>,
25}
26
27pub struct Group {
29 pub factor: SpectralAtom,
30 pub placed: Vec<Line>,
31 pub summed: Vec<Line>,
32}
33
34pub enum LanePlan {
35 Grouped { groups: Vec<Group>, bins: usize },
36 Sweep { atoms: usize, samples: usize },
37}
38
39pub struct Sampled {
40 pub sum: Box<SpectralSum>,
41 pub rule: Rule,
42 pub lanes: Vec<LanePlan>,
43}
44
45pub enum Plan {
47 Spectrum(Box<SpectralSum>),
48 Lines(Box<LinePlan>),
49 Sampled(Box<Sampled>),
50 Point {
52 written: Box<ClosedForm>,
53 nodes: usize,
54 width: usize,
55 },
56 Added(Vec<Plan>),
58}
59
60pub fn of(
62 sum: &SpectralSum,
63 rate: u32,
64 horizon: Horizon,
65 profile: &Profile,
66 len: usize,
67) -> Result<Plan, CollapseError> {
68 let ceiling = profile.ceiling(rate);
69 if sum.var == Var::F {
70 return Ok(Plan::Spectrum(Box::new(sum.clone())));
71 }
72 if let Some(found) = line_plan(sum, rate, horizon, profile, len, ceiling)? {
73 return Ok(Plan::Lines(Box::new(found)));
74 }
75 let truncated = truncate::spectral_sum(sum, Audible::of(profile, rate))?;
76 let rule = match () {
77 () if atoms::band_limited(&truncated, ceiling, profile) => Rule::BandLimited,
78 () if atoms::windowed(&truncated) => Rule::CroppedPair,
79 () => Rule::PointSampled,
80 };
81 let lanes = truncated
82 .lanes
83 .iter()
84 .map(|lane| lane_plan(lane, rate, horizon, len))
85 .collect();
86 Ok(Plan::Sampled(Box::new(Sampled {
87 sum: Box::new(truncated),
88 rule,
89 lanes,
90 })))
91}
92
93pub fn of_written(
95 form: &ClosedForm,
96 rate: u32,
97 horizon: Horizon,
98 profile: &Profile,
99 len: usize,
100) -> Result<Plan, CollapseError> {
101 let Some(addends) = addends(form) else {
102 return point_plan(form, rate, profile);
103 };
104 let mut parts = Vec::with_capacity(addends.len());
105 for addend in &addends {
106 match of_term(addend, rate, horizon, profile, len)? {
107 Plan::Added(inner) => parts.extend(inner),
108 part => parts.push(part),
109 }
110 }
111 match parts.iter().all(|part| matches!(part, Plan::Point { .. })) {
113 true => point_plan(form, rate, profile),
114 false => Ok(Plan::Added(parts)),
115 }
116}
117
118fn of_term(
121 form: &ClosedForm,
122 rate: u32,
123 horizon: Horizon,
124 profile: &Profile,
125 len: usize,
126) -> Result<Plan, CollapseError> {
127 let Ok(sum) = normalize_closed_form(form) else {
128 return of_written(form, rate, horizon, profile, len);
129 };
130 match of(&sum, rate, horizon, profile, len) {
131 Err(nested @ CollapseError::NestedSeries { .. }) => Err(nested),
132 Err(_) => of_written(form, rate, horizon, profile, len),
133 held => held,
134 }
135}
136
137fn addends(form: &ClosedForm) -> Option<Vec<ClosedForm>> {
138 if let Body::Add(parts) = &form.body {
139 return Some(
140 parts
141 .iter()
142 .map(|part| ClosedForm {
143 var: form.var,
144 body: (*part.body).clone(),
145 origin: part.origin,
146 })
147 .collect(),
148 );
149 }
150 let under = linear_over(&form.body)?;
151 let held = addends(&ClosedForm {
152 body: under,
153 ..form.clone()
154 })?;
155 Some(
156 held.into_iter()
157 .map(|addend| ClosedForm {
158 body: map_children(&form.body, |_| {
159 Part::new(addend.origin, addend.body.clone())
160 }),
161 ..addend
162 })
163 .collect(),
164 )
165}
166
167fn linear_over(f: &Body) -> Option<Body> {
169 match f {
170 Body::Crop { of, .. }
171 | Body::Shift { of, .. }
172 | Body::Deriv { of, .. }
173 | Body::Channel(of, _) => Some((*of.body).clone()),
174 _ => None,
175 }
176}
177
178fn point_plan(form: &ClosedForm, rate: u32, profile: &Profile) -> Result<Plan, CollapseError> {
179 let written = ClosedForm {
180 body: truncate::written(&form.body, Audible::of(profile, rate))?,
181 ..form.clone()
182 };
183 let width = point::width_of(&written.body, &point::NoRefs).max(1);
184 Ok(Plan::Point {
185 nodes: point_nodes(&written.body),
186 written: Box::new(written),
187 width,
188 })
189}
190
191pub fn point_nodes(f: &Body) -> usize {
193 let width = point::width_of(f, &point::NoRefs).max(1);
194 (0..width).map(|c| nodes_at(f, c)).sum()
195}
196
197fn nodes_at(f: &Body, component: usize) -> usize {
198 let branch = |part: &Part, c: usize| nodes_at(&part.body, c);
199 match f {
200 Body::Join(parts) => {
201 let widths: Vec<usize> = parts
202 .iter()
203 .map(|p| point::width_of(&p.body, &point::NoRefs))
204 .collect();
205 match point::lane_of(&widths, component) {
206 Some((at, inner)) => 1 + branch(&parts[at], inner),
207 None => 1,
208 }
209 }
210 Body::Channel(of, k) => 1 + branch(of, usize::from(*k)),
211 other => {
212 1 + sva_formula::closed_form::children(other)
213 .iter()
214 .map(|part| branch(part, component))
215 .sum::<usize>()
216 }
217 }
218}
219
220fn line_plan(
222 sum: &SpectralSum,
223 rate: u32,
224 horizon: Horizon,
225 profile: &Profile,
226 len: usize,
227 ceiling: f64,
228) -> Result<Option<LinePlan>, CollapseError> {
229 let mut per_lane = Vec::with_capacity(sum.lanes.len());
230 let mut grids: Vec<f64> = Vec::new();
231 let mut tail: Option<f64> = None;
232 for lane in &sum.lanes {
233 match lines::of_lane(lane, ceiling, profile.floor(ceiling)) {
234 Some(found) => {
235 if let Some(left) = found.tail_db {
236 tail = Some(tail.map_or(left, |held: f64| held.max(left)));
237 }
238 grids.extend(found.grids);
239 per_lane.push(found.lines);
240 }
241 None => return Ok(None),
242 }
243 }
244
245 let (mut kept, mut dropped): (Vec<Vec<Line>>, Vec<Vec<Line>>) = (Vec::new(), Vec::new());
246 for found in &per_lane {
247 let (here, gone): (Vec<Line>, Vec<Line>) = found.iter().partition(|l| l.hz.abs() < ceiling);
248 kept.push(here);
249 dropped.push(gone);
250 }
251 if kept.iter().all(Vec::is_empty) && !dropped.iter().all(Vec::is_empty) {
253 let lowest = dropped
254 .iter()
255 .flatten()
256 .map(|l| l.hz.abs())
257 .fold(f64::INFINITY, f64::min);
258 return Err(CollapseError::EmptyBand { ceiling, lowest });
259 }
260
261 let bins = lines::grid(&kept, &grids, rate, horizon.span(), len);
262 let split: Vec<(Vec<Line>, Vec<Line>)> = kept
263 .iter()
264 .map(|k| match direct_is_cheaper(k.len(), bins, len) {
265 true => (Vec::new(), k.clone()),
266 false => lines::split(k, bins, rate),
267 })
268 .collect();
269 Ok(Some(LinePlan {
270 placed: split.iter().map(|(on, _)| on.clone()).collect(),
271 summed: split.into_iter().map(|(_, off)| off).collect(),
272 dropped,
273 bins,
274 tail_db: tail,
275 }))
276}
277
278fn lane_plan(lane: &Lane, rate: u32, horizon: Horizon, len: usize) -> LanePlan {
279 let bins = lines::bins(horizon.span(), rate);
280 let sweep = LanePlan::Sweep {
281 atoms: lane.atoms.len(),
282 samples: super::span::nonzero(lane, horizon, rate, len)
283 .iter()
284 .map(|(from, to)| to - from)
285 .sum(),
286 };
287 let Some(found) = lines::grouped(lane) else {
288 return sweep;
289 };
290 let groups: Vec<Group> = found
291 .into_iter()
292 .map(|(factor, held)| {
293 let (placed, summed) = match direct_is_cheaper(held.len(), bins, len) {
294 true => (Vec::new(), held),
295 false => lines::split(&held, bins, rate),
296 };
297 Group {
298 factor,
299 placed,
300 summed,
301 }
302 })
303 .collect();
304 let grouped = LanePlan::Grouped { groups, bins };
305 match grouped.flops(len) <= sweep.flops(len) {
306 true => grouped,
307 false => sweep,
308 }
309}
310
311pub fn transform_flops(n: usize) -> u128 {
314 let stages = |m: usize| m as u128 * (m.max(2).trailing_zeros() as u128);
315 match n.is_power_of_two() {
316 true => stages(n),
317 false => 3 * stages((2 * n - 1).next_power_of_two()),
318 }
319}
320
321impl LinePlan {
322 pub fn flops(&self, len: usize) -> u128 {
323 let transforms: u128 = self
324 .placed
325 .iter()
326 .filter(|p| !p.is_empty())
327 .map(|_| transform_flops(self.bins))
328 .sum();
329 let direct: u128 = self
330 .summed
331 .iter()
332 .map(|s| s.len() as u128 * len as u128)
333 .sum();
334 transforms + direct
335 }
336
337 pub fn rule(&self) -> Rule {
338 match (lines::distinct(&self.placed), lines::distinct(&self.summed)) {
339 (_, 0) => Rule::LineSpectrumExact,
340 (0, _) => Rule::LineSpectrumSummed,
341 _ => Rule::LineSpectrumMixed,
342 }
343 }
344}
345
346impl LanePlan {
347 pub fn flops(&self, len: usize) -> u128 {
348 match self {
349 LanePlan::Grouped { groups, bins } => groups
350 .iter()
351 .map(|g| {
352 let placed = match g.placed.is_empty() {
353 true => 0,
354 false => transform_flops(*bins),
355 };
356 placed + (g.summed.len() as u128 + 1) * len as u128
357 })
358 .sum(),
359 LanePlan::Sweep { atoms, samples } => *atoms as u128 * *samples as u128,
360 }
361 }
362
363 fn swept_flops(&self, len: u128) -> u128 {
364 let terms: u128 = match self {
365 LanePlan::Grouped { groups, .. } => groups
366 .iter()
367 .map(|g| (g.placed.len() + g.summed.len() + 1) as u128)
368 .sum(),
369 LanePlan::Sweep { atoms, .. } => *atoms as u128,
370 };
371 terms * len
372 }
373}
374
375impl Plan {
376 pub fn flops(&self, len: usize) -> u128 {
377 match self {
378 Plan::Spectrum(_) => transform_flops(len),
379 Plan::Lines(found) => found.flops(len),
380 Plan::Sampled(held) => held.lanes.iter().map(|lane| lane.flops(len)).sum(),
381 Plan::Point { nodes, .. } => *nodes as u128 * len as u128,
382 Plan::Added(parts) => parts.iter().map(|part| part.flops(len)).sum(),
383 }
384 }
385
386 pub fn alias_flops(&self, len: usize) -> u128 {
387 self.alias_flops_at(len, super::ALIAS_OVERSAMPLE)
388 }
389
390 pub fn alias_flops_at(&self, len: usize, oversample: usize) -> u128 {
393 let reference = (len * oversample) as u128;
394 match self {
395 Plan::Point { nodes, .. } => *nodes as u128 * reference,
396 Plan::Sampled(held) if held.rule == Rule::PointSampled => held
397 .lanes
398 .iter()
399 .map(|lane| lane.swept_flops(reference))
400 .sum(),
401 Plan::Added(parts) => parts
402 .iter()
403 .map(|part| part.alias_flops_at(len, oversample))
404 .sum(),
405 Plan::Spectrum(_) | Plan::Lines(_) | Plan::Sampled(_) => 0,
406 }
407 }
408
409 pub fn rule(&self) -> Rule {
410 match self {
411 Plan::Spectrum(_) => Rule::InverseSpectrum,
412 Plan::Lines(found) => found.rule(),
413 Plan::Sampled(held) => held.rule,
414 Plan::Point { .. } => Rule::PointSampled,
415 Plan::Added(_) => Rule::Added,
416 }
417 }
418}