forge_audio/

graph.rs

//! Audio processing DAG — directed acyclic graph with topological sort.
//!
//! Nodes are audio processors (decks, FX, groups, master).
//! Edges define signal flow. Graph is evaluated in topological order.
//! Mutation via RCU: clone → modify → atomic swap (arc-swap).

use std::collections::{HashMap, VecDeque};

// ---------------------------------------------------------------------------
// Node identity
// ---------------------------------------------------------------------------

/// Unique identifier for a node in the audio graph.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct NodeId(pub u32);

/// What kind of node this is (for display / routing logic).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NodeKind {
    Deck,
    Effect,
    Group,
    Master,
    Cue,
}

// ---------------------------------------------------------------------------
// AudioNode trait
// ---------------------------------------------------------------------------

/// Trait for any audio-processing node in the graph.
///
/// All processing is f64 internally. Conversion to f32 happens only at
/// the hardware output boundary.
pub trait AudioNode: Send + Sync {
    /// Process `frames` of audio.
    /// `inputs` contains one buffer per incoming edge (may be empty for source nodes).
    /// Write output into `output`. Both are interleaved stereo f64.
    fn process(&mut self, inputs: &[&[f64]], output: &mut [f64], frames: usize);

    /// The node's unique ID.
    fn node_id(&self) -> NodeId;

    /// What kind of node this is.
    fn kind(&self) -> NodeKind;

    /// Number of output channels (default: 2 for stereo).
    fn channels(&self) -> usize { 2 }
}

// ---------------------------------------------------------------------------
// Graph edge
// ---------------------------------------------------------------------------

/// A directed edge: audio flows from `from` to `to`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Edge {
    pub from: NodeId,
    pub to: NodeId,
}

// ---------------------------------------------------------------------------
// AudioGraph
// ---------------------------------------------------------------------------

/// The audio processing graph. Immutable once built — mutate via clone + swap.
pub struct AudioGraph {
    /// Nodes keyed by ID. Order doesn't matter — eval_order defines execution.
    nodes: HashMap<NodeId, Box<dyn AudioNode>>,
    /// All edges in the graph.
    edges: Vec<Edge>,
    /// Topologically sorted execution order. Computed once on build.
    eval_order: Vec<NodeId>,
    /// Pre-allocated output buffers per node (frames * channels * f64).
    buffers: HashMap<NodeId, Vec<f64>>,
    /// Max frames this graph was allocated for.
    max_frames: usize,
}

impl AudioGraph {
    /// Build a new graph from nodes and edges. Runs topological sort.
    ///
    /// Returns Err if the graph contains a cycle that can't be broken.
    pub fn build(
        nodes: Vec<Box<dyn AudioNode>>,
        edges: Vec<Edge>,
        max_frames: usize,
    ) -> Result<Self, GraphError> {
        let node_map: HashMap<NodeId, Box<dyn AudioNode>> = nodes
            .into_iter()
            .map(|n| (n.node_id(), n))
            .collect();

        let eval_order = topological_sort(&node_map, &edges)?;

        // Pre-allocate output buffers for every node.
        let mut buffers = HashMap::new();
        for (id, node) in &node_map {
            buffers.insert(*id, vec![0.0f64; max_frames * node.channels()]);
        }

        Ok(Self { nodes: node_map, edges, eval_order, buffers, max_frames })
    }

    /// Process one block of audio through the entire graph.
    ///
    /// REAL-TIME SAFE: no allocations, no locks, no I/O.
    /// Caller must ensure `frames <= max_frames`.
    pub fn process(&mut self, frames: usize) {
        debug_assert!(frames <= self.max_frames);

        // Clear all buffers.
        for buf in self.buffers.values_mut() {
            for s in buf.iter_mut() { *s = 0.0; }
        }

        // Process nodes in topological order.
        for &node_id in &self.eval_order {
            // Gather input buffers from upstream nodes.
            let input_ids: Vec<NodeId> = self.edges.iter()
                .filter(|e| e.to == node_id)
                .map(|e| e.from)
                .collect();

            // SAFETY: We need to borrow input buffers immutably and the
            // current node's buffer mutably. Since they're different keys
            // in the HashMap, we use raw pointers to avoid borrow conflicts.
            let input_ptrs: Vec<*const [f64]> = input_ids.iter()
                .filter_map(|id| self.buffers.get(id).map(|b| b.as_slice() as *const [f64]))
                .collect();

            let inputs: Vec<&[f64]> = unsafe {
                input_ptrs.iter().map(|p| &**p).collect()
            };

            let output_buf = self.buffers.get_mut(&node_id).unwrap();
            let node = self.nodes.get_mut(&node_id).unwrap();
            node.process(&inputs, output_buf, frames);
        }
    }

    /// Get the output buffer of a specific node (e.g., Master for final output).
    pub fn output(&self, node_id: NodeId) -> Option<&[f64]> {
        self.buffers.get(&node_id).map(|b| b.as_slice())
    }

    /// List all node IDs in evaluation order.
    pub fn eval_order(&self) -> &[NodeId] {
        &self.eval_order
    }

    /// Get the number of nodes.
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    /// Get the number of edges.
    pub fn edge_count(&self) -> usize {
        self.edges.len()
    }

    /// Get a reference to a node by ID.
    pub fn node(&self, id: NodeId) -> Option<&dyn AudioNode> {
        self.nodes.get(&id).map(|n| n.as_ref())
    }
}

// ---------------------------------------------------------------------------
// Topological sort — Kahn's algorithm
// ---------------------------------------------------------------------------

/// Errors that can occur during graph construction.
#[derive(Debug, Clone)]
pub enum GraphError {
    /// The graph contains a cycle that couldn't be auto-broken.
    Cycle(Vec<NodeId>),
    /// An edge references a node that doesn't exist.
    MissingNode(NodeId),
}

impl std::fmt::Display for GraphError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            GraphError::Cycle(ids) => write!(f, "Graph cycle detected involving {:?}", ids),
            GraphError::MissingNode(id) => write!(f, "Edge references missing node {:?}", id),
        }
    }
}

impl std::error::Error for GraphError {}

/// Kahn's algorithm for topological sorting.
///
/// Returns the node IDs in an order such that for every edge U→V,
/// U appears before V.
fn topological_sort(
    nodes: &HashMap<NodeId, Box<dyn AudioNode>>,
    edges: &[Edge],
) -> Result<Vec<NodeId>, GraphError> {
    // Validate edges reference existing nodes.
    for edge in edges {
        if !nodes.contains_key(&edge.from) {
            return Err(GraphError::MissingNode(edge.from));
        }
        if !nodes.contains_key(&edge.to) {
            return Err(GraphError::MissingNode(edge.to));
        }
    }

    // Build in-degree map and adjacency list.
    let mut in_degree: HashMap<NodeId, usize> = nodes.keys().map(|&id| (id, 0)).collect();
    let mut adjacency: HashMap<NodeId, Vec<NodeId>> = nodes.keys().map(|&id| (id, Vec::new())).collect();

    for edge in edges {
        *in_degree.entry(edge.to).or_insert(0) += 1;
        adjacency.entry(edge.from).or_default().push(edge.to);
    }

    // Seed the queue with zero-in-degree nodes (sources).
    let mut queue: VecDeque<NodeId> = in_degree.iter()
        .filter(|(_, &deg)| deg == 0)
        .map(|(&id, _)| id)
        .collect();

    let mut sorted = Vec::with_capacity(nodes.len());

    while let Some(node_id) = queue.pop_front() {
        sorted.push(node_id);
        if let Some(neighbors) = adjacency.get(&node_id) {
            for &neighbor in neighbors {
                if let Some(deg) = in_degree.get_mut(&neighbor) {
                    *deg -= 1;
                    if *deg == 0 {
                        queue.push_back(neighbor);
                    }
                }
            }
        }
    }

    if sorted.len() != nodes.len() {
        // Cycle detected — collect remaining nodes.
        let remaining: Vec<NodeId> = nodes.keys()
            .filter(|id| !sorted.contains(id))
            .copied()
            .collect();
        return Err(GraphError::Cycle(remaining));
    }

    Ok(sorted)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    /// Minimal test node that passes input through to output.
    struct PassthroughNode {
        id: NodeId,
        kind: NodeKind,
    }

    impl AudioNode for PassthroughNode {
        fn process(&mut self, inputs: &[&[f64]], output: &mut [f64], frames: usize) {
            // Sum all inputs.
            let len = frames * self.channels();
            for input in inputs {
                for i in 0..len.min(input.len()).min(output.len()) {
                    output[i] += input[i];
                }
            }
        }
        fn node_id(&self) -> NodeId { self.id }
        fn kind(&self) -> NodeKind { self.kind }
    }

    /// Source node that generates a constant value.
    struct ConstNode {
        id: NodeId,
        value: f64,
    }

    impl AudioNode for ConstNode {
        fn process(&mut self, _inputs: &[&[f64]], output: &mut [f64], frames: usize) {
            let len = frames * self.channels();
            for i in 0..len.min(output.len()) {
                output[i] = self.value;
            }
        }
        fn node_id(&self) -> NodeId { self.id }
        fn kind(&self) -> NodeKind { NodeKind::Deck }
    }

    fn make_node(id: u32, kind: NodeKind) -> Box<dyn AudioNode> {
        Box::new(PassthroughNode { id: NodeId(id), kind })
    }

    fn make_const(id: u32, value: f64) -> Box<dyn AudioNode> {
        Box::new(ConstNode { id: NodeId(id), value })
    }

    #[test]
    fn empty_graph() {
        let g = AudioGraph::build(vec![], vec![], 64).unwrap();
        assert_eq!(g.node_count(), 0);
        assert_eq!(g.edge_count(), 0);
    }

    #[test]
    fn single_node() {
        let nodes: Vec<Box<dyn AudioNode>> = vec![make_const(1, 0.5)];
        let mut g = AudioGraph::build(nodes, vec![], 64).unwrap();
        g.process(64);
        let out = g.output(NodeId(1)).unwrap();
        assert!((out[0] - 0.5).abs() < 1e-10);
    }

    #[test]
    fn linear_chain() {
        // Deck(1) → FX(2) → Master(3)
        let nodes: Vec<Box<dyn AudioNode>> = vec![
            make_const(1, 1.0),
            make_node(2, NodeKind::Effect),
            make_node(3, NodeKind::Master),
        ];
        let edges = vec![
            Edge { from: NodeId(1), to: NodeId(2) },
            Edge { from: NodeId(2), to: NodeId(3) },
        ];
        let mut g = AudioGraph::build(nodes, edges, 64).unwrap();
        assert_eq!(g.eval_order().len(), 3);

        // Deck must come before FX, FX before Master.
        let order = g.eval_order();
        let pos = |id: u32| order.iter().position(|n| n.0 == id).unwrap();
        assert!(pos(1) < pos(2));
        assert!(pos(2) < pos(3));

        g.process(64);
        let out = g.output(NodeId(3)).unwrap();
        // Const(1.0) passes through passthrough nodes unchanged.
        assert!((out[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn diamond_graph() {
        // Deck(1) → FX(2) → Master(4)
        //         ↘ FX(3) ↗
        let nodes: Vec<Box<dyn AudioNode>> = vec![
            make_const(1, 0.5),
            make_node(2, NodeKind::Effect),
            make_node(3, NodeKind::Effect),
            make_node(4, NodeKind::Master),
        ];
        let edges = vec![
            Edge { from: NodeId(1), to: NodeId(2) },
            Edge { from: NodeId(1), to: NodeId(3) },
            Edge { from: NodeId(2), to: NodeId(4) },
            Edge { from: NodeId(3), to: NodeId(4) },
        ];
        let mut g = AudioGraph::build(nodes, edges, 64).unwrap();
        g.process(64);
        let out = g.output(NodeId(4)).unwrap();
        // Master receives 0.5 from each FX path = 1.0.
        assert!((out[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn four_decks_to_master() {
        // Deck A(1), B(2), C(3), D(4) → Group L(5) / Group R(6) → Master(7)
        let nodes: Vec<Box<dyn AudioNode>> = vec![
            make_const(1, 0.25), // Deck A
            make_const(2, 0.25), // Deck B
            make_const(3, 0.25), // Deck C
            make_const(4, 0.25), // Deck D
            make_node(5, NodeKind::Group),  // Group L
            make_node(6, NodeKind::Group),  // Group R
            make_node(7, NodeKind::Master), // Master
        ];
        let edges = vec![
            Edge { from: NodeId(1), to: NodeId(5) }, // A → L
            Edge { from: NodeId(3), to: NodeId(5) }, // C → L
            Edge { from: NodeId(2), to: NodeId(6) }, // B → R
            Edge { from: NodeId(4), to: NodeId(6) }, // D → R
            Edge { from: NodeId(5), to: NodeId(7) }, // L → Master
            Edge { from: NodeId(6), to: NodeId(7) }, // R → Master
        ];
        let mut g = AudioGraph::build(nodes, edges, 64).unwrap();
        g.process(64);
        let out = g.output(NodeId(7)).unwrap();
        // All 4 decks at 0.25 sum to 1.0 at master.
        assert!((out[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn cycle_detected() {
        // A → B → A (cycle)
        let nodes: Vec<Box<dyn AudioNode>> = vec![
            make_node(1, NodeKind::Effect),
            make_node(2, NodeKind::Effect),
        ];
        let edges = vec![
            Edge { from: NodeId(1), to: NodeId(2) },
            Edge { from: NodeId(2), to: NodeId(1) },
        ];
        let result = AudioGraph::build(nodes, edges, 64);
        assert!(result.is_err());
        if let Err(GraphError::Cycle(ids)) = result {
            assert_eq!(ids.len(), 2);
        }
    }

    #[test]
    fn missing_node_in_edge() {
        let nodes: Vec<Box<dyn AudioNode>> = vec![make_node(1, NodeKind::Deck)];
        let edges = vec![Edge { from: NodeId(1), to: NodeId(99) }];
        let result = AudioGraph::build(nodes, edges, 64);
        assert!(matches!(result, Err(GraphError::MissingNode(NodeId(99)))));
    }

    #[test]
    fn eval_order_respects_topology() {
        // 1 → 3, 2 → 3, 3 → 4
        let nodes: Vec<Box<dyn AudioNode>> = vec![
            make_node(1, NodeKind::Deck),
            make_node(2, NodeKind::Deck),
            make_node(3, NodeKind::Group),
            make_node(4, NodeKind::Master),
        ];
        let edges = vec![
            Edge { from: NodeId(1), to: NodeId(3) },
            Edge { from: NodeId(2), to: NodeId(3) },
            Edge { from: NodeId(3), to: NodeId(4) },
        ];
        let g = AudioGraph::build(nodes, edges, 64).unwrap();
        let order = g.eval_order();
        let pos = |id: u32| order.iter().position(|n| n.0 == id).unwrap();
        assert!(pos(1) < pos(3));
        assert!(pos(2) < pos(3));
        assert!(pos(3) < pos(4));
    }
}