Allow registering MIR-passes through compiler plugins

Author: oli-obk

mk/crates.mk

@@ -106,7 +106,7 @@ DEPS_rustc_passes := syntax rustc core rustc_front
 DEPS_rustc_mir := rustc rustc_front syntax
 DEPS_rustc_resolve := arena rustc rustc_front log syntax
 DEPS_rustc_platform_intrinsics := rustc rustc_llvm
-DEPS_rustc_plugin := rustc rustc_metadata syntax
+DEPS_rustc_plugin := rustc rustc_metadata syntax rustc_mir
 DEPS_rustc_privacy := rustc rustc_front log syntax
 DEPS_rustc_trans := arena flate getopts graphviz libc rustc rustc_back rustc_mir \
                     log syntax serialize rustc_llvm rustc_front rustc_platform_intrinsics

src/librustc/mir/mir_map.rs

@@ -10,8 +10,19 @@
 
 use util::nodemap::NodeMap;
 use mir::repr::Mir;
+use mir::transform::MirPass;
+use middle::ty;
 
 pub struct MirMap<'tcx> {
     pub map: NodeMap<Mir<'tcx>>,
 }
 
+impl<'tcx> MirMap<'tcx> {
+    pub fn run_passes(&mut self, passes: &mut [Box<MirPass>], tcx: &ty::ctxt<'tcx>) {
+        for (_, ref mut mir) in &mut self.map {
+            for pass in &mut *passes {
+                pass.run_on_mir(mir, tcx)
+            }
+        }
+    }
+}

src/librustc/session/mod.rs

@@ -13,6 +13,7 @@ use middle::cstore::CrateStore;
 use middle::dependency_format;
 use session::search_paths::PathKind;
 use util::nodemap::{NodeMap, FnvHashMap};
+use mir::transform::MirPass;
 
 use syntax::ast::{NodeId, NodeIdAssigner, Name};
 use syntax::codemap::{Span, MultiSpan};
@@ -59,6 +60,7 @@ pub struct Session {
     pub lint_store: RefCell<lint::LintStore>,
     pub lints: RefCell<NodeMap<Vec<(lint::LintId, Span, String)>>>,
     pub plugin_llvm_passes: RefCell<Vec<String>>,
+    pub plugin_mir_passes: RefCell<Vec<Box<MirPass>>>,
     pub plugin_attributes: RefCell<Vec<(String, AttributeType)>>,
     pub crate_types: RefCell<Vec<config::CrateType>>,
     pub dependency_formats: RefCell<dependency_format::Dependencies>,
@@ -475,6 +477,7 @@ pub fn build_session_(sopts: config::Options,
         lint_store: RefCell::new(lint::LintStore::new()),
         lints: RefCell::new(NodeMap()),
         plugin_llvm_passes: RefCell::new(Vec::new()),
+        plugin_mir_passes: RefCell::new(Vec::new()),
         plugin_attributes: RefCell::new(Vec::new()),
         crate_types: RefCell::new(Vec::new()),
         dependency_formats: RefCell::new(FnvHashMap()),

src/librustc_driver/driver.rs

@@ -545,7 +545,7 @@ pub fn phase_2_configure_and_expand(sess: &Session,
     });
 
     let Registry { syntax_exts, early_lint_passes, late_lint_passes, lint_groups,
-                   llvm_passes, attributes, .. } = registry;
+                   llvm_passes, attributes, mir_passes, .. } = registry;
 
     try!(sess.track_errors(|| {
         let mut ls = sess.lint_store.borrow_mut();
@@ -561,6 +561,7 @@ pub fn phase_2_configure_and_expand(sess: &Session,
         }
 
         *sess.plugin_llvm_passes.borrow_mut() = llvm_passes;
+        *sess.plugin_mir_passes.borrow_mut() = mir_passes;
         *sess.plugin_attributes.borrow_mut() = attributes.clone();
     }));
 
@@ -843,11 +844,15 @@ pub fn phase_3_run_analysis_passes<'tcx, F, R>(sess: &'tcx Session,
              "match checking",
              || middle::check_match::check_crate(tcx));
 
-        let mir_map =
+        let mut mir_map =
             time(time_passes,
                  "MIR dump",
                  || mir::mir_map::build_mir_for_crate(tcx));
 
+        time(time_passes,
+             "MIR passes",
+             || mir_map.run_passes(&mut sess.plugin_mir_passes.borrow_mut(), tcx));
+
         time(time_passes,
              "liveness checking",
              || middle::liveness::check_crate(tcx));

src/librustc_plugin/lib.rs

@@ -71,6 +71,7 @@
 extern crate rustc;
 extern crate rustc_front;
 extern crate rustc_metadata;
+extern crate rustc_mir;
 
 pub use self::registry::Registry;
 

src/librustc_plugin/registry.rs

@@ -13,6 +13,8 @@
 use rustc::lint::{EarlyLintPassObject, LateLintPassObject, LintId, Lint};
 use rustc::session::Session;
 
+use rustc::mir::transform::MirPass;
+
 use syntax::ext::base::{SyntaxExtension, NamedSyntaxExtension, NormalTT};
 use syntax::ext::base::{IdentTT, MultiModifier, MultiDecorator};
 use syntax::ext::base::{MacroExpanderFn, MacroRulesTT};
@@ -53,6 +55,9 @@ pub struct Registry<'a> {
     #[doc(hidden)]
     pub late_lint_passes: Vec<LateLintPassObject>,
 
+    #[doc(hidden)]
+    pub mir_passes: Vec<Box<MirPass>>,
+
     #[doc(hidden)]
     pub lint_groups: HashMap<&'static str, Vec<LintId>>,
 
@@ -76,6 +81,7 @@ impl<'a> Registry<'a> {
             lint_groups: HashMap::new(),
             llvm_passes: vec!(),
             attributes: vec!(),
+            mir_passes: Vec::new(),
         }
     }
 
@@ -134,6 +140,11 @@ impl<'a> Registry<'a> {
         self.lint_groups.insert(name, to.into_iter().map(|x| LintId::of(x)).collect());
     }
 
+    /// Register a MIR pass
+    pub fn register_mir_pass(&mut self, pass: Box<MirPass>) {
+        self.mir_passes.push(pass);
+    }
+
     /// Register an LLVM pass.
     ///
     /// Registration with LLVM itself is handled through static C++ objects with

src/test/auxiliary/dummy_mir_pass.rs

@@ -0,0 +1,52 @@
+// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// force-host
+
+#![feature(plugin_registrar, rustc_private)]
+#![feature(box_syntax)]
+
+#[macro_use] extern crate rustc;
+extern crate rustc_front;
+extern crate rustc_plugin;
+extern crate syntax;
+
+use rustc::mir::transform::MirPass;
+use rustc::mir::repr::{Mir, Literal};
+use rustc::mir::visit::MutVisitor;
+use rustc::middle::ty;
+use rustc::middle::const_eval::ConstVal;
+use rustc::lint::{LateContext, LintContext, LintPass, LateLintPass, LateLintPassObject, LintArray};
+use rustc_plugin::Registry;
+use rustc_front::hir;
+use syntax::attr;
+
+struct Pass;
+
+impl MirPass for Pass {
+    fn run_on_mir<'tcx>(&mut self, mir: &mut Mir<'tcx>, tcx: &ty::ctxt<'tcx>) {
+        Visitor.visit_mir(mir)
+    }
+}
+
+struct Visitor;
+
+impl<'tcx> MutVisitor<'tcx> for Visitor {
+    fn visit_literal(&mut self, literal: &mut Literal<'tcx>) {
+        if let Literal::Value { value: ConstVal::Int(ref mut i @ 11) } = *literal {
+            *i = 42;
+        }
+    }
+}
+
+#[plugin_registrar]
+pub fn plugin_registrar(reg: &mut Registry) {
+    reg.register_mir_pass(box Pass);
+}

src/test/run-pass-fulldeps/mir-pass.rs

@@ -0,0 +1,24 @@
+// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// aux-build:dummy_mir_pass.rs
+// ignore-stage1
+
+#![feature(plugin, rustc_attrs)]
+#![plugin(dummy_mir_pass)]
+
+#[rustc_mir]
+fn math() -> i32 {
+    11
+}
+
+pub fn main() {
+    assert_eq!(math(), 42);
+}