main.rs

#![feature(rustc_private)]
#![feature(decl_macro)]

extern crate getopts;
extern crate rustc;
extern crate rustc_data_structures;
extern crate rustc_driver;
extern crate rustc_errors;
extern crate rustc_mir;
extern crate syntax;

use rustc::hir::def_id::DefId;
use rustc::mir::Mir;
use rustc::session::Session;
use rustc::session::config::{self, ErrorOutputType, Input};
use rustc::ty::TyCtxt;
use rustc_driver::{Compilation, CompilerCalls,
                   RustcDefaultCalls};
use rustc_driver::driver::CompileController;
use rustc_mir::transform::{add_call_guards,
                           add_validation, copy_prop,
                           deaggregator, dump_mir,
                           elaborate_drops, erase_regions,
                           generator, inline, instcombine,
                           no_landing_pads, run_passes,
                           simplify, simplify_branches,
                           MirPass, MirSource,
                           lower_128bit};
use std::path::PathBuf;
use syntax::ast;

mod pass;
use pass::MyMirPass;

struct AdventMirCompilerCalls {
    default: RustcDefaultCalls,
}

impl Default for AdventMirCompilerCalls {
    fn default() -> Self {
        Self {
            default: RustcDefaultCalls,
        }
    }
}

impl<'a> CompilerCalls<'a> for AdventMirCompilerCalls {
    fn early_callback(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        descriptions: &rustc_errors::registry::Registry,
        output: ErrorOutputType,
    ) -> Compilation {
        self.default.early_callback(
            matches,
            sopts,
            cfg,
            descriptions,
            output,
        )
    }

    fn no_input(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
        descriptions: &rustc_errors::registry::Registry,
    ) -> Option<(Input, Option<PathBuf>)> {
        self.default.no_input(
            matches,
            sopts,
            cfg,
            odir,
            ofile,
            descriptions,
        )
    }

    fn late_callback(
        &mut self,
        matches: &getopts::Matches,
        sess: &Session,
        crate_stores: &rustc::middle::cstore::CrateStore,
        input: &Input,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
    ) -> Compilation {
        self.default.late_callback(
            matches,
            sess,
            crate_stores,
            input,
            odir,
            ofile,
        )
    }

    fn build_controller(
        &mut self,
        sess: &Session,
        matches: &getopts::Matches,
    ) -> CompileController<'a> {
        let mut controller =
            self.default.build_controller(sess, matches);
        controller.provide = Box::new(|providers| {
            providers.optimized_mir = new_optimized_mir;
        });
        controller
    }
}

fn new_optimized_mir<'a, 'tcx>(
    tcx: TyCtxt<'a, 'tcx, 'tcx>,
    def_id: DefId,
) -> &'tcx Mir<'tcx> {
    let _ = tcx.mir_borrowck(def_id);
    let _ = tcx.borrowck(def_id);

    let mut mir = tcx.mir_validated(def_id).steal();
    run_passes![tcx, mir, def_id, 2;
        no_landing_pads::NoLandingPads,
        simplify_branches::SimplifyBranches::new("initial"),

        // These next passes must be executed together
        add_call_guards::CriticalCallEdges,
        elaborate_drops::ElaborateDrops,
        no_landing_pads::NoLandingPads,
        // AddValidation needs to run after ElaborateDrops and before EraseRegions, and it needs
        // an AllCallEdges pass right before it.
        add_call_guards::AllCallEdges,
        add_validation::AddValidation,
        simplify::SimplifyCfg::new("elaborate-drops"),
        // No lifetime analysis based on borrowing can be done from here on out.

        // From here on out, regions are gone.
        erase_regions::EraseRegions,

        lower_128bit::Lower128Bit,

        MyMirPass,

        // Optimizations begin.
        inline::Inline,
        instcombine::InstCombine,
        deaggregator::Deaggregator,
        copy_prop::CopyPropagation,
        simplify::SimplifyLocals,

        generator::StateTransform,
        add_call_guards::CriticalCallEdges,
        dump_mir::Marker("PreTrans"),
    ];

    tcx.alloc_mir(mir)
}

fn main() {
    let sysroot = std::process::Command::new("rustc")
        .args(&["--print", "sysroot"])
        .output()
        .unwrap()
        .stdout;
    let sysroot = std::str::from_utf8(&sysroot)
        .unwrap()
        .trim_right()
        .to_string();

    rustc_driver::in_rustc_thread(|| {
        let args: Vec<_> = std::env::args()
            .chain(vec!["--sysroot".to_string(), sysroot])
            .collect();

        rustc_driver::run_compiler(
            &args,
            &mut AdventMirCompilerCalls::default(),
            None,
            None,
        ).0
            .expect("run_compiler failed");
    }).expect("in_rustc_thread failed");
}

pass.rs

use rustc::middle::const_val::ConstVal;
use rustc::mir::{BinOp, Literal, Local, LocalDecl,
                 Location, Mir, Operand, Rvalue};
use rustc::mir::visit::{MutVisitor, Visitor};
use rustc::ty::{Const, TyCtxt};
use rustc_data_structures::indexed_vec::IndexVec;
use rustc_mir::transform::{MirPass, MirSource};
use std::borrow::Cow;

pub struct MyMirPass;

impl MirPass for MyMirPass {
    fn name<'a>(&'a self) -> Cow<'a, str> {
        "MyMirPass".into()
    }

    fn run_pass<'a, 'tcx>(
        &self,
        tcx: TyCtxt<'a, 'tcx, 'tcx>,
        src: MirSource,
        mir: &mut Mir<'tcx>,
    ) {
        SimplifyOp::visit_mir(
            &mut SimplifyOp {
                tcx,
                local_decls: mir.local_decls.clone(),
            },
            mir,
        );
    }
}

struct SimplifyOp<'tcx: 'a, 'a> {
    local_decls: IndexVec<Local, LocalDecl<'tcx>>,
    tcx: TyCtxt<'a, 'tcx, 'tcx>,
}

impl<'tcx: 'a, 'a> MutVisitor<'tcx>
    for SimplifyOp<'tcx, 'a> {
    fn visit_rvalue(
        &mut self,
        rvalue: &mut Rvalue<'tcx>,
        location: Location,
    ) {
        let mut replace = None;
        match *rvalue {
            Rvalue::BinaryOp(
                BinOp::Add,
                ref l,
                Operand::Constant(ref c),
            ) if l.ty(&self.local_decls, self.tcx)
                .is_primitive_ty() =>
            {
                if let Literal::Value {
                    value:
                        &Const {
                            ty: _,
                            val: ConstVal::Integral(ref val),
                        },
                } = c.literal
                {
                    if val.to_u128_unchecked() == 0 {
                        replace = Some(l.clone());
                    }
                }
            }
            _ => {}
        }

        if let Some(replace) = replace {
            *rvalue = Rvalue::Use(replace);
        }

        self.super_rvalue(rvalue, location);
    }
}