Subkey fix and refactor (#418)

* Fix bug and refactoring

- Fix panic in some cases
- Extract the generation as a reusable function
- Improve scoring calculation
- Add tests
- Sort the results
- Fix the case when called with a desired pattern either empty or of 1 char
- Add a 'even more' paranoiac mode
- Add benchmarks

Fix #417

* Fixed compilation for nightly

substrate/subkey/Cargo.toml

@@ -7,3 +7,6 @@ authors = ["Parity Technologies <admin@parity.io>"]
 ed25519 = { version = "*", path = "../substrate/ed25519" }
 substrate-primitives = { version = "*", path = "../substrate/primitives" }
 rand = "0.4"
+
+[features]
+bench = []

substrate/subkey/README.adoc

@@ -0,0 +1,22 @@
+= Subkey
+
+A key generation utility with vanity address support.
+
+Usage:
+
+	subkey <search string> <number of keys>
+
+Sample use:
+
+	$ subkey
+or
+	$ subkey polka
+or
+	$ subkey polka 3
+
+
+Result:
+
+	5CxS39ykKsmPetYQjTqW6aJXkSChnuvPdziA8uphuPaCyRZ1: 406ac59ccbb8358f7c95b726d3ccb039afe35e2dd62045189d1abae8d7805b8a (54%)
+	5CujMhFmChyq3AMUwMasfbqSpZYpbFfZS5UQ7zUn2d63CGBo: 5b6ac59ccbb8358f7c95b726d3ccb039afe35e2dd62045189d1abae8d7805b8a (69%)
+	5EfdN3zChABKsXT9bEg33zqPsBu4YCu1h7yoovvjtsUMqyFU: c46ac59ccbb8358f7c95b726d3ccb039afe35e2dd62045189d1abae8d7805b8a (69%)

substrate/subkey/src/main.rs

@@ -1,3 +1,6 @@
+#![cfg_attr(feature = "bench", feature(test))]
+#[cfg(feature = "bench")]
+extern crate test;
 extern crate ed25519;
 extern crate substrate_primitives;
 extern crate rand;
@@ -6,6 +9,7 @@ use rand::{OsRng, Rng};
 use std::env::args;
 use ed25519::Pair;
 use substrate_primitives::hexdisplay::HexDisplay;
+use std::cmp;
 
 fn good_waypoint(done: u64) -> u64 {
 	match done {
@@ -26,39 +30,60 @@ fn next_seed(mut seed: [u8; 32]) -> [u8; 32] {
 	return seed;
 }
 
-fn main() {
-	if args().len() != 2 {
-		println!("Usage: subkey <search string>");
-		return;
-	}
-	let desired = args().last().unwrap();
-	let score = |s: &str| {
-		for truncate in 0..desired.len() - 1 {
-			let snip_size = desired.len() - truncate;
-			let truncated = &desired[0..snip_size];
-			if let Some(pos) = s.find(truncated) {
-				return (31 - pos) + (snip_size * 32);
-			}
+/// A structure used to carry both Pair and seed.
+/// This should usually NOT been used. If unsure, use Pair.
+pub struct KeyPair {
+	pub pair: Pair,
+	pub seed: [u8; 32],
+	pub score: usize,
+}
+
+/// Calculate the score of a key based on the desired
+/// input.
+fn calculate_score(_desired: &str, key: &str) -> usize {
+	for truncate in 0.._desired.len() {
+		let snip_size = _desired.len() - truncate;
+		let truncated = &_desired[0..snip_size];
+		if let Some(pos) = key.find(truncated) {
+			let score = cmp::min(100, (51 - pos) + (snip_size * 50 / _desired.len()));
+			return score;
 		}
-		0
-	};
-	let top = 30 + (desired.len() * 32);
+	}
+	0
+}
+
+pub fn generate_key(_desired: &str, _amount: usize, paranoiac: bool) -> Result<Vec<KeyPair>, &str> {
+	println!("Generating {} keys with pattern '{}'", _amount, &_desired);
+
+	let top = 30 + (_desired.len() * 32);
 	let mut best = 0;
 	let mut seed = [0u8; 32];
 	let mut done = 0;
+	let mut res = vec![];
+
+	OsRng::new().unwrap().fill_bytes(&mut seed[..]);
+
 	loop {
+		if res.len() >= _amount { break; }
+
 		// reset to a new random seed at beginning and regularly after for paranoia.
-		if done % 100000 == 0 {
+		if paranoiac || done % 100000 == 0 {
 			OsRng::new().unwrap().fill_bytes(&mut seed[..]);
 		}
 
 		let p = Pair::from_seed(&seed);
 		let ss58 = p.public().to_ss58check();
-		let s = score(&ss58);
-		if s > best {
-			println!("{}: {} ({}% complete)", ss58, HexDisplay::from(&seed), s * 100 / top);
-			best = s;
+		let score = calculate_score(&_desired, &ss58);
+		if score > best || _desired.len() < 2 {
+			best = score;
+			let keypair = KeyPair {
+				pair: p,
+				seed: seed.clone(),
+				score: score,
+			};
+			res.push(keypair);
 			if best == top {
+				println!("best: {} == top: {}", best, top);
 				break;
 			}
 		}
@@ -66,7 +91,85 @@ fn main() {
 		done += 1;
 
 		if done % good_waypoint(done) == 0 {
-			println!("{} keys searched", done);
+			println!("Stopping after {} keys searched", done);
+			break;
 		}
 	}
+	res.sort_unstable_by(|a, b| b.score.cmp(&a.score));
+	Ok(res)
+}
+
+fn main() {
+	let desired: String = args().nth(1).unwrap_or_default();
+	let amount_of_keys: String = args().nth(2).unwrap_or_else(|| String::from("1"));
+	let amount_of_keys: usize = amount_of_keys.parse::<usize>().expect("Failed to parse number");
+
+	let keys = generate_key(&desired, amount_of_keys, true).expect("Key generation failed");
+	for key in keys {
+		println!("{} - {} ({}%)",
+			key.pair.public().to_ss58check(),
+			HexDisplay::from(&key.seed),
+			key.score);
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	#[cfg(feature = "bench")]
+	use test::Bencher;
+
+	#[test]
+	fn test_generation_no_args() {
+		assert!(generate_key("",1, false).unwrap().len() == 1);
+	}
+
+	#[test]
+	fn test_generation_with_single_char() {
+		assert!(generate_key("j", 1, false).unwrap().len() == 1);
+	}
+
+	#[test]
+	fn test_generation_with_args() {
+		assert!(generate_key("polka", 2, false).unwrap().len() == 2);
+	}
+
+	#[test]
+	fn test_score_1_char_100() {
+		let score = calculate_score("j", "5jolkadotwHY5k9GpdTgpqs9xjuNvtv8EcwCFpEeyEf3KHim");
+		assert!(score  == 100, format!("Wrong score, we found {}", score));
+	}
+
+	#[test]
+	fn test_score_100() {
+		let score = calculate_score("Polkadot", "5PolkadotwHY5k9GpdTgpqs9xjuNvtv8EcwCFpEeyEf3KHim");
+		assert!( score == 100, format!("Wrong score, we found {}", score));
+	}
+
+	#[test]
+	fn test_score_50_2() {
+		// 50% for the position + 50% for the size
+		assert!(calculate_score("Polkadot", "5PolkXXXXwHY5k9GpdTgpqs9xjuNvtv8EcwCFpEeyEf3KHim") == 75);
+	}
+
+	#[test]
+	fn test_score_0() {
+		assert!(calculate_score("Polkadot", "5GUWv4bLCchGUHJrzULXnh4JgXsMpTKRnjuXTY7Qo1Kh9uYK") == 0);
+	}
+
+	#[cfg(feature = "bench")]
+	#[bench]
+    fn bench_paranoiac(b: &mut Bencher) {
+        b.iter(|| {
+			generate_key("polka", 3, true)
+		});
+    }
+
+	#[cfg(feature = "bench")]
+    #[bench]
+    fn bench_not_paranoiac(b: &mut Bencher) {
+        b.iter(|| {
+			generate_key("polka", 3, false)
+		});
+    }
 }