1
\$\begingroup\$

This is my solution to the third exercise (and last) at the end of chapter 8 of The Rust Programming Language book (https://doc.rust-lang.org/book/ch08-03-hash-maps.html#summary).

There are only 3 types of commands: ADD, GET and EXIT:

  • ADD syntax: ADD name TO department (where "name" and "department" are variables)
  • GET syntax: GET department (we can also use ALL instead of department variable to get the employees of all the departments)
  • EXIT: exit the program
use std::collections::HashMap;
use std::io;

// get the employees from all departments
fn get_all_employees(dict: &HashMap<String, Vec<String>>) -> Vec<String> {
    let mut employees = Vec::new();
    for (_, v) in dict.iter() {
        for employee in v {
            employees.push(employee.to_string());
        }
    }   
    employees
}

// get the employees from a department 
fn get_dept_employees(dept: &str, dict: &HashMap<String, Vec<String>>) -> Vec<String> {
    let employees = dict.get(dept);
    match employees {
        Some(v) => v.to_vec(),
        None => {
            println!("No such department.");
            vec![]
        },
    }
}

// execute an ADD command 
fn exec_add(cmd_tokens: &Vec<&str>, dict: &mut HashMap<String, Vec<String>>) {

    if cmd_tokens.len() != 4 {
        println!("Invalid command.");
        return;
    }

    // add employee name to department (create it if non-existent)
    let dept = dict.entry(cmd_tokens[3].to_string()).or_insert(Vec::new());
    dept.push(cmd_tokens[1].to_string());
    println!("'{}' added to '{}' department.", cmd_tokens[1], cmd_tokens[3]);
}

// execute a GET command
fn exec_get(cmd_tokens: &Vec<&str>, dict: &mut HashMap<String, Vec<String>>) {

    if cmd_tokens.len() != 2 {
        println!("Invalid command.");
        return;
    }

    let mut employees: Vec<String> = Vec::new();

    match cmd_tokens.get(1) {
        Some(&"ALL") => employees = get_all_employees(&dict),
        Some(dept) => employees = get_dept_employees(dept, &dict),
        _ => println!("Invalid command."),
    }

    employees.sort();
    for employee in employees {
        println!("{}", employee);
    }
}

fn main() {
    let mut dict: HashMap<String, Vec<String>> = HashMap::new();
    loop {

        // read command from stdin 
        let mut cmd = String::new();
        println!("Enter database command: ");
        io::stdin()
            .read_line(&mut cmd)
            .expect("Failed to read from stdin");
        cmd = cmd.trim().to_string();
        
        // parse command and execute
        let tokens: Vec<&str> = cmd.split_whitespace().collect();
        if let Some(action) = tokens.get(0) {
            match action {
                &"ADD" => exec_add(&tokens, &mut dict),
                &"GET" => exec_get(&tokens, &mut dict),
                &"EXIT" => break,
                _ => println!("Invalid command."),
            }
        };
    }
    println!("Bye.");
}

Any suggestion would be appreciated!

\$\endgroup\$

1 Answer 1

1
\$\begingroup\$

The get_all_employees function can be simplified using iterators

fn get_all_employees(dict: &HashMap<String, Vec<String>>) -> Vec<String> {
  dict
    .values()
    .map(|v| v.iter())
    .flatten()
    .map(Clone::clone)
    .collect()
}

Similarly, get_dept_employees can be entirely removed, in favor of a better design: is the given department is not present, you print "No such department." and return right away in exec_get.

fn exec_get(cmd_tokens: &Vec<&str>, dict: &mut HashMap<String, Vec<String>>) {
  [...]
  let employees = match cmd_tokens.get(1) {
    Some(&"ALL") => get_all_employees(dict),
    Some(dept) => match dict.get(dept) {
      Some(x) => x,
      None => {
        println!("No such department.");
        return;
      }
    },
    None => println!("Invalid command."),
  }
  [...]
}

There are further improvements possible with this function:

  • make it return a Result<(), &'static str> which contains an error message in case something went wrong;
  • cmd_tokens.get(1) could be replaced with cmd_tokens[1] because we have already checked that cmd_tokens.len() == 2;
  • we only need a regular borrow of HashMap<String, Vec<String>>, not a mutable one;
  • taking a &Vec<_> as an argument can be replaced with &[_].
fn exec_get(cmd_tokens: &[&str], dict: &HashMap<String, Vec<String>>) -> Result<(), &'static str> {
  if cmd_tokens.len() == 2 {
    let mut employees = match cmd_tokens[1] {
      "ALL" => get_all_employees(dict),
      dept => dict
        .get(dept)
        .ok_or("No such department.")?
        .to_vec(),
    };
    employees.sort();
    for employee in employees {
      println!("{}", employee);
    }
    Ok(())
  } else {
    Err("Invalid command.")
  }
}

But then we realize the match statement is really just an if statement

fn exec_get(cmd_tokens: &[&str], dict: &HashMap<String, Vec<String>>) -> Result<(), &'static str> {
  [...]
    let mut employees = if cmd_tokens[1] == "ALL" {
      get_all_employees(dict)
    } else {
      dict
        .get(cmd_tokens[1])
        .ok_or("No such department.")?
        .to_vec()
    };
  [...]
}

The same is applied to exec_add:

fn exec_add(cmd_tokens: &[&str], dict: &mut HashMap<String, Vec<String>>) -> Result<(), &'static str> {
  if cmd_tokens.len() == 4 {
    dict
      .entry(cmd_tokens[3].to_string())
      .or_default() // Because `Vec::default()` is like `Vec::new()`
      .push(cmd_token[1].to_string());
    println!("'{}' added to the department '{}'", cmd_tokens[1], cmd_tokens[3]);
    Ok(())
  } else {
    Err("Invalid command.")
  }
}

Finally let's adapt main to these changes

fn main () {
  let mut dict = HashMap::new(); // No need for type annotations here
  let mut command = String::new(); // Allocate only once this buffer
  loop {
    print!("Enter database command: ");
    io::stdin()
      .read_line(&mut command)
      .expect("Failed to read from stdin");
    let tokens = command
      .trim()
      .split_whitespace()
      .collect::<Vec<_>>();
    if let Some(&action) = tokens.get(0) { // Dereference `action` here so you don't have to write `&"ADD"`
      if let Err(msg) = match action {
        "ADD" => exec_add(&tokens, &mut dict),
        "GET" => exec_get(&tokens, &dict),
        "EXIT" => break,
        _ => Err("Invalid command."),
      } {
        println!("Error: {}", msg);
      }
    };
    command.clear(); // Clear buffer (less expensive than re-allocate).
  }
  println!("Bye.");
}

The full code is available at the playground.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.