UMassCTF 2021 - Chains: Dumping the first 90000000 values for Collatz's function

collatz-dump.rs

use std::fs::File;
use std::io::prelude::*;
use std::io::BufWriter;

/* Needs lru = "0.6.5" in Cargo.toml */
use lru::LruCache;

const CACHE_SIZE: usize = 1024 * 2000;

fn collatz(mut n: u32, cache: &mut LruCache<u32, i32>) -> i32 {
	let mut i: i32 = 0;

	while n != 1 {

		if let Some(e) = cache.get(&n) {
			return *e + i;
		}

		if n & 1 == 0 {
			n = n >> 1;
		} else {
			n = n * 3 + 1;
		}

		i += 1;
	}

	return i;
}

fn main() -> std::io::Result<()> {
	let mut cache = LruCache::new(CACHE_SIZE);
	let file = File::create("collatz.out")?;
	let mut writer = BufWriter::new(file);
	let mut res: i32;

	let desired_outputs = [105, 110, 113, 116, 136, 160, 181, 185, 199, 202, 211, 247, 266, 271, 273, 303, 318, 341, 352] ;

	for i in 1..900_000_001 {

		/* Show progress */
		if (i % 100_000) == 0 {
			print!("{}\r", i);
		}

		/* Run collatz and update cache */
		res = collatz(i, &mut cache);
		cache.put(i, res);

		/* Update file if collatz's output is found in `desired_outputs` */
		if let Ok(_) = desired_outputs.binary_search(&res) {
			writer.write(format!("{} {}\n", i, res).as_bytes())?;
		}
	}

	println!("Cache entries used: {}/{}", cache.len(), CACHE_SIZE);

	Ok(())
}