day14.rs

use std::io;
use std::fs;
use std::str::FromStr;
use std::convert::TryInto;
use std::collections::HashMap;

#[derive(Debug, Clone)]
enum Instruction {
    Mask(u64, u64, u64),
    Set(u64, u64),
}

impl TryInto<Instruction> for String {
    type Error = io::Error;
    
    fn try_into(self) -> io::Result<Instruction> {
        if self.starts_with("mask") {
            let mut and1_mask = 0u64;
            let mut and2_mask = 0u64;
            let mut or_mask = 0u64;

            for c in self.as_bytes() {
                if c == &b'0' || c == &b'1' || c == &b'X' {
                    and1_mask = and1_mask << 1;
                    and2_mask = and2_mask << 1;
                    or_mask = or_mask << 1;

                    if c == &b'1' {
                        or_mask = or_mask | 1
                    }

                    if c == &b'X' {
                        and1_mask = and1_mask | 1
                    } else {
                        and2_mask = and2_mask | 1;
                    }
                }
            }

            Ok(Instruction::Mask(and1_mask, and2_mask, or_mask))
        } else {
            let mut parts = self.trim_start_matches("mem[").split("] = ");
            parts.next().and_then(|addr_str|
                u64::from_str(addr_str).ok().and_then(|addr|
                    parts.next().and_then(|value_str|
                        u64::from_str(value_str).ok().map(|value| Instruction::Set(addr, value))
                    )                        
                )
            ).ok_or(io::Error::new(io::ErrorKind::Other, "Can not parse mem set"))
        }
    }
}

fn evaluate_with_value_decoder<T, K>(source: T, result: &mut HashMap<u64, u64>) -> io::Result<()> 
    where T: IntoIterator<Item=K>,
          K: TryInto<Instruction, Error = io::Error> {

    let mut and_mask = 0u64;
    
    let mut or_mask = 0u64;
    
    for line in source {
        match line.try_into()? {
            Instruction::Mask(and1, _, or) => {
                and_mask = and1;
                or_mask = or;
            },
            Instruction::Set(address, value) => {
                result.insert(address, (value & and_mask) | or_mask);
            },
        }
    }

    Ok(())
}

struct Submasks {
    submask: u64,
    mask: u64,
    empty_seen: bool,
}

impl Iterator for Submasks {
    type Item = u64;

    fn next(&mut self) -> Option<Self::Item> {
        match self.submask {
            0 => match self.empty_seen {
                false => {
                    self.empty_seen = true;

                    Some(0)
                },
                _ => None
            },
            _ => {
                let value = self.submask;

                self.submask = (self.submask - 1) & self.mask;

                Some(value)
            }
        }
    }
}

fn submasks(mask: u64) -> Submasks {
    Submasks {
        submask: mask,
        mask: mask,
        empty_seen: false,
    }
}

fn evaluate_with_address_decoder<T, K>(source: T, result: &mut HashMap<u64, u64>) -> io::Result<()> 
    where T: IntoIterator<Item = K>,
          K: TryInto<Instruction, Error = io::Error> {

    let mut and1_mask = 0u64;

    let mut and2_mask = 0u64;
    
    let mut or_mask = 0u64;
    
    for line in source {
        match line.try_into()? {
            Instruction::Mask(and1, and2, or) => {
                and1_mask = and1;
                and2_mask = and2;
                or_mask = or;
            },
            Instruction::Set(address, value) => {
                let stable_address = (address | or_mask) & and2_mask;
                
                for floating_mask in submasks(and1_mask) {
                    result.insert(stable_address | floating_mask, value);
                }
            },
        }
    }

    Ok(())
}

#[test]
fn test_evaluate_with_value_decoder() {
    let mut result = HashMap::new();

    evaluate_with_value_decoder(
        vec![
            "mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X".to_string(),
            "mem[8] = 11".to_string(),
            "mem[7] = 101".to_string(),
            "mem[8] = 0".to_string(),
        ],
        &mut result
    ).unwrap();

    assert_eq!(Some(&101), result.get(&7));
    assert_eq!(Some(&64), result.get(&8));
}

#[test]
fn test_evaluate_with_address_decoder() {
    let mut result = HashMap::new();

    evaluate_with_address_decoder(
        vec![
            "mask = 000000000000000000000000000000X1001X".to_string(),
            "mem[42] = 100".to_string(),
            "mask = 00000000000000000000000000000000X0XX".to_string(),
            "mem[26] = 1".to_string(),
        ],
        &mut result
    ).unwrap();

    assert_eq!(208, result.values().sum::<u64>());
}

#[allow(dead_code)]
pub fn solve_p1(input: &str) -> io::Result<u64> {
    let mut result = HashMap::new();
    evaluate_with_value_decoder(fs::read_to_string(input)?.lines().map(|s| s.to_string()), &mut result)?;
    
    Ok(result.values().sum::<u64>())
}

#[allow(dead_code)]
pub fn solve_p2(input: &str) -> io::Result<u64> {
    let mut result = HashMap::new();
    evaluate_with_address_decoder(fs::read_to_string(input)?.lines().map(|s| s.to_string()), &mut result)?;
    Ok(result.values().sum::<u64>())
}