sgf_render/
query.rs

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}