From 1444cabc081f67ef30fbb86d6088754ec76163ac Mon Sep 17 00:00:00 2001
From: Oliver Schneider <git-spam-no-reply9815368754983@oli-obk.de>
Date: Tue, 5 Jul 2016 13:04:46 +0200
Subject: [PATCH] make the memory limit configurable

---
 src/bin/miri.rs        | 34 ++++++++++++++++++++++++++++++++--
 src/interpreter/mod.rs |  7 ++++---
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/src/bin/miri.rs b/src/bin/miri.rs
index 4340729d7e02..087da20ac8be 100644
--- a/src/bin/miri.rs
+++ b/src/bin/miri.rs
@@ -4,6 +4,7 @@ extern crate getopts;
 extern crate miri;
 extern crate rustc;
 extern crate rustc_driver;
+extern crate rustc_plugin;
 extern crate env_logger;
 extern crate log_settings;
 extern crate syntax;
@@ -13,6 +14,7 @@ use miri::{eval_main, run_mir_passes};
 use rustc::session::Session;
 use rustc::mir::mir_map::MirMap;
 use rustc_driver::{driver, CompilerCalls, Compilation};
+use syntax::ast::MetaItemKind;
 
 struct MiriCompilerCalls;
 
@@ -23,7 +25,9 @@ impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
         _: &getopts::Matches
     ) -> driver::CompileController<'a> {
         let mut control = driver::CompileController::basic();
-
+        control.after_hir_lowering.callback = Box::new(|state| {
+            state.session.plugin_attributes.borrow_mut().push(("miri".to_owned(), syntax::feature_gate::AttributeType::Whitelisted));
+        });
         control.after_analysis.stop = Compilation::Stop;
         control.after_analysis.callback = Box::new(|state| {
             state.session.abort_if_errors();
@@ -33,9 +37,35 @@ impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
             let (node_id, _) = state.session.entry_fn.borrow()
                 .expect("no main or start function found");
 
+            let krate = state.hir_crate.as_ref().unwrap();
+            let mut memory_size = 100*1024*1024; // 100MB
+            fn extract_str(lit: &syntax::ast::Lit) -> syntax::parse::token::InternedString {
+                match lit.node {
+                    syntax::ast::LitKind::Str(ref s, _) => s.clone(),
+                    _ => panic!("attribute values need to be strings"),
+                }
+            }
+            for attr in krate.attrs.iter() {
+                match attr.node.value.node {
+                    MetaItemKind::List(ref name, _) if name != "miri" => {}
+                    MetaItemKind::List(_, ref items) => for item in items {
+                        match item.node {
+                            MetaItemKind::NameValue(ref name, ref value) => {
+                                match &**name {
+                                    "memory_size" => memory_size = extract_str(value).parse::<u64>().expect("not a number"),
+                                    _ => state.session.span_err(item.span, "unknown miri attribute"),
+                                }
+                            }
+                            _ => state.session.span_err(item.span, "miri attributes need to be of key = value kind"),
+                        }
+                    },
+                    _ => {},
+                }
+            }
+
             let mut mir_map = MirMap { map: mir_map.map.clone() };
             run_mir_passes(tcx, &mut mir_map);
-            eval_main(tcx, &mir_map, node_id);
+            eval_main(tcx, &mir_map, node_id, memory_size);
 
             state.session.abort_if_errors();
         });
diff --git a/src/interpreter/mod.rs b/src/interpreter/mod.rs
index 66d147e2d8b4..ea1f82f5f62a 100644
--- a/src/interpreter/mod.rs
+++ b/src/interpreter/mod.rs
@@ -133,12 +133,12 @@ enum ConstantKind {
 }
 
 impl<'a, 'tcx> EvalContext<'a, 'tcx> {
-    pub fn new(tcx: TyCtxt<'a, 'tcx, 'tcx>, mir_map: &'a MirMap<'tcx>) -> Self {
+    pub fn new(tcx: TyCtxt<'a, 'tcx, 'tcx>, mir_map: &'a MirMap<'tcx>, memory_size: u64) -> Self {
         EvalContext {
             tcx: tcx,
             mir_map: mir_map,
             mir_cache: RefCell::new(DefIdMap()),
-            memory: Memory::new(&tcx.data_layout, 100*1024*1024 /* 100MB */),
+            memory: Memory::new(&tcx.data_layout, memory_size),
             statics: HashMap::new(),
             stack: Vec::new(),
         }
@@ -928,10 +928,11 @@ pub fn eval_main<'a, 'tcx: 'a>(
     tcx: TyCtxt<'a, 'tcx, 'tcx>,
     mir_map: &'a MirMap<'tcx>,
     node_id: ast::NodeId,
+    memory_size: u64,
 ) {
     let mir = mir_map.map.get(&node_id).expect("no mir for main function");
     let def_id = tcx.map.local_def_id(node_id);
-    let mut ecx = EvalContext::new(tcx, mir_map);
+    let mut ecx = EvalContext::new(tcx, mir_map, memory_size);
     let substs = tcx.mk_substs(subst::Substs::empty());
     let return_ptr = ecx.alloc_ret_ptr(mir.return_ty, substs)
                         .expect("should at least be able to allocate space for the main function's return value")