diff --git a/tools/test-drive/src/main.rs b/tools/test-drive/src/main.rs
index b1dd168..1195b51 100644
--- a/tools/test-drive/src/main.rs
+++ b/tools/test-drive/src/main.rs
@@ -6,12 +6,14 @@
 
 mod sanity_checks;
 
+extern crate rustc_driver;
+extern crate rustc_interface;
 extern crate rustc_middle;
 extern crate rustc_smir;
 extern crate stable_mir;
 
 use rustc_middle::ty::TyCtxt;
-use rustc_smir::rustc_internal;
+use rustc_smir::{run, rustc_internal};
 use stable_mir::CompilerError;
 use std::ops::ControlFlow;
 use std::panic::{catch_unwind, AssertUnwindSafe};
@@ -39,7 +41,7 @@ type TestResult = Result<(), String>;
 fn main() -> ExitCode {
     let args = std::env::args();
     let (smir_args, rustc_args): (Vec<String>, _) = args.partition(|arg| arg.starts_with("--smir"));
-    let callback = if smir_args.contains(&CHECK_ARG.to_string()) {
+    let result = if smir_args.contains(&CHECK_ARG.to_string()) {
         VERBOSE.store(
             smir_args.contains(&VERBOSE_ARG.to_string()),
             Ordering::Relaxed,
@@ -48,11 +50,10 @@ fn main() -> ExitCode {
             smir_args.contains(&FIXME_ARG.to_string()),
             Ordering::Relaxed,
         );
-        test_stable_mir
+        run!(rustc_args, tcx, test_stable_mir(tcx))
     } else {
-        |_: TyCtxt| ControlFlow::<()>::Continue(())
+        run!(rustc_args, ControlFlow::<()>::Continue(()))
     };
-    let result = rustc_internal::StableMir::new(rustc_args, callback).run();
     if result.is_ok() || matches!(result, Err(CompilerError::Skipped)) {
         ExitCode::SUCCESS
     } else {
diff --git a/tools/test-drive/src/sanity_checks.rs b/tools/test-drive/src/sanity_checks.rs
index f5f12a8..cbc0946 100644
--- a/tools/test-drive/src/sanity_checks.rs
+++ b/tools/test-drive/src/sanity_checks.rs
@@ -98,8 +98,8 @@ pub fn test_crates() -> TestResult {
 
 /// Visit all local types, statements and terminator to ensure nothing crashes.
 fn check_body(body: stable_mir::mir::Body) {
-    for bb in body.blocks {
-        for stable_mir::mir::Statement { kind, .. } in bb.statements {
+    for bb in &body.blocks {
+        for stable_mir::mir::Statement { kind, .. } in &bb.statements {
             black_box(matches!(kind, stable_mir::mir::StatementKind::Assign(..)));
         }
         black_box(matches!(
@@ -108,7 +108,7 @@ fn check_body(body: stable_mir::mir::Body) {
         ));
     }
 
-    for local in body.locals {
+    for local in body.locals() {
         black_box(matches!(local.ty.kind(), stable_mir::ty::TyKind::Alias(..)));
     }
 }