'Learning' Animal Guesser Program

Question

I've written a program that uses a binary tree to essentially 'learn' about animals given the name and questions about them. Learn is probably a little strong of a word, all it does is store the questions using bincode (pre-release 2.0.0-rc.2) and ask them later to figure out what the user is talking about.

#![warn(clippy::pedantic, clippy::nursery, rust_2018_idioms)]

use std::{fs::File, io::Seek, path::PathBuf};

use bincode::{Decode, Encode};

#[derive(Encode, Decode, Debug)]
enum Btree {
    Branch {
        question: String,
        // No
        left: Box<Btree>,
        // Yes
        right: Box<Btree>,
    },
    Leaf {
        animal: String,
    },
}

impl Btree {
    fn get_current(&self, current: &[bool]) -> &Self {
        let mut tree = self;
        for &i in current {
            // Should always be a branch
            if let Self::Branch {
                question: _,
                left,
                right,
            } = tree
            {
                if i {
                    tree = right;
                } else {
                    tree = left;
                }
            } else {
                unreachable!()
            }
        }
        tree
    }

    fn get_current_mut(&mut self, current: &[bool]) -> &mut Self {
        let mut tree = self;
        for &i in current {
            // Should always be a branch
            if let Self::Branch {
                question: _,
                left,
                right,
            } = tree
            {
                if i {
                    tree = right;
                } else {
                    tree = left;
                }
            } else {
                unreachable!()
            }
        }
        tree
    }
}

fn create_db(db_path: PathBuf) {
    println!("Database not found, creating now");
    let mut db_file = std::fs::File::create(db_path).expect("Couldn't create database file");
    let animal1: String =
        promptly::prompt("Please enter the name of an animal").expect("Couldn't read answer");
    let animal2: String =
        promptly::prompt("Please enter the name of another animal").expect("Couldn't read answer");
    let question: String = promptly::prompt(format!(
        "Please write a yes/no question to differentiate between a(n) {animal1} and a(n) {animal2}"
    ))
    .expect("Couldn't read question");
    let animal_yesno: bool = promptly::prompt(format!(
        "Does 'yes' in your question correspond to a {animal1}?"
    ))
    .expect("Couldn't read answer");
    let btree = if animal_yesno {
        Btree::Branch {
            question,
            left: Box::new(Btree::Leaf { animal: animal2 }),
            right: Box::new(Btree::Leaf { animal: animal1 }),
        }
    } else {
        Btree::Branch {
            question,
            left: Box::new(Btree::Leaf { animal: animal1 }),
            right: Box::new(Btree::Leaf { animal: animal2 }),
        }
    };
    bincode::encode_into_std_write(btree, &mut db_file, bincode::config::standard())
        .expect("Couldn't save database");
}

fn read_db() -> Option<(Btree, File)> {
    let db_path: PathBuf =
        promptly::prompt("Enter the database path").expect("Couldn't read database path");
    let mut db_file = if let Ok(db_file) = std::fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open(&db_path)
    {
        db_file
    } else {
        create_db(db_path);
        return None;
    };
    Some((
        bincode::decode_from_std_read(&mut db_file, bincode::config::standard())
            .expect("Couldn't parse database"),
        db_file,
    ))
}

fn main() {
    let (mut db, mut db_file) = match read_db() {
        Some(dbs) => dbs,
        None => return,
    };
    let mut current_db = Vec::new();
    loop {
        match db.get_current(&current_db) {
            Btree::Branch {
                question,
                left: _,
                right: _,
            } => {
                let ans: bool = promptly::prompt(question).expect("Couldn't read answer");
                if ans {
                    current_db.push(true);
                } else {
                    current_db.push(false);
                }
            }
            Btree::Leaf { animal } => {
                let ans: bool = promptly::prompt(format!("Is your animal a(n) {animal}?"))
                    .expect("Couldn't read answer");
                if ans {
                    println!("Yay!");
                    break;
                }
                let new_animal: String = promptly::prompt("Please enter the name of your animal")
                    .expect("Couldn't read answer");
                let new_question: String = promptly::prompt(format!("Please write a yes/no question to differentiate between a(n) {new_animal} and a(n) {animal}")).expect("Couldn't read question");
                let new_animal_yesno: bool = promptly::prompt(format!(
                    "Does 'yes' in your question correspond to a(n) {new_animal}?"
                ))
                .expect("Couldn't read answer");
                let animal = animal.clone();
                let parent = db.get_current_mut(&current_db[0..current_db.len() - 1]);
                let parent = if let Btree::Branch {
                    question: _,
                    left,
                    right,
                } = parent
                {
                    if *current_db.last().unwrap() {
                        right
                    } else {
                        left
                    }
                } else {
                    unreachable!();
                };
                *parent = if new_animal_yesno {
                    Box::new(Btree::Branch {
                        question: new_question,
                        left: Box::new(Btree::Leaf { animal }),
                        right: Box::new(Btree::Leaf { animal: new_animal }),
                    })
                } else {
                    Box::new(Btree::Branch {
                        question: new_question,
                        left: Box::new(Btree::Leaf { animal: new_animal }),
                        right: Box::new(Btree::Leaf { animal }),
                    })
                };
                db_file.rewind().expect("Couldn't truncate database");
                bincode::encode_into_std_write(db, &mut db_file, bincode::config::standard())
                    .expect("Couldn't save database");
                println!("Database changes saved");
                break;
            }
        }
    }
}

I'm hoping to be able to find a better way to navigate down the tree instead of using a Vec<bool>. If possible, I'm also hoping for a better way to get around the can't have an immutable and mutable reference at the same time error that I tried to fix with let animal = animal.clone(). Of course, any other improvements are also welcome.

Answer 1

When it comes to tree data structures in Rust, you have three options:

Firstly, you can write your algorithm recursively. So you could write a function something like:

  fn play_game(node: &mut Btree) {
         // either:
         play_game(node.left);
         // or:
         play_game(node.right);
         // or:
         *node = Btree {
            // replace the node
          }
     }

Secondly, you can use indexes to keep track your location in the tree. This is essentially what you've done with current_db which keep track of the location of the current node. You can go a little further and always use indexes, even with the data structure itself.

enum Node {
    Branch {
        question: String,
        left: usize,
        right: usize,
    },
    Leaf {
        animal: String,
    },
}

In this case left and right would be indexes into a Vec<Node> on a basic tree object. Then just keep track of simpler usize indexes to know which node you are working with.

Thirdly, you can use RefCell. Instead of Box<Btree> use Box<RefCell<Btree>>. RefCell has interior mutability, you can modify the interior of the RefCell even if you don't have a mutable reference to the RefCell itself. This means that you can navigate and manipulate the parts of the tree more easily at the cost of extra runtime checking.

Personally, I favor the first option when my algorithm is naturally suited to running recursively on trees. I favor the second option when I need to freely manipulate a tree. I pretty much never take the third option.

As for the animal.clone() issue:

If you rewrite your game playing algorithm in a recursive fashion, it should be possible to avoid it. Basically, you need to write the algorithm in such a way that Rust realize that you are replacing the same Btree that you taking the animal String from, so its okay to move it.

But in the event that proves tricky, there is an alternative:

let animal = std::mem::replace(animal, String::new());

This will replace the original animal name with a temporary empty string (which will not allocate and thus be cheaper than cloning the string) and gives you the original animal String for whatever use you might put it to.