mcai-graph 0.3.0

Graph definition for Media Cloud AI Workflows representation
Documentation
use crate::{Link, LinkType};
use mcai_types::Coordinates;
use std::{cell::RefCell, rc::Rc};

pub type SharedNode = Rc<RefCell<Node>>;

#[derive(Clone, Debug, PartialEq)]
pub struct Dimensions {
  width: usize,
  height: usize,
}

impl Dimensions {
  pub fn new(width: usize, height: usize) -> Self {
    Self { width, height }
  }

  pub fn width(&self) -> usize {
    self.width
  }

  pub fn height(&self) -> usize {
    self.height
  }
}

#[derive(Clone, Debug, PartialEq)]
pub struct Node {
  id: u32,
  coordinates: Coordinates,
  dimensions: Dimensions,
  parents: Vec<SharedNode>,
  requires: Vec<SharedNode>,
}

impl Node {
  pub fn new(id: u32, coordinates: Coordinates, dimensions: Dimensions) -> Self {
    Node {
      id,
      coordinates,
      dimensions,
      parents: vec![],
      requires: vec![],
    }
  }

  pub fn into_rc(self) -> SharedNode {
    Rc::new(RefCell::new(self))
  }

  pub fn id(&self) -> u32 {
    self.id
  }

  pub fn coordinates(&self) -> Coordinates {
    self.coordinates.clone()
  }

  pub fn dimensions(&self) -> Dimensions {
    self.dimensions.clone()
  }

  pub fn parents(&self) -> Vec<SharedNode> {
    self.parents.clone()
  }

  pub fn required(&self) -> Vec<SharedNode> {
    self.requires.clone()
  }

  pub fn get_input_coordinates(&self) -> Coordinates {
    let mut coordinates = self.coordinates.clone();
    coordinates.x += (self.dimensions.width / 2) as isize;
    coordinates
  }

  pub fn get_output_coordinates(&self) -> Coordinates {
    let mut coordinates = self.coordinates.clone();
    coordinates.x += (self.dimensions.width / 2) as isize;
    coordinates.y += (self.dimensions.height) as isize;
    coordinates
  }

  pub fn get_links(&self, link_type: &LinkType) -> Vec<Link> {
    match link_type {
      LinkType::Parentage => &self.parents,
      LinkType::Requirement => &self.requires,
    }
    .iter()
    .map(|node| (self, node.borrow(), *link_type))
    .map(Link::from)
    .collect::<Vec<Link>>()
  }

  pub fn set_parents(&mut self, parents: Vec<SharedNode>) {
    self.parents = parents;
  }

  pub fn set_required_nodes(&mut self, required_to_start: Vec<SharedNode>) {
    self.requires = required_to_start;
  }

  pub fn set_coordinates(&mut self, coordinates: Coordinates) {
    self.coordinates = coordinates;
  }

  pub fn add_parent(&mut self, parent: SharedNode) {
    if !self.parents.contains(&parent) {
      self.parents.push(parent);
    }
  }

  pub fn add_required(&mut self, required: SharedNode) {
    if !self.requires.contains(&required) {
      self.requires.push(required);
    }
  }

  pub fn remove_parent(&mut self, parent_id: u32) {
    self
      .parents
      .retain(|parent| parent.borrow().id != parent_id);
  }

  pub fn remove_required(&mut self, required_id: u32) {
    self
      .requires
      .retain(|required| required.borrow().id != required_id);
  }
}