1use std::collections::BTreeSet;
4
5use crate::{DistanceTrace, Lattice, Symbol};
6
7const DOT_BEST_PATH_COLOR: &str = "#9a5b38";
8const DOT_DEFAULT_NODE_COLOR: &str = "#495057";
9const DOT_MUTED_EDGE_COLOR: &str = "#868e96";
10
11#[derive(Clone, Copy, Debug)]
13pub struct LatticeDotData<'a> {
14 pub left_input: &'a str,
16 pub right_input: &'a str,
18 pub left_lattice: &'a Lattice,
20 pub right_lattice: &'a Lattice,
22 pub distance: usize,
24 pub trace: Option<&'a DistanceTrace>,
26 pub trace_error: Option<&'a str>,
28}
29
30pub fn romaji_lattice_dot(data: &LatticeDotData<'_>) -> String {
32 lattice_pair_dot("moine_romaji_lattice", data)
33}
34
35pub fn pinyin_lattice_dot(data: &LatticeDotData<'_>) -> String {
37 lattice_pair_dot("moine_pinyin_lattice", data)
38}
39
40fn lattice_pair_dot(graph_name: &str, data: &LatticeDotData<'_>) -> String {
41 let left_symbols = data.trace.as_ref().map(|trace| trace.left_symbols());
42 let right_symbols = data.trace.as_ref().map(|trace| trace.right_symbols());
43 let left_best_arcs = left_symbols
44 .as_deref()
45 .map(|symbols| best_arc_keys(data.left_lattice, symbols))
46 .unwrap_or_default();
47 let right_best_arcs = right_symbols
48 .as_deref()
49 .map(|symbols| best_arc_keys(data.right_lattice, symbols))
50 .unwrap_or_default();
51 let left_best_nodes = best_nodes(&left_best_arcs);
52 let right_best_nodes = best_nodes(&right_best_arcs);
53
54 let best_path_label = match (&left_symbols, &right_symbols, data.trace_error) {
55 (Some(left), Some(right), _) => format!(
56 "best_left={}\\nbest_right={}",
57 dot_escape(&symbols_to_string(left)),
58 dot_escape(&symbols_to_string(right))
59 ),
60 (_, _, Some(error)) => format!("best path unavailable: {}", dot_escape(error)),
61 _ => "best path unavailable".to_string(),
62 };
63 let graph_label = format!("distance={}\\n{}", data.distance, best_path_label);
64
65 let mut dot = String::new();
66 dot.push_str(&format!("digraph {graph_name} {{\n"));
67 dot.push_str(" rankdir=LR;\n");
68 dot.push_str(" graph [fontname=\"Helvetica\", labelloc=\"t\", label=\"");
69 dot.push_str(&graph_label);
70 dot.push_str("\"];\n");
71 dot.push_str(&format!(
72 " node [fontname=\"Helvetica\", shape=circle, width=0.48, fixedsize=true, color=\"{DOT_DEFAULT_NODE_COLOR}\"];\n",
73 ));
74 dot.push_str(&format!(
75 " edge [fontname=\"Helvetica\", color=\"{DOT_DEFAULT_NODE_COLOR}\", arrowsize=0.7];\n\n"
76 ));
77
78 append_lattice_cluster(
79 &mut dot,
80 "right",
81 "RIGHT",
82 data.right_input,
83 data.right_lattice,
84 &right_best_arcs,
85 &right_best_nodes,
86 );
87 dot.push('\n');
88 append_lattice_cluster(
89 &mut dot,
90 "left",
91 "LEFT",
92 data.left_input,
93 data.left_lattice,
94 &left_best_arcs,
95 &left_best_nodes,
96 );
97 dot.push_str("}\n");
98 dot
99}
100
101pub fn symbols_to_string(symbols: &[Symbol]) -> String {
103 symbols
104 .iter()
105 .map(|&symbol| char::from_u32(symbol).unwrap_or(char::REPLACEMENT_CHARACTER))
106 .collect()
107}
108
109fn append_lattice_cluster(
110 dot: &mut String,
111 prefix: &str,
112 lane_label: &str,
113 input: &str,
114 lattice: &Lattice,
115 best_arcs: &BTreeSet<ArcKey>,
116 best_nodes: &BTreeSet<usize>,
117) {
118 dot.push_str(&format!(" subgraph cluster_{prefix} {{\n"));
119 dot.push_str(" style=\"rounded\";\n");
120 dot.push_str(" color=\"#ced4da\";\n");
121 dot.push_str(" label=\"");
122 dot.push_str(&format!("{lane_label}\\ninput={}", dot_escape(input)));
123 dot.push_str("\";\n");
124 for node in 0..lattice.node_count() {
125 let label = if node == lattice.start() {
126 "BOS".to_string()
127 } else if node == lattice.end() {
128 "EOS".to_string()
129 } else {
130 node.to_string()
131 };
132 let shape = if node == lattice.end() {
133 "doublecircle"
134 } else {
135 "circle"
136 };
137 let color = if best_nodes.contains(&node) {
138 DOT_BEST_PATH_COLOR
139 } else {
140 DOT_DEFAULT_NODE_COLOR
141 };
142 let penwidth = if best_nodes.contains(&node) {
143 "2.4"
144 } else {
145 "1.2"
146 };
147 dot.push_str(&format!(
148 " {prefix}_{node} [label=\"{}\", shape={shape}, color=\"{color}\", penwidth={penwidth}];\n",
149 dot_escape(&label)
150 ));
151 }
152 for arc in lattice.arcs() {
153 let key = arc_key(arc.src, arc.dst, arc.symbol);
154 let is_best = best_arcs.contains(&key);
155 let color = if is_best {
156 DOT_BEST_PATH_COLOR
157 } else {
158 DOT_MUTED_EDGE_COLOR
159 };
160 let penwidth = if is_best { "3.0" } else { "1.1" };
161 dot.push_str(&format!(
162 " {prefix}_{} -> {prefix}_{} [label=\"{}\", color=\"{color}\", fontcolor=\"{color}\", penwidth={penwidth}];\n",
163 arc.src,
164 arc.dst,
165 dot_escape(&symbol_to_string(arc.symbol))
166 ));
167 }
168 dot.push_str(" }\n");
169}
170
171type ArcKey = (usize, usize, Symbol);
172
173fn arc_key(src: usize, dst: usize, symbol: Symbol) -> ArcKey {
174 (src, dst, symbol)
175}
176
177fn best_arc_keys(lattice: &Lattice, symbols: &[Symbol]) -> BTreeSet<ArcKey> {
178 let mut path = Vec::new();
179 if find_arc_path(lattice, lattice.start(), symbols, 0, &mut path) {
180 path.into_iter().collect()
181 } else {
182 BTreeSet::new()
183 }
184}
185
186fn find_arc_path(
187 lattice: &Lattice,
188 node: usize,
189 symbols: &[Symbol],
190 symbol_idx: usize,
191 path: &mut Vec<ArcKey>,
192) -> bool {
193 if symbol_idx == symbols.len() {
194 return node == lattice.end();
195 }
196
197 for arc in lattice.outgoing_arcs(node) {
198 if arc.symbol != symbols[symbol_idx] {
199 continue;
200 }
201 path.push(arc_key(arc.src, arc.dst, arc.symbol));
202 if find_arc_path(lattice, arc.dst, symbols, symbol_idx + 1, path) {
203 return true;
204 }
205 path.pop();
206 }
207 false
208}
209
210fn best_nodes(best_arcs: &BTreeSet<ArcKey>) -> BTreeSet<usize> {
211 let mut nodes = BTreeSet::new();
212 for &(src, dst, _) in best_arcs {
213 nodes.insert(src);
214 nodes.insert(dst);
215 }
216 nodes
217}
218
219fn symbol_to_string(symbol: Symbol) -> String {
220 char::from_u32(symbol)
221 .unwrap_or(char::REPLACEMENT_CHARACTER)
222 .to_string()
223}
224
225fn dot_escape(value: &str) -> String {
226 let mut escaped = String::with_capacity(value.len());
227 for ch in value.chars() {
228 match ch {
229 '"' => escaped.push_str("\\\""),
230 '\\' => escaped.push_str("\\\\"),
231 '\n' => escaped.push_str("\\n"),
232 '\r' => {}
233 _ => escaped.push(ch),
234 }
235 }
236 escaped
237}
238
239#[cfg(test)]
240mod tests {
241 use super::*;
242 use crate::{try_distance_with_trace, Lattice};
243
244 #[test]
245 fn renders_best_path_in_dot() {
246 let left_lattice = Lattice::from_paths(["insat"]);
247 let right_lattice = Lattice::from_paths(["insatu", "zzzzz"]);
248 let trace = try_distance_with_trace(&left_lattice, &right_lattice).unwrap();
249 let dot = romaji_lattice_dot(&LatticeDotData {
250 left_input: "いんさt",
251 right_input: "印刷",
252 left_lattice: &left_lattice,
253 right_lattice: &right_lattice,
254 distance: trace.distance,
255 trace: Some(&trace),
256 trace_error: None,
257 });
258
259 assert!(dot.contains("digraph moine_romaji_lattice"));
260 assert!(dot.contains("best_left=insat"));
261 assert!(dot.contains("best_right=insatu"));
262 assert!(dot.contains("label=\"u\", color=\"#9a5b38\""));
263 }
264}