Improved value maximizer

src/agents/minmaxtree.rs

@@ -1,300 +1,51 @@
-use std::{
+use std::{iter, num::NonZeroU8};
-	fmt::{Debug, Display},
-	iter,
-	num::NonZeroU8,
-	thread,
-};
 
 use anyhow::Result;
 use itertools::Itertools;
-use rayon::iter::{ParallelBridge, ParallelIterator};
 
-use super::{MaximizerAgent, MinimizerAgent, RandomAgent};
+use super::{
+	util::{fill_partials, TreeCoords},
+	MaximizerAgent, MinimizerAgent, RandomAgent,
+};
 use crate::{
+	agents::util::{find_partials, free_chars},
 	board::{Board, PlayerAction, TreeElement},
 	util::Symb,
 };
 
 pub struct MinMaxTree {}
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+fn find_best_numbers_v1<'a, F>(
-enum TreeDir {
-	Right,
-	Left,
-	This,
-}
-
-#[derive(Clone, Copy)]
-struct TreeCoords {
-	len: usize,
-	coords: [TreeDir; 4],
-	inversion: [bool; 4],
-}
-
-impl Display for TreeCoords {
-	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-		if self.get_inversion() {
-			write!(f, "-")?
-		} else {
-			write!(f, "+")?
-		}
-
-		for c in self.coords {
-			match c {
-				TreeDir::Left => write!(f, "L")?,
-				TreeDir::Right => write!(f, "R")?,
-				TreeDir::This => break,
-			}
-		}
-
-		Ok(())
-	}
-}
-
-impl Debug for TreeCoords {
-	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-		Display::fmt(self, f)
-	}
-}
-
-#[allow(dead_code)]
-impl TreeCoords {
-	pub fn new() -> Self {
-		Self {
-			len: 0,
-			coords: [TreeDir::This; 4],
-			inversion: [false; 4],
-		}
-	}
-
-	pub fn push(&mut self, dir: TreeDir, invert: bool) {
-		if self.len == 4 || dir == TreeDir::This {
-			return;
-		}
-
-		self.coords[self.len] = dir;
-		self.inversion[self.len] = invert;
-		self.len += 1;
-	}
-
-	pub fn pop(&mut self) -> Option<(TreeDir, bool)> {
-		if self.len == 0 {
-			return None;
-		}
-
-		self.len -= 1;
-		let dir = self.coords[self.len];
-		let inv = self.inversion[self.len];
-		self.coords[self.len] = TreeDir::This;
-		self.inversion[self.len] = false;
-		Some((dir, inv))
-	}
-
-	pub fn get_inversion(&self) -> bool {
-		if self.len == 0 {
-			false
-		} else {
-			self.inversion[self.len - 1]
-		}
-	}
-
-	pub fn get_from<'a>(&self, mut tree: &'a TreeElement) -> Option<&'a TreeElement> {
-		for i in 0..self.len {
-			match &self.coords[i] {
-				TreeDir::Left => {
-					if let Some(t) = tree.left() {
-						tree = t
-					} else {
-						return None;
-					}
-				}
-
-				TreeDir::Right => {
-					if let Some(t) = tree.right() {
-						tree = t
-					} else {
-						return None;
-					}
-				}
-
-				TreeDir::This => return Some(tree),
-			}
-		}
-
-		Some(tree)
-	}
-
-	pub fn get_from_mut<'a>(&self, mut tree: &'a mut TreeElement) -> Option<&'a mut TreeElement> {
-		for i in 0..self.len {
-			match &self.coords[i] {
-				TreeDir::Left => {
-					if let Some(t) = tree.left_mut() {
-						tree = t
-					} else {
-						return None;
-					}
-				}
-
-				TreeDir::Right => {
-					if let Some(t) = tree.right_mut() {
-						tree = t
-					} else {
-						return None;
-					}
-				}
-
-				TreeDir::This => return Some(tree),
-			}
-		}
-
-		Some(tree)
-	}
-}
-
-/// Count the number of free spaces in partials we want to minimize
-fn count_min_slots(tree: &TreeElement, partials: &[TreeCoords]) -> usize {
-	partials
-		.iter()
-		.filter(|x| x.get_inversion())
-		.map(|x| match x.get_from(tree) {
-			Some(TreeElement::Partial(s)) => s.chars().filter(|x| *x == '_').count(),
-			_ => unreachable!(),
-		})
-		.sum()
-}
-
-/// Count the number of free spaces in partials we want to maximize
-fn count_max_slots(tree: &TreeElement, partials: &[TreeCoords]) -> usize {
-	partials
-		.iter()
-		.filter(|x| !x.get_inversion())
-		.map(|x| match x.get_from(tree) {
-			Some(TreeElement::Partial(s)) => s.chars().filter(|x| *x == '_').count(),
-			_ => unreachable!(),
-		})
-		.sum()
-}
-
-/// Find the coordinates of all partials in the given tree
-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.get_inversion(),
-									Some(TreeElement::Mul { .. }) => current_coords.get_inversion(),
-									Some(TreeElement::Sub { .. }) => {
-										!current_coords.get_inversion()
-									}
-									Some(TreeElement::Div { .. }) => {
-										!current_coords.get_inversion()
-									}
-									_ => 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.get_inversion()),
-			TreeElement::Neg { .. } => {
-				current_coords.push(TreeDir::Right, !current_coords.get_inversion())
-			}
-		}
-	}
-}
-
-fn fill_maxs(
 	tree: &TreeElement,
 	partials: &[TreeCoords],
-	mut numbers: impl Iterator<Item = Symb>,
+	numbers: impl Iterator<Item = &'a Symb>,
-) -> TreeElement {
+	minimize: bool,
-	let mut tmp_tree = tree.clone();
-	for p in partials.iter().filter(|x| !x.get_inversion()) {
-		let x = p.get_from_mut(&mut tmp_tree).unwrap();
 
-		let x_str = match x {
+	// Returns true if we want to maximize the given partial,
-			TreeElement::Partial(s) => s,
+	// and false if we want to fix it.
-			_ => unreachable!(),
+	filter: F,
-		};
+) -> Vec<Symb>
-		let mut new_str = String::new();
+where
-		for c in x_str.chars() {
+	F: Fn(&&TreeCoords) -> bool,
-			if c == '_' {
+{
-				new_str.push_str(&format!("{}", numbers.next().unwrap()))
+	// Fill maximizer slots with arbitrary numbers
-			} else {
+	let min_tree_base = fill_partials(
-				new_str.push(c);
-			}
-		}
-		*x = TreeElement::Number(new_str.parse().unwrap())
-	}
-
-	tmp_tree
-}
-
-fn fill_mins(
-	tree: &TreeElement,
-	partials: &[TreeCoords],
-	mut numbers: impl Iterator<Item = Symb>,
-) -> TreeElement {
-	let mut tmp_tree = tree.clone();
-	for p in partials.iter().filter(|x| x.get_inversion()) {
-		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::Number(new_str.parse().unwrap())
-	}
-
-	tmp_tree
-}
-
-fn find_best_maxs(tree: &TreeElement, partials: &[TreeCoords], maxs: &[Symb]) -> Vec<Symb> {
-	// Fill maximizer slots in arbitrary order
-	let min_tree_base = fill_mins(
 		tree,
-		partials,
+		partials.iter().filter(|x| !filter(x)),
-		iter::repeat(Symb::Number(NonZeroU8::new(5).unwrap())),
+		iter::repeat(&Symb::Number(NonZeroU8::new(5).unwrap())),
 	);
 
-	let trees: Vec<(f32, Vec<&Symb>)> = maxs
+	let partials_to_optimize: Vec<TreeCoords> = partials.iter().filter(filter).cloned().collect();
-		.iter()
+	let n_empty = free_chars(tree, partials_to_optimize.iter()).len();
-		.permutations(maxs.len())
+	println!("{:?}", n_empty);
+
+	let trees: Vec<(f32, Vec<&Symb>)> = numbers
+		.permutations(n_empty)
 		.unique()
-		.par_bridge()
+		.filter_map(move |l| {
-		.filter_map(|l| {
 			let mut i = l.iter();
 			let mut tmp_tree = min_tree_base.clone();
-			for p in partials.iter().filter(|x| !x.get_inversion()) {
+			for p in &partials_to_optimize {
 				let x = p.get_from_mut(&mut tmp_tree).unwrap();
 
 				let x_str = match x {
@@ -312,82 +63,128 @@ fn find_best_maxs(tree: &TreeElement, partials: &[TreeCoords], maxs: &[Symb]) ->
 				*x = TreeElement::Number(new_str.parse().unwrap())
 			}
 
+			println!("{:?}", tmp_tree);
 			tmp_tree.evaluate().map(|x| (x, l))
 		})
 		.collect();
 
-	let mut max_list: Option<Vec<&Symb>> = None;
+	let mut best_list: Option<Vec<&Symb>> = None;
 	let mut best_value: Option<f32> = None;
 
 	for (x, list) in trees {
 		if let Some(m) = best_value {
-			if m < x {
+			if (minimize && x < m) || (!minimize && x > m) {
 				best_value = Some(x);
-				max_list = Some(list);
+				best_list = Some(list);
 			}
 		} else {
 			best_value = Some(x);
-			max_list = Some(list);
+			best_list = Some(list);
 		}
 	}
 
-	max_list.unwrap().into_iter().cloned().collect()
+	best_list.unwrap().into_iter().cloned().collect()
 }
 
-fn find_best_mins(tree: &TreeElement, partials: &[TreeCoords], mins: &[Symb]) -> Vec<Symb> {
+fn find_best_numbers(
-	// Fill maximizer slots in arbitrary order
+	tree: &TreeElement,
-	let min_tree_base = fill_maxs(
+	partials: &[TreeCoords],
+
+	// The numbers we're allowed to add, sorted in ascending order
+	available_numbers: &[Symb],
+) -> TreeElement {
+	// Fill all empty slots with fives
+	let tree_filled = fill_partials(
 		tree,
-		partials,
+		partials.iter(),
-		iter::repeat(Symb::Number(NonZeroU8::new(5).unwrap())),
+		iter::repeat(&Symb::Number(NonZeroU8::new(5).unwrap())),
 	);
 
-	let trees: Vec<(f32, Vec<&Symb>)> = mins
+	let base = tree_filled.evaluate().unwrap();
-		.iter()
-		.permutations(mins.len())
-		.unique()
-		.par_bridge()
-		.filter_map(|l| {
-			let mut i = l.iter();
-			let mut tmp_tree = min_tree_base.clone();
-			for p in partials.iter().filter(|x| x.get_inversion()) {
-				let x = p.get_from_mut(&mut tmp_tree).unwrap();
 
-				let x_str = match x {
+	// Test each slot:
-					TreeElement::Partial(s) => s,
+	// Increase its value by 1, and record its effect on the
+	// expression's total value.
+	// This isn't a perfect metric, but it's pretty good.
+	let mut slots: Vec<(usize, &TreeCoords, usize, f32)> = free_chars(tree, partials.iter())
+		.into_iter()
+		.enumerate()
+		.map(|(i_slot, (c, i))| {
+			let mut new_tree = tree_filled.clone();
+			let p = c.get_from_mut(&mut new_tree).unwrap();
+			match p {
+				TreeElement::Partial(s) => s.replace_range(i..i + 1, "6"),
 				_ => unreachable!(),
-				};
-				let mut new_str = String::new();
-				for c in x_str.chars() {
-					if c == '_' {
-						new_str.push_str(&format!("{}", i.next().unwrap()))
-					} else {
-						new_str.push(c);
 			}
-				}
+			// This shouldn't ever be None.
-				*x = TreeElement::Number(new_str.parse().unwrap())
+			(i_slot, c, i, new_tree.evaluate().unwrap() - base)
-			}
-
-			tmp_tree.evaluate().map(|x| (x, l))
 		})
 		.collect();
 
-	let mut min_list: Option<Vec<&Symb>> = None;
+	// Sort by least to most influence
-	let mut best_value: Option<f32> = None;
+	slots.sort_by(|a, b| a.3.partial_cmp(&b.3).unwrap());
 
-	for (x, list) in trees {
+	let all_symbols = {
-		if let Some(m) = best_value {
+		// We need this many from the bottom, and this many from the top.
-			if m < x {
+		let neg_count = slots.iter().filter(|(_, _, _, x)| *x <= 0.0).count();
-				best_value = Some(x);
+		let pos_count = slots.iter().filter(|(_, _, _, x)| *x > 0.0).count();
-				min_list = Some(list);
+
+		let mut a_iter = available_numbers
+			.iter()
+			.take(neg_count)
+			.chain(available_numbers.iter().rev().take(pos_count).rev());
+
+		let mut g = slots
+			// Group slots with equal weights
+			// and count the number of elements in each group
+			.iter()
+			.group_by(|x| x.3)
+			.into_iter()
+			.map(|(_, x)| x.count())
+			// Generate the digits we should try for each group of
+			// equal-weight slots
+			.map(|s| {
+				(0..s)
+					.map(|_| a_iter.next().unwrap().clone())
+					.permutations(s)
+					.unique()
+					.collect_vec()
+			})
+			// Now, covert this to an array of all cartesian products
+			// of this set of sets
+			.multi_cartesian_product()
+			.map(|x| x.iter().flatten().cloned().collect_vec())
+			.map(|v| slots.iter().zip(v).collect_vec())
+			.collect_vec();
+
+		// Sort these vectors so the order of values
+		// matches the order of empty slots
+		g.iter_mut()
+			.for_each(|v| v.sort_by(|(a, _), (b, _)| a.0.partial_cmp(&b.0).unwrap()));
+		g.into_iter()
+			.map(|v| v.into_iter().map(|(_, s)| s).collect_vec())
+	};
+
+	let mut best_tree = None;
+	let mut best_value = None;
+	for i in all_symbols {
+		let tmp_tree = fill_partials(&tree, partials.iter(), i.iter());
+		let val = tmp_tree.evaluate();
+
+		if let Some(val) = val {
+			if let Some(best) = best_value {
+				if val > best {
+					best_value = Some(val);
+					best_tree = Some(tmp_tree)
 				}
 			} else {
-			best_value = Some(x);
+				best_value = Some(val);
-			min_list = Some(list);
+				best_tree = Some(tmp_tree)
+			}
 		}
 	}
 
-	min_list.unwrap().into_iter().cloned().collect()
+	best_tree.unwrap()
 }
 
 impl MinMaxTree {}
@@ -397,9 +194,6 @@ impl MinimizerAgent for MinMaxTree {
 		let tree = board.to_tree();
 		let partials = find_partials(&tree);
 
-		let max_slots = count_max_slots(&tree, &partials);
-		let min_slots = count_min_slots(&tree, &partials);
-
 		let available_numbers = (0..=9)
 			.map(|x| match x {
 				0 => Symb::Zero,
@@ -408,32 +202,13 @@ impl MinimizerAgent for MinMaxTree {
 			.filter(|x| !board.contains(*x))
 			.collect::<Vec<_>>();
 
-		if available_numbers.len() < max_slots {
+		// For the code below, we must guarantee that
+		// that is, min_slots + max_slots <= available_numbers.len
+		if available_numbers.len() < free_chars(&tree, partials.iter()).len() {
 			return RandomAgent {}.step_max(board);
 		}
 
-		// Assume these won't ever overlap
+		let t = find_best_numbers(&tree, &partials, &available_numbers);
-		// (that is, min_slots + max_slots <= available_numbers.len)
-		let mins: Vec<Symb> = available_numbers[0..min_slots].to_vec();
-		let maxs: Vec<Symb> = available_numbers[available_numbers.len() - max_slots..]
-			.iter()
-			.copied()
-			.rev()
-			.collect();
-
-		let t = tree.clone();
-		let p = partials.clone();
-		let ha = thread::spawn(move || find_best_mins(&t, &p, &mins[..]));
-
-		let t = tree.clone();
-		let p = partials.clone();
-		let hb = thread::spawn(move || find_best_maxs(&t, &p, &maxs[..]));
-
-		let best_min_list = ha.join().unwrap();
-		let best_max_list = hb.join().unwrap();
-
-		let t = fill_mins(&tree, &partials, best_min_list.into_iter());
-		let t = fill_maxs(&t, &partials, best_max_list.into_iter());
 
 		println!("{:?}", t);
 		RandomAgent {}.step_max(board)

src/agents/mod.rs

@@ -2,6 +2,7 @@ mod diffuse;
 mod minmaxtree;
 mod player;
 mod random;
+pub mod util;
 
 pub use diffuse::DiffuseAgent;
 pub use minmaxtree::MinMaxTree;

src/agents/util/mod.rs

@@ -0,0 +1,6 @@
+/// Common helper functions that may be used by agents.
+mod partials;
+mod treecoords;
+
+pub use partials::*;
+pub use treecoords::*;

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()
+}

src/agents/util/treecoords.rs

@@ -0,0 +1,137 @@
+use std::fmt::{Debug, Display};
+
+use crate::board::TreeElement;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum TreeDir {
+	Right,
+	Left,
+	This,
+}
+
+#[derive(Clone, Copy)]
+pub struct TreeCoords {
+	len: usize,
+	coords: [TreeDir; 4],
+	inversion: [bool; 4],
+}
+
+impl Display for TreeCoords {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		if self.is_inverted() {
+			write!(f, "-")?
+		} else {
+			write!(f, "+")?
+		}
+
+		for c in self.coords {
+			match c {
+				TreeDir::Left => write!(f, "L")?,
+				TreeDir::Right => write!(f, "R")?,
+				TreeDir::This => break,
+			}
+		}
+
+		Ok(())
+	}
+}
+
+impl Debug for TreeCoords {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		Display::fmt(self, f)
+	}
+}
+
+#[allow(dead_code)]
+impl TreeCoords {
+	pub fn new() -> Self {
+		Self {
+			len: 0,
+			coords: [TreeDir::This; 4],
+			inversion: [false; 4],
+		}
+	}
+
+	pub fn push(&mut self, dir: TreeDir, invert: bool) {
+		if self.len == 4 || dir == TreeDir::This {
+			return;
+		}
+
+		self.coords[self.len] = dir;
+		self.inversion[self.len] = invert;
+		self.len += 1;
+	}
+
+	pub fn pop(&mut self) -> Option<(TreeDir, bool)> {
+		if self.len == 0 {
+			return None;
+		}
+
+		self.len -= 1;
+		let dir = self.coords[self.len];
+		let inv = self.inversion[self.len];
+		self.coords[self.len] = TreeDir::This;
+		self.inversion[self.len] = false;
+		Some((dir, inv))
+	}
+
+	pub fn is_inverted(&self) -> bool {
+		if self.len == 0 {
+			false
+		} else {
+			self.inversion[self.len - 1]
+		}
+	}
+
+	pub fn get_from<'a>(&self, mut tree: &'a TreeElement) -> Option<&'a TreeElement> {
+		for i in 0..self.len {
+			match &self.coords[i] {
+				TreeDir::Left => {
+					if let Some(t) = tree.left() {
+						tree = t
+					} else {
+						return None;
+					}
+				}
+
+				TreeDir::Right => {
+					if let Some(t) = tree.right() {
+						tree = t
+					} else {
+						return None;
+					}
+				}
+
+				TreeDir::This => return Some(tree),
+			}
+		}
+
+		Some(tree)
+	}
+
+	pub fn get_from_mut<'a>(&self, mut tree: &'a mut TreeElement) -> Option<&'a mut TreeElement> {
+		for i in 0..self.len {
+			match &self.coords[i] {
+				TreeDir::Left => {
+					if let Some(t) = tree.left_mut() {
+						tree = t
+					} else {
+						return None;
+					}
+				}
+
+				TreeDir::Right => {
+					if let Some(t) = tree.right_mut() {
+						tree = t
+					} else {
+						return None;
+					}
+				}
+
+				TreeDir::This => return Some(tree),
+			}
+		}
+
+		Some(tree)
+	}
+}