jonhoo / roget Goto Github PK

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum Correctness {
	Wrong = 0,
	Misplaced = 1,
	Correct = 2,
}


impl Correctness {
	
	
	// Performs equality check
	// Example:
	// Input:
	//     answer = 01101 00101 01000 00010 01111
	//     guess  = 01101 01111 01011 00010 01010
	//     result = 11111 00000 00000 11111 00000
	fn check_for_eq(answer: u32, guess: u32) -> u32 {
		let mut result = !(answer ^ guess);
		result &= (result >> 1) & (result >> 2) & (result >> 3) & (result >> 4) & 0b00001_00001_00001_00001_00001u32;
		result |= (result << 1) | (result << 2) | (result << 3) | (result << 4);
		result
	}
	
	pub fn compute(answer: &str, guess: &str) -> [Correctness; 5] {
		const TO_NUMBER: u8 = 'a' as u8 - 1;
		
		
		let answer = answer.as_bytes();
		let guess = guess.as_bytes();
		
		// Assuming answer and guess are lowercase and use only alphabetic characters
		// we can rappresent them using 5bits per letter (in this case 'a' = 1)
		let mut answer = (answer[0] - TO_NUMBER) as u32 | 
			(((answer[1] - TO_NUMBER) as u32) << 5) | 
			(((answer[2] - TO_NUMBER) as u32) << 10) | 
			(((answer[3] - TO_NUMBER) as u32) << 15) | 
			(((answer[4] - TO_NUMBER) as u32) << 20);
		
		let mut guess = (guess[0] - TO_NUMBER) as u32 | 
			(((guess[1] - TO_NUMBER) as u32) << 5) | 
			(((guess[2] - TO_NUMBER) as u32) << 10) | 
			(((guess[3] - TO_NUMBER) as u32) << 15) | 
			(((guess[4] - TO_NUMBER) as u32) << 20);
		
		let green = Self::check_for_eq(answer, guess);
		
		// Removing green letters:
		// setting used letters from guess to 11111
		guess |= green;
		// setting used letters from answer to 00000
		answer &= !green;
		
		// To detect yellow letters we can shift guess 4 times and check the equality
		let mut yellow = 0u32;
		for _ in 0..4 {
			guess = (guess >> 5) | (guess << 20);
			let n = Self::check_for_eq(answer, guess);
			// setting used letters from guess to 11111
			guess |= n;
			// setting used letters from answer to 00000
			answer &= !n;
			// yellow must be moved with guess to avoid overwrites
			yellow = n | (yellow >> 5) | (yellow << 20);
		}
		// moving yellow to its original position
		yellow = (yellow >> 5) | (yellow << 20);
		
		let result = (green & 0b00010_00010_00010_00010_00010u32) | yellow & 0b00001_00001_00001_00001_00001;
		
		// Safety: ((result >> (n * 5)) & 0b11111) as u8, where: 0 <= n < 5
		//             - Can only give 0, 1 or 2 as result
		//             - Enum is #[repr(u8)]
		[
			unsafe { std::mem::transmute((result & 0b11111) as u8) },
			unsafe { std::mem::transmute(((result >> 5) & 0b11111) as u8) },
			unsafe { std::mem::transmute(((result >> 10) & 0b11111) as u8) },
			unsafe { std::mem::transmute(((result >> 15) & 0b11111) as u8) },
			unsafe { std::mem::transmute(((result >> 20) & 0b11111) as u8) },
		]
	}
	
}

fn main() {
	let b0 = benchmark(|| Correctness::compute("abcde", "fghij"));
	let b1 = benchmark(|| jonhoo::Correctness::compute("abcde", "fghij"));
	println!("{} ns/op", b0);
	println!("{} ns/op", b1);
	println!("b1/b0 = {}", b1 / b0);
}

// ns/op
fn benchmark<T, F: Fn() -> T>(f: F) -> f64 {
	const ITERATIONS: usize = 1000000000;
	let now = Instant::now();
	for _ in 0..ITERATIONS {
		execute_it_pls(f());
	}
	let el = now.elapsed().as_millis();
	el as f64 / (ITERATIONS as f64 / 1000000.0)
}

fn execute_it_pls<T>(dummy: T) -> T {
	let ptr = (&dummy) as *const _;
	unsafe { asm!("/* {0} */", in(reg) ptr) }
	dummy
}