tree_morph/
helpers.rs

1//! Various selector functions for different use cases.
2//!
3//! A selector function accepts an iterator over ([`Rule`], [`Update`]) pairs and returns the
4//! selected [`Update`], or `None`.
5//!
6//! [`morph`](crate::morph) will call the given selector function when there is an ambiguity in
7//! which rule to apply. That is, when more than one rule from the same group returns [`Some(...)`]
8//! for a given sub-tree.
9
10use std::{collections::VecDeque, fmt::Display, io::Write};
11
12use crate::prelude::{Rule, Update};
13use multipeek::multipeek;
14use uniplate::Uniplate;
15
16/// Returns the first [`Update`] if the iterator only yields one, otherwise calls `select`.
17///
18/// See the module-level documentation for more information.
19pub(crate) fn one_or_select<T, M, R>(
20    select: impl Fn(&T, &mut dyn Iterator<Item = (&R, Update<T, M>)>) -> Option<Update<T, M>>,
21    t: &T,
22    rs: &mut dyn Iterator<Item = (&R, Update<T, M>)>,
23) -> Option<Update<T, M>>
24where
25    T: Uniplate,
26    R: Rule<T, M>,
27{
28    let mut rs = multipeek(rs);
29    if rs.peek_nth(1).is_none() {
30        return rs.next().map(|(_, u)| u);
31    }
32    select(t, &mut rs)
33}
34
35/// A uniform type for selector functions such as `select_first`.
36pub type SelectorFn<T, M, R> =
37    fn(&T, &mut dyn Iterator<Item = (&R, Update<T, M>)>) -> Option<Update<T, M>>;
38
39/// Returns the first available [`Update`] if there is one, otherwise returns `None`.
40///
41/// This is a good default selection strategy, especially when you expect only one possible
42/// rule to apply to any one term.
43///
44/// See the module-level documentation for more information.
45pub fn select_first<T, M, R>(
46    _: &T,
47    rs: &mut dyn Iterator<Item = (&R, Update<T, M>)>,
48) -> Option<Update<T, M>>
49where
50    T: Uniplate,
51    R: Rule<T, M>,
52{
53    rs.next().map(|(_, u)| u)
54}
55
56/// Panics when called by the engine, printing the original subtree and all applicable rules
57/// and their results.
58///
59/// This is useful when you always expect no more than one rule to be applicable, as the
60/// engine will only call the selector function when there is an ambiguity in which to apply.
61///
62/// See the module-level documentation for more information.
63pub fn select_panic<T, M, R>(
64    t: &T,
65    rs: &mut dyn Iterator<Item = (&R, Update<T, M>)>,
66) -> Option<Update<T, M>>
67where
68    T: Uniplate + std::fmt::Debug,
69    R: Rule<T, M> + std::fmt::Debug,
70{
71    let rules = rs.map(|(r, _)| r).collect::<Vec<_>>();
72    // Since `one_or_select` only calls the selector if there is more than one rule,
73    // at this point there is guaranteed to be more than one rule.
74    panic!("Multiple rules applicable to expression {t:?}\n{rules:?}",);
75}
76
77/// Selects an [`Update`] based on user input through stdin.
78///
79/// See the module-level documentation for more information.
80pub fn select_user_input<T, M, R>(
81    t: &T,
82    rs: &mut dyn Iterator<Item = (&R, Update<T, M>)>,
83) -> Option<Update<T, M>>
84where
85    T: Uniplate + Display,
86    R: Rule<T, M> + Display,
87{
88    let mut choices: Vec<_> = rs.collect();
89
90    let rules = choices
91        .iter()
92        .enumerate()
93        .map(
94            |(
95                i,
96                (
97                    r,
98                    Update {
99                        new_subtree: new_tree,
100                        ..
101                    },
102                ),
103            )| {
104                format!(
105                    "{}. {}
106   ~> {}",
107                    i + 1,
108                    r,
109                    new_tree
110                )
111            },
112        )
113        .collect::<Vec<_>>()
114        .join("\n\n");
115
116    loop {
117        print!(
118            "--- Current Expression ---
119{t}
120
121--- Rules ---
122{rules}
123
124---
125q   No change
126<n> Apply rule n
127
128:",
129        );
130        std::io::stdout().flush().unwrap(); // Print the : on same line
131
132        let mut line = String::new();
133        std::io::stdin().read_line(&mut line).unwrap();
134
135        match line.trim() {
136            "q" => return None,
137            n => {
138                if let Ok(n) = n.parse::<usize>()
139                    && n > 0
140                    && n <= choices.len()
141                {
142                    let ret = choices.swap_remove(n - 1).1;
143                    return Some(ret);
144                }
145            }
146        }
147    }
148}
149
150/// Selects a random [`Update`] from the iterator.
151///
152/// See the module-level documentation for more information.
153pub fn select_random<T, M, R>(
154    _: &T,
155    rs: &mut dyn Iterator<Item = (&R, Update<T, M>)>,
156) -> Option<Update<T, M>>
157where
158    T: Uniplate,
159    R: Rule<T, M>,
160{
161    use rand::seq::IteratorRandom;
162    let mut rng = rand::rng();
163    rs.choose(&mut rng).map(|(_, u)| u)
164}
165
166/// Selects the [`Update`] which results in the smallest subtree.
167///
168/// Subtree size is determined by maximum depth.
169/// Among trees with the same depth, the first in the iterator order is selected.
170///
171/// See the module-level documentation for more information.
172pub fn select_smallest_subtree<T, M, R>(
173    _: &T,
174    rs: &mut dyn Iterator<Item = (&R, Update<T, M>)>,
175) -> Option<Update<T, M>>
176where
177    T: Uniplate,
178    R: Rule<T, M>,
179{
180    rs.min_by_key(|(_, u)| {
181        u.new_subtree.cata(&|_, cs: VecDeque<i32>| {
182            // Max subtree height + 1
183            cs.iter().max().unwrap_or(&0) + 1
184        })
185    })
186    .map(|(_, u)| u)
187}