Generic A*-Algorithm without clone

Question

I wrote a generic implementation of A* with the focus of not having to clone nodes or actions. Thanks to some very good advice I received on Stack Overflow, I have arrived at the following final result:

lib.rs

extern crate typed_arena;

use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::hash::{Hash, Hasher};

use typed_arena::Arena;

/// Extension of SearchTree that features action cost
pub trait SearchTreeState<A>: Eq + Hash {
    /// Get a list of available actions for this tree state and their corresponding costs
    fn available_actions(&self) -> Vec<(A, f64)>;
    
    /// Apply an action previously returned by available_actions
    /// and return the new search tree state
    fn apply_action(&self, act: &A) -> Self;
}

/// Node in the expanded search tree for uniform cost search with heuristic
struct HucsNode<A, T>
where
    T: SearchTreeState<A> + Hash
{
    cost: f64,
    heuristic_cost: f64,
    parent_index: usize,
    action: Option<A>,
    tree: T,
}

impl<A, T> PartialEq for HucsNode<A, T>
where
    T: SearchTreeState<A> + Hash
{
    fn eq(&self, other: &HucsNode<A, T>) -> bool {
        // Nodes are equal if their tree states as this is what we are
        // interested in when comparing them in the map
        return self.tree == other.tree
    }
}

impl<A, T> Eq for HucsNode<A, T>
where
    T: SearchTreeState<A> + Hash
{
}

impl<A, T> Hash for HucsNode<A, T>
where
    T: SearchTreeState<A> + Hash
{
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher
    {
        // Nodes are just hashed by their tree states since that also
        // defines their equality
        self.tree.hash(state);
    }
}

impl<A, T> PartialOrd for HucsNode<A, T>
where
    T: SearchTreeState<A> + Hash
{
    fn partial_cmp(&self, other: &HucsNode<A, T>) -> Option<Ordering> {
        // Must never return None here as used by Ord
        Some(self.cmp(other))
    }
}

impl<A, T> Ord for HucsNode<A, T>
where
    T: SearchTreeState<A> + Hash
{
    fn cmp(&self, other: &HucsNode<A, T>) -> Ordering {
        // Nodes are ordered by their summed costs for the open list
        let self_cost = self.cost + self.heuristic_cost;
        let other_cost = other.cost + other.heuristic_cost;
        
        // Flip for min-heap, PartialOrd should never return None
        return other_cost.partial_cmp(&self_cost).unwrap();
    }
}

/// A wrapper for a borrowed hashable thing
struct BackedHashWrapper<'a, T: 'a + Hash + Eq> {
    source: &'a T
}

impl<'a, T> Hash for BackedHashWrapper<'a, T>
where
    T: 'a + Eq + Hash
{
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher
    {
        self.source.hash(state);
    }
}

impl<'a, T> PartialEq for BackedHashWrapper<'a, T>
where
    T: 'a + Eq + Hash
{
    fn eq(&self, other: &BackedHashWrapper<T>) -> bool {
        self.source == other.source
    }
}

impl<'a, T> Eq for BackedHashWrapper<'a, T>
where
    T: 'a + Eq + Hash
{
}

/// The initial capacity of both the arena of old nodes and the hash map of old nodes
const INITIAL_OLD_LIST_CAPACITY: usize = 1_000_000;

/// Perform a uniform cost search with a valid heuristic function on a search tree.
/// Returns a sequence of actions if a state is found that satisfies the predicate or 
/// None if the search terminates beforehand.
pub fn hucs<A, T: SearchTreeState<A> + Hash> (
    tree: T,
    predicate: &Fn(&T) -> bool,
    heuristic: &Fn(&T) -> f64,
) -> Option<Vec<A>> {

    // Min heap of the nodes in the expanded tree, ordered by actual cost to get to node
    // + heuristic cost to goal
    let mut open_list = BinaryHeap::new();
    
    // Push the initial node onto the tree
    open_list.push(HucsNode {
        action: None,
        // Parent index of root should never be read
        parent_index: usize::max_value(),
        cost: 0.0,
        heuristic_cost: heuristic(&tree),
        tree: tree,
    });
    
    let mut found_leaf = None;
    
    let old_nodes_arena = Arena::with_capacity(INITIAL_OLD_LIST_CAPACITY);
    
    // Destroy hash_map after this scope so items in the arena are no longer immutably borrowed
    {
        // Contains hashes and references to old_nodes_arena
        let mut hash_map = HashMap::with_capacity(INITIAL_OLD_LIST_CAPACITY);
        let mut current_node_index = 0 as usize;
    
        'outer: while let Some(current_node) = open_list.pop() {
        
            if predicate(&current_node.tree) {
                found_leaf = Some(current_node);
                break 'outer;
            }
            
            // Temporarily wrap the current node so we can check against the hash map
            match hash_map.get(&BackedHashWrapper{ source: &current_node }) {
                // Skip if we already had a better or equal path to this state with less cost
                Some(old_cost) => if *old_cost <= current_node.cost { continue 'outer; },
                None           => {}
            }
            
            for (action, action_cost) in current_node.tree.available_actions() {
            
                let new_tree = current_node.tree.apply_action(&action);
                let new_cost = current_node.cost + action_cost;
                
                let new_node = HucsNode {
                    action: Some(action),
                    cost: new_cost,
                    parent_index: current_node_index,
                    heuristic_cost: heuristic(&new_tree),
                    tree: new_tree,
                };
                
                open_list.push(new_node);
            }
            
            // Add the current node to the arena of old nodes
            let current_node_ref = old_nodes_arena.alloc(current_node);
            
            // Add a wrapper to the hash map for the current node
            hash_map.insert(BackedHashWrapper {
                source: current_node_ref
            }, current_node_ref.cost);
            
            current_node_index += 1;
        }
    }
    
    return found_leaf.map(|leaf| {form_action_sequence(leaf, old_nodes_arena.into_vec())});
}

/// Restore the sequence of actions that was used to get to this node by climbing the tree of expanded nodes
fn form_action_sequence<A, T: SearchTreeState<A> + Hash>(
    leaf: HucsNode<A, T>,
    mut older_nodes: Vec<HucsNode<A, T>>,
) -> Vec<A> {

    let mut action_vector = Vec::new() as Vec<A>;
    
    let mut current = leaf;
    
    while let Some(action) = current.action {
        action_vector.insert(0, action);
        
        // Safe to swap as nodes' parents are always before them
        current = older_nodes.swap_remove(current.parent_index);
    }

    return action_vector;
}

I used the following program to test the algorithm, where AddGame tries to find the smallest combination of the given summands that add up to the goal (in a rather simple way).

main.rs

extern crate hucs;

use std::hash::{Hash, Hasher};
use std::time::Instant;

use hucs::{hucs, SearchTreeState};

fn main() {

    // Search that should fail
    run_game(11);

    for i in 1_000_000..1_000_020 {
        run_game(i);
    }
}

fn run_game(goal: u64) -> usize {

    let start = Instant::now();

    let game = AddGame::new(vec!(6, 16, 2017), goal);
    
    let predicate = game.make_predicate();
    let heuristic = game.make_heuristic();
    
    let count = match hucs(game, &*predicate, &*heuristic) {
        Some(solution) => solution.len(),
        None           => 0
    };
    
    let duration = start.elapsed();
    
    println!("{:>10}; {:>10}; {}.{:03}", goal, count, duration.as_secs(), duration.subsec_nanos() / 1000000);
    
    return count;
}

struct AddGame {
    summands: Vec<u64>,
    goal: u64,
    sum: u64,
}

/// AddGame is a basic test for the Uniform Cost Search where the combination
/// of the summands is wanted that has the least number of summands and adds up to
/// the given goal (a summand may be used multiple times)
impl AddGame {

    fn new(summands: Vec<u64>, goal: u64) -> Self {
        return AddGame {
            summands: summands,
            goal: goal,
            sum: 0,
        }
    }
    
    fn make_predicate(&self) -> Box<Fn(&AddGame) -> bool> {
        // Goal state is reached when the sum is exactly goal
        Box::new(move |game: &AddGame| game.sum == game.goal)
    }

    fn make_heuristic(&self) -> Box<Fn(&AddGame) -> f64> {
        // Heuristic score is remaining difference to goal divided by maximum summand
        let max_summand = self.summands.iter().cloned().max().unwrap();
        Box::new(move |game: &AddGame| (((game.goal - game.sum) / max_summand) as f64))
    }
}

impl Eq for AddGame {}

impl PartialEq for AddGame {
    fn eq(&self, other: &AddGame) -> bool {
        // Just compare by sum, used summands may differ
        // and goal should be equal if both trees originated
        // from the same root
        self.sum == other.sum
    }
}

impl Hash for AddGame
{
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher
    {
        // Just hash the state as only the state is relevant
        // for Eq
        self.sum.hash(state);
    }
}

impl SearchTreeState<u64> for AddGame {
    fn available_actions(&self) -> Vec<(u64, f64)> {
        let difference = self.goal - self.sum;
        
        return self.summands.iter()
            .cloned()
            // Do not exceed goal
            .filter(|s| *s <= difference)
            // Cost of applying an action is always 1 so we find minimum number of actions
            .map(|s| (s, 1.0))
            .collect();
    }
    
    fn apply_action(&self, act: &u64) -> Self {
        return AddGame {
            // Remove summands lower than the applied summands to order the action of adding summands,
            // otherwise have to go through a lot of different combinations of the same summand
            summands: self.summands.iter().cloned().filter(|s| *s >= *act).collect(),
            goal: self.goal,
            sum: self.sum + act
        }
    }
}

Since I am a complete beginner, I would be very interested in your feedback regarding both the idiomaticness of the code and the performance, though I think the latter is already quite good.

I used typed-arena 1.3.0.

Answer 1

The code looks pretty good! I can tell that you've paid attention to the current Rust idiomatic formatting, which is nice to see. The code is also pretty short and readable.

Overall

1. Write some tests!!11!1!!1 The code I show below compiles and I think it has the same results, but I don't know for sure.

2. Clippy has three classes of warnings for this code:

   1. warning: you seem to be trying to use match for destructuring a single pattern. Consider using `if let`

      Basically, you want to avoid this pattern

      match foo {
          One => interesting(),
          Two => {},
      }
      

      Instead, do

      if let One = foo {
          interesting();
      }
      

   2. warning: you should put `Foo` between ticks in the documentation

      This formats the documentation better.

   3. warning: unneeded return statement

      It's idiomatic to not have explicit returns.

main.rs

1. Probably want to derive Debug for almost all types.

2. In many cases, can derive Eq instead of the manual impl.

3. Some of the closures don't need move.

4. Closures don't need a type on the closure argument; they can inferred.

5. Recommend calling cloned after the filter, that way you aren't cloning needlessly. In this case, cloning numbers should see no difference.

6. I might just dereference the Copy types in the closure argument to avoid the cloned call and any dereferencing in the bodies.

7. Can compare a reference to a value to another reference to a value (e.g. &i32 with &i32), there's no need to dereference all the way to the value.

8. The pattern of

   match foo {
       Some(x) => x.bar(),
       None => default,
   }
   

   Can be replaced with map_or:

   let count = hucs(game, &*predicate, &*heuristic).map_or(0, |solution| solution.len());
   

extern crate hucs;

use std::hash::{Hash, Hasher};
use std::time::Instant;

use hucs::{hucs, SearchTreeState};

fn main() {
    // Search that should fail
    run_game(11);

    for i in 1_000_000..1_000_020 {
        run_game(i);
    }
}

fn run_game(goal: u64) -> usize {
    let start = Instant::now();

    let game = AddGame::new(vec!(6, 16, 2017), goal);

    let predicate = game.make_predicate();
    let heuristic = game.make_heuristic();

    let count = hucs(game, &*predicate, &*heuristic).map_or(0, |solution| solution.len());

    let duration = start.elapsed();

    println!("{:>10}; {:>10}; {}.{:03}", goal, count, duration.as_secs(), duration.subsec_nanos() / 1000000);

    count
}

#[derive(Debug, Eq)]
struct AddGame {
    summands: Vec<u64>,
    goal: u64,
    sum: u64,
}

/// `AddGame` is a basic test for the Uniform Cost Search where the combination
/// of the summands is wanted that has the least number of summands and adds up to
/// the given goal (a summand may be used multiple times)
impl AddGame {
    fn new(summands: Vec<u64>, goal: u64) -> Self {
        AddGame {
            summands: summands,
            goal: goal,
            sum: 0,
        }
    }

    fn make_predicate(&self) -> Box<Fn(&AddGame) -> bool> {
        // Goal state is reached when the sum is exactly goal
        Box::new(|game| game.sum == game.goal)
    }

    fn make_heuristic(&self) -> Box<Fn(&AddGame) -> f64> {
        // Heuristic score is remaining difference to goal divided by maximum summand
        let max_summand = self.summands.iter().cloned().max().unwrap();
        Box::new(move |game| (((game.goal - game.sum) / max_summand) as f64))
    }
}

impl PartialEq for AddGame {
    fn eq(&self, other: &AddGame) -> bool {
        // Just compare by sum, used summands may differ
        // and goal should be equal if both trees originated
        // from the same root
        self.sum == other.sum
    }
}

impl Hash for AddGame {
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher
    {
        // Just hash the state as only the state is relevant
        // for Eq
        self.sum.hash(state);
    }
}

impl SearchTreeState<u64> for AddGame {
    fn available_actions(&self) -> Vec<(u64, f64)> {
        let difference = self.goal - self.sum;

        self.summands.iter()
            // Do not exceed goal
            .filter(|&&s| s <= difference)
            // Cost of applying an action is always 1 so we find minimum number of actions
            .map(|&s| (s, 1.0))
            .collect()
    }

    fn apply_action(&self, act: &u64) -> Self {
        AddGame {
            // Remove summands lower than the applied summands to order the action of adding summands,
            // otherwise have to go through a lot of different combinations of the same summand
            summands: self.summands.iter().filter(|&s| s >= act).cloned().collect(),
            goal: self.goal,
            sum: self.sum + act
        }
    }
}

lib.rs

1. Don't need + Hash on every T: SearchTreeState<A> as that trait already has Hash as a supertrait.

2. Introduce a function on HucsNode to compute the total cost and use it in the cmp implementation. This avoids the value for self and other from diverging.

3. BackedHashWrapper can use derive all of the traits.

4. In fact, BackedHashWrapper doesn't provide any value; you can just store the reference in the hashmap directly.

5. I don't like the name hucs; would spell it out fully.

6. Why take a trait object reference (&Fn(...)) instead of a generic that implements the traits?

7. Why not combine the predicate and heuristic into a single trait; it seems like they are related and always together.

8. Naming a variable hash_map isn't descriptive. The type tells us that. Instead, describe what it's used for (instead of using a comment for that!)

9. In fact, you don't even need a hash map, a hash set will do, since you can get the cost from the key! Now the name is really bad ;-)

10. No need to specify as usize; type inference will take care of it.

11. Don't need the 'outer loop label since you aren't breaking or continuiing out of nested loops. Just break and continue are fine.

12. No need for as Vec<A>; type inference handles it.

13. Inserting at the beginning vector seems inefficient, as it has to move everything down each insert, (O(N^2)). Instead, insert at the end and then reverse. (O(N))

extern crate typed_arena;

use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashSet};
use std::hash::{Hash, Hasher};

use typed_arena::Arena;

/// Extension of `SearchTree` that features action cost
pub trait SearchTreeState<A>: Eq + Hash {
    /// Get a list of available actions for this tree state and their corresponding costs
    fn available_actions(&self) -> Vec<(A, f64)>;

    /// Apply an action previously returned by available_actions
    /// and return the new search tree state
    fn apply_action(&self, act: &A) -> Self;
}

/// Node in the expanded search tree for uniform cost search with heuristic
#[derive(Debug)]
struct HucsNode<A, T>
where
    T: SearchTreeState<A>
{
    cost: f64,
    heuristic_cost: f64,
    parent_index: usize,
    action: Option<A>,
    tree: T,
}

impl<A, T> HucsNode<A, T>
where
    T: SearchTreeState<A>
{
    fn total_cost(&self) -> f64 {
        self.cost + self.heuristic_cost
    }
}

impl<A, T> PartialEq for HucsNode<A, T>
where
    T: SearchTreeState<A>
{
    fn eq(&self, other: &HucsNode<A, T>) -> bool {
        // Nodes are equal if their tree states as this is what we are
        // interested in when comparing them in the map
        self.tree == other.tree
    }
}

impl<A, T> Eq for HucsNode<A, T>
where
    T: SearchTreeState<A>
{
}

impl<A, T> Hash for HucsNode<A, T>
where
    T: SearchTreeState<A>
{
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher
    {
        // Nodes are just hashed by their tree states since that also
        // defines their equality
        self.tree.hash(state);
    }
}

impl<A, T> PartialOrd for HucsNode<A, T>
where
    T: SearchTreeState<A>
{
    fn partial_cmp(&self, other: &HucsNode<A, T>) -> Option<Ordering> {
        // Must never return None here as used by Ord
        Some(self.cmp(other))
    }
}

impl<A, T> Ord for HucsNode<A, T>
where
    T: SearchTreeState<A>
{
    fn cmp(&self, other: &HucsNode<A, T>) -> Ordering {
        // Flip for min-heap, PartialOrd should never return None
        other.total_cost().partial_cmp(&self.total_cost()).unwrap()
    }
}

/// The initial capacity of both the arena of old nodes and the hash map of old nodes
const INITIAL_OLD_LIST_CAPACITY: usize = 1_000_000;

/// Perform a uniform cost search with a valid heuristic function on a search tree.
/// Returns a sequence of actions if a state is found that satisfies the predicate or
/// None if the search terminates beforehand.
pub fn hucs<A, T, P, H>(tree: T, predicate: P, heuristic: H) -> Option<Vec<A>>
where
    T: SearchTreeState<A>,
    P: Fn(&T) -> bool,
    H: Fn(&T) -> f64,
{
    // Min heap of the nodes in the expanded tree, ordered by actual cost to get to node
    // + heuristic cost to goal
    let mut open_list = BinaryHeap::new();

    // Push the initial node onto the tree
    open_list.push(HucsNode {
        action: None,
        // Parent index of root should never be read
        parent_index: usize::max_value(),
        cost: 0.0,
        heuristic_cost: heuristic(&tree),
        tree: tree,
    });

    let mut found_leaf = None;

    let old_nodes_arena = Arena::with_capacity(INITIAL_OLD_LIST_CAPACITY);

    // Destroy hash_map after this scope so items in the arena are no longer immutably borrowed
    {
        // Contains hashes and references to old_nodes_arena
        let mut hash_map: HashSet<&_> = HashSet::with_capacity(INITIAL_OLD_LIST_CAPACITY);
        let mut current_node_index = 0;

        while let Some(current_node) = open_list.pop() {
            if predicate(&current_node.tree) {
                found_leaf = Some(current_node);
                break;
            }

            // Check against the hash map
            if let Some(old_node) = hash_map.get(&current_node) {
                // Skip if we already had a better or equal path to this state with less cost
                if old_node.cost <= current_node.cost { continue; }
            }

            for (action, action_cost) in current_node.tree.available_actions() {
                let new_tree = current_node.tree.apply_action(&action);
                let new_cost = current_node.cost + action_cost;

                let new_node = HucsNode {
                    action: Some(action),
                    cost: new_cost,
                    parent_index: current_node_index,
                    heuristic_cost: heuristic(&new_tree),
                    tree: new_tree,
                };

                open_list.push(new_node);
            }

            // Add the current node to the arena of old nodes
            hash_map.insert(old_nodes_arena.alloc(current_node));

            current_node_index += 1;
        }
    }

    found_leaf.map(|leaf| form_action_sequence(leaf, old_nodes_arena.into_vec()))
}

/// Restore the sequence of actions that was used to get to this node by climbing the tree of expanded nodes
fn form_action_sequence<A, T: SearchTreeState<A>>(
    leaf: HucsNode<A, T>,
    mut older_nodes: Vec<HucsNode<A, T>>,
) -> Vec<A> {

    let mut action_vector = Vec::new();

    let mut current = leaf;

    while let Some(action) = current.action {
        action_vector.push(action);

        // Safe to swap as nodes' parents are always before them
        current = older_nodes.swap_remove(current.parent_index);
    }

    action_vector.reverse();

    action_vector
}