From b2fe3a59a93d5c20a690417ed3c0a19c7fdbcc52 Mon Sep 17 00:00:00 2001
From: Guanqun Lu <guanqun.lu@gmail.com>
Date: Mon, 10 Sep 2018 14:49:47 +0800
Subject: [PATCH] enhance the environmental macro

---
 substrate/environmental/src/lib.rs     | 32 +++++++++++++++++++++++---
 substrate/environmental/with_std.rs    |  1 +
 substrate/environmental/without_std.rs |  1 +
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/substrate/environmental/src/lib.rs b/substrate/environmental/src/lib.rs
index f6f0e94c74522..1ec8ddd034272 100644
--- a/substrate/environmental/src/lib.rs
+++ b/substrate/environmental/src/lib.rs
@@ -217,6 +217,32 @@ macro_rules! environmental {
 			}
 		}
 	};
+	($name:ident<$traittype:ident> : trait $t:ident <$concretetype:ty>) => {
+		#[allow(non_camel_case_types, dead_code)]
+		struct $name <H: $traittype> { _private_field: $crate::imp::PhantomData<H> }
+
+		thread_local_impl!(static GLOBAL: $crate::imp::RefCell<Option<*mut ($t<$concretetype> + 'static)>>
+			= $crate::imp::RefCell::new(None));
+
+		impl<H: $traittype> $name<H> {
+			#[allow(unused_imports)]
+			pub fn using<R, F: FnOnce() -> R>(
+				protected: &mut $t<H>,
+				f: F
+			) -> R {
+				let lifetime_extended = unsafe {
+					$crate::imp::transmute::<&mut $t<H>, &mut ($t<$concretetype> + 'static)>(protected)
+				};
+				$crate::using(&GLOBAL, lifetime_extended, f)
+			}
+
+			pub fn with<R, F: for<'a> FnOnce(&'a mut ($t<$concretetype> + 'a)) -> R>(
+				f: F
+			) -> Option<R> {
+				$crate::with(&GLOBAL, |x| f(x))
+			}
+		}
+	};
 	($name:ident : trait $t:ident <>) => { environmental! { $name : trait @$t [] } };
 	($name:ident : trait $t:ident < $($args:ty),* $(,)* >) => { environmental! { $name : trait @$t [$($args,)*] } };
 	($name:ident : trait $t:ident) => { environmental! { $name : trait @$t [] } };
@@ -345,11 +371,11 @@ mod tests {
 
 		let numbers = vec![1, 2, 3];
 		let mut numbers = &numbers[..];
-		let out = foo::using(&mut numbers, || {
-			foo::with(|x| x.mul_and_add() )
+		let out = foo::<ConcretePlus>::using(&mut numbers, || {
+			foo::<ConcretePlus>::with(|x| x.mul_and_add() )
 		}).unwrap();
 
 		assert_eq!(out, 6 + 42);
-		environmental!(foo: trait Multiplier<ConcretePlus>);
+		environmental!(foo<Plus>: trait Multiplier<ConcretePlus>);
 	}
 }
diff --git a/substrate/environmental/with_std.rs b/substrate/environmental/with_std.rs
index b8b413ab42252..63ec5a71af317 100644
--- a/substrate/environmental/with_std.rs
+++ b/substrate/environmental/with_std.rs
@@ -20,6 +20,7 @@ pub mod imp {
 	pub use std::thread::LocalKey;
 	pub use std::mem::transmute;
 	pub use std::mem::replace;
+	pub use std::marker::PhantomData;
 }
 
 #[doc(hidden)]
diff --git a/substrate/environmental/without_std.rs b/substrate/environmental/without_std.rs
index 7c64ab86ddd1a..0576419d1d3f5 100644
--- a/substrate/environmental/without_std.rs
+++ b/substrate/environmental/without_std.rs
@@ -19,6 +19,7 @@ pub mod imp {
 	pub use core::cell::RefCell;
 	pub use core::mem::transmute;
 	pub use core::mem::replace;
+	pub use core::marker::PhantomData;
 
 	// This code is a simplified version of [`LocalKey`] and it's wasm32 specialization: [`statik::Key`].
 	// [`LocalKey`]: https://github.com/alexcrichton/rust/blob/98931165a23a1c2860d99759385f45d6807c8982/src/libstd/thread/local.rs#L89