advent/2021/src/lib.rs

use std::io::Read;
use std::collections::VecDeque;
use std::str::FromStr;
use std::fs::File;
use std::fmt::{Display, Formatter};
use std::ops::{Index, IndexMut};

use color_eyre::{eyre};
use eyre::eyre;


pub fn load(filename: &str) -> eyre::Result<String> {
    let mut file = File::open(filename)?;
    let mut data = String::new();
    file.read_to_string(&mut data)?;
    Ok(data)
}


pub struct LineParser<'a, T: FromStr> {
    it: std::str::Lines<'a>,
    convert: fn(&str) -> Result<T, <T as FromStr>::Err>
}


impl<'a, T: FromStr> Iterator for LineParser<'a, T> {
    type Item = Result<T, <T as FromStr>::Err>;

    fn next(&mut self) -> Option<Self::Item> {
        self.it.next().map(|s|
            (self.convert)(s)
        )
    }
}


// define a trait so that we can attach it to String
pub trait ParseLines {
    fn parse_lines<T: FromStr>(&self) -> LineParser<'_, T>;
}


impl ParseLines for str {
    fn parse_lines<T: FromStr>(&self) -> LineParser<'_, T> {
        LineParser {
            it: self.lines(),
            convert: str::parse::<T>
        }
    }
}


impl ParseLines for String {
    fn parse_lines<T: FromStr>(&self) -> LineParser<'_, T> {
        self[..].parse_lines::<T>()
    }
}


pub trait IterExt: Iterator {
    fn take_pair(&mut self) -> eyre::Result<(Self::Item, Self::Item)> {
        let a = match self.next() {
            Some(v) => v,
            None => return Err(eyre!("Not enough values to unpack")),
        };
        let b = match self.next() {
            Some(v) => v,
            None => return Err(eyre!("Not enough values to unpack")),
        };
        Ok((a, b))
    }

    fn max_n(&mut self, n: usize) -> eyre::Result<VecDeque<Self::Item>>
        where Self: Sized, Self::Item: Ord,
    {
        let mut items = VecDeque::with_capacity(n);
        while let Some(item) = self.next() {
            for i in 0..=items.len() {
                if i < items.len() && item > items[i] {
                    if items.len() == n {
                        items.pop_back();
                    }
                    items.insert(i, item);
                    break;
                }
                else if i == items.len() && i < n {
                    items.push_back(item);
                    break;
                }
            }
        }

        if items.len() < n {
            Err(eyre!("Not enough values to take the maximum {}", n))
        }
        else {
            Ok(items)
        }
    }
}

impl<I: Iterator> IterExt for I {}


// 2d container, functions kind of like a "vec of vecs" but all data is stored in a single vec for memory contiguity
pub struct Vec2<T> {
    data: Vec<T>,
    columns: usize,
}


pub struct NeighborCoords {
    num_rows: usize,
    num_cols: usize,
    start_row: usize,
    start_col: usize,
    i: usize,
}

impl Iterator for NeighborCoords {
    type Item = (usize, usize);
    fn next(&mut self) -> Option<Self::Item> {
        if self.i > 8 {
            return None;
        }

        let (rel_row, rel_col) = (self.i / 3, self.i % 3);
        self.i += 1;
        if rel_row == 1 && rel_col == 1 {
            return self.next();
        }
        
        let row_offset = self.start_row + rel_row;
        let col_offset = self.start_col + rel_col;
        // the "neighbor window" is off by 1, so we make the comparisons off by 1 as well
        if row_offset > 0 && row_offset <= self.num_rows
           && col_offset > 0 && col_offset <= self.num_cols
        {
            Some((row_offset - 1, col_offset - 1))
        }
        else {
            self.next()
        }
    }
}


#[allow(dead_code)]
impl<T> Vec2<T> {
    pub fn new(columns: usize) -> Vec2<T> {
        Vec2 {data: Vec::<T>::new(), columns}
    }

    pub fn with_capacity(rows: usize, columns: usize) -> Vec2<T> {
        Vec2 {data: Vec::<T>::with_capacity(rows * columns), columns}
    }

    pub fn push<S: IntoIterator<Item=T>>(&mut self, row: S) {
        self.data.extend(row.into_iter());
    }

    pub fn row_count(&self) -> usize {
        self.data.len() / self.columns
    }

    pub fn col_count(&self) -> usize {
        self.columns
    }

    pub fn coords(&self, i: usize) -> (usize, usize) {
        // convert underlying index into 2d coordinates
        (i / self.columns, i % self.columns)
    }

    pub fn offset(&self, row: usize, col: usize) -> usize {
        // and the reverse
        self.columns * row + col
    }

    pub fn rows(&self) -> impl Iterator<Item=&[T]> {
        self.data.chunks_exact(self.columns)
    }

    pub fn rows_mut(&mut self) -> impl Iterator<Item=&mut [T]> {
        self.data.chunks_exact_mut(self.columns)
    }
    
    pub fn values(&self) -> impl Iterator<Item=&T> {
        self.data.iter()
    }

    pub fn values_mut(&mut self) -> impl Iterator<Item=&mut T> {
        self.data.iter_mut()
    }

    pub fn neighbor_coords(&self, row: usize, col: usize) -> NeighborCoords {
        NeighborCoords {
            num_rows: self.row_count(), 
            num_cols: self.col_count(),
            start_row: row,
            start_col: col,
            i: 0,
        }
    }

    pub fn get(&self, row: usize, col: usize) -> &T {
        &self.data[self.offset(row, col)]
    }

    pub fn get_mut(&mut self, row: usize, col: usize) -> &mut T {
        self.data.index_mut(self.offset(row, col))
    }

    pub fn set(&mut self, row: usize, col: usize, val: T) {
        let existing = self.data.index_mut(self.offset(row, col));
        *existing = val;
    }
}

#[allow(dead_code)]
impl<T: Copy> Vec2<T> {
    pub fn fill(&mut self, value: T) {
        for i in 0..self.data.len() {
            self.data[i] = value;
        }
    }
    
    pub fn fill_to_cap(&mut self, value: T) {
        self.fill(value);
        for _i in self.data.len()..self.data.capacity() {
            self.data.push(value);
        }
    }
}


impl<T> Index<usize> for Vec2<T> {
    type Output = [T];
    fn index(&self, index: usize) -> &Self::Output {
        let start = self.columns * index;
        &self.data[start..(start + self.columns)]
    }
}


impl<T> IndexMut<usize> for Vec2<T> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        let start = self.columns * index;
        &mut self.data[start..(start + self.columns)]
    }
}


impl<T: Display> Display for Vec2<T> {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        for row in self.rows() {
            write!(f, "[")?;
            for (i, v) in row.iter().enumerate() {
                if i > 0 {write!(f, ", ")?;}
                write!(f, "{}", v)?;
            }
            write!(f, "]\n")?;
        }
        Ok(())
    }
}