1use std::io::{stdout, Write};
2
3use sgf_parse::{go::Prop, SgfNode};
4
5use crate::errors::QueryError;
6use crate::sgf_traversal::{variation_nodes, variation_roots, SgfTraversalNode};
7use crate::{QueryArgs, QueryMode};
8
9pub fn query(sgf: &str, query_args: &QueryArgs) -> Result<(), QueryError> {
10 let collection = sgf_parse::go::parse(sgf)?;
11 match query_args.mode() {
12 QueryMode::Default => write_query_text(&collection, stdout())?,
13 QueryMode::LastGame => println!("{}", query_game_index(&collection)?),
14 QueryMode::LastVariation => println!(
15 "{}",
16 query_variation_index(&collection, query_args.game_number)?
17 ),
18 QueryMode::LastNode => println!(
19 "{}",
20 query_node_index(&collection, query_args.game_number, query_args.variation)?
21 ),
22 }
23 Ok(())
24}
25
26fn query_game_index(collection: &[SgfNode<Prop>]) -> Result<usize, QueryError> {
27 match collection.len() {
28 0 => Err(QueryError::GameNotFound),
29 n => Ok(n - 1),
30 }
31}
32
33fn query_variation_index(
34 collection: &[SgfNode<Prop>],
35 game_number: u64,
36) -> Result<u64, QueryError> {
37 let sgf_node = collection
38 .get(game_number as usize)
39 .ok_or(QueryError::GameNotFound)?;
40 let node = variation_roots(sgf_node)
41 .last()
42 .ok_or(QueryError::VariationNotFound)?;
43 Ok(node.variation)
44}
45
46fn query_node_index(
47 collection: &[SgfNode<Prop>],
48 game_number: u64,
49 variation: u64,
50) -> Result<usize, QueryError> {
51 let sgf_node = collection
52 .get(game_number as usize)
53 .ok_or(QueryError::GameNotFound)?;
54 let count = variation_nodes(sgf_node, variation)
55 .map_err(|_| QueryError::VariationNotFound)?
56 .count();
57 match count {
58 0 => Err(QueryError::VariationNotFound),
59 n => Ok(n - 1),
60 }
61}
62
63fn write_query_text(
64 collection: &[SgfNode<Prop>],
65 mut writer: impl Write,
66) -> Result<(), QueryError> {
67 for (game_num, node) in collection.iter().enumerate() {
68 writeln!(writer, "Game #{game_num}")?;
69 write_game_text(node, &mut writer)?;
70 if game_num < collection.len() - 1 {
71 writeln!(writer)?;
72 }
73 }
74 Ok(())
75}
76
77fn write_game_text<W: Write>(sgf_node: &SgfNode<Prop>, writer: &mut W) -> Result<(), QueryError> {
78 for fork_node in variation_roots(sgf_node) {
79 let SgfTraversalNode {
80 sgf_node,
81 variation_node_number,
82 variation,
83 parent_variation: _,
84 branch_number: _,
85 is_variation_root: _,
86 branches,
87 } = fork_node;
88 let branch_diagram_for_line = {
89 let s: Vec<&str> = branches
90 .iter()
91 .enumerate()
92 .map(|(i, b)| {
93 if i < branches.len() - 1 {
94 match b {
95 true => "│ ",
96 false => " ",
97 }
98 } else {
99 match b {
100 true => "├── ",
101 false => "└── ",
102 }
103 }
104 })
105 .collect();
106 s.join("")
107 };
108 let last_move = std::iter::successors(Some(sgf_node), |n| n.children().next()).count() - 1
109 + variation_node_number as usize;
110 writeln!(
111 writer,
112 "{branch_diagram_for_line}v{variation}, {variation_node_number}-{last_move}",
113 )?;
114 }
115 Ok(())
116}
117
118#[cfg(test)]
119mod tests {
120 use sgf_parse::{go::Prop, SgfNode};
121
122 use super::{
123 query_game_index, query_node_index, query_variation_index, write_query_text, QueryError,
124 };
125
126 static TEST_DATA: &str = "\
127(;GM[1]FF[4]
128CA[UTF-8]
129AP[Quarry:0.2.0]
130SZ[9]
131KM[6.5]
132PB[Black]
133PW[White]
134;B[ee];W[ec]
135(;B[eg];W[eh];B[dh]
136(;W[cd];B[gd];W[fh])
137(;W[gd];B[fh]))
138(;B[fc];W[fb]
139(;B[fd]
140(;W[gb]
141(;B[eg])
142(;B[dg])
143(;B[ce]))
144(;W[eg]
145(;B[gb])
146(;B[fg];W[fh])))
147(;B[gc])
148(;B[eg]))
149)
150
151(;GM[1]FF[4]
152CA[UTF-8]
153AP[Quarry:0.2.0]
154SZ[9]
155KM[6.5]
156PB[Black]
157PW[White]
158;B[de];W[fe]
159(;B[ge])
160(;B[fg])
161)";
162
163 fn get_collection() -> Vec<SgfNode<Prop>> {
164 sgf_parse::go::parse(TEST_DATA).unwrap()
165 }
166
167 #[test]
168 fn query_diagram() {
169 let expected = "\
170Game #0
171v0, 0-8
172├── v0, 3-8
173│ ├── v0, 6-8
174│ └── v1, 6-7
175└── v2, 3-7
176 ├── v2, 5-7
177 │ ├── v2, 6-7
178 │ │ ├── v2, 7-7
179 │ │ ├── v3, 7-7
180 │ │ └── v4, 7-7
181 │ └── v5, 6-7
182 │ ├── v5, 7-7
183 │ └── v6, 7-8
184 ├── v7, 5-5
185 └── v8, 5-5
186
187Game #1
188v0, 0-3
189├── v0, 3-3
190└── v1, 3-3
191";
192
193 let mut output = vec![];
194 write_query_text(&get_collection(), &mut output).unwrap();
195 let output = String::from_utf8(output).unwrap();
196 assert_eq!(output, expected);
197 }
198
199 #[test]
200 fn game_index() {
201 let result = query_game_index(&get_collection()).unwrap();
202 assert_eq!(result, 1);
203 }
204
205 #[test]
206 fn game_index_error() {
207 let result = query_game_index(&[]);
208 assert!(matches!(result, Err(QueryError::GameNotFound)));
209 }
210
211 #[test]
212 fn variation_index() {
213 assert_eq!(query_variation_index(&get_collection(), 0).unwrap(), 8);
214 assert_eq!(query_variation_index(&get_collection(), 1).unwrap(), 1);
215 }
216
217 #[test]
218 fn variation_index_error() {
219 let result = query_variation_index(&get_collection(), 2);
220 assert!(matches!(result, Err(QueryError::GameNotFound)));
221 }
222
223 #[test]
224 fn node_index() {
225 assert_eq!(query_node_index(&get_collection(), 0, 0).unwrap(), 8);
226 assert_eq!(query_node_index(&get_collection(), 0, 1).unwrap(), 7);
227 assert_eq!(query_node_index(&get_collection(), 0, 2).unwrap(), 7);
228 assert_eq!(query_node_index(&get_collection(), 0, 3).unwrap(), 7);
229 assert_eq!(query_node_index(&get_collection(), 0, 4).unwrap(), 7);
230 assert_eq!(query_node_index(&get_collection(), 0, 5).unwrap(), 7);
231 assert_eq!(query_node_index(&get_collection(), 0, 6).unwrap(), 8);
232 assert_eq!(query_node_index(&get_collection(), 0, 7).unwrap(), 5);
233 assert_eq!(query_node_index(&get_collection(), 0, 8).unwrap(), 5);
234 assert_eq!(query_node_index(&get_collection(), 1, 0).unwrap(), 3);
235 assert_eq!(query_node_index(&get_collection(), 1, 1).unwrap(), 3);
236 }
237
238 #[test]
239 fn node_index_error() {
240 let result = query_node_index(&get_collection(), 0, 9);
241 assert!(matches!(result, Err(QueryError::VariationNotFound)));
242 }
243}