Skip to content

Commit

Permalink
Merge pull request #658 from egraphs-good/ajpal-svg2
Browse files Browse the repository at this point in the history 
Color edges in SVGs
  • Loading branch information
@ajpal
ajpal authored Nov 12, 2024
2 parents ba456ef + 0432704 commit 0dc0be2
Show file tree
Hide file tree
Showing 5 changed files with 207 additions and 107 deletions.
150 changes: 133 additions & 17 deletions src/rvsdg/rvsdg2svg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use std::fmt;
use std::iter::once;

use bril_rs::ConstOps;
use indexmap::IndexMap;

use super::{BasicExpr, Id, Operand, RvsdgBody, RvsdgFunction, RvsdgProgram};

Expand All @@ -22,6 +23,38 @@ pub(crate) struct Region {
edges: Vec<Edge>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) enum Color {
Teal,
SlateBlue,
Purple,
MediumVioletRed,
DarkOrange,
ForestGreen,
Maroon,
Crimson,
Goldenrod,
Black,
}

impl Color {
fn from_usize(i: usize) -> Self {
match i % 10 {
0 => Color::Teal,
1 => Color::SlateBlue,
2 => Color::Purple,
3 => Color::MediumVioletRed,
4 => Color::DarkOrange,
5 => Color::ForestGreen,
6 => Color::Maroon,
7 => Color::Crimson,
8 => Color::Goldenrod,
9 => Color::Black,
_ => unreachable!(),
}
}
}

#[derive(Debug)]
enum Node {
Unit(String, usize, usize),
Expand All @@ -34,6 +67,28 @@ enum Node {
// `None` refers to the region, `Some(i)` refers to the node at index `i`.
// The second number is the index of the port that is being referred to.
type Edge = ((Option<Id>, usize), (Option<Id>, usize));
type ColoredEdge = ((Option<Id>, usize), (Option<Id>, usize), Color);

fn color_edges(
edges: Vec<Edge>,
inputs: &[Color],
edge_colors: &mut IndexMap<Edge, Color>,
) -> Vec<ColoredEdge> {
edges
.iter()
.enumerate()
.map(|(idx, edge)| {
let ((x, i), (y, j)) = edge;
let color = match (x, edge_colors.get(edge)) {
(None, _) => inputs[*i],
(_, None) => Color::from_usize(idx),
(_, Some(c)) => *c,
};
edge_colors.insert(*edge, color);
((*x, *i), (*y, *j), color)
})
.collect()
}

struct Size {
width: f32,
Expand Down Expand Up @@ -133,7 +188,7 @@ fn port(x: f32, y: f32, color: &str) -> Xml {
}

impl Node {
fn to_xml(&self) -> (Size, Xml) {
fn to_xml(&self, inputs: &[Color], edge_colors: &mut IndexMap<Edge, Color>) -> (Size, Xml) {
match self {
Node::Unit(text, _, _) => {
let size = Size {
Expand Down Expand Up @@ -177,7 +232,11 @@ impl Node {
(size, group)
}
Node::Match(name, branches) => {
let children: Vec<_> = branches.iter().map(|t| t.1.to_xml(false)).collect();
let inputs = &inputs[1..]; // first input is the predicate, which we should ignore
let children: Vec<_> = branches
.iter()
.map(|t| t.1.to_xml(false, inputs, edge_colors))
.collect();
let size = Size {
width: REGION_SPACING * (children.len() + 1) as f32
+ children.iter().map(|t| t.0.width).sum::<f32>(),
Expand Down Expand Up @@ -265,7 +324,7 @@ impl Node {
(size, group)
}
Node::Loop(region) => {
let (s, mut xml) = region.to_xml(true);
let (s, mut xml) = region.to_xml(true, inputs, edge_colors);
let size = Size {
width: s.width + REGION_SPACING * 2.0,
height: s.height + FONT_SIZE + REGION_SPACING * 2.0,
Expand Down Expand Up @@ -319,11 +378,51 @@ impl Node {
}

impl Region {
fn to_xml(&self, in_loop: bool) -> (Size, Xml) {
let mut edges = self.edges.clone();
fn to_xml(
&self,
in_loop: bool,
inputs: &[Color],
edge_colors: &mut IndexMap<Edge, Color>,
) -> (Size, Xml) {
let mut edges: Vec<Edge> = self.edges.clone();
edges.sort();

let mut children: BTreeMap<_, _> = self.nodes.iter().map(|t| (t.0, t.1.to_xml())).collect();
let mut children: BTreeMap<_, _> = self
.nodes
.iter()
.map(|t| {
// Edges that are inputs to this node
let mut input_edges: Vec<&Edge> = edges
.iter()
.filter(|((_, _), (x, _))| match x {
Some(i) => i == t.0,
None => false,
})
.collect();
input_edges.sort_by_key(|(_, (_, x))| x);

let inputs: Vec<Color> = input_edges
.iter()
.enumerate()
.map(|(idx, edge)| {
let ((x, i), _) = edge;
let color = match (x, edge_colors.get(*edge)) {
(None, None) => inputs[*i],
(None, Some(c)) => {
assert!(inputs[*i] == *c);
*c
}
(Some(_), None) => Color::from_usize(idx),
(Some(_), Some(c)) => *c,
};
edge_colors.insert(**edge, color);
color
})
.collect();

(t.0, t.1.to_xml(&inputs, edge_colors))
})
.collect();

let mut layers: Vec<Vec<Id>> = vec![];
let mut to_order: BTreeSet<Id> = self.nodes.keys().copied().collect();
Expand Down Expand Up @@ -384,7 +483,9 @@ impl Region {
assert_eq!(w, size.width);
assert_eq!(h, size.height);

let edges = Xml::group(edges.iter().map(|((a, i), (b, j))| {
let colored_edges = color_edges(edges, inputs, edge_colors);

let colored_edges = Xml::group(colored_edges.iter().map(|((a, i), (b, j), edge_color)| {
let (a_x, a_y) = match a {
None => (blend(size.width, self.srcs, *i), 0.0),
Some(a) => (
Expand Down Expand Up @@ -449,7 +550,7 @@ impl Region {
"path",
[
("fill", "transparent"),
("stroke", "black"),
("stroke", &format!("{:?}", edge_color)),
("stroke-linecap", "round"),
("stroke-width", &format!("{}", STROKE_WIDTH)),
("d", &path_string),
Expand Down Expand Up @@ -483,13 +584,16 @@ impl Region {
"",
);

(size, Xml::group([background, edges, nodes, srcs, dsts]))
(
size,
Xml::group([background, colored_edges, nodes, srcs, dsts]),
)
}
}

impl Region {
fn to_svg(&self) -> String {
let (size, xml) = self.to_xml(false);
fn to_svg(&self, inputs: &[Color], edge_colors: &mut IndexMap<Edge, Color>) -> String {
let (size, xml) = self.to_xml(false, inputs, edge_colors);
let svg = Xml::new(
"svg",
[
Expand Down Expand Up @@ -650,7 +754,16 @@ impl RvsdgProgram {
height += spacing;
}

let (size, mut xml) = function.to_region().to_xml(false);
let colors: Vec<Color> = function
.args
.iter()
.enumerate()
.map(|(i, _)| Color::from_usize(i + 5))
.collect();
let mut edge_colors = IndexMap::new(); // fresh edge_colors map for each function
let (size, mut xml) = function
.to_region()
.to_xml(false, &colors, &mut edge_colors);
// assert that it doesn't have a transform yet
assert!(!xml.attributes.contains_key("transform"));
xml.attributes
Expand Down Expand Up @@ -679,7 +792,13 @@ impl RvsdgProgram {

impl RvsdgFunction {
pub(crate) fn to_svg(&self) -> String {
self.to_region().to_svg()
let colors: Vec<Color> = self
.args
.iter()
.enumerate()
.map(|(i, _)| Color::from_usize(i + 5))
.collect();
self.to_region().to_svg(&colors, &mut IndexMap::new())
}

pub(crate) fn to_region(&self) -> Region {
Expand Down Expand Up @@ -761,10 +880,7 @@ mod tests {
Type::Int,
)),
],
results: vec![
(RvsdgType::Bril(Type::Int), Operand::Project(0, 10)),
(RvsdgType::PrintState, Operand::Arg(2)),
],
results: vec![(RvsdgType::Bril(Type::Int), Operand::Project(0, 10))],
}
.to_svg();

Expand Down
11 changes: 5 additions & 6 deletions src/rvsdg/snapshots/eggcc__rvsdg__rvsdg2svg__add_rvsdg_svg.snap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,35 +21,35 @@ expression: svg
<path
d="M 325 0V 15Q 325 25 335 25H 556.6667Q 566.6667 25 566.6667 35V 350"
fill="transparent"
stroke="black"
stroke="ForestGreen"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 100 150V 170Q 100 180 110 180H 373.33334Q 383.33334 180 383.33334 190V 200"
fill="transparent"
stroke="black"
stroke="Teal"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 250 150V 160Q 250 170 260 170H 406.66666Q 416.66666 170 416.66666 180V 200"
fill="transparent"
stroke="black"
stroke="SlateBlue"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 400 300V 315Q 400 325 410 325H 523.3333Q 533.3333 325 533.3333 335V 350"
fill="transparent"
stroke="black"
stroke="Teal"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 550 450V 465Q 550 475 540 475H 335Q 325 475 325 485V 500"
fill="transparent"
stroke="black"
stroke="DarkOrange"
stroke-linecap="round"
stroke-width="2">
</path>
Expand Down Expand Up @@ -226,4 +226,3 @@ expression: svg
</g>
</g>
</svg>

22 changes: 11 additions & 11 deletions src/rvsdg/snapshots/eggcc__rvsdg__rvsdg2svg__diamond_rvsdg_svg.snap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,35 +21,35 @@ expression: svg
<path
d="M 341.66666 0V 10Q 341.66666 20 331.66666 20H 93.33333Q 83.33333 20 83.33333 30V 50"
fill="transparent"
stroke="black"
stroke="ForestGreen"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 341.66666 0V 10Q 341.66666 20 331.66666 20H 126.666664Q 116.666664 20 116.666664 30V 50"
fill="transparent"
stroke="black"
stroke="ForestGreen"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 683.3333 0V 20Q 683.3333 30 693.3333 30H 706.6667Q 716.6667 30 716.6667 40V 200"
fill="transparent"
stroke="black"
stroke="Maroon"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 100 150V 165Q 100 175 110 175H 448.33334Q 458.33334 175 458.33334 185V 200"
fill="transparent"
stroke="black"
stroke="Teal"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 587.5 700V 715Q 587.5 725 577.5 725H 522.5Q 512.5 725 512.5 735V 750"
fill="transparent"
stroke="black"
stroke="DarkOrange"
stroke-linecap="round"
stroke-width="2">
</path>
Expand Down Expand Up @@ -132,21 +132,21 @@ expression: svg
<path
d="M 175 0V 15Q 175 25 185 25H 256.66666Q 266.66666 25 266.66666 35V 200"
fill="transparent"
stroke="black"
stroke="Maroon"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 100 150V 165Q 100 175 110 175H 223.33333Q 233.33333 175 233.33333 185V 200"
fill="transparent"
stroke="black"
stroke="Teal"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 250 300V 315Q 250 325 240 325H 185Q 175 325 175 335V 350"
fill="transparent"
stroke="black"
stroke="Purple"
stroke-linecap="round"
stroke-width="2">
</path>
Expand Down Expand Up @@ -272,21 +272,21 @@ expression: svg
<path
d="M 175 0V 15Q 175 25 185 25H 256.66666Q 266.66666 25 266.66666 35V 200"
fill="transparent"
stroke="black"
stroke="Maroon"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 100 150V 165Q 100 175 110 175H 223.33333Q 233.33333 175 233.33333 185V 200"
fill="transparent"
stroke="black"
stroke="Teal"
stroke-linecap="round"
stroke-width="2">
</path>
<path
d="M 250 300V 315Q 250 325 240 325H 185Q 175 325 175 335V 350"
fill="transparent"
stroke="black"
stroke="Purple"
stroke-linecap="round"
stroke-width="2">
</path>
Expand Down
Loading

0 comments on commit 0dc0be2

Please sign in to comment.