1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

use std::collections::HashSet;

type Frequency = i32;

fn parse_frequency_changes(
    input_string: &str,
) -> impl Iterator<Item = Result<Frequency, String>> + Clone + '_ {
    input_string.lines().enumerate().map(|(line_index, line)| {
        line.parse::<Frequency>().map_err(|error| {
            format!(
                "Invalid input on line {}: {}",
                line_index + 1,
                error.to_string()
            )
        })
    })
}

pub fn part1(input_string: &str) -> Result<Frequency, String> {
    parse_frequency_changes(input_string).sum::<Result<_, _>>()
}

pub fn part2(input_string: &str) -> Result<Frequency, String> {
    const MAX_ITERATIONS: usize = 1_000_000;

    let mut frequency: Frequency = 0;
    let mut seen_frequencies = HashSet::new();

    for change in parse_frequency_changes(input_string)
        .cycle()
        .take(MAX_ITERATIONS)
    {
        if seen_frequencies.insert(frequency) {
            frequency = frequency.checked_add(change?).ok_or("Too high frequency")?;
        } else {
            return Ok(frequency);
        }
    }

    Err(format!(
        "Frequency not repeated after {} iterations",
        MAX_ITERATIONS
    ))
}

#[test]
pub fn tests_part1() {
    assert_eq!(Ok(3), part1("+1\n-2\n+3\n+1"));
    assert_eq!(Ok(3), part1("+1\n+1\n+1"));
    assert_eq!(Ok(0), part1("+1\n+1\n-2"));
    assert_eq!(Ok(-6), part1("-1\n-2\n-3"));
    assert_eq!(Ok(477), part1(include_str!("day01_input.txt")));
}

#[test]
fn tests_part2() {
    assert_eq!(Ok(0), part2("+1\n-1"));
    assert_eq!(Ok(10), part2("+3\n+3\n+4\n-2\n-4"));
    assert_eq!(Ok(5), part2("-6\n+3\n+8\n+5\n-6"));
    assert_eq!(Ok(14), part2("+7\n+7\n-2\n-7\n-4"));
    assert_eq!(Ok(390), part2(include_str!("day01_input.txt")));
}