SwiftUI Snippet: Orbit Animation

import simd
import SwiftUI

// A view that renders an iCloud-style Orbit animation.
//
// It consists of many colorful bubbles orbiting a central icon. Over time,
// some of them scale up and reveal one of multiple symbols. As they scale up,
// other bubbles are pushed aside.
//
// This view uses three different animation techniques simultaneously:
//
// 1. A simple physics simulation is used to calculate the position of the
//    individual bubbles as they push each other out of the way.
// 2. Manually-evaluated keyframe animations are used to describe how
//    individual circles scale up and down and transition their contents.
// 3. The rotation of the circles is calculated directly from the rotational
//    speed and the elapsed time.
struct OrbitView: View {
  // The animated properties of the orbit.
  struct Track {
    // The scale of the orbit.
    var scale: Double = 1

    // The speed at which the animation rotates.
    var speed: Angle = .degrees(30)
  }

  // Set up a simulation.
  @State var simulation = {
    let nodes = Simulation.Node.concentricPattern

    return Simulation(
      nodes: nodes,
      forces: [
        Collision(),
        // Weakly pin every node to its initial position.
        Pin(pins: nodes.enumerated().map { i, node in
          .init(i, to: node.position, strength: 0.05)
        }),
      ])
  }()

  // The time at which the animation started.
  @State var startDate = Date.now

  // Scale up the orbit and increase its velocity over the course of 3s.
  let orbitTimeline = KeyframeTimeline(initialValue: Track()) {
    KeyframeTrack(\.scale) {
      MoveKeyframe(0.3)
      LinearKeyframe(1, duration: 3, timingCurve: .easeInOut)
    }

    KeyframeTrack(\.speed) {
      MoveKeyframe(.degrees(0))
      LinearKeyframe(.degrees(55), duration: 3, timingCurve: .easeInOut)
    }
  }

  // Define the bubbles that will scale up during the animation. They are
  // offset by 1/3 of their animation duration.
  static let bubbles: [Bubble] = {
    // Randomly offset the base-indices for some variety.
    let offset = Int.random(in: -20 ..< 30)

    return [
      Bubble(index: 37 + offset, image: Image(systemName: "swift"), offset: 0.0 * Bubble.timeline.duration),
      Bubble(index: 48 + offset, image: Image(systemName: "iphone"), offset: 0.333 * Bubble.timeline.duration),
      Bubble(index: 56 + offset, image: Image(systemName: "apple.terminal"), offset: 0.666 * Bubble.timeline.duration),
    ]
  }()

  var body: some View {
    // Run the simulation every time the `TimelineView` updates.
    TimelineView(.animation) { animationContext in
      // The relative time since the appearance of the view.
      let t = animationContext.date.timeIntervalSince(startDate)

      // Delay any bubble animations by 6s.
      let offsetTime = max(0, t - 6)

      // Delay the orbit timeline by 2s and resolve it.
      let orbitValues = orbitTimeline.value(time: t - 2)

      // The rotation of the orbit can be calculated directly.
      let angle = Angle.degrees(orbitValues.speed.degrees * t)

      Canvas { context, size in
        let scale = 0.9 * (min(size.width, size.height) / 2)
        // Center the drawing context on the midpoint of the canvas.
        context.translateBy(x: size.width / 2, y: size.height / 2)
        context.scaleBy(x: scale, y: scale)

        context.rotate(by: angle)
        context.scaleBy(x: orbitValues.scale, y: orbitValues.scale)

        // Draw every node.
        for (i, node) in simulation.nodes.enumerated() {
          var bubbleContext = context
          // Counter-rotate around the center of the bubble to have
          // the icons remain upright.
          bubbleContext.translateBy(x: node.frame.midX, y: node.frame.midY)
          bubbleContext.rotate(by: -angle)
          bubbleContext.translateBy(x: -node.frame.midX, y: -node.frame.midY)

          // Calculate the base color of the bubble based on its 2d
          // position in the (unrotated) orbit.
          //
          // 336pts refers to the maximum diameter (2 * 168) used in
          // `concentricPattern`.
          let bubbleColor = Bubble.color(at: UnitPoint(node.position))

          // If the current node refers to a bubble, use its timeline
          // to interpolate the base color – otherwise use 0.
          let bubble = Self.bubbles.first { $0.index == i }

          if let bubble {
            let bubbleValues = bubble.values(at: offsetTime)

            // Fade the color towards blue as the bubble expands.
            bubbleContext.fill(
              Path(ellipseIn: node.frame.inset(by: 0.02)),
              with: .style(bubbleColor.mix(with: .darkBlue, by: bubbleValues.colorBlend))
            )

            // Draw the icons, blur and fade them.
            var t = context.resolve(bubble.image)
            t.shading = .style(.white)

            let scale = 0.6 * node.frame.height / max(t.size.width, t.size.height)

            bubbleContext.opacity = bubbleValues.opacity
            bubbleContext.addFilter(.blur(radius: bubbleValues.blurRadius))
            bubbleContext.draw(t, in: CGRect(center: node.location, size: t.size.scaled(by: scale)))
          } else {
            bubbleContext.fill(
              Path(ellipseIn: node.frame.inset(by: 0.02)),
              with: .color(bubbleColor)
            )
          }
        }
      }
      .onChange(of: animationContext.date) { a, b in
        // Update the scale of each of the bubbles on every tick.
        // This bridges the keyframe animation and the physics
        // simulation.
        for bubble in Self.bubbles {
          simulation.nodes[bubble.index].scale = bubble.values(at: offsetTime).scale
        }

        // Advance the animation every frame.
        //
        // TODO: Take the time since the last update into account.
        simulation.tick()
      }
    }
    .onAppear {
      startDate = Date.now
    }
    .overlay {
      let t = Transaction(animation: .default)
      AsyncImage(url: URL(string: "https://github.com/robb.png"), transaction: t) { result in
        result.image?
          .resizable()
          .aspectRatio(contentMode: .fit)
      }
      .containerRelativeFrame([.horizontal, .vertical]) { length, _ in length / 3 }
      .background(.purple.mix(with: .black, by: 0.4))
      .clipShape(.circle)
    }
  }
}

struct Bubble {
  // The animated properties of the bubble.
  struct Track {
    // The blur radius of the bubble's icon.
    var blurRadius: CGFloat = 8

    // The opacity of the bubble's icon.
    var opacity: Double = 0

    // The scale of the bubble.
    var scale: Double = 1

    // The amount by which the bubble's base color is mixed with blue.
    var colorBlend: Double = 0
  }

  static var timeline = KeyframeTimeline(initialValue: Track()) {
    KeyframeTrack(\.opacity) {
      CubicKeyframe(1, duration: 0.31)
      LinearKeyframe(1, duration: 2)
      CubicKeyframe(0, duration: 1.5)
    }

    KeyframeTrack(\.colorBlend) {
      CubicKeyframe(1, duration: 0.31)
      LinearKeyframe(1, duration: 2)
      SpringKeyframe(0, duration: 0.5)
      LinearKeyframe(0, duration: 3.3)
    }

    KeyframeTrack(\.blurRadius) {
      CubicKeyframe(0, duration: 0.31)
      LinearKeyframe(0, duration: 2)
      CubicKeyframe(8, duration: 1.5)
    }

    KeyframeTrack(\.scale) {
      CubicKeyframe(4, duration: 0.4)
      LinearKeyframe(4, duration: 2)
      SpringKeyframe(1)
    }
  }

  static func color(at position: UnitPoint) -> Color {
    let x = Color.orange.mix(with: .blue, by: position.x + 0.5)
    let y = Color.pink.mix(with: .mint, by: position.y + 0.5)
    return x.mix(with: y, by: 0.5)
  }

  // The index of the node this bubble represents.
  var index: Int

  // The icon drawn inside the bubble.
  var image: Image

  // The relative time offset at which this bubble animates. Used to stagger
  // the bubbles.
  var offset: TimeInterval

  // Resolve the timeline, taking the bubble's offset into account.
  func values(at time: TimeInterval) -> Track {
    guard time >= 0 else { return Track() }

    return Self.timeline.value(wrappingTime: time - offset)
  }
}

// A simple semi-implicit Euler integrator.
struct Simulation {
  // A node in the physics simulation is a circle of a given radius.
  struct Node {
    var position: SIMD2<Double>

    var velocity: SIMD2<Double>

    var initialRadius: Double

    // We want to animate the bubbles scaling up later, so we add a
    // dedicated scale parameter here.
    var scale: Double = 1

    // The scaled radius, this value is considered for collision detection
    // and drawing.
    var radius: Double {
      initialRadius * scale
    }

    var frame: CGRect {
      let r = radius

      return CGRect(center: CGPoint(x: position.x, y: position.y), size: CGSize(width: r * 2, height: r * 2))
    }

    init(position: CGPoint, radius: Double = 5) {
      self.init(x: position.x, y: position.y, radius: radius)
    }

    init(x: Double, y: Double, radius: Double = 5) {
      self.position = [x, y]
      self.velocity = [0, 0]
      self.initialRadius = radius
    }
  }

  var nodes: [Node]

  var forces: [any Force]

  // Velocities decay by (1 - this) every tick.
  var decay: Double = 0.6

  // Advances the simulation by one step.
  //
  // TODO: Take the time since the previous tick into account to make the simulation more resilient.
  //
  // Search for Verlet Integration for more info on how to write a more stable simulation.
  mutating func tick() {
    for i in forces.indices {
      forces[i].apply(to: &nodes)
    }

    for i in nodes.indices {
      nodes[i].position += nodes[i].velocity
      nodes[i].velocity *= decay
    }
  }
}

protocol Force {
  // Applies a force to a list of nodes by updating their velocities.
  mutating func apply(to nodes: inout [Simulation.Node])
}

// Moves a node by adding a velocity that points towards a fixed point.
//
// This is used so nodes return to their original position.
struct Pin: Force {
  struct Constant {
    // The index of the node.
    var index: Int

    // The target position.
    var position: SIMD2<Double>

    // The strength of the movement – if 1, the node with move to `position`
    // on the next tick.
    var strength: Double

    init(_ index: Int, to position: SIMD2<Double>, strength: Double) {
      self.index = index
      self.position = position
      self.strength = strength
    }
  }

  var pins: [Constant]

  func apply(to nodes: inout [Simulation.Node]) {
    for pin in pins {
      let d = pin.position - nodes[pin.index].position
      nodes[pin.index].velocity += d * pin.strength
    }
  }
}

// Moves nodes apart that touch each other.
struct Collision: Force {
  func apply(to nodes: inout [Simulation.Node]) {
    // Add all nodes to a quad tree to quickly find their neighbors.
    let quadTree = QuadTree(contentsOf: nodes)

    for (i, node) in nodes.enumerated() {
      let r = node.radius + 0.2
      let neighborhood = node.frame.insetBy(dx: -r, dy: -r)

      // Apply a force to all neighbors with a center within ±30pts, the
      // quad tree makes this lookup efficient.
      for j in quadTree.indices(in: neighborhood) {
        // Only consider the case of i being of higher index to avoid
        // doing the work twice.
        guard i > j else { continue }

        let node = nodes[i]
        let other = nodes[j]

        let delta = node.position - other.position
        // Calculate the distance between centers.
        let distance = length(delta)
        // Calculate the sum of the radii.
        let combinedRadii = node.radius + other.radius

        // If the distance is less then the sum of radii, the nodes
        // overlap.
        if distance <= combinedRadii {
          // Calculate the amount of overlap and the normalized
          // direction vector.
          let overlap = combinedRadii - distance
          let direction = delta / distance

          // Weight the velocity based on the area of the nodes. This
          // moves bigger nodes less, implying a higher mass.
          let weight = pow(other.radius, 2) / pow(combinedRadii, 2)

          // Apply two opposing velocities, weighted by their mass.
          nodes[i].velocity += direction * overlap * weight
          nodes[j].velocity -= direction * overlap * (1 - weight)
        }
      }
    }
  }
}

extension Simulation.Node {
  // A pattern of 4 x 24 nodes, arranged in concentric rings.
  //
  // The outermost ring represents the unit circle.
  static var concentricPattern: [Self] {
    let indices = (0 ..< 24).map(Double.init)
    let offset = Angle.degrees(360 / (2 * 24))
    let step = Angle.degrees(360 / 24)

    var nodes: [Self] = []
    nodes.append(contentsOf: indices.map { i in
      Self(position: CGPoint(theta: step * i, radius: 0.65), radius: 0.04)
    })
    nodes.append(contentsOf: indices.map { i in
      Self(position: CGPoint(theta: step * i + offset, radius: 0.78), radius: 0.05)
    })
    nodes.append(contentsOf: indices.map { i in
      Self(position: CGPoint(theta: step * i, radius: 0.9), radius: 0.045)
    })
    nodes.append(contentsOf: indices.map { i in
      Self(position: CGPoint(theta: step * i + offset, radius: 1), radius: 0.055)
    })
    return nodes
  }
}

#Preview {
  OrbitView()
}

#Preview("Quad Tree") {
  @Previewable @State var quadTree = QuadTree<CGPoint>(bounds: CGRect(x: 0, y: 0, width: 350, height: 350))

  VStack {
    Canvas { context, size in
      let subrect = CGRect(x: 240, y: 50, width: 80, height: 200)
      context.stroke(Path(subrect), with: .color(.red))

      var copy = context
      copy.opacity = 0.3
      copy.drawLayer { context in
        for node in quadTree.nodes {
          context.stroke(Path(node.bounds), with: .color(.purple))
        }
      }

      for point in quadTree.elements {
        let frame = CGRect(origin: point.offset(x: -2.5, y: -2.5), size: CGSize(width: 5, height: 5))
        context.fill(Path(ellipseIn: frame), with: .color(.blue))
      }

      for point in quadTree.elements(in: subrect) {
        let frame = CGRect(origin: point.offset(x: -4, y: -4), size: CGSize(width: 8, height: 8))
        context.stroke(Path(ellipseIn: frame), with: .color(.red))
      }
    }
    .frame(width: 350, height: 350)

    HStack {
      let insert = {
        quadTree.append(
          CGPoint(
            theta: Angle.degrees(.random(in: -180 ... 180)),
            radius: CGFloat.random(in: 120 ... 140)
          )
          .offset(x: 175, y: 175)
        )
      }

      Button("Insert Point") {
        insert()
      }

      Button("Insert 10 Points") {
        for _ in 0 ..< 10 { insert() }
      }
    }
    .buttonStyle(.borderedProminent)
  }
  .padding()
}

extension Simulation.Node: QuadTreeElement {
  var location: CGPoint {
    CGPoint(x: position.x, y: position.y)
  }
}

extension CGRect {
  init(center: CGPoint, size: CGSize) {
    self.init(x: center.x - size.width / 2, y: center.y - size.height / 2, width: size.width, height: size.height)
  }

  var center: CGPoint {
    .init(x: midX, y: midY)
  }

  func inset(by fraction: CGFloat) -> Self {
    insetBy(dx: width * fraction , dy: height * fraction)
  }
}

extension CGSize {
  func scaled(by multiplier: Double) -> CGSize {
    .init(width: width * multiplier, height: height * multiplier)
  }
}

extension CGPoint {
  init(theta: Angle, radius: CGFloat) {
    self.init()
    x = radius * cos(theta.radians)
    y = radius * sin(theta.radians)
  }
}

extension UnitPoint {
  init(_ point: SIMD2<Double>) {
    self.init(x: point.x, y: point.y)
  }
}

extension KeyframeTimeline {
  public func value(wrappingTime time: TimeInterval) -> Value {
    value(time: fmod(time, duration))
  }
}

extension Color {
  static var darkBlue: Self { Color.blue.mix(with: .black, by: 0.1) }
}

import SwiftUI

public protocol QuadTreeElement {
  var location: CGPoint { get }
}

public struct QuadTree<Element: QuadTreeElement> {
  public typealias Index = Int

  public struct Node {
    enum Storage {
      indirect case children(InlineArray<4, Node>)
      case elements(ContiguousArray<(Array<Element>.Index, CGPoint)>)
    }

    public var bounds: CGRect

    var storage: Storage = .elements([])

    public var elements: (some RandomAccessCollection<(Index, CGPoint)>)? {
      if case let .elements(elements) = storage {
        elements
      } else {
        nil
      }
    }

    @discardableResult
    mutating func insert(_ element: Array<Element>.Index, at location: CGPoint) -> Bool {
      guard bounds.contains(location) else { return false }

      switch storage {
      case .children(var children):
        defer { self.storage = .children(children) }

        return children.indices.contains { children[$0].insert(element, at: location) }
      case .elements(var elements):
        elements.append((element, location))

        self.storage = .elements(elements)

        subdivideIfNeeded()

        return true
      }
    }

    private mutating func subdivideIfNeeded() {
      guard case let .elements(elements) = storage else { return }

      guard elements.count > 4 else { return }

      let q = bounds.quadrants

      var children: InlineArray<4, Node> = [
        Node(bounds: q[0]), Node(bounds: q[1]),
        Node(bounds: q[2]), Node(bounds: q[3]),
      ]

      for (element, location) in elements {
        let _ = children.indices.contains { children[$0].insert(element, at: location) }
      }

      self.storage = .children(children)
    }

    func elements(in rect: CGRect, accumulator: inout [Array<Element>.Index]) {
      guard rect.intersects(bounds) else { return }

      switch storage {
      case .children(let children):
        for i in children.indices {
          children[i].elements(in: rect, accumulator: &accumulator)
        }
      case .elements(let elements):
        for (element, location) in elements where rect.contains(location) {
          accumulator.append(element)
        }
      }
    }
  }

  struct NodeSequence: Sequence {
    struct Iterator: IteratorProtocol {
      var nodes: [QuadTree.Node]

      mutating func next() -> QuadTree.Node? {
        guard !nodes.isEmpty else { return nil }

        let next = nodes.removeFirst()

        if case let .children(children) = next.storage {
          nodes.append(contentsOf: children.indices.map { children[$0] })
        }

        return next
      }
    }

    var tree: QuadTree

    func makeIterator() -> Iterator {
      Iterator(nodes: [tree.root])
    }
  }

  var root: Node

  var elements: [Element] = []

  public var nodes: some Sequence<Node> {
    NodeSequence(tree: self)
  }

  public init(bounds: CGRect) {
    self.root = Node(bounds: bounds)
  }

  public init(contentsOf array: [Element]) {
    self.elements = array

    let x = array.map(\.location.x)
    let y = array.map(\.location.y)

    let minX = x.min() ?? 0
    let minY = y.min() ?? 0
    let maxX = x.max() ?? 0
    let maxY = y.max() ?? 0

    self.root = .init(bounds: CGRect(x: minX, y: minY, width: maxX - minX, height: maxY - minY))

    for (i, element) in zip(elements.indices, elements) {
      root.insert(i, at: element.location)
    }
  }

  public mutating func append(_ element: Element) {
    elements.append(element)
    root.insert(elements.endIndex - 1, at: element.location)
  }

  public func indices(in rect: CGRect) -> [Index] {
    var indices: [Index] = []
    root.elements(in: rect, accumulator: &indices)

    return indices
  }

  public func elements(in rect: CGRect) -> [Element] {
    indices(in: rect).map { elements[$0] }
  }
}

extension CGRect {
  var quadrants: InlineArray<4, CGRect> {
    let q = CGSize(width: width / 2, height: height / 2)

    return [
      CGRect(origin: origin, size: q),
      CGRect(origin: origin.offset(x: q.width), size: q),
      CGRect(origin: origin.offset(x: q.width, y: q.height), size: q),
      CGRect(origin: origin.offset(y: q.height), size: q),
    ]
  }
}

extension CGPoint {
  func offset(x: CGFloat = 0, y: CGFloat = 0) -> CGPoint {
    CGPoint(x: self.x + x, y: self.y + y)
  }
}

extension CGPoint: QuadTreeElement {
  public var location: CGPoint { self }
}