Missing nubmers

main.rs

use std::fs;
use std::io;
use std::io::BufRead;
use std::collections::HashMap;

type Prefixes = Vec<u64>;
type Counters = Vec<u64>;

struct ArgsContext {
    filename: String,
    mask_size: u8
}

struct AppContextAfterFirstPass {
    args_context: ArgsContext,
    range_size: u8,
}

fn create_bitmask(size: u8, offset: u8) -> u64 {
    let unshifted = 2u64.pow(size as _) - 1;
    unshifted << offset
}

impl ArgsContext {
    fn from_env_args() -> ArgsContext {
        let mut args = std::env::args();
        args.next().expect("First value of args is the name of the executable");

        let filename = args.next().expect("Need at least one argument");
        let mask_size = if let Some(mask_size) = args.next() {
            mask_size.parse().unwrap()
        } else {
            1 // default value
        };

        assert!(args.next().is_none(), "Too many arguments");

        ArgsContext {
            filename: filename,
            mask_size: mask_size,
        }
    }

    fn read_numbers(&self) -> Box<Iterator<Item=u64>> {
        let file = fs::File::open(&self.filename).unwrap();
        let reader = io::BufReader::new(file);
        Box::new(reader.lines().map(|line| line.unwrap().parse().unwrap()))
    }

    fn create_counters_vec(&self) -> Counters {
        vec![0; 2usize.pow(self.mask_size as _)]
    }

    fn first_pass(self) -> (AppContextAfterFirstPass, Prefixes) {
        let mut highest_number = 0;
        let mut counts = self.create_counters_vec();
        let bitmask = create_bitmask(self.mask_size, 0);
        for number in self.read_numbers() {
            if highest_number < number {
                highest_number = number;
            }

            let first_byte = number & bitmask;
            counts[first_byte as usize] += 1;
        }

        let range_size = {
            let mut mask = std::u64::MAX;
            let mut range_size = 0;
            while mask & highest_number != 0 {
                mask = mask << 1;
                range_size += 1;
            }
            range_size
        };

        let expected_count = 2u64.pow((range_size - self.mask_size) as u32);
        let prefixes = counts.iter().enumerate().filter_map(|(idx, count)| {
            if *count < expected_count {
                Some(idx as u64)
            } else {
                None
            }
        }).collect::<Vec<_>>();

        let app_context = AppContextAfterFirstPass {
            args_context: self,
            range_size: range_size,
        };

        (app_context, prefixes)
    }
}

impl AppContextAfterFirstPass {
    fn make_pass(&self, offset: u8, prefixes: Prefixes) -> Prefixes {
        let mut prefixes_counters = HashMap::new();
        for prefix in prefixes {
            prefixes_counters.insert(prefix, self.args_context.create_counters_vec());
        }
        let bitmask = create_bitmask(self.args_context.mask_size, offset);
        let prefix_bitmask = create_bitmask(offset, 0);
        let _ = self.range_size;
        let _ = offset;
        for number in self.args_context.read_numbers() {
            let prefix = number & prefix_bitmask;
            if let Some(counts) = prefixes_counters.get_mut(&prefix) {
                let relevant_byte = (number & bitmask) >> offset;
                counts[relevant_byte as usize] += 1;
            }
        }

        let expected_count = {
            let remaining_bits = self.range_size as isize - (self.args_context.mask_size + offset) as isize;
            if 0 < remaining_bits {
                2u64.pow(remaining_bits as u32)
            } else {
                1 // no remaining bits - only one option
            }
        };

        let prefixes = prefixes_counters.iter().flat_map(|(prefix, counts)| {
            counts.iter().enumerate().filter_map(move |(idx, count)| {
                if *count < expected_count {
                    Some(prefix + ((idx as u64) << offset))
                } else {
                    None
                }
            })
        }).collect::<Vec<_>>();
        prefixes
    }
}

fn main() {
    let args_context = ArgsContext::from_env_args();
    let (app_context, mut prefixes) = args_context.first_pass();

    // inclusive ranges and step_by() are unstable, so I'll just iterate manually
    let mut offset = app_context.args_context.mask_size;
    while offset < app_context.range_size {
        prefixes = app_context.make_pass(offset, prefixes);
        offset += app_context.args_context.mask_size;
    }

    prefixes.sort();
    for number in prefixes {
        println!("{}", number);
    }
}

test_missing_numbers.py

from random import sample, shuffle

from plumbum import local

import pytest

@pytest.mark.parametrize('range_size', [16])
@pytest.mark.parametrize('missin_numbers', [4])
@pytest.mark.parametrize('mask_size', [8])
def test_missing_numbers(range_size, missin_numbers, mask_size):
    number_of_numbers = 2 ** range_size
    missin_numbers = set(sample(range(number_of_numbers), missin_numbers))
    numbers = list(range(number_of_numbers))
    shuffle(numbers)
    with local.path('numbers.txt').open('w') as f:
        for number in numbers:
            if number not in missin_numbers:
                f.write('%s\n' % number)

    result = {int(line) for line in local['cargo']('run', '-q', 'numbers.txt', mask_size).splitlines()}

    assert missin_numbers == result