From abb61d99ad8cc3c65a576d2cea4f6d9c52887880 Mon Sep 17 00:00:00 2001
From: J Bailey <jj2baile@uwaterloo.ca>
Date: Tue, 28 Apr 2015 19:24:16 -0400
Subject: [PATCH] Extend the nullable pointer optimization to captured vars of
 closures

---
 src/librustc_trans/trans/adt.rs               | 16 +++++++
 .../run-pass/nullable-pointer-opt-closures.rs | 43 +++++++++++++++++++
 2 files changed, 59 insertions(+)
 create mode 100644 src/test/run-pass/nullable-pointer-opt-closures.rs

diff --git a/src/librustc_trans/trans/adt.rs b/src/librustc_trans/trans/adt.rs
index cecf35001c4..001de615fb1 100644
--- a/src/librustc_trans/trans/adt.rs
+++ b/src/librustc_trans/trans/adt.rs
@@ -440,6 +440,22 @@ fn find_discr_field_candidate<'tcx>(tcx: &ty::ctxt<'tcx>,
             None
         },
 
+        // Perhaps one of the upvars of this struct is non-zero
+        // Let's recurse and find out!
+        ty::ty_closure(def_id, substs) => {
+            let typer = NormalizingClosureTyper::new(tcx);
+            let upvars = typer.closure_upvars(def_id, substs).unwrap();
+            let upvar_types = upvars.iter().map(|u| u.ty).collect::<Vec<_>>();
+
+            for (j, &ty) in upvar_types.iter().enumerate() {
+                if let Some(mut fpath) = find_discr_field_candidate(tcx, ty, path.clone()) {
+                    fpath.push(j);
+                    return Some(fpath);
+                }
+            }
+            None
+        },
+
         // Can we use one of the fields in this tuple?
         ty::ty_tup(ref tys) => {
             for (j, &ty) in tys.iter().enumerate() {
diff --git a/src/test/run-pass/nullable-pointer-opt-closures.rs b/src/test/run-pass/nullable-pointer-opt-closures.rs
new file mode 100644
index 00000000000..ac5634e6cdd
--- /dev/null
+++ b/src/test/run-pass/nullable-pointer-opt-closures.rs
@@ -0,0 +1,43 @@
+// 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.
+
+
+use std::mem;
+
+pub fn main() {
+    // By Ref Capture
+    let a = 10i32;
+    let b = Some(|| println!("{}", a));
+    // When we capture by reference we can use any of the
+    // captures as the discriminant since they're all
+    // behind a pointer.
+    assert_eq!(mem::size_of_val(&b), mem::size_of::<usize>());
+
+    // By Value Capture
+    let a = Box::new(12i32);
+    let b = Some(move || println!("{}", a));
+    // We captured `a` by value and since it's a `Box` we can use it
+    // as the discriminant.
+    assert_eq!(mem::size_of_val(&b), mem::size_of::<Box<i32>>());
+
+    // By Value Capture - Transitive case
+    let a = "Hello".to_string(); // String -> Vec -> Unique -> NonZero
+    let b = Some(move || println!("{}", a));
+    // We captured `a` by value and since down the chain it contains
+    // a `NonZero` field, we can use it as the discriminant.
+    assert_eq!(mem::size_of_val(&b), mem::size_of::<String>());
+
+    // By Value - No Optimization
+    let a = 14i32;
+    let b = Some(move || println!("{}", a));
+    // We captured `a` by value but we can't use it as the discriminant
+    // thus we end up with an extra field for the discriminant
+    assert_eq!(mem::size_of_val(&b), mem::size_of::<(i32, i32)>());
+}