Improved value maximizer

src/agents/util/partials.rs

@@ -0,0 +1,96 @@
+use super::{TreeCoords, TreeDir};
+use crate::{board::TreeElement, util::Symb};
+
+/// Find the coordinates of all partials in the given tree
+pub fn find_partials(tree: &TreeElement) -> Vec<TreeCoords> {
+	let mut partials = Vec::new();
+	let mut current_coords = TreeCoords::new();
+
+	loop {
+		let t = current_coords.get_from(tree).unwrap();
+		match t {
+			TreeElement::Number(_) | TreeElement::Partial(_) => {
+				if let TreeElement::Partial(_) = t {
+					partials.push(current_coords);
+				}
+
+				loop {
+					match current_coords.pop() {
+						Some((TreeDir::Left, _)) => {
+							current_coords.push(
+								TreeDir::Right,
+								match current_coords.get_from(tree) {
+									Some(TreeElement::Add { .. }) => current_coords.is_inverted(),
+									Some(TreeElement::Mul { .. }) => current_coords.is_inverted(),
+									Some(TreeElement::Sub { .. }) => !current_coords.is_inverted(),
+									Some(TreeElement::Div { .. }) => !current_coords.is_inverted(),
+									_ => unreachable!(),
+								},
+							);
+							break;
+						}
+						Some((TreeDir::Right, _)) => {}
+						Some((TreeDir::This, _)) => unreachable!(),
+						None => return partials,
+					}
+				}
+			}
+			TreeElement::Div { .. }
+			| TreeElement::Mul { .. }
+			| TreeElement::Sub { .. }
+			| TreeElement::Add { .. } => current_coords.push(TreeDir::Left, current_coords.is_inverted()),
+			TreeElement::Neg { .. } => {
+				current_coords.push(TreeDir::Right, !current_coords.is_inverted())
+			}
+		}
+	}
+}
+
+/// Fill empty slots in the given partials, in order.
+/// Will panic if we run out of numbers to fill with.
+///
+/// Returns a new tree with filled partials.
+pub fn fill_partials<'a>(
+	tree: &'a TreeElement,
+	partials: impl Iterator<Item = &'a TreeCoords>,
+	mut numbers: impl Iterator<Item = &'a Symb>,
+) -> TreeElement {
+	let mut tmp_tree = tree.clone();
+	for p in partials {
+		let x = p.get_from_mut(&mut tmp_tree).unwrap();
+
+		let x_str = match x {
+			TreeElement::Partial(s) => s,
+			_ => unreachable!(),
+		};
+		let mut new_str = String::new();
+		for c in x_str.chars() {
+			if c == '_' {
+				new_str.push_str(&format!("{}", numbers.next().unwrap()))
+			} else {
+				new_str.push(c);
+			}
+		}
+		*x = TreeElement::Partial(new_str)
+	}
+
+	tmp_tree
+}
+
+/// Find all empty slots in the given partials
+/// Returns (coords of partial, index of slot in string)
+pub fn free_chars<'a>(
+	tree: &'a TreeElement,
+	partials: impl Iterator<Item = &'a TreeCoords>,
+) -> Vec<(&TreeCoords, usize)> {
+	partials
+		.flat_map(|x| match x.get_from(tree) {
+			Some(TreeElement::Partial(s)) => {
+				s.chars()
+					.enumerate()
+					.filter_map(move |(i, c)| if c == '_' { Some((x, i)) } else { None })
+			}
+			_ => unreachable!(),
+		})
+		.collect()
+}