cyborg/ramen/src/command.rs

use super::{Graph, GraphError, GraphResult};
use crate::node::{Edge, Node, NodeKind, SlotIndex};
use std::fmt::Debug;

pub trait Command: Debug {
    fn undo(&self, graph: &Graph) -> GraphResult<DynCommand>;
    fn apply(&self, graph: &mut Graph) -> GraphResult<()>;
}

pub type DynCommand = Box<dyn Command>;

pub struct CommandHistory {
    commands: Vec<(DynCommand, DynCommand)>,
    cursor: usize,
}

impl CommandHistory {
    pub fn new() -> Self {
        Self {
            commands: Vec::new(),
            cursor: 0,
        }
    }

    pub fn push(&mut self, graph: &mut Graph, cmd: DynCommand) -> GraphResult<()> {
        let undo = cmd.undo(graph)?;
        cmd.apply(graph)?;

        // Discard redo history on new push
        if self.cursor < self.commands.len() {
            self.commands.truncate(self.cursor);
        }

        self.commands.push((cmd, undo));
        self.cursor += 1;

        Ok(())
    }

    pub fn undo(&mut self, graph: &mut Graph) -> GraphResult<bool> {
        if self.cursor > 0 {
            self.cursor -= 1;
            self.commands[self.cursor].1.apply(graph)?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    pub fn redo(&mut self, graph: &mut Graph) -> GraphResult<bool> {
        if self.cursor < self.commands.len() {
            self.commands[self.cursor].0.apply(graph)?;
            self.cursor += 1;
            Ok(true)
        } else {
            Ok(false)
        }
    }
}

#[derive(Clone, Debug)]
pub struct CreateNode {
    pub kind: NodeKind,
    pub pos: glam::Vec2,
    pub inputs: Vec<Option<SlotIndex>>,
}

impl Command for CreateNode {
    fn undo(&self, graph: &Graph) -> GraphResult<DynCommand> {
        let target = graph.nodes.vacant_key();
        let undo = DeleteNode { target };
        Ok(Box::new(undo))
    }

    fn apply(&self, graph: &mut Graph) -> GraphResult<()> {
        let node = Node {
            kind: self.kind,
            pos: self.pos,
            inputs: self.inputs.clone(),
        };

        graph.nodes.insert(node);

        Ok(())
    }
}

impl From<Node> for CreateNode {
    fn from(node: Node) -> Self {
        Self {
            kind: node.kind,
            pos: node.pos,
            inputs: node.inputs,
        }
    }
}

#[derive(Clone, Debug)]
pub struct DeleteNode {
    pub target: usize,
}

impl Command for DeleteNode {
    fn undo(&self, graph: &Graph) -> GraphResult<DynCommand> {
        let target = graph.get_node(self.target)?;

        let undo = CreateNode {
            kind: target.kind,
            pos: target.pos,
            inputs: target.inputs.clone(),
        };

        Ok(Box::new(undo))
    }

    fn apply(&self, graph: &mut Graph) -> GraphResult<()> {
        graph
            .nodes
            .try_remove(self.target)
            .ok_or(GraphError::InvalidReference)
            .map(|_| ())
    }
}

#[derive(Clone, Debug)]
pub struct MoveNode {
    pub target: usize,
    pub to: glam::Vec2,
    pub relative: bool,
}

impl Command for MoveNode {
    fn undo(&self, graph: &Graph) -> GraphResult<DynCommand> {
        let target = graph.get_node(self.target)?;

        let undo = Self {
            target: self.target,
            to: target.pos,
            relative: false,
        };

        Ok(Box::new(undo))
    }

    fn apply(&self, graph: &mut Graph) -> GraphResult<()> {
        let target = graph.get_node_mut(self.target)?;
        
        if self.relative {
            target.pos += self.to;
        } else {
            target.pos = self.to;
        }

        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct SetEdge {
    pub edge: Edge,
}

impl Command for SetEdge {
    fn undo(&self, graph: &Graph) -> GraphResult<DynCommand> {
        let target = graph.get_node(self.edge.output.node)?;

        let old_input = target
            .inputs
            .get(self.edge.output.slot)
            .ok_or(GraphError::InvalidReference)?;

        Ok(match old_input {
            Some(old_input) => Box::new(SetEdge {
                edge: Edge {
                    input: *old_input,
                    output: self.edge.output,
                },
            }),
            None => Box::new(DeleteEdge {
                input: self.edge.output,
            }),
        })
    }

    fn apply(&self, graph: &mut Graph) -> GraphResult<()> {
        let target = graph.get_node_mut(self.edge.output.node)?;

        let input = target
            .inputs
            .get_mut(self.edge.output.slot)
            .ok_or(GraphError::InvalidReference)?;

        *input = Some(self.edge.input);

        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct DeleteEdge {
    pub input: SlotIndex,
}

impl Command for DeleteEdge {
    fn undo(&self, graph: &Graph) -> GraphResult<DynCommand> {
        let target = graph.get_node(self.input.node)?;

        let old_input = target
            .inputs
            .get(self.input.slot)
            .ok_or(GraphError::InvalidReference)?
            .ok_or(GraphError::MissingEdge)?;

        let undo = SetEdge {
            edge: Edge {
                input: old_input,
                output: self.input,
            },
        };

        Ok(Box::new(undo))
    }

    fn apply(&self, graph: &mut Graph) -> GraphResult<()> {
        let target = graph.get_node_mut(self.input.node)?;

        let input = target
            .inputs
            .get_mut(self.input.slot)
            .ok_or(GraphError::InvalidReference)?;

        if input.is_none() {
            Err(GraphError::MissingEdge)
        } else {
            *input = None;
            Ok(())
        }
    }
}

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

    impl CommandHistory {
        pub fn push_multi(
            &mut self,
            graph: &mut Graph,
            cmds: impl IntoIterator<Item = DynCommand>,
        ) -> GraphResult<()> {
            for cmd in cmds {
                println!("{:#?}", cmd);
                self.push(graph, cmd)?;
                println!("{:#?}", graph);
            }

            Ok(())
        }

        pub fn rewind(&mut self, graph: &mut Graph) -> GraphResult<()> {
            while self.cursor > 0 {
                self.cursor -= 1;
                let undo = &self.commands[self.cursor].1;
                println!("{:#?}", undo);
                undo.apply(graph)?;
                println!("{:#?}", graph);
            }

            Ok(())
        }
    }

    fn three_node_add_cmds() -> Vec<DynCommand> {
        vec![
            Box::new(CreateNode::from(Node::literal(1.0))),
            Box::new(CreateNode::from(Node::literal(2.0))),
            Box::new(CreateNode::from(Node {
                kind: NodeKind::Op(AtomOp::Add),
                pos: glam::Vec2::ZERO,
                inputs: vec![None; 2],
            })),
            Box::new(SetEdge {
                edge: Edge {
                    input: SlotIndex { node: 0, slot: 0 },
                    output: SlotIndex { node: 2, slot: 0 },
                },
            }),
            Box::new(SetEdge {
                edge: Edge {
                    input: SlotIndex { node: 1, slot: 0 },
                    output: SlotIndex { node: 2, slot: 1 },
                },
            }),
        ]
    }

    #[test]
    fn create_node() -> GraphResult<()> {
        let mut graph = Graph::new();

        let cmd = CreateNode {
            kind: NodeKind::Lit(0.0),
            pos: glam::Vec2::ZERO,
            inputs: vec![],
        };

        let undo = cmd.undo(&graph)?;

        cmd.apply(&mut graph)?;
        assert_eq!(graph.nodes.len(), 1);

        undo.apply(&mut graph)?;
        assert_eq!(graph.nodes.len(), 0);

        Ok(())
    }

    #[test]
    fn three_node_add() -> GraphResult<()> {
        let mut graph = Graph::new();
        let cmds = three_node_add_cmds();

        for cmd in cmds.into_iter() {
            println!("{:#?}", cmd);
            cmd.apply(&mut graph)?;
            println!("{:#?}", graph);
        }

        graph.topological_order()?;

        Ok(())
    }

    #[test]
    fn three_node_add_rewind() -> GraphResult<()> {
        let mut graph = Graph::new();
        let mut history = CommandHistory::new();
        let cmds = three_node_add_cmds();
        history.push_multi(&mut graph, cmds)?;
        graph.topological_order()?;
        history.rewind(&mut graph)?;
        assert_eq!(graph.nodes.len(), 0);
        Ok(())
    }
}