Merge pull request #4513 from mockydevs/fix/3579-nested-blend-memory

Blend: nested blends keep to a step budget (#3579)

crates/doc/src/blend.rs

@@ -1059,13 +1059,146 @@ pub fn steps_knockout_shows(steps: &[Arc<Node>]) -> bool {
     !steps.iter().all(|c| paints_opaquely(c))
 }
 
-/// Evaluate a blend: the keys and generated steps in paint order.
+/// Most steps a blend makes in all, the steps of the blends inside its keys and steps included.
+/// Each step of a blend whose keys are blends is a copy of one of them, which makes its own steps
+/// when drawn, so nesting multiplies the step counts capped pair by pair: five levels of 15 steps
+/// are 1.3 million steps from a 23 KB file, 3.4 GB to draw (#3579). Past it, a blend makes
+/// proportionally fewer steps between each pair of keys ([`fit_steps`]).
+pub const MAX_BLEND_STEPS: usize = 100_000;
+
+/// The steps between each pair of keys, as the blend's spacing asks for them.
+fn requested_steps(keys: &[Arc<Node>], spec: &BlendSpec, rail: Option<&Rail>) -> Vec<usize> {
+    keys.windows(2)
+        .enumerate()
+        .map(|(i, w)| match w {
+            [a, b] => blend_step_count(a, b, spec.spacing, rail.map_or_else(|| center_of(a).distance(center_of(b)), |r| r.len(i))),
+            _ => 0,
+        })
+        .collect()
+}
+
+/// Most steps `spacing` makes between two keys, whatever their size and distance
+/// ([`blend_step_count`]'s clamps).
+fn most_steps(spacing: BlendSpacing) -> usize {
+    match spacing {
+        BlendSpacing::Steps(n) => n.clamp(1, 1000) as usize,
+        BlendSpacing::Distance(_) => 1000,
+        BlendSpacing::SmoothColor => 256,
+    }
+}
+
+/// `want` steps between each pair of keys, fewer where they would make more than
+/// [`MAX_BLEND_STEPS`] in all with `keys` steps already made by the keys and `each[i]` made by
+/// one step of pair `i` (itself and what the blends in it make): every pair keeps the same share
+/// of its steps, rounded down, and none when the keys alone fill the budget.
+fn fit_steps(want: &[usize], each: &[usize], keys: usize) -> Vec<usize> {
+    let room = MAX_BLEND_STEPS.saturating_sub(keys) as u128;
+    let demand = want.iter().zip(each).map(|(&n, &c)| n as u128 * c as u128).fold(0u128, u128::saturating_add);
+    if demand <= room {
+        return want.to_vec();
+    }
+    // `demand` is more than `room`, so it isn't 0, and each share is below `n`.
+    want.iter().map(|&n| usize::try_from(n as u128 * room / demand).unwrap_or(0)).collect()
+}
+
+/// Counts the steps blends make in all ([`MAX_BLEND_STEPS`]), each object counted once by its
+/// address while the count runs.
+#[derive(Default)]
+struct StepCount {
+    known: std::collections::HashMap<(usize, bool), usize>,
+}
+
+impl StepCount {
+    /// Steps the blends in `n` make in all when drawn, each keeping to [`MAX_BLEND_STEPS`].
+    /// `copy`: `n` is part of a blend's step, a copy moved and scaled into place. A blend spaced
+    /// by distance (or by distance where the colours match) makes more steps scaled up, so a
+    /// copy counts the most its spacing can make between each pair, an upper bound for any scale.
+    fn node(&mut self, n: &Node, copy: bool) -> usize {
+        let at = (std::ptr::from_ref(n) as usize, copy);
+        if let Some(&known) = self.known.get(&at) {
+            return known;
+        }
+        let count = match &n.kind {
+            NodeKind::Blend { children, spec } if copy => self.most(children, spec),
+            NodeKind::Blend { children, spec } => self.blend(children, spec, Rail::new(children, spec).as_ref()).1,
+            // Each instance draws the source again.
+            NodeKind::Repeat(r) => r.instances(None).saturating_mul(self.nodes(&r.source, copy)),
+            _ => self.nodes(n.children().map_or(&[], Vec::as_slice), copy),
+        };
+        self.known.insert(at, count);
+        count
+    }
+
+    fn nodes(&mut self, nodes: &[Arc<Node>], copy: bool) -> usize {
+        nodes.iter().map(|c| self.node(c, copy)).fold(0, usize::saturating_add)
+    }
+
+    /// What one step between keys `a` and `b` makes: itself and the steps of the blends in it.
+    /// Groups pair their children as [`Lerp::new`] does; anything else is a copy of the nearer key.
+    fn step(&mut self, a: &Node, b: &Node) -> usize {
+        let inner = match (&a.kind, &b.kind) {
+            (NodeKind::Group { children: ca, clip: k1 }, NodeKind::Group { children: cb, clip: k2 })
+                if k1 == k2 && !(*k1 && ca.len() != cb.len()) =>
+            {
+                (0..ca.len().max(cb.len()))
+                    .map(|i| match (ca.get(i), cb.get(i)) {
+                        (Some(x), Some(y)) => self.step(x, y).saturating_sub(1),
+                        (Some(x), None) | (None, Some(x)) => self.node(x, true),
+                        (None, None) => 0,
+                    })
+                    .fold(0, usize::saturating_add)
+            }
+            _ => self.node(a, true).max(self.node(b, true)),
+        };
+        inner.saturating_add(1)
+    }
+
+    /// What each step between each pair of `keys` makes ([`Self::step`]).
+    fn steps(&mut self, keys: &[Arc<Node>]) -> Vec<usize> {
+        keys.windows(2)
+            .map(|w| match w {
+                [a, b] => self.step(a, b),
+                _ => 1,
+            })
+            .collect()
+    }
+
+    /// The steps between each pair of keys a blend makes ([`fit_steps`]), and the steps it makes
+    /// in all, its keys' included.
+    fn blend(&mut self, keys: &[Arc<Node>], spec: &BlendSpec, rail: Option<&Rail>) -> (Vec<usize>, usize) {
+        let inside = self.nodes(keys, false);
+        let each = self.steps(keys);
+        let kept = fit_steps(&requested_steps(keys, spec, rail), &each, inside);
+        let total = kept.iter().zip(&each).map(|(&n, &c)| n.saturating_mul(c)).fold(inside, usize::saturating_add);
+        (kept, total)
+    }
+
+    /// The most steps a copy of a blend makes in all, at any scale: what [`Self::blend`] counts
+    /// with [`most_steps`] between each pair, before rounding.
+    fn most(&mut self, keys: &[Arc<Node>], spec: &BlendSpec) -> usize {
+        let inside = self.nodes(keys, true);
+        let each = self.steps(keys);
+        let demand = each.iter().map(|&c| c.saturating_mul(most_steps(spec.spacing))).fold(0, usize::saturating_add);
+        inside.saturating_add(demand.min(MAX_BLEND_STEPS.saturating_sub(inside)))
+    }
+}
+
+/// How many steps `keys` blended with `spec` make between each pair of keys: as many as the
+/// spacing asks for, fewer when they and the steps of the blends nested in them would pass
+/// [`MAX_BLEND_STEPS`] ([`blend_expand`] makes these).
+pub fn blend_steps(keys: &[Arc<Node>], spec: &BlendSpec) -> Vec<usize> {
+    StepCount::default().blend(keys, spec, Rail::new(keys, spec).as_ref()).0
+}
+
+/// Evaluate a blend: the keys and generated steps in paint order. The steps keep to
+/// [`MAX_BLEND_STEPS`] in all ([`blend_steps`]).
 pub fn blend_expand(keys: &[Arc<Node>], spec: &BlendSpec) -> Vec<Node> {
     let k = keys.len();
     if k < 2 {
         return keys.iter().map(|n| (**n).clone()).collect();
     }
     let rail = Rail::new(keys, spec);
+    let (counts, _) = StepCount::default().blend(keys, spec, rail.as_ref());
     let mut out = Vec::new();
     for (i, a) in keys.iter().enumerate() {
         let mut key = (**a).clone();
@@ -1074,8 +1207,7 @@ pub fn blend_expand(keys: &[Arc<Node>], spec: &BlendSpec) -> Vec<Node> {
         }
         out.push(key);
         let Some(b) = keys.get(i + 1) else { break };
-        let len = rail.as_ref().map_or_else(|| center_of(a).distance(center_of(b)), |r| r.len(i));
-        let n = blend_step_count(a, b, spec.spacing, len);
+        let n = counts.get(i).copied().unwrap_or(0);
         let lerp = Lerp::new(a, b, (spec.start(i), spec.start(i + 1)));
         for j in 1..=n {
             // Even position of the step, then where the easing puts its shape and its colours.
@@ -1301,4 +1433,128 @@ mod tests {
         mid.walk(&mut |n| ids.push(n.id));
         assert!(ids.iter().all(|i| *i == NodeId(0)), "{ids:?}");
     }
+
+    /// A blend of `a` and `b` spaced by `spacing`.
+    fn blend_of(a: Node, b: Node, spacing: BlendSpacing) -> Node {
+        let spec = BlendSpec { spacing, ..Default::default() };
+        Node::new(NodeId(1), NodeKind::Blend { children: vec![Arc::new(a), Arc::new(b)], spec: Box::new(spec) })
+    }
+
+    /// A blend's keys and options.
+    fn blend_parts(n: &Node) -> Option<(&Vec<Arc<Node>>, &BlendSpec)> {
+        match &n.kind {
+            NodeKind::Blend { children, spec } => Some((children, spec)),
+            _ => None,
+        }
+    }
+
+    /// A square of side `side` at (`x`, `y`).
+    fn square(x: f64, y: f64, side: f64) -> Node {
+        Node::path(NodeId(1), shapes::rectangle(Rect::new(x, y, x + side, y + side)), Appearance::default())
+    }
+
+    /// `levels` of blends of 15 steps whose keys are the blends a level down (#3579's input).
+    fn nested(levels: u32, x: f64) -> Node {
+        if levels == 0 {
+            return square(x, 0.0, 4.0);
+        }
+        let w = 40.0 * f64::from(levels);
+        blend_of(nested(levels - 1, x), nested(levels - 1, x + w), BlendSpacing::Steps(15))
+    }
+
+    /// The steps drawing `n` makes in all, each blend evaluated as the renderer does: one level at
+    /// a time, the blends in its keys and steps evaluated again.
+    fn drawn_steps(n: &Node) -> usize {
+        match &n.kind {
+            NodeKind::Blend { children, spec } => {
+                let out = blend_expand(children, spec);
+                out.len() - children.len() + out.iter().map(drawn_steps).sum::<usize>()
+            }
+            _ => n.children().map_or(0, |c| c.iter().map(|c| drawn_steps(c)).sum()),
+        }
+    }
+
+    /// Steps `levels` of [`nested`] blends of 15 steps make in all: each level its 15 steps,
+    /// and again the steps of the level below in each key and each step.
+    fn nested_steps(levels: u32) -> usize {
+        (17usize.pow(levels) - 1) * 15 / 16
+    }
+
+    /// #3579: blends whose keys are blends draw their steps again in every step above them. Five
+    /// levels of 15 steps made 1.3 million steps from a 23 KB file. Now a blend makes fewer steps
+    /// where they would pass `MAX_BLEND_STEPS`, and none where its keys already fill it, so each
+    /// level past that adds only what its keys draw.
+    #[test]
+    fn nested_blends_keep_to_the_step_budget() {
+        assert_eq!(nested_steps(5), 1_331_115);
+        let four = nested_steps(4);
+        assert!(four < MAX_BLEND_STEPS && 2 * four > MAX_BLEND_STEPS);
+        let five = nested(5, 0.0);
+        let (children, spec) = blend_parts(&five).unwrap();
+        assert_eq!(blend_steps(children, spec), vec![0]);
+        assert_eq!(drawn_steps(&five), 2 * four);
+        assert_eq!(drawn_steps(&nested(6, 0.0)), 4 * four);
+        // A blend whose keys leave room keeps the same share of its steps between each pair.
+        let mut three = nested(3, 0.0);
+        if let NodeKind::Blend { spec, .. } = &mut three.kind {
+            spec.spacing = BlendSpacing::Steps(1000);
+        }
+        let (children, spec) = blend_parts(&three).unwrap();
+        let inside = 2 * nested_steps(2);
+        let each = 1 + nested_steps(2);
+        let kept = (MAX_BLEND_STEPS - inside) * 1000 / (1000 * each);
+        assert_eq!(blend_steps(children, spec), vec![kept]);
+        let drawn = drawn_steps(&three);
+        assert_eq!(drawn, inside + kept * each);
+        assert!(drawn <= MAX_BLEND_STEPS && drawn > MAX_BLEND_STEPS - each, "{drawn}");
+    }
+
+    /// Nested blends within the budget make every step their spacing asks for, as before.
+    #[test]
+    fn nested_blends_within_the_budget_draw_every_step() {
+        for levels in 1..=4 {
+            let n = nested(levels, 0.0);
+            let (children, spec) = blend_parts(&n).unwrap();
+            assert_eq!(blend_steps(children, spec), vec![15]);
+            assert_eq!(drawn_steps(&n), nested_steps(levels), "{levels}");
+        }
+        // The keys, then 15 steps, then the last key, each a blend whose keys moved with it.
+        let two = nested(2, 0.0);
+        let (children, spec) = blend_parts(&two).unwrap();
+        let out = blend_expand(children, spec);
+        assert_eq!(out.len(), 17);
+        assert_eq!(out[0], *children[0]);
+        assert!(out.iter().all(|s| matches!(s.kind, NodeKind::Blend { .. })));
+        assert_eq!(out[8].geometric_bounds(), children[0].geometric_bounds().map(|b| b + Vec2::new(40.0, 0.0)));
+    }
+
+    /// A blend spaced by distance makes more steps scaled up, so a step's copy of one counts the
+    /// most its spacing can make. Here a tiny blend of one step is blended with a large blend of
+    /// one step: copies of the tiny one, scaled up toward the large one's size, make up to 1000
+    /// steps each, which neither key's own count shows.
+    #[test]
+    fn scaled_copies_of_blends_spaced_by_distance_keep_to_the_budget() {
+        let tiny = blend_of(square(0.0, 0.0, 1.0), square(2.0, 0.0, 1.0), BlendSpacing::Distance(1.0));
+        let large = blend_of(square(0.0, 100.0, 1000.0), square(2000.0, 100.0, 1000.0), BlendSpacing::Steps(1));
+        assert_eq!((drawn_steps(&tiny), drawn_steps(&large)), (1, 1));
+        let n = blend_of(tiny, large, BlendSpacing::Steps(1000));
+        let (children, spec) = blend_parts(&n).unwrap();
+        let steps = blend_steps(children, spec);
+        assert!(steps[0] < 1000, "{steps:?}");
+        let drawn = drawn_steps(&n);
+        assert!(drawn <= MAX_BLEND_STEPS && drawn > 1000, "{drawn}");
+    }
+
+    /// A flat blend of many keys shares the budget evenly between its pairs.
+    #[test]
+    fn many_keys_share_the_step_budget() {
+        let keys: Vec<Arc<Node>> = (0..300).map(|i| key(f64::from(i) * 20.0)).collect();
+        let spec = BlendSpec { spacing: BlendSpacing::Steps(1000), ..Default::default() };
+        let steps = blend_steps(&keys, &spec);
+        assert!(steps.iter().all(|&n| n == MAX_BLEND_STEPS * 1000 / (299 * 1000)), "{:?}", &steps[..3]);
+        assert_eq!(blend_expand(&keys, &spec).len(), 300 + steps.iter().sum::<usize>());
+        // Within the budget, every step stays.
+        assert_eq!(blend_steps(&keys[..50], &spec), vec![1000; 49]);
+        assert_eq!(fit_steps(&[10, 20], &[1, 1], MAX_BLEND_STEPS), vec![0, 0]);
+    }
 }