nesfuzz/src/input.rs

use crate::{Rng, DISABLE_START_PRESSES_AFTER, MUTATION_RATE, MUTATION_RATE_SOFT_RESET};
use std::fs::File;
use std::io;
use std::io::BufRead;

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub enum ConsoleAction {
    None,
    SoftReset,
}

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct FrameInput {
    pub console_action: ConsoleAction,
    pub player_1_input: u8,
    pub player_2_input: u8,
}

#[derive(Clone, Eq, PartialEq, Hash)]
pub struct FuzzingInput {
    frames: Vec<FrameInput>,
    disable_start_after: usize,
    pub(crate) mutated: bool,
}

impl FuzzingInput {
    pub fn disable_start_after(&mut self, frames: usize) {
        self.disable_start_after = frames;
    }

    pub fn mutate_from(&mut self, rng: &mut Rng, frame: usize, frames_to_consider: usize) {
        self.mutated = true;
        for frame_num in frame..(frame + frames_to_consider) {
            while frame_num >= self.frames.len() {
                self.frames.push(FrameInput {
                    console_action: ConsoleAction::None,
                    player_1_input: rng.next() as u8,
                    player_2_input: 0,
                });
            }

            let check = (rng.next() % frames_to_consider as u32) as f64;

            if (check / frames_to_consider as f64) < MUTATION_RATE {
                if frame_num < self.frames.len() {
                    self.frames[frame_num].player_1_input = rng.next() as u8;
                }
            }

            if (check / frames_to_consider as f64) < MUTATION_RATE_SOFT_RESET {
                if frame_num < self.frames.len() {
                    self.frames[frame_num].console_action = ConsoleAction::SoftReset
                }
            }

            if frame_num > DISABLE_START_PRESSES_AFTER {
                self.frames[frame_num].player_1_input &= 0b11110011;
            }
        }
    }

    // Get the fuzzing input for the give frame...
    pub fn get_frame_input(&self, frame: usize) -> Option<&FrameInput> {
        self.frames.get(frame)
    }

    // Hacky code to parse an fm2 movie/input file into our initial fuzzing input
    pub fn load(fm2: &str) -> FuzzingInput {
        let file = File::open(fm2).unwrap();
        let mut frames = vec![];
        let lines = io::BufReader::new(file).lines();
        for (_line_num, line) in lines.enumerate() {
            if let Ok(lip) = line {
                let ip = lip.as_bytes();
                if ip.is_empty() {
                    continue;
                }
                if ip[0] as char == '|' {
                    let mut frame_input = FrameInput {
                        console_action: ConsoleAction::None,
                        player_1_input: 0,
                        player_2_input: 0,
                    };
                    if ip[1] as char == '1' {
                        frame_input.console_action = ConsoleAction::SoftReset;
                    }
                    // RLDUTSBA
                    if ip[3] as char == 'R' {
                        frame_input.player_1_input |= 128
                    }
                    if ip[4] as char == 'L' {
                        frame_input.player_1_input |= 64
                    }
                    if ip[5] as char == 'D' {
                        frame_input.player_1_input |= 32
                    }
                    if ip[6] as char == 'U' {
                        frame_input.player_1_input |= 16
                    }
                    if ip[7] as char == 'T' {
                        frame_input.player_1_input |= 8
                    }
                    if ip[8] as char == 'S' {
                        frame_input.player_1_input |= 4
                    }
                    if ip[9] as char == 'B' {
                        frame_input.player_1_input |= 2
                    }
                    if ip[10] as char == 'A' {
                        frame_input.player_1_input |= 1
                    }
                    // RLDUTSBA
                    if ip[12] as char == 'R' {
                        frame_input.player_2_input |= 128
                    }
                    if ip[13] as char == 'L' {
                        frame_input.player_2_input |= 64
                    }
                    if ip[14] as char == 'D' {
                        frame_input.player_2_input |= 32
                    }
                    if ip[15] as char == 'U' {
                        frame_input.player_2_input |= 16
                    }
                    if ip[16] as char == 'T' {
                        frame_input.player_2_input |= 8
                    }
                    if ip[17] as char == 'S' {
                        frame_input.player_2_input |= 4
                    }
                    if ip[18] as char == 'B' {
                        frame_input.player_2_input |= 2
                    }
                    if ip[19] as char == 'A' {
                        frame_input.player_2_input |= 1
                    }
                    frames.push(frame_input);
                }
            }
        }

        FuzzingInput {
            frames,
            disable_start_after: 0xFFFFFFFF,
            mutated: false,
        }
    }
}