Skip to main content

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