Introduce ChalkDatabase::with_program helper

Author: tmandry

chalk-solve/src/solve/test/bench.rs

@@ -19,26 +19,22 @@ fn run_bench(
     bencher: &mut Bencher,
     expected: &str,
 ) {
-    let mut db = ChalkDatabase::default();
-
-    db.query_mut(ProgramText)
-        .set((), Arc::new(text.to_string()));
-    db.query_mut(ProgramSolverChoice).set((), solver_choice);
-
-    let program = db.lowered_program().unwrap();
-    let env = db.environment().unwrap();
-    ir::tls::set_current_program(&program, || {
-        let goal = parse_and_lower_goal(&program, goal_text).unwrap();
-        let peeled_goal = goal.into_peeled_goal();
-
-        // Execute once to get an expected result.
-        let result = solver_choice.solve_root_goal(&env, &peeled_goal);
-
-        // Check expectation.
-        assert_result(&result, expected);
-
-        // Then do it many times to measure time.
-        bencher.iter(|| solver_choice.solve_root_goal(&env, &peeled_goal));
+    ChalkDatabase::with_program(Arc::new(program_text.to_string()), solver_choice, |db| {
+        let program = db.lowered_program().unwrap();
+        let env = db.environment().unwrap();
+        ir::tls::set_current_program(&program, || {
+            let goal = parse_and_lower_goal(&program, goal_text).unwrap();
+            let peeled_goal = goal.into_peeled_goal();
+
+            // Execute once to get an expected result.
+            let result = solver_choice.solve_root_goal(&env, &peeled_goal);
+
+            // Check expectation.
+            assert_result(&result, expected);
+
+            // Then do it many times to measure time.
+            bencher.iter(|| solver_choice.solve_root_goal(&env, &peeled_goal));
+        });
     });
 }
 

src/bin/chalki.rs

@@ -18,15 +18,14 @@ use std::process::exit;
 use std::sync::Arc;
 
 use chalk::db::ChalkDatabase;
-use chalk::query::{LoweringDatabase, ProgramSolverChoice, ProgramText};
+use chalk::query::LoweringDatabase;
 use chalk::rust_ir;
 use chalk::rust_ir::lowering::*;
 use chalk_engine::fallible::NoSolution;
 use chalk_solve::ext::*;
 use chalk_solve::solve::SolverChoice;
 use docopt::Docopt;
 use rustyline::error::ReadlineError;
-use salsa::Database;
 
 const USAGE: &'static str = "
 chalk repl
@@ -77,15 +76,11 @@ impl Program {
     ///
     /// [`SolverChoice`]: struct.solve.SolverChoice.html
     fn new(text: String, solver_choice: SolverChoice) -> Result<Program> {
-        let mut db = ChalkDatabase::default();
-
-        db.query_mut(ProgramText)
-            .set((), Arc::new(text.to_string()));
-        db.query_mut(ProgramSolverChoice).set((), solver_choice);
-
-        let ir = db.checked_program().unwrap();
-        let env = db.environment().unwrap();
-        Ok(Program { text, ir, env })
+        ChalkDatabase::with_program(Arc::new(text.clone()), solver_choice, |db| {
+            let ir = db.checked_program().unwrap();
+            let env = db.environment().unwrap();
+            Ok(Program { text, ir, env })
+        })
     }
 }
 

src/db.rs

@@ -1,16 +1,34 @@
-use crate::query;
+use crate::query::{self, ProgramSolverChoice, ProgramText};
+use chalk_solve::solve::SolverChoice;
+use salsa::Database;
+use std::sync::Arc;
 
 #[derive(Default)]
 pub struct ChalkDatabase {
     runtime: salsa::Runtime<ChalkDatabase>,
 }
 
-impl salsa::Database for ChalkDatabase {
+impl Database for ChalkDatabase {
     fn salsa_runtime(&self) -> &salsa::Runtime<ChalkDatabase> {
         &self.runtime
     }
 }
 
+impl ChalkDatabase {
+    pub fn with_program<F: FnOnce(&mut ChalkDatabase) -> R, R>(
+        program_text: Arc<String>,
+        solver_choice: SolverChoice,
+        f: F,
+    ) -> R {
+        let mut db = ChalkDatabase::default();
+
+        db.query_mut(ProgramText).set((), program_text);
+        db.query_mut(ProgramSolverChoice).set((), solver_choice);
+
+        f(&mut db)
+    }
+}
+
 salsa::database_storage! {
     pub struct DatabaseStorage for ChalkDatabase {
         impl query::LoweringDatabase {

src/test_util.rs

@@ -3,7 +3,6 @@
 use crate::db::ChalkDatabase;
 use crate::errors::Result;
 use crate::query::LoweringDatabase;
-use crate::query::{ProgramSolverChoice, ProgramText};
 use crate::rust_ir::lowering::LowerGoal;
 use crate::rust_ir::Program;
 use chalk_ir::Goal;
@@ -12,7 +11,6 @@ use chalk_parse;
 use chalk_solve::solve::SolverChoice;
 use diff;
 use itertools::Itertools;
-use salsa::Database;
 use std::fmt::Write;
 use std::prelude::v1::Result as StdResult;
 use std::sync::Arc;
@@ -21,27 +19,19 @@ pub fn parse_and_lower_program(
     text: &str,
     solver_choice: SolverChoice,
 ) -> StdResult<Arc<Program>, String> {
-    let mut db = ChalkDatabase::default();
-
-    db.query_mut(ProgramText)
-        .set((), Arc::new(text.to_string()));
-    db.query_mut(ProgramSolverChoice).set((), solver_choice);
-
-    db.checked_program()
+    ChalkDatabase::with_program(Arc::new(text.to_string()), solver_choice, |db| {
+        db.checked_program()
+    })
 }
 
 pub fn parse_and_lower_program_with_env(
     text: &str,
     solver_choice: SolverChoice,
 ) -> StdResult<(Arc<Program>, Arc<ProgramEnvironment>), String> {
-    let mut db = ChalkDatabase::default();
-
-    db.query_mut(ProgramText)
-        .set((), Arc::new(text.to_string()));
-    db.query_mut(ProgramSolverChoice).set((), solver_choice);
-
-    db.checked_program()
-        .and_then(|program| Ok((program, db.environment()?)))
+    ChalkDatabase::with_program(Arc::new(text.to_string()), solver_choice, |db| {
+        db.checked_program()
+            .and_then(|program| Ok((program, db.environment()?)))
+    })
 }
 
 pub fn parse_and_lower_goal(program: &Program, text: &str) -> Result<Box<Goal>> {