<commit_before>#![cfg(windows)]

extern crate winapi;
extern crate kernel32;

use self::kernel32::*;
use self::winapi::*;
use std::ptr;
use std::iter::repeat;
use std::cmp::min;
use screen::Key;
use screen::Key::*;

use std::thread;
use std::sync::mpsc::Receiver;
use std::sync::mpsc;
use ::NEWLINE;

macro_rules! win32 {
    ($funcall:expr) => (
        if unsafe { $funcall } == 0 {
            panic!("Win32 call failed");
        }
    );
}

pub struct Screen {
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    start_line: u16,
    original_console_mode: DWORD,
    original_colors: WORD,
    input: Receiver<u16>,
    conout: HANDLE,
}

impl Screen {
    pub fn open_screen(desired_rows: u16) -> Screen {
        let mut orig_mode;
        let conin: HANDLE;
        let conout: HANDLE;
        unsafe {
            const OPEN_EXISTING: DWORD = 3;
            let rw_access = GENERIC_READ | GENERIC_WRITE;
            conin = CreateFileA("CONIN$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            conout = CreateFileA("CONOUT$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            orig_mode = ::std::mem::uninitialized();
        }

        if conin == INVALID_HANDLE_VALUE || conout == INVALID_HANDLE_VALUE {
            panic!("Unable to open console");
        }

        let (cols, rows) = Screen::winsize(conout).unwrap();

        win32!(GetConsoleMode(conin, &mut orig_mode));
        let new_mode = orig_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
        win32!(SetConsoleMode(conin, new_mode));

        let rx = Screen::spawn_input_thread(conin as usize);
        let initial_pos = Screen::get_cursor_pos(conout);
        let visible_choices = min(desired_rows, rows - 1);
        let start_line = get_start_line(rows, visible_choices, initial_pos);
        let original_colors = Screen::get_original_colors(conout);
        let (column, _) = initial_pos;
        if column > 0 {
            Self::write_to(conout, NEWLINE);
        }
        for _ in 0..visible_choices {
            Self::write_to(conout, NEWLINE);
        }
        Screen {
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line + Self::get_buffer_offset(conout),
            original_console_mode: orig_mode,
            original_colors: original_colors,
            input: rx,
            conout: conout,
        }
    }

    // We have to take the conin handle as a usize instead of a *mut c_void in order to avoid a
    // lecture from the compiler about how the latter type cannot be safely sent between threads.
    // I'm not sure if a better solution exists at this time.
    fn spawn_input_thread(conin: usize) -> Receiver<u16> {
        let (tx, rx) = mpsc::channel();

        thread::spawn(move || {
            loop {
                let conin = conin as *mut c_void;
                let mut buf: Vec<u16> = repeat(0u16).take(0x1000).collect();
                let mut chars_read: DWORD = 0;
                win32!(ReadFile(conin, buf.as_mut_ptr() as LPVOID, 1, &mut chars_read as LPDWORD, ptr::null_mut()));
                for i in 0..chars_read as usize {
                    tx.send(buf[i]).unwrap();
                }
            }
        });

        rx
    }

    fn move_cursor(&mut self, line: u16, column: u16) {
        win32!(SetConsoleCursorPosition(self.conout, COORD { X: column as i16, Y: line as i16}));
    }

    pub fn move_cursor_to_prompt_line(&mut self, col: u16) {
        let start_line = self.start_line;
        self.move_cursor(start_line, col);
    }

    pub fn blank_screen(&mut self) {
        let blank_line = repeat(' ').take((self.width - 1) as usize).collect::<String>();
        let start_line = self.start_line;
        self.move_cursor(start_line, 0);
        for _ in 0..self.visible_choices {
            self.write(&blank_line);
            self.write(NEWLINE);
        }
        self.write(&blank_line);
        self.move_cursor(start_line, 0);
    }

    pub fn show_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: TRUE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn hide_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: FALSE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn write(&mut self, s: &str) {
        Self::write_to(self.conout, s);
    }

    fn write_to(conout: HANDLE, s: &str) {
        let mut chars_written: DWORD = 0;
        let utf16 = s.encode_utf16().collect::<Vec<u16>>();
        let chars_to_write = utf16.len() as DWORD;
        win32!(WriteConsoleW(conout, utf16.as_ptr() as PVOID, chars_to_write, &mut chars_written as LPDWORD, ptr::null_mut()));
    }

    pub fn write_red_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const WHITE_ON_RED: WORD = BACKGROUND_RED as WORD;
        self.set_colors(WHITE_ON_RED);
        self.write(s);
        self.set_colors(orig);
    }

    pub fn write_red(&mut self, s: &str) {
        let orig = self.original_colors;
        const RED_ON_BLACK: WORD = FOREGROUND_RED as WORD;
        self.set_colors(RED_ON_BLACK);
        self.write(s);
        self.set_colors(orig);
    }

    pub fn write_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const BLACK_ON_WHITE: WORD = (BACKGROUND_RED | BACKGROUND_GREEN | BACKGROUND_BLUE) as WORD;
        self.set_colors(BLACK_ON_WHITE);
        self.write(s);
        self.set_colors(orig);
    }

    fn set_colors(&mut self, colors: WORD) {
        win32!(SetConsoleTextAttribute(self.conout, colors));
    }

    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut ret = Vec::new();
        while let Ok(byte) = self.input.try_recv() {
            ret.push(Screen::translate_byte(byte));
        }
        if ret.is_empty() {
            let byte = self.input.recv().unwrap();
            ret.push(Screen::translate_byte(byte));
        }
        ret
    }

    fn translate_byte(byte: u16) -> Key {
        if byte == '\r' as u16 {
            Enter
        } else if byte == 9 {
            Tab
        } else if byte == 127 {
            Backspace
        } else if byte & 96 == 0 {
            Control(((byte + 96u16) as u8) as char)
        } else {
            Char((byte as u8) as char)
        }
    }

    fn get_cursor_pos(handle: HANDLE) -> (u16, u16) {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        let cursor_pos = buffer_info.dwCursorPosition;
        (cursor_pos.X as u16, cursor_pos.Y as u16)
    }

    fn get_original_colors(handle: HANDLE) -> WORD {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        buffer_info.wAttributes
    }

    fn winsize(conout: HANDLE) -> Option<(u16, u16)> {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        let result = unsafe { GetConsoleScreenBufferInfo(conout, &mut buffer_info) };
        if result != 0 {
            // This code specifically computes the size of the window,
            // *not* the size of the buffer (which is easily available
            // from dwSize). I got the algorithm from:
            //
            // http://stackoverflow.com/a/12642749
            let left = buffer_info.srWindow.Left;
            let top = buffer_info.srWindow.Top;
            let right = buffer_info.srWindow.Right;
            let bottom = buffer_info.srWindow.Bottom;
            let cols = right - left + 1;
            let rows = bottom - top + 1;
            Some((cols as u16, rows as u16))
        } else {
            None
        }
    }

    fn get_buffer_offset(conout: HANDLE) -> u16 {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(conout, &mut buffer_info));
        buffer_info.srWindow.Top as u16
    }
}

fn get_start_line(rows: u16, visible_choices: u16, initial_pos: (u16, u16)) -> u16 {
    let bottom_most_line = rows - visible_choices - 1;
    let (initial_x, initial_y) = initial_pos;
    let line_under_cursor = if initial_x == 0 { initial_y } else { initial_y + 1 };
    if line_under_cursor + 1 + visible_choices > rows {
        bottom_most_line
    } else {
        line_under_cursor
    }
}

#[cfg(test)]
mod tests {
    use super::kernel32;
    use super::winapi::STD_OUTPUT_HANDLE;
    use super::{Screen, get_start_line};

    #[test]
    fn winsize_test() {
        // AppVeyor builds run without a console, making this test impossible.
        if option_env!("APPVEYOR").is_some() {
            // TODO: It should be made obvious from the output that this test was skipped
            return;
        }
        let conout = unsafe { kernel32::GetStdHandle(STD_OUTPUT_HANDLE) };
        let (cols, rows) = Screen::winsize(conout).expect("Failed to get window size!");
        // We don't know the window size a priori, but we can at least
        // assert that it is within some kind of sensible range.
        assert!(cols > 40);
        assert!(rows > 40);
        assert!(cols < 1000);
        assert!(rows < 1000);
    }

    #[test]
    fn start_line_test() {
        assert_eq!(5, get_start_line(100, 20, (0, 5)));
        assert_eq!(6, get_start_line(100, 20, (1, 5)));
        assert_eq!(79, get_start_line(100, 20, (0, 100)));
        assert_eq!(0, get_start_line(15, 14, ((0, 5))));
        assert_eq!(79, get_start_line(100, 20, (50, 100)));
    }
}
<commit_msg>Migrate Windows input to ReadConsoleInput<commit_after>#![cfg(windows)]

extern crate winapi;
extern crate kernel32;

use self::kernel32::*;
use self::winapi::*;
use std::ptr;
use std::iter::repeat;
use std::cmp::min;
use screen::Key;
use screen::Key::*;

use std::thread;
use std::sync::mpsc::Receiver;
use std::sync::mpsc;
use ::NEWLINE;

macro_rules! win32 {
    ($funcall:expr) => (
        if unsafe { $funcall } == 0 {
            panic!("Win32 call failed");
        }
    );
}

pub struct Screen {
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    start_line: u16,
    original_console_mode: DWORD,
    original_colors: WORD,
    input: Receiver<INPUT_RECORD>,
    conout: HANDLE,
}

impl Screen {
    pub fn open_screen(desired_rows: u16) -> Screen {
        let mut orig_mode;
        let conin: HANDLE;
        let conout: HANDLE;
        unsafe {
            const OPEN_EXISTING: DWORD = 3;
            let rw_access = GENERIC_READ | GENERIC_WRITE;
            conin = CreateFileA("CONIN$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            conout = CreateFileA("CONOUT$\0".as_ptr() as *const i8, rw_access, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, 0, ptr::null_mut());
            orig_mode = ::std::mem::uninitialized();
        }

        if conin == INVALID_HANDLE_VALUE || conout == INVALID_HANDLE_VALUE {
            panic!("Unable to open console");
        }

        let (cols, rows) = Screen::winsize(conout).unwrap();

        win32!(GetConsoleMode(conin, &mut orig_mode));
        let new_mode = orig_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
        win32!(SetConsoleMode(conin, new_mode));

        let rx = Screen::spawn_input_thread(conin as usize);
        let initial_pos = Screen::get_cursor_pos(conout);
        let visible_choices = min(desired_rows, rows - 1);
        let start_line = get_start_line(rows, visible_choices, initial_pos);
        let original_colors = Screen::get_original_colors(conout);
        let (column, _) = initial_pos;
        if column > 0 {
            Self::write_to(conout, NEWLINE);
        }
        for _ in 0..visible_choices {
            Self::write_to(conout, NEWLINE);
        }
        Screen {
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line + Self::get_buffer_offset(conout),
            original_console_mode: orig_mode,
            original_colors: original_colors,
            input: rx,
            conout: conout,
        }
    }

    // We have to take the conin handle as a usize instead of a *mut c_void in order to avoid a
    // lecture from the compiler about how the latter type cannot be safely sent between threads.
    // I'm not sure if a better solution exists at this time.
    fn spawn_input_thread(conin: usize) -> Receiver<INPUT_RECORD> {
        let (tx, rx) = mpsc::channel();

        thread::spawn(move || {
            loop {
                let conin = conin as *mut c_void;
                let mut input_record = INPUT_RECORD{EventType: 0 as WORD, Event: [0, 0, 0, 0]};
                let mut events_read: DWORD = 0;
                win32!(ReadConsoleInputW(conin, &mut input_record as PINPUT_RECORD, 1, &mut events_read as LPDWORD));
                if events_read > 0 {
                    tx.send(input_record).unwrap();
                }
            }
        });

        rx
    }

    fn move_cursor(&mut self, line: u16, column: u16) {
        win32!(SetConsoleCursorPosition(self.conout, COORD { X: column as i16, Y: line as i16}));
    }

    pub fn move_cursor_to_prompt_line(&mut self, col: u16) {
        let start_line = self.start_line;
        self.move_cursor(start_line, col);
    }

    pub fn blank_screen(&mut self) {
        let blank_line = repeat(' ').take((self.width - 1) as usize).collect::<String>();
        let start_line = self.start_line;
        self.move_cursor(start_line, 0);
        for _ in 0..self.visible_choices {
            self.write(&blank_line);
            self.write(NEWLINE);
        }
        self.write(&blank_line);
        self.move_cursor(start_line, 0);
    }

    pub fn show_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: TRUE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn hide_cursor(&mut self) {
        let cursor_info = CONSOLE_CURSOR_INFO { dwSize: 100, bVisible: FALSE };
        win32!(SetConsoleCursorInfo(self.conout, &cursor_info));
    }

    pub fn write(&mut self, s: &str) {
        Self::write_to(self.conout, s);
    }

    fn write_to(conout: HANDLE, s: &str) {
        let mut chars_written: DWORD = 0;
        let utf16 = s.encode_utf16().collect::<Vec<u16>>();
        let chars_to_write = utf16.len() as DWORD;
        win32!(WriteConsoleW(conout, utf16.as_ptr() as PVOID, chars_to_write, &mut chars_written as LPDWORD, ptr::null_mut()));
    }

    pub fn write_red_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const WHITE_ON_RED: WORD = BACKGROUND_RED as WORD;
        self.set_colors(WHITE_ON_RED);
        self.write(s);
        self.set_colors(orig);
    }

    pub fn write_red(&mut self, s: &str) {
        let orig = self.original_colors;
        const RED_ON_BLACK: WORD = FOREGROUND_RED as WORD;
        self.set_colors(RED_ON_BLACK);
        self.write(s);
        self.set_colors(orig);
    }

    pub fn write_inverted(&mut self, s: &str) {
        let orig = self.original_colors;
        const BLACK_ON_WHITE: WORD = (BACKGROUND_RED | BACKGROUND_GREEN | BACKGROUND_BLUE) as WORD;
        self.set_colors(BLACK_ON_WHITE);
        self.write(s);
        self.set_colors(orig);
    }

    fn set_colors(&mut self, colors: WORD) {
        win32!(SetConsoleTextAttribute(self.conout, colors));
    }

    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut ret = Vec::new();
        while let Ok(event) = self.input.try_recv() {
            ret.push(Screen::translate_event(event));
        }
        if ret.is_empty() {
            let event = self.input.recv().unwrap();
            ret.push(Screen::translate_event(event));
        }
        ret
    }

    fn translate_event(event: INPUT_RECORD) -> Key {
        if event.EventType != KEY_EVENT {
            return Nothing;
        }

        let key_event = unsafe { event.KeyEvent() };
        if key_event.bKeyDown == FALSE {
            return Nothing;
        }
        let byte = key_event.UnicodeChar;
        if byte == '\r' as u16 {
            Enter
        } else if byte == 9 {
            Tab
        } else if byte == 127 {
            Backspace
        } else if byte & 96 == 0 {
            Control(((byte + 96u16) as u8) as char)
        } else {
            Char((byte as u8) as char)
        }
    }

    fn get_cursor_pos(handle: HANDLE) -> (u16, u16) {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        let cursor_pos = buffer_info.dwCursorPosition;
        (cursor_pos.X as u16, cursor_pos.Y as u16)
    }

    fn get_original_colors(handle: HANDLE) -> WORD {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(handle, &mut buffer_info));
        buffer_info.wAttributes
    }

    fn winsize(conout: HANDLE) -> Option<(u16, u16)> {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        let result = unsafe { GetConsoleScreenBufferInfo(conout, &mut buffer_info) };
        if result != 0 {
            // This code specifically computes the size of the window,
            // *not* the size of the buffer (which is easily available
            // from dwSize). I got the algorithm from:
            //
            // http://stackoverflow.com/a/12642749
            let left = buffer_info.srWindow.Left;
            let top = buffer_info.srWindow.Top;
            let right = buffer_info.srWindow.Right;
            let bottom = buffer_info.srWindow.Bottom;
            let cols = right - left + 1;
            let rows = bottom - top + 1;
            Some((cols as u16, rows as u16))
        } else {
            None
        }
    }

    fn get_buffer_offset(conout: HANDLE) -> u16 {
        let mut buffer_info = unsafe { ::std::mem::uninitialized() };
        win32!(GetConsoleScreenBufferInfo(conout, &mut buffer_info));
        buffer_info.srWindow.Top as u16
    }
}

fn get_start_line(rows: u16, visible_choices: u16, initial_pos: (u16, u16)) -> u16 {
    let bottom_most_line = rows - visible_choices - 1;
    let (initial_x, initial_y) = initial_pos;
    let line_under_cursor = if initial_x == 0 { initial_y } else { initial_y + 1 };
    if line_under_cursor + 1 + visible_choices > rows {
        bottom_most_line
    } else {
        line_under_cursor
    }
}

#[cfg(test)]
mod tests {
    use super::kernel32;
    use super::winapi::STD_OUTPUT_HANDLE;
    use super::{Screen, get_start_line};

    #[test]
    fn winsize_test() {
        // AppVeyor builds run without a console, making this test impossible.
        if option_env!("APPVEYOR").is_some() {
            // TODO: It should be made obvious from the output that this test was skipped
            return;
        }
        let conout = unsafe { kernel32::GetStdHandle(STD_OUTPUT_HANDLE) };
        let (cols, rows) = Screen::winsize(conout).expect("Failed to get window size!");
        // We don't know the window size a priori, but we can at least
        // assert that it is within some kind of sensible range.
        assert!(cols > 40);
        assert!(rows > 40);
        assert!(cols < 1000);
        assert!(rows < 1000);
    }

    #[test]
    fn start_line_test() {
        assert_eq!(5, get_start_line(100, 20, (0, 5)));
        assert_eq!(6, get_start_line(100, 20, (1, 5)));
        assert_eq!(79, get_start_line(100, 20, (0, 100)));
        assert_eq!(0, get_start_line(15, 14, ((0, 5))));
        assert_eq!(79, get_start_line(100, 20, (50, 100)));
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::prelude::*;

type Vector3 = (i32, i32, i32);

#[derive(Debug, Clone, Eq, PartialEq)]
struct MoonState {
  position: Vector3,
  velocity: Vector3
}

impl MoonState {
  fn potential_energy(&self) -> i32 {
    norm(self.position)
  }

  fn kinetic_enery(&self) -> i32 {
    norm(self.velocity)
  }

  fn total_energy(&self) -> i32 {
    self.potential_energy() * self.kinetic_enery()
  }
}

fn get_input() -> std::io::Result<String> {
  let mut file = File::open("12.txt")?;
  let mut contents = String::new();
  file.read_to_string(&mut contents)?;

  Ok(contents)
}

fn parse_input(input: &str) -> Vec<MoonState> {
  input.lines()
  .filter_map(|line| line.get(1..line.len() - 1))
  .filter_map(|line| {
    let mut coords = line.split(", ")
      .filter_map(|token| {
        match token.split('=').nth(1) {
          Some(x) => x.parse::<i32>().ok(),
          _ => None
        }
      });

    match (coords.nth(0), coords.nth(0), coords.nth(0)) {
      (Some(x), Some(y), Some(z)) => Some((x, y, z)),
      _ => None
    }
  })
  .map(|position| MoonState {
    position,
    velocity: (0, 0, 0)
  })
  .collect()
}

fn norm((x, y, z): Vector3) -> i32 {
  x.abs() + y.abs() + z.abs()
}

fn simulate_tick(moons: &mut Vec<MoonState>) {
  // Apply gravity

  fn calc_acceleration(x: i32, y: i32) -> i32 {
    if x < y { 1 }
    else if x > y { -1 }
    else { 0 }
  }

  let n = moons.len();
  let pairs = (0..n)
    .flat_map(|x| {
      (0..n)
      .map(move |y| (x, y))
    })
    .filter(|&(x, y)| x < y);

  for (i, j) in pairs {
    let moon1 = &moons[i];
    let moon2 = &moons[j];

    let acceleration = (
      calc_acceleration(moon1.position.0, moon2.position.0),
      calc_acceleration(moon1.position.1, moon2.position.1),
      calc_acceleration(moon1.position.2, moon2.position.2)
    );

    let moon1 = &mut moons[i];

    moon1.velocity = (
      moon1.velocity.0 + acceleration.0,
      moon1.velocity.1 + acceleration.1,
      moon1.velocity.2 + acceleration.2
    );

    let moon2 = &mut moons[j];

    moon2.velocity = (
      moon2.velocity.0 - acceleration.0,
      moon2.velocity.1 - acceleration.1,
      moon2.velocity.2 - acceleration.2
    );
  }

  // Apply velocity

  for moon in moons {
    moon.position = (
      moon.position.0 + moon.velocity.0,
      moon.position.1 + moon.velocity.1,
      moon.position.2 + moon.velocity.2
    );
  }
}

fn main() {
  let input = get_input().unwrap();
  let initial_moons = parse_input(&input);
  let mut moons = initial_moons.clone();

  for _ in 0..1000 {
    simulate_tick(&mut moons);
  }

  let total_energy = moons.iter()
    .map(|moon| moon.total_energy())
    .sum::<i32>();

  println!("Part 1: {}", total_energy);
}
<commit_msg>Add day 12 part 2 solution<commit_after>use std::fs::File;
use std::io::prelude::*;
use std::fmt;

type Vector3 = (i32, i32, i32);

#[derive(Clone)]
struct MoonState {
  position: Vector3,
  velocity: Vector3
}

impl fmt::Debug for MoonState {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "{:?} - {:?}", self.position, self.velocity)
  }
}

impl MoonState {
  fn potential_energy(&self) -> i32 {
    norm(self.position)
  }

  fn kinetic_enery(&self) -> i32 {
    norm(self.velocity)
  }

  fn total_energy(&self) -> i32 {
    self.potential_energy() * self.kinetic_enery()
  }
}

fn get_input() -> std::io::Result<String> {
  let mut file = File::open("12.txt")?;
  let mut contents = String::new();
  file.read_to_string(&mut contents)?;

  Ok(contents)
}

fn parse_input(input: &str) -> Vec<MoonState> {
  input.lines()
  .filter_map(|line| line.get(1..line.len() - 1))
  .filter_map(|line| {
    let mut coords = line.split(", ")
      .filter_map(|token| {
        match token.split('=').nth(1) {
          Some(x) => x.parse::<i32>().ok(),
          _ => None
        }
      });

    match (coords.nth(0), coords.nth(0), coords.nth(0)) {
      (Some(x), Some(y), Some(z)) => Some((x, y, z)),
      _ => None
    }
  })
  .map(|position| MoonState {
    position,
    velocity: (0, 0, 0)
  })
  .collect()
}

fn norm((x, y, z): Vector3) -> i32 {
  x.abs() + y.abs() + z.abs()
}

fn gcd(a: i64, b: i64) -> i64 {
  let (mut a, mut b) = (a, b);
  let mut h;

  if a == 0 {
    return b.abs();
  } else if b == 0 {
    return a.abs();
  }

  loop {
    h = a % b;
    a = b;
    b = h;

    if b == 0 {
      break a.abs();
    }
  }
}

fn lcm(a: i64, b: i64) -> i64 {
  if a == 0 && b == 0 {
    0
  } else {
    a.abs() / gcd(a, b) * b.abs()
  }
}

fn multiple_lcm(numbers: &[i64]) -> Option<i64> {
  numbers.iter()
  .fold(None, |acc, &x| match acc {
    None => Some(x),
    Some(acc) => Some(lcm(acc, x))
  })
}

fn calc_acceleration(pos1: Vector3, pos2: Vector3) -> Vector3 {
  fn get_sign(x: i32, y: i32) -> i32 {
    if x < y { 1 }
    else if x > y { -1 }
    else { 0 }
  }

  (
    get_sign(pos1.0, pos2.0),
    get_sign(pos1.1, pos2.1),
    get_sign(pos1.2, pos2.2)
  )
}

fn simulate_tick(moons: &mut Vec<MoonState>) {
  // Apply gravity

  let n = moons.len();
  let pairs = (0..n)
    .flat_map(|x| {
      (0..n)
      .map(move |y| (x, y))
    })
    .filter(|&(x, y)| x < y);

  for (i, j) in pairs {
    let moon1 = &moons[i];
    let moon2 = &moons[j];
    let acceleration = calc_acceleration(moon1.position, moon2.position);

    let moon1 = &mut moons[i];

    moon1.velocity = (
      moon1.velocity.0 + acceleration.0,
      moon1.velocity.1 + acceleration.1,
      moon1.velocity.2 + acceleration.2
    );

    let moon2 = &mut moons[j];

    moon2.velocity = (
      moon2.velocity.0 - acceleration.0,
      moon2.velocity.1 - acceleration.1,
      moon2.velocity.2 - acceleration.2
    );
  }

  // Apply velocity

  for moon in moons.iter_mut() {
    moon.position = (
      moon.position.0 + moon.velocity.0,
      moon.position.1 + moon.velocity.1,
      moon.position.2 + moon.velocity.2
    );
  }
}

fn main() {
  let input = get_input().unwrap();
  let initial_moons = parse_input(&input);
  let mut moons = initial_moons.clone();

  for _ in 0..1000 {
    simulate_tick(&mut moons);
  }

  let total_energy = moons.iter()
    .map(|moon| moon.total_energy())
    .sum::<i32>();

  println!("Part 1: {}", total_energy);

  let mut moons = initial_moons.clone();
  let (mut xf, mut yf, mut zf) = (false, false, false);
  let (mut xp, mut yp, mut zp) = (0, 0, 0);

  fn count_var<F, T>(
    counter: &mut i64,
    finished: &mut bool,
    moons: &Vec<MoonState>,
    compare: &Vec<MoonState>,
    predicate: F
  )
  where
    T: PartialEq,
    F: Fn(&MoonState) -> T
  {
    if !*finished {
      *counter += 1;

      let component_equals = moons.iter().map(&predicate)
        .zip(compare.iter().map(&predicate))
        .all(|(x, y)| x == y);

      if component_equals {
        *finished = true;
      }
    }
  }

  loop {
    simulate_tick(&mut moons);

    count_var(&mut xp, &mut xf, &moons, &initial_moons, |moon| [moon.position.0, moon.velocity.0]);
    count_var(&mut yp, &mut yf, &moons, &initial_moons, |moon| [moon.position.1, moon.velocity.1]);
    count_var(&mut zp, &mut zf, &moons, &initial_moons, |moon| [moon.position.2, moon.velocity.2]);

    if xf && yf && zf {
      break;
    }
  }

  println!("Part 2: {}", multiple_lcm(&[xp, yp, zp]).unwrap());
}
<|endoftext|>


<commit_before>//! Utility methods and constructors for Lua objects

use std::fmt::Display;
use std::convert::From;
use rlua::{self, Lua, Table, UserData, AnyUserData, Value,
           FromLua, ToLua, Function};
use super::signal;
use super::class::Class;
use super::property::Property;

/// All Lua objects can be cast to this.
#[derive(Clone, Debug)]
pub struct Object<'lua> {
    pub object: AnyUserData<'lua>
}

impl <'lua> From<AnyUserData<'lua>> for Object<'lua> {
    fn from(object: AnyUserData<'lua>) -> Self {
        Object { object }
    }
}

/// Construct a new object, used when using the default Objectable::new.
pub struct ObjectBuilder<'lua>{
    lua: &'lua Lua,
    object: Object<'lua>
}

impl <'lua> ObjectBuilder<'lua> {
    pub fn add_to_meta(self, new_meta: Table<'lua>) -> rlua::Result<Self> {
        let meta = self.object.table()?.get_metatable()
            .expect("Object had no meta table");
        for entry in new_meta.pairs::<rlua::Value, rlua::Value>() {
            let (key, value) = entry?;
            meta.set(key, value)?;
        }
        self.object.table()?.set_metatable(Some(meta));
        Ok(self)
    }

    #[allow(dead_code)]
    pub fn add_to_signals(self, name: String, func: rlua::Function)
                          -> rlua::Result<Self> {
        signal::connect_signal(self.lua, self.object.clone(), name, &[func])?;
        Ok(self)
    }

    pub fn handle_constructor_argument(self, args: Table) -> rlua::Result<Self> {
        let meta = self.object.table()?.get_metatable()
            .expect("Object had no meta table");
        let class = meta.get::<_, Table>("__class")?;
        let props = class.get::<_, Vec<Property>>("properties")?;

        // Handle all table entries that correspond to known properties,
        // silently ignore all other keys
        for pair in args.pairs() {
            let (key, value): (Value, Value) = pair?;
            if let rlua::Value::String(key) = key {
                if let Ok(key) = key.to_str() {
                    // Find the property
                    for prop in props.iter() {
                        if prop.name == key {
                            // Property exists and has a cb_new callback
                            if let Some(ref new) = prop.cb_new {
                                let _: () = new.bind(self.object.table()?.clone())?.call(value)?;
                                break;
                            }
                        }
                    }
                }
            }
        }
        Ok(self)
    }

    pub fn build(self) -> Object<'lua> {
        self.object
    }
}

/// Trait implemented by all objects that represent OO lua objects.
///
/// This trait allows casting an object gotten back from the Lua runtime
/// into a concrete object so that Rust can do things with it.
///
/// You can't do anything to the object until it has been converted into a
/// canonical form using this trait.
pub trait Objectable<'lua, T, S: UserData + Default + Display + Clone> {
    fn cast(obj: Object<'lua>) -> rlua::Result<T> {
        //let data = obj.table()?.get::<_, AnyUserData>("userdata")?;
        if obj.object.is::<S>()? {
            Ok(Self::_wrap(obj))
        } else {
            use rlua::Error::RuntimeError;
            Err(RuntimeError("Could not cast object to concrete type".into()))
        }
    }

    /// Given the internal object table, constructs a concrete object out
    /// of it.
    ///
    /// This is only used internally, which is why it's prefixed with a "_"
    /// Please do not use it outside of object.rs.
    fn _wrap(table: Object<'lua>) -> T;

    fn state(&self) -> rlua::Result<S> {
        Ok(self.get_object()?.clone())
    }

    fn set_state(&mut self, data: S) -> rlua::Result<()> {
        *self.get_object_mut()? = data;
        Ok(())
    }

    /// Gets the internal table for the concrete object.
    /// Used internally by cast, though there's nothing wrong with it being
    /// used outside of internal object use.
    //fn get_table(&self) -> Table<'lua>;

    fn get_object(&self) -> rlua::Result<S>;
    fn get_object_mut(&mut self) -> rlua::Result<::std::cell::RefMut<S>>;

    fn allocate(lua: &'lua Lua, class: Class) -> rlua::Result<ObjectBuilder<'lua>>
    {
        // This is the format of the lua objects in Way Cooler:
        /*
        -- This is the weird wrapper table, binding the user data and the object table together
        {
            -- Actual user data object, e.g the Rust object
            userdata = <user_data_here>,
            -- This is the object data table, e.g the thing that in Awesome
            -- is linked to the user data with lua_setuservalue
            data = {},
            -- That metatable that is allocated in the new macro for objects.
            -- It can access both the userdata and the lua data fields
            __metatable = {}
        }
        */
        let object = match S::default().to_lua(lua)? {
            rlua::Value::UserData(data) => data,
            _ => unreachable!()
        };
        // TODO Increment the instance count
        let wrapper_table = lua.create_table()?;
        //wrapper_table.set("userdata", object)?;
        let data_table = lua.create_table()?;
        wrapper_table.set("data", data_table)?;
        let meta = lua.create_table()?;
        meta.set("__class", class)?;
        meta.set("properties", Vec::<Property>::new().to_lua(lua)?)?;
        meta.set("signals", lua.create_table()?)?;
        meta.set("connect_signal",
                 lua.create_function(connect_signal)?)?;
        meta.set("disconnect_signal",
                 lua.create_function(disconnect_signal)?)?;
        meta.set("emit_signal", lua.create_function(emit_signal)?)?;
        // TODO Remove these index functions? They are defined on the user data now I think
        // But maybe this is a fallback option still?
        meta.set("__index", lua.create_function(table_default_index)?)?;
        // NOTE TODO Very sure that I don't need this, but why I need above?
        //meta.set("__newindex", lua.create_function(default_newindex)?)?;
        meta.set("__tostring", lua.create_function(|_, data: AnyUserData| {
            Ok(format!("{}", data.borrow::<S>()?.clone()))
        })?)?;
        wrapper_table.set_metatable(Some(meta));
        object.set_user_value(wrapper_table)?;
        // TODO Emit new signal event
        let object = Object { object };
        Ok(ObjectBuilder { object, lua })
    }
}

impl <'lua> ToLua<'lua> for Object<'lua> {
    fn to_lua(self, _: &'lua Lua) -> rlua::Result<Value<'lua>> {
        Ok(Value::UserData(self.object))
    }
}

impl <'lua> Object<'lua> {
    pub fn signals(&self) -> rlua::Result<rlua::Table<'lua>> {
        self.table()?.get::<_, Table>("signals")
        //let obj: Object<'lua> = self.object.clone().into();
        //obj.table()?.get::<_, Table>("signals")
    }

    pub fn table(&self) -> rlua::Result<Table<'lua>> {
        self.object.get_user_value::<Table<'lua>>()
    }
}

pub fn table_default_index<'lua>(lua: &'lua Lua,
                         (obj_table, index): (Table<'lua>, Value<'lua>))
                         -> rlua::Result<Value<'lua>> {
    // Look up in metatable first
    let meta = obj_table.get_metatable()
        .expect("Object had no metatable");
    if meta.get::<_, Table>("__class").is_ok() {
        if let Ok(val) = meta.raw_get::<_, Value>(index.clone()) {
            match val {
                Value::Nil => {},
                val => return Ok(val)
            }
        }
    }
    let index = String::from_lua(index, lua)?;
    match index.as_str() {
        "valid" => {
            // TODO What do?
            unimplemented!()
        },
        "data" => {
            obj_table.to_lua(lua)
        },
        index => {
            // Try see if there is a property of the class with the name
            if let Ok(class) = meta.get::<_, Table>("__class") {
                let props = class.get::<_, Vec<Property>>("properties")?;
                for prop in props {
                    if prop.name.as_str() == index {
                        // Property exists and has an index callback
                        if let Some(index) = prop.cb_index {
                            // TODO what do?
                            unimplemented!()
                           // return index.call(obj)
                        }
                    }
                }
                match class.get::<_, Function>("__index_miss_handler") {
                    Ok(function) => {
                        return function.bind(obj_table)?.call(index)
                    },
                    Err(_) => {}
                }
            }
            // TODO property miss handler if index doesn't exst
            Ok(rlua::Value::Nil)
        }
    }
}

/// Default indexing of an Awesome object.
///
/// Automatically looks up contents in meta table, so instead of overriding this
/// it's easier to just add the required data in the meta table.
pub fn default_index<'lua>(lua: &'lua Lua,
                           (obj, index): (AnyUserData<'lua>, Value<'lua>))
                           -> rlua::Result<Value<'lua>> {
    // Look up in metatable first
    let obj: Object<'lua> = obj.into();
    let obj_table = obj.table()?;
    let meta = obj_table.get_metatable()
        .expect("Object had no metatable");
    if meta.get::<_, Table>("__class").is_ok() {
        if let Ok(val) = meta.raw_get::<_, Value>(index.clone()) {
            match val {
                Value::Nil => {},
                val => return Ok(val)
            }
        }
    }
    let index = String::from_lua(index, lua)?;
    match index.as_str() {
        "valid" => {
            Ok(Value::Boolean(
                if let Ok(class) = meta.get::<_, Table>("__class") {
                    let class: Class = class.into();
                    class.checker()?
                        .map(|checker| checker(obj))
                        .unwrap_or(true)
                } else {
                    false
                }))
        },
        "data" => {
            obj_table.to_lua(lua)
        },
        index => {
            // Try see if there is a property of the class with the name
            if let Ok(class) = meta.get::<_, Table>("__class") {
                let props = class.get::<_, Vec<Property>>("properties")?;
                for prop in props {
                    if prop.name.as_str() == index {
                        // Property exists and has an index callback
                        if let Some(index) = prop.cb_index {
                            return index.call(obj)
                        }
                    }
                }
                match class.get::<_, Function>("__index_miss_handler") {
                    Ok(function) => {
                        return function.bind(obj_table)?.call(index)
                    },
                    Err(_) => {}
                }
            }
            // TODO property miss handler if index doesn't exst
            Ok(rlua::Value::Nil)
        }
    }
}

/// Default new indexing (assignment) of an Awesome object.
///
/// Automatically looks up contents in meta table, so instead of overriding this
/// it's easier to just add the required data in the meta table.
pub fn default_newindex<'lua>(_: &'lua Lua,
                              (obj, index, val):
                              (AnyUserData<'lua>, String, Value<'lua>))
                              -> rlua::Result<Value<'lua>> {
    // Look up in metatable first
    let obj: Object<'lua> = obj.into();
    let obj_table = obj.table()?;
    if let Some(meta) = obj_table.get_metatable() {
        if let Ok(val) = meta.raw_get::<_, Value>(index.clone()) {
            match val {
                Value::Nil => {},
                val => return Ok(val)
            }
        }
        let class = meta.get::<_, Table>("__class")?;
        let props = class.get::<_, Vec<Property>>("properties")?;
        for prop in props {
            if prop.name.as_str() == index {
                // Property exists and has a newindex callback
                if let Some(newindex) = prop.cb_newindex {
                    return newindex.bind(obj.clone())?.call(val)
                }
            }
        }
        match class.get::<_, Function>("__newindex_miss_handler") {
            Ok(function) => {
                return function.bind(obj_table)?.call((index, val))
            },
            Err(_) => {}
        }
        // TODO property miss handler if index doesn't exst
    }
    Ok(Value::Nil)
}

fn connect_signal(lua: &Lua, (obj, signal, func): (AnyUserData, String, Function))
                  -> rlua::Result<()> {
    signal::connect_signal(lua, obj.into(), signal, &[func])
}

fn disconnect_signal(lua: &Lua, (obj, signal): (AnyUserData, String))
                     -> rlua::Result<()> {
    signal::disconnect_signal(lua, obj.into(), signal)
}

fn emit_signal(lua: &Lua, (obj, signal, args): (AnyUserData, String, Value))
               -> rlua::Result<()> {
    signal::emit_signal(lua, obj.into(), signal, args)
}
<commit_msg>Fixed __index on class<commit_after>//! Utility methods and constructors for Lua objects

use std::fmt::Display;
use std::convert::From;
use rlua::{self, Lua, Table, UserData, AnyUserData, Value,
           FromLua, ToLua, Function};
use super::signal;
use super::class::Class;
use super::property::Property;

/// All Lua objects can be cast to this.
#[derive(Clone, Debug)]
pub struct Object<'lua> {
    pub object: AnyUserData<'lua>
}

impl <'lua> From<AnyUserData<'lua>> for Object<'lua> {
    fn from(object: AnyUserData<'lua>) -> Self {
        Object { object }
    }
}

/// Construct a new object, used when using the default Objectable::new.
pub struct ObjectBuilder<'lua>{
    lua: &'lua Lua,
    object: Object<'lua>
}

impl <'lua> ObjectBuilder<'lua> {
    pub fn add_to_meta(self, new_meta: Table<'lua>) -> rlua::Result<Self> {
        let meta = self.object.table()?.get_metatable()
            .expect("Object had no meta table");
        for entry in new_meta.pairs::<rlua::Value, rlua::Value>() {
            let (key, value) = entry?;
            meta.set(key, value)?;
        }
        self.object.table()?.set_metatable(Some(meta));
        Ok(self)
    }

    #[allow(dead_code)]
    pub fn add_to_signals(self, name: String, func: rlua::Function)
                          -> rlua::Result<Self> {
        signal::connect_signal(self.lua, self.object.clone(), name, &[func])?;
        Ok(self)
    }

    pub fn handle_constructor_argument(self, args: Table) -> rlua::Result<Self> {
        let meta = self.object.table()?.get_metatable()
            .expect("Object had no meta table");
        let class = meta.get::<_, Table>("__class")?;
        let props = class.get::<_, Vec<Property>>("properties")?;

        // Handle all table entries that correspond to known properties,
        // silently ignore all other keys
        for pair in args.pairs() {
            let (key, value): (Value, Value) = pair?;
            if let rlua::Value::String(key) = key {
                if let Ok(key) = key.to_str() {
                    // Find the property
                    for prop in props.iter() {
                        if prop.name == key {
                            // Property exists and has a cb_new callback
                            if let Some(ref new) = prop.cb_new {
                                let _: () = new.bind(self.object.table()?.clone())?.call(value)?;
                                break;
                            }
                        }
                    }
                }
            }
        }
        Ok(self)
    }

    pub fn build(self) -> Object<'lua> {
        self.object
    }
}

/// Trait implemented by all objects that represent OO lua objects.
///
/// This trait allows casting an object gotten back from the Lua runtime
/// into a concrete object so that Rust can do things with it.
///
/// You can't do anything to the object until it has been converted into a
/// canonical form using this trait.
pub trait Objectable<'lua, T, S: UserData + Default + Display + Clone> {
    fn cast(obj: Object<'lua>) -> rlua::Result<T> {
        //let data = obj.table()?.get::<_, AnyUserData>("userdata")?;
        if obj.object.is::<S>()? {
            Ok(Self::_wrap(obj))
        } else {
            use rlua::Error::RuntimeError;
            Err(RuntimeError("Could not cast object to concrete type".into()))
        }
    }

    /// Given the internal object table, constructs a concrete object out
    /// of it.
    ///
    /// This is only used internally, which is why it's prefixed with a "_"
    /// Please do not use it outside of object.rs.
    fn _wrap(table: Object<'lua>) -> T;

    fn state(&self) -> rlua::Result<S> {
        Ok(self.get_object()?.clone())
    }

    fn set_state(&mut self, data: S) -> rlua::Result<()> {
        *self.get_object_mut()? = data;
        Ok(())
    }

    /// Gets the internal table for the concrete object.
    /// Used internally by cast, though there's nothing wrong with it being
    /// used outside of internal object use.
    //fn get_table(&self) -> Table<'lua>;

    fn get_object(&self) -> rlua::Result<S>;
    fn get_object_mut(&mut self) -> rlua::Result<::std::cell::RefMut<S>>;

    fn allocate(lua: &'lua Lua, class: Class) -> rlua::Result<ObjectBuilder<'lua>>
    {
        // This is the format of the lua objects in Way Cooler:
        /*
        -- This is the weird wrapper table, binding the user data and the object table together
        {
            -- Actual user data object, e.g the Rust object
            userdata = <user_data_here>,
            -- This is the object data table, e.g the thing that in Awesome
            -- is linked to the user data with lua_setuservalue
            data = {},
            -- That metatable that is allocated in the new macro for objects.
            -- It can access both the userdata and the lua data fields
            __metatable = {}
        }
        */
        let object = match S::default().to_lua(lua)? {
            rlua::Value::UserData(data) => data,
            _ => unreachable!()
        };
        // TODO Increment the instance count
        let wrapper_table = lua.create_table()?;
        //wrapper_table.set("userdata", object)?;
        let data_table = lua.create_table()?;
        wrapper_table.set("data", data_table)?;
        let meta = lua.create_table()?;
        meta.set("__class", class)?;
        meta.set("properties", Vec::<Property>::new().to_lua(lua)?)?;
        meta.set("signals", lua.create_table()?)?;
        meta.set("connect_signal",
                 lua.create_function(connect_signal)?)?;
        meta.set("disconnect_signal",
                 lua.create_function(disconnect_signal)?)?;
        meta.set("emit_signal", lua.create_function(emit_signal)?)?;
        meta.set("__index", meta.clone())?;
        meta.set("__tostring", lua.create_function(|_, data: AnyUserData| {
            Ok(format!("{}", data.borrow::<S>()?.clone()))
        })?)?;
        wrapper_table.set_metatable(Some(meta));
        object.set_user_value(wrapper_table)?;
        // TODO Emit new signal event
        let object = Object { object };
        Ok(ObjectBuilder { object, lua })
    }
}

impl <'lua> ToLua<'lua> for Object<'lua> {
    fn to_lua(self, _: &'lua Lua) -> rlua::Result<Value<'lua>> {
        Ok(Value::UserData(self.object))
    }
}

impl <'lua> Object<'lua> {
    pub fn signals(&self) -> rlua::Result<rlua::Table<'lua>> {
        self.table()?.get::<_, Table>("signals")
        //let obj: Object<'lua> = self.object.clone().into();
        //obj.table()?.get::<_, Table>("signals")
    }

    pub fn table(&self) -> rlua::Result<Table<'lua>> {
        self.object.get_user_value::<Table<'lua>>()
    }
}

/// Default indexing of an Awesome object.
///
/// Automatically looks up contents in meta table, so instead of overriding this
/// it's easier to just add the required data in the meta table.
pub fn default_index<'lua>(lua: &'lua Lua,
                           (obj, index): (AnyUserData<'lua>, Value<'lua>))
                           -> rlua::Result<Value<'lua>> {
    // Look up in metatable first
    let obj: Object<'lua> = obj.into();
    let obj_table = obj.table()?;
    let meta = obj_table.get_metatable()
        .expect("Object had no metatable");
    if meta.get::<_, Table>("__class").is_ok() {
        if let Ok(val) = meta.raw_get::<_, Value>(index.clone()) {
            match val {
                Value::Nil => {},
                val => return Ok(val)
            }
        }
    }
    let index = String::from_lua(index, lua)?;
    match index.as_str() {
        "valid" => {
            Ok(Value::Boolean(
                if let Ok(class) = meta.get::<_, Table>("__class") {
                    let class: Class = class.into();
                    class.checker()?
                        .map(|checker| checker(obj))
                        .unwrap_or(true)
                } else {
                    false
                }))
        },
        "data" => {
            obj_table.to_lua(lua)
        },
        index => {
            // Try see if there is a property of the class with the name
            if let Ok(class) = meta.get::<_, Table>("__class") {
                let props = class.get::<_, Vec<Property>>("properties")?;
                for prop in props {
                    if prop.name.as_str() == index {
                        // Property exists and has an index callback
                        if let Some(index) = prop.cb_index {
                            return index.call(obj)
                        }
                    }
                }
                match class.get::<_, Function>("__index_miss_handler") {
                    Ok(function) => {
                        return function.bind(obj_table)?.call(index)
                    },
                    Err(_) => {}
                }
            }
            // TODO property miss handler if index doesn't exst
            Ok(rlua::Value::Nil)
        }
    }
}

/// Default new indexing (assignment) of an Awesome object.
///
/// Automatically looks up contents in meta table, so instead of overriding this
/// it's easier to just add the required data in the meta table.
pub fn default_newindex<'lua>(_: &'lua Lua,
                              (obj, index, val):
                              (AnyUserData<'lua>, String, Value<'lua>))
                              -> rlua::Result<Value<'lua>> {
    // Look up in metatable first
    let obj: Object<'lua> = obj.into();
    let obj_table = obj.table()?;
    if let Some(meta) = obj_table.get_metatable() {
        if let Ok(val) = meta.raw_get::<_, Value>(index.clone()) {
            match val {
                Value::Nil => {},
                val => return Ok(val)
            }
        }
        let class = meta.get::<_, Table>("__class")?;
        let props = class.get::<_, Vec<Property>>("properties")?;
        for prop in props {
            if prop.name.as_str() == index {
                // Property exists and has a newindex callback
                if let Some(newindex) = prop.cb_newindex {
                    return newindex.bind(obj.clone())?.call(val)
                }
            }
        }
        match class.get::<_, Function>("__newindex_miss_handler") {
            Ok(function) => {
                return function.bind(obj_table)?.call((index, val))
            },
            Err(_) => {}
        }
        // TODO property miss handler if index doesn't exst
    }
    Ok(Value::Nil)
}

fn connect_signal(lua: &Lua, (obj, signal, func): (AnyUserData, String, Function))
                  -> rlua::Result<()> {
    signal::connect_signal(lua, obj.into(), signal, &[func])
}

fn disconnect_signal(lua: &Lua, (obj, signal): (AnyUserData, String))
                     -> rlua::Result<()> {
    signal::disconnect_signal(lua, obj.into(), signal)
}

fn emit_signal(lua: &Lua, (obj, signal, args): (AnyUserData, String, Value))
               -> rlua::Result<()> {
    signal::emit_signal(lua, obj.into(), signal, args)
}
<|endoftext|>


<commit_before>use std::io::{Result, Write};

use analysis::rust_type::bounds_rust_type;
use case::CaseExt;
use analysis;
use library;
use env::Env;
use nameutil;
use super::child_properties;
use super::function;
use super::general;
use super::properties;
use super::signal;
use super::trait_impls;
use super::trampoline;

pub fn generate(
    w: &mut Write,
    env: &Env,
    analysis: &analysis::object::Info,
    generate_display_trait: bool,
) -> Result<()> {
    try!(general::start_comments(w, &env.config));
    try!(general::uses(w, env, &analysis.imports));

    try!(general::define_object_type(
        w,
        env,
        &analysis.name,
        &analysis.c_type,
        &analysis.c_class_type.as_ref().map(|s| &s[..]),
        &analysis.rust_class_type.as_ref().map(|s| &s[..]),
        &analysis.get_type,
        analysis.is_interface,
        &analysis.supertypes,
    ));

    if need_generate_inherent(analysis) {
        try!(writeln!(w));
        try!(write!(w, "impl {} {{", analysis.name));
        for func_analysis in &analysis.constructors() {
            try!(function::generate(w, env, func_analysis, false, false, 1));
        }

        if !need_generate_trait(analysis) {
            for func_analysis in &analysis.methods() {
                try!(function::generate(w, env, func_analysis, false, false, 1));
            }

            for property in &analysis.properties {
                try!(properties::generate(w, env, property, false, false, 1));
            }

            for child_property in &analysis.child_properties {
                try!(child_properties::generate(
                    w,
                    env,
                    child_property,
                    false,
                    false,
                    1,
                ));
            }
        }

        for func_analysis in &analysis.functions() {
            try!(function::generate(w, env, func_analysis, false, false, 1));
        }

        if !need_generate_trait(analysis) {
            for signal_analysis in analysis
                .signals
                .iter()
                .chain(analysis.notify_signals.iter())
            {
                try!(signal::generate(
                    w,
                    env,
                    signal_analysis,
                    &analysis.trampolines,
                    false,
                    false,
                    1,
                ));
            }
        }

        try!(writeln!(w, "}}"));

        try!(general::declare_default_from_new(
            w,
            env,
            &analysis.name,
            &analysis.functions
        ));
    }

    try!(trait_impls::generate(
        w,
        &analysis.name,
        &analysis.functions,
        &analysis.specials,
        if need_generate_trait(analysis) {
            Some(&analysis.trait_name)
        } else {
            None
        },
    ));

    if analysis.concurrency != library::Concurrency::None {
        try!(writeln!(w));
    }

    if analysis.properties.iter().any(|property| property.construct || property.construct_only) {
        let mut methods = vec![];
        let mut properties = vec![];
        writeln!(w, "#[cfg(any(feature = \"builders\", feature = \"dox\"))]
pub struct {}Builder {{", analysis.name)?;
        for property in &analysis.properties {
            if (!property.is_get && property.construct) || property.construct_only {
                match bounds_rust_type(env, property.typ) {
                    Ok(type_string) => {
                        let attribute_type =
                            match type_string.as_str() {
                                "str" => "String",
                                _ => &type_string,
                            };
                        let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
                        if let Some(version) = property.version {
                            writeln!(w, "    #[cfg(any(feature = \"{}\", feature = \"dox\"))]", version.to_feature())?;
                        }
                        writeln!(w, "    {}: Option<{}>,", name, attribute_type)?;
                        let prefix =
                            if let Some(version) = property.version {
                                format!("    #[cfg(any(feature = \"{}\", feature = \"dox\"))]\n", version.to_feature())
                            }
                            else {
                                String::new()
                            };
                        let (type_string, conversion) =
                            match type_string.as_str() {
                                "str" => ("&str", ".to_string()"),
                                _ => (&*type_string, ""),
                            };
                        methods.push(format!("\n{}    pub fn {name}(mut self, {name}: {}) -> Self {{
        self.{name} = Some({name}{});
        self
    }}", prefix, type_string, conversion, name=name));
                        properties.push(property);
                    },
                    Err(_) => writeln!(w, "    //{}: /*Unknown type*/,", property.name)?,
                }
            }
        }
        writeln!(w, "}}\n

#[cfg(any(feature = \"builders\", feature = \"dox\"))]
impl {}Builder {{
    pub fn new() -> Self {{
        Self {{", analysis.name)?;
        for property in &properties {
            if let Some(version) = property.version {
                writeln!(w, "            #[cfg(any(feature = \"{}\", feature = \"dox\"))]", version.to_feature())?;
            }
            let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
            writeln!(w, "            {}: None,", name)?;
        }
        writeln!(w, "        }}
    }}

    pub fn build(self) -> {} {{
        let mut n_properties = 0;
        let mut property_names: Vec<CString> = vec![];
        let mut names = vec![];
        let mut values = vec![];", analysis.name)?;
        for property in &properties {
            let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
            if let Some(version) = property.version {
                writeln!(w,
"        #[cfg(any(feature = \"{}\", feature = \"dox\"))]
        {{", version.to_feature())?;
            }
            writeln!(w,
"        if let Some({property}) = self.{property} {{
            property_names.push(CString::new(\"{}\").unwrap());
            names.push(property_names[property_names.len() - 1].as_ptr());
            let property = {property}.to_value();
            values.push(property.into_raw());
            n_properties += 1;
        }}", property.name, property=name)?;
            if property.version.is_some() {
                writeln!(w, "        }}")?;
            }
        }
        writeln!(w,
"        unsafe {{
            crate::Object::from_glib_none(gobject_sys::g_object_new_with_properties(
                {}::static_type().to_glib(), n_properties, names.as_mut_ptr(), values.as_ptr())
            as *mut _).downcast().expect(\"downcast\")
        }}
    }}", analysis.name)?;
        for method in methods {
            writeln!(w, "{}", method)?;
        }
        writeln!(w, "}}")?;
    }

    match analysis.concurrency {
        library::Concurrency::Send | library::Concurrency::SendSync => {
            try!(writeln!(w, "unsafe impl Send for {} {{}}", analysis.name));
        }
        library::Concurrency::SendUnique => {
            if env.namespaces.is_glib_crate {
                try!(writeln!(w, "unsafe impl ::SendUnique for {} {{", analysis.name));
            } else {
                try!(writeln!(w, "unsafe impl glib::SendUnique for {} {{", analysis.name));
            }

            try!(writeln!(w, "    fn is_unique(&self) -> bool {{"));
            try!(writeln!(w, "        self.ref_count() == 1"));
            try!(writeln!(w, "    }}"));

            try!(writeln!(w, "}}"));
        },
        _ => (),
    }

    if let library::Concurrency::SendSync = analysis.concurrency {
        try!(writeln!(w, "unsafe impl Sync for {} {{}}", analysis.name));
    }

    if !analysis.final_type {
        try!(writeln!(w));
        try!(writeln!(w, "pub const NONE_{}: Option<&{}> = None;", analysis.name.to_snake().to_uppercase(), analysis.name));
    }

    if need_generate_trait(analysis) {
        try!(writeln!(w));
        try!(generate_trait(w, env, analysis));
    }

    if !analysis.trampolines.is_empty() {
        for trampoline in &analysis.trampolines {
            try!(trampoline::generate(
                w,
                env,
                trampoline,
                need_generate_trait(analysis),
                &analysis.name,
            ));
        }
    }

    if generate_display_trait {
        try!(writeln!(
            w,
            "\nimpl fmt::Display for {} {{",
            analysis.name,
        ));
        // Generate Display trait implementation.
        try!(writeln!(w, "\tfn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{\n\
                            \t\twrite!(f, \"{}\")\n\
                          \t}}\n\
                        }}", analysis.name));
    }

    Ok(())
}

fn generate_trait(
    w: &mut Write,
    env: &Env,
    analysis: &analysis::object::Info,
) -> Result<()> {
    try!(write!(w, "pub trait {}: 'static {{", analysis.trait_name));

    for func_analysis in &analysis.methods() {
        try!(function::generate(w, env, func_analysis, true, true, 1));
    }
    for property in &analysis.properties {
        try!(properties::generate(w, env, property, true, true, 1));
    }
    for child_property in &analysis.child_properties {
        try!(child_properties::generate(
            w,
            env,
            child_property,
            true,
            true,
            1,
        ));
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        try!(signal::generate(
            w,
            env,
            signal_analysis,
            &analysis.trampolines,
            true,
            true,
            1,
        ));
    }
    try!(writeln!(w, "}}"));

    try!(writeln!(w));
    try!(write!(
        w,
        "impl<O: IsA<{}>> {} for O {{",
        analysis.name,
        analysis.trait_name,
    ));

    for func_analysis in &analysis.methods() {
        try!(function::generate(w, env, func_analysis, true, false, 1));
    }
    for property in &analysis.properties {
        try!(properties::generate(w, env, property, true, false, 1));
    }
    for child_property in &analysis.child_properties {
        try!(child_properties::generate(
            w,
            env,
            child_property,
            true,
            false,
            1,
        ));
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        try!(signal::generate(
            w,
            env,
            signal_analysis,
            &analysis.trampolines,
            true,
            false,
            1,
        ));
    }
    try!(writeln!(w, "}}"));

    Ok(())
}

fn need_generate_inherent(analysis: &analysis::object::Info) -> bool {
    analysis.has_constructors || analysis.has_functions || !need_generate_trait(analysis)
}

fn need_generate_trait(analysis: &analysis::object::Info) -> bool {
    analysis.generate_trait
}

pub fn generate_reexports(
    env: &Env,
    analysis: &analysis::object::Info,
    module_name: &str,
    contents: &mut Vec<String>,
    traits: &mut Vec<String>,
) {
    let mut cfgs: Vec<String> = Vec::new();
    if let Some(cfg) = general::cfg_condition_string(&analysis.cfg_condition, false, 0) {
        cfgs.push(cfg);
    }
    if let Some(cfg) = general::version_condition_string(env, analysis.version, false, 0) {
        cfgs.push(cfg);
    }
    contents.push("".to_owned());
    contents.extend_from_slice(&cfgs);
    contents.push(format!("mod {};", module_name));
    contents.extend_from_slice(&cfgs);

    let none_type = if !analysis.final_type {
        format!(", NONE_{}", analysis.name.to_snake().to_uppercase())
    } else {
        String::new()
    };

    if let Some(ref class_name) = analysis.rust_class_type {
        contents.push(format!("pub use self::{}::{{{}, {}{}}};", module_name, analysis.name, class_name, none_type));
    } else {
        contents.push(format!("pub use self::{}::{{{}{}}};", module_name, analysis.name, none_type));
    }
    if need_generate_trait(analysis) {
        contents.extend_from_slice(&cfgs);
        contents.push(format!(
            "pub use self::{}::{};",
            module_name,
            analysis.trait_name
        ));
        for cfg in &cfgs {
            traits.push(format!("\t{}", cfg));
        }
        traits.push(format!("\tpub use super::{};", analysis.trait_name));
    }

    let mut generate_builder = false;
    for property in &analysis.properties {
        if (!property.is_get && property.construct) || property.construct_only {
            generate_builder = true;
            break;
        }
    }
    if generate_builder {
        contents.extend_from_slice(&cfgs);
        contents.push(format!("#[cfg(any(feature = \"builders\", feature = \"dox\"))]
pub use self::{}::{}Builder;", module_name, analysis.name));
    }
}
<commit_msg>Simplify generated code<commit_after>use std::io::{Result, Write};

use analysis::rust_type::bounds_rust_type;
use case::CaseExt;
use analysis;
use library;
use env::Env;
use nameutil;
use super::child_properties;
use super::function;
use super::general;
use super::properties;
use super::signal;
use super::trait_impls;
use super::trampoline;

pub fn generate(
    w: &mut Write,
    env: &Env,
    analysis: &analysis::object::Info,
    generate_display_trait: bool,
) -> Result<()> {
    try!(general::start_comments(w, &env.config));
    try!(general::uses(w, env, &analysis.imports));

    try!(general::define_object_type(
        w,
        env,
        &analysis.name,
        &analysis.c_type,
        &analysis.c_class_type.as_ref().map(|s| &s[..]),
        &analysis.rust_class_type.as_ref().map(|s| &s[..]),
        &analysis.get_type,
        analysis.is_interface,
        &analysis.supertypes,
    ));

    if need_generate_inherent(analysis) {
        try!(writeln!(w));
        try!(write!(w, "impl {} {{", analysis.name));
        for func_analysis in &analysis.constructors() {
            try!(function::generate(w, env, func_analysis, false, false, 1));
        }

        if !need_generate_trait(analysis) {
            for func_analysis in &analysis.methods() {
                try!(function::generate(w, env, func_analysis, false, false, 1));
            }

            for property in &analysis.properties {
                try!(properties::generate(w, env, property, false, false, 1));
            }

            for child_property in &analysis.child_properties {
                try!(child_properties::generate(
                    w,
                    env,
                    child_property,
                    false,
                    false,
                    1,
                ));
            }
        }

        for func_analysis in &analysis.functions() {
            try!(function::generate(w, env, func_analysis, false, false, 1));
        }

        if !need_generate_trait(analysis) {
            for signal_analysis in analysis
                .signals
                .iter()
                .chain(analysis.notify_signals.iter())
            {
                try!(signal::generate(
                    w,
                    env,
                    signal_analysis,
                    &analysis.trampolines,
                    false,
                    false,
                    1,
                ));
            }
        }

        try!(writeln!(w, "}}"));

        try!(general::declare_default_from_new(
            w,
            env,
            &analysis.name,
            &analysis.functions
        ));
    }

    try!(trait_impls::generate(
        w,
        &analysis.name,
        &analysis.functions,
        &analysis.specials,
        if need_generate_trait(analysis) {
            Some(&analysis.trait_name)
        } else {
            None
        },
    ));

    if analysis.concurrency != library::Concurrency::None {
        try!(writeln!(w));
    }

    if analysis.properties.iter().any(|property| property.construct || property.construct_only) {
        let mut methods = vec![];
        let mut properties = vec![];
        writeln!(w, "#[cfg(any(feature = \"builders\", feature = \"dox\"))]
pub struct {}Builder {{", analysis.name)?;
        for property in &analysis.properties {
            if (!property.is_get && property.construct) || property.construct_only {
                match bounds_rust_type(env, property.typ) {
                    Ok(type_string) => {
                        let attribute_type =
                            match type_string.as_str() {
                                "str" => "String",
                                _ => &type_string,
                            };
                        let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
                        if let Some(version) = property.version {
                            writeln!(w, "    #[cfg(any(feature = \"{}\", feature = \"dox\"))]", version.to_feature())?;
                        }
                        writeln!(w, "    {}: Option<{}>,", name, attribute_type)?;
                        let prefix =
                            if let Some(version) = property.version {
                                format!("    #[cfg(any(feature = \"{}\", feature = \"dox\"))]\n", version.to_feature())
                            }
                            else {
                                String::new()
                            };
                        let (type_string, conversion) =
                            match type_string.as_str() {
                                "str" => ("&str", ".to_string()"),
                                _ => (&*type_string, ""),
                            };
                        methods.push(format!("\n{}    pub fn {name}(mut self, {name}: {}) -> Self {{
        self.{name} = Some({name}{});
        self
    }}", prefix, type_string, conversion, name=name));
                        properties.push(property);
                    },
                    Err(_) => writeln!(w, "    //{}: /*Unknown type*/,", property.name)?,
                }
            }
        }
        writeln!(w, "}}\n

#[cfg(any(feature = \"builders\", feature = \"dox\"))]
impl {}Builder {{
    pub fn new() -> Self {{
        Self {{", analysis.name)?;
        for property in &properties {
            if let Some(version) = property.version {
                writeln!(w, "            #[cfg(any(feature = \"{}\", feature = \"dox\"))]", version.to_feature())?;
            }
            let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
            writeln!(w, "            {}: None,", name)?;
        }
        writeln!(w, "        }}
    }}

    pub fn build(self) -> {} {{
        let mut properties: Vec<(&str, &dyn ToValue)> = vec![];", analysis.name)?;
        for property in &properties {
            let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
            if let Some(version) = property.version {
                writeln!(w,
"        #[cfg(any(feature = \"{}\", feature = \"dox\"))]
        {{", version.to_feature())?;
            }
            writeln!(w,
"        if let Some(ref {property}) = self.{property} {{
            properties.push((\"{}\", {property}));
        }}", property = name)?;
            if property.version.is_some() {
                writeln!(w, "        }}")?;
            }
        }
        writeln!(w,
"        crate::Object::new({}::static_type(), &properties).expect(\"object new\").downcast().expect(\"downcast\")
    }}", analysis.name)?;
        for method in methods {
            writeln!(w, "{}", method)?;
        }
        writeln!(w, "}}")?;
    }

    match analysis.concurrency {
        library::Concurrency::Send | library::Concurrency::SendSync => {
            try!(writeln!(w, "unsafe impl Send for {} {{}}", analysis.name));
        }
        library::Concurrency::SendUnique => {
            if env.namespaces.is_glib_crate {
                try!(writeln!(w, "unsafe impl ::SendUnique for {} {{", analysis.name));
            } else {
                try!(writeln!(w, "unsafe impl glib::SendUnique for {} {{", analysis.name));
            }

            try!(writeln!(w, "    fn is_unique(&self) -> bool {{"));
            try!(writeln!(w, "        self.ref_count() == 1"));
            try!(writeln!(w, "    }}"));

            try!(writeln!(w, "}}"));
        },
        _ => (),
    }

    if let library::Concurrency::SendSync = analysis.concurrency {
        try!(writeln!(w, "unsafe impl Sync for {} {{}}", analysis.name));
    }

    if !analysis.final_type {
        try!(writeln!(w));
        try!(writeln!(w, "pub const NONE_{}: Option<&{}> = None;", analysis.name.to_snake().to_uppercase(), analysis.name));
    }

    if need_generate_trait(analysis) {
        try!(writeln!(w));
        try!(generate_trait(w, env, analysis));
    }

    if !analysis.trampolines.is_empty() {
        for trampoline in &analysis.trampolines {
            try!(trampoline::generate(
                w,
                env,
                trampoline,
                need_generate_trait(analysis),
                &analysis.name,
            ));
        }
    }

    if generate_display_trait {
        try!(writeln!(
            w,
            "\nimpl fmt::Display for {} {{",
            analysis.name,
        ));
        // Generate Display trait implementation.
        try!(writeln!(w, "\tfn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{\n\
                            \t\twrite!(f, \"{}\")\n\
                          \t}}\n\
                        }}", analysis.name));
    }

    Ok(())
}

fn generate_trait(
    w: &mut Write,
    env: &Env,
    analysis: &analysis::object::Info,
) -> Result<()> {
    try!(write!(w, "pub trait {}: 'static {{", analysis.trait_name));

    for func_analysis in &analysis.methods() {
        try!(function::generate(w, env, func_analysis, true, true, 1));
    }
    for property in &analysis.properties {
        try!(properties::generate(w, env, property, true, true, 1));
    }
    for child_property in &analysis.child_properties {
        try!(child_properties::generate(
            w,
            env,
            child_property,
            true,
            true,
            1,
        ));
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        try!(signal::generate(
            w,
            env,
            signal_analysis,
            &analysis.trampolines,
            true,
            true,
            1,
        ));
    }
    try!(writeln!(w, "}}"));

    try!(writeln!(w));
    try!(write!(
        w,
        "impl<O: IsA<{}>> {} for O {{",
        analysis.name,
        analysis.trait_name,
    ));

    for func_analysis in &analysis.methods() {
        try!(function::generate(w, env, func_analysis, true, false, 1));
    }
    for property in &analysis.properties {
        try!(properties::generate(w, env, property, true, false, 1));
    }
    for child_property in &analysis.child_properties {
        try!(child_properties::generate(
            w,
            env,
            child_property,
            true,
            false,
            1,
        ));
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        try!(signal::generate(
            w,
            env,
            signal_analysis,
            &analysis.trampolines,
            true,
            false,
            1,
        ));
    }
    try!(writeln!(w, "}}"));

    Ok(())
}

fn need_generate_inherent(analysis: &analysis::object::Info) -> bool {
    analysis.has_constructors || analysis.has_functions || !need_generate_trait(analysis)
}

fn need_generate_trait(analysis: &analysis::object::Info) -> bool {
    analysis.generate_trait
}

pub fn generate_reexports(
    env: &Env,
    analysis: &analysis::object::Info,
    module_name: &str,
    contents: &mut Vec<String>,
    traits: &mut Vec<String>,
) {
    let mut cfgs: Vec<String> = Vec::new();
    if let Some(cfg) = general::cfg_condition_string(&analysis.cfg_condition, false, 0) {
        cfgs.push(cfg);
    }
    if let Some(cfg) = general::version_condition_string(env, analysis.version, false, 0) {
        cfgs.push(cfg);
    }
    contents.push("".to_owned());
    contents.extend_from_slice(&cfgs);
    contents.push(format!("mod {};", module_name));
    contents.extend_from_slice(&cfgs);

    let none_type = if !analysis.final_type {
        format!(", NONE_{}", analysis.name.to_snake().to_uppercase())
    } else {
        String::new()
    };

    if let Some(ref class_name) = analysis.rust_class_type {
        contents.push(format!("pub use self::{}::{{{}, {}{}}};", module_name, analysis.name, class_name, none_type));
    } else {
        contents.push(format!("pub use self::{}::{{{}{}}};", module_name, analysis.name, none_type));
    }
    if need_generate_trait(analysis) {
        contents.extend_from_slice(&cfgs);
        contents.push(format!(
            "pub use self::{}::{};",
            module_name,
            analysis.trait_name
        ));
        for cfg in &cfgs {
            traits.push(format!("\t{}", cfg));
        }
        traits.push(format!("\tpub use super::{};", analysis.trait_name));
    }

    let mut generate_builder = false;
    for property in &analysis.properties {
        if (!property.is_get && property.construct) || property.construct_only {
            generate_builder = true;
            break;
        }
    }
    if generate_builder {
        contents.extend_from_slice(&cfgs);
        contents.push(format!("#[cfg(any(feature = \"builders\", feature = \"dox\"))]
pub use self::{}::{}Builder;", module_name, analysis.name));
    }
}
<|endoftext|>


<commit_before>//! Types for the Global Descriptor Table and segment selectors.

use bit_field::BitField;
use core::fmt;
use structures::tss::TaskStateSegment;
use PrivilegeLevel;

/// Specifies which element to load into a segment from
/// descriptor tables (i.e., is a index to LDT or GDT table
/// with some additional flags).
///
/// See Intel 3a, Section 3.4.2 "Segment Selectors"
pub struct SegmentSelector(pub u16);

impl SegmentSelector {
    /// Creates a new SegmentSelector
    ///
    /// # Arguments
    ///  * `index`: index in GDT or LDT array (not the offset)
    ///  * `rpl`: the requested privilege level
    pub const fn new(index: u16, rpl: PrivilegeLevel) -> SegmentSelector {
        SegmentSelector(index << 3 | (rpl as u16))
    }

    /// Returns the GDT index.
    pub fn index(&self) -> u16 {
        self.0 >> 3
    }

    /// Returns the requested privilege level.
    pub fn rpl(&self) -> PrivilegeLevel {
        PrivilegeLevel::from_u16(self.0.get_bits(0..2))
    }
}

impl fmt::Debug for SegmentSelector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut s = f.debug_struct("SegmentSelector");
        s.field("index", &self.index());
        s.field("rpl", &self.rpl());
        s.finish()
    }
}

pub struct GlobalDescriptorTable {
    table: [u64; 8],
    next_free: usize,
}

impl GlobalDescriptorTable {
    pub fn new() -> GlobalDescriptorTable {
        GlobalDescriptorTable {
            table: [0; 8],
            next_free: 1,
        }
    }

    pub fn add_entry(&mut self, entry: Descriptor) -> SegmentSelector {
        let index = match entry {
            Descriptor::UserSegment(value) => self.push(value),
            Descriptor::SystemSegment(value_low, value_high) => {
                let index = self.push(value_low);
                self.push(value_high);
                index
            }
        };
        SegmentSelector::new(index as u16, PrivilegeLevel::Ring0)
    }

    pub fn load(&'static self) {
        use core::mem::size_of;
        use instructions::tables::{lgdt, DescriptorTablePointer};

        let ptr = DescriptorTablePointer {
            base: self.table.as_ptr() as u64,
            limit: (self.table.len() * size_of::<u64>() - 1) as u16,
        };

        unsafe { lgdt(&ptr) };
    }

    fn push(&mut self, value: u64) -> usize {
        if self.next_free < self.table.len() {
            let index = self.next_free;
            self.table[index] = value;
            self.next_free += 1;
            index
        } else {
            panic!("GDT full");
        }
    }
}

pub enum Descriptor {
    UserSegment(u64),
    SystemSegment(u64, u64),
}

bitflags! {
    pub struct DescriptorFlags: u64 {
        const CONFORMING        = 1 << 42;
        const EXECUTABLE        = 1 << 43;
        const USER_SEGMENT      = 1 << 44;
        const PRESENT           = 1 << 47;
        const LONG_MODE         = 1 << 53;
    }
}

impl Descriptor {
    pub fn kernel_code_segment() -> Descriptor {
        use self::DescriptorFlags as Flags;

        let flags = Flags::USER_SEGMENT | Flags::PRESENT | Flags::EXECUTABLE | Flags::LONG_MODE;
        Descriptor::UserSegment(flags.bits())
    }

    pub fn tss_segment(tss: &'static TaskStateSegment) -> Descriptor {
        use self::DescriptorFlags as Flags;
        use bit_field::BitField;
        use core::mem::size_of;

        let ptr = tss as *const _ as u64;

        let mut low = Flags::PRESENT.bits();
        // base
        low.set_bits(16..40, ptr.get_bits(0..24));
        low.set_bits(56..64, ptr.get_bits(24..32));
        // limit (the `-1` in needed since the bound is inclusive)
        low.set_bits(0..16, (size_of::<TaskStateSegment>() - 1) as u64);
        // type (0b1001 = available 64-bit tss)
        low.set_bits(40..44, 0b1001);

        let mut high = 0;
        high.set_bits(0..32, ptr.get_bits(32..64));

        Descriptor::SystemSegment(low, high)
    }
}
<commit_msg>Make SegmentSelector repr(transparent) and derive Copy and Eq<commit_after>//! Types for the Global Descriptor Table and segment selectors.

use bit_field::BitField;
use core::fmt;
use structures::tss::TaskStateSegment;
use PrivilegeLevel;

/// Specifies which element to load into a segment from
/// descriptor tables (i.e., is a index to LDT or GDT table
/// with some additional flags).
///
/// See Intel 3a, Section 3.4.2 "Segment Selectors"
#[derive(Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
pub struct SegmentSelector(pub u16);

impl SegmentSelector {
    /// Creates a new SegmentSelector
    ///
    /// # Arguments
    ///  * `index`: index in GDT or LDT array (not the offset)
    ///  * `rpl`: the requested privilege level
    pub const fn new(index: u16, rpl: PrivilegeLevel) -> SegmentSelector {
        SegmentSelector(index << 3 | (rpl as u16))
    }

    /// Returns the GDT index.
    pub fn index(&self) -> u16 {
        self.0 >> 3
    }

    /// Returns the requested privilege level.
    pub fn rpl(&self) -> PrivilegeLevel {
        PrivilegeLevel::from_u16(self.0.get_bits(0..2))
    }
}

impl fmt::Debug for SegmentSelector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut s = f.debug_struct("SegmentSelector");
        s.field("index", &self.index());
        s.field("rpl", &self.rpl());
        s.finish()
    }
}

pub struct GlobalDescriptorTable {
    table: [u64; 8],
    next_free: usize,
}

impl GlobalDescriptorTable {
    pub fn new() -> GlobalDescriptorTable {
        GlobalDescriptorTable {
            table: [0; 8],
            next_free: 1,
        }
    }

    pub fn add_entry(&mut self, entry: Descriptor) -> SegmentSelector {
        let index = match entry {
            Descriptor::UserSegment(value) => self.push(value),
            Descriptor::SystemSegment(value_low, value_high) => {
                let index = self.push(value_low);
                self.push(value_high);
                index
            }
        };
        SegmentSelector::new(index as u16, PrivilegeLevel::Ring0)
    }

    pub fn load(&'static self) {
        use core::mem::size_of;
        use instructions::tables::{lgdt, DescriptorTablePointer};

        let ptr = DescriptorTablePointer {
            base: self.table.as_ptr() as u64,
            limit: (self.table.len() * size_of::<u64>() - 1) as u16,
        };

        unsafe { lgdt(&ptr) };
    }

    fn push(&mut self, value: u64) -> usize {
        if self.next_free < self.table.len() {
            let index = self.next_free;
            self.table[index] = value;
            self.next_free += 1;
            index
        } else {
            panic!("GDT full");
        }
    }
}

pub enum Descriptor {
    UserSegment(u64),
    SystemSegment(u64, u64),
}

bitflags! {
    pub struct DescriptorFlags: u64 {
        const CONFORMING        = 1 << 42;
        const EXECUTABLE        = 1 << 43;
        const USER_SEGMENT      = 1 << 44;
        const PRESENT           = 1 << 47;
        const LONG_MODE         = 1 << 53;
    }
}

impl Descriptor {
    pub fn kernel_code_segment() -> Descriptor {
        use self::DescriptorFlags as Flags;

        let flags = Flags::USER_SEGMENT | Flags::PRESENT | Flags::EXECUTABLE | Flags::LONG_MODE;
        Descriptor::UserSegment(flags.bits())
    }

    pub fn tss_segment(tss: &'static TaskStateSegment) -> Descriptor {
        use self::DescriptorFlags as Flags;
        use bit_field::BitField;
        use core::mem::size_of;

        let ptr = tss as *const _ as u64;

        let mut low = Flags::PRESENT.bits();
        // base
        low.set_bits(16..40, ptr.get_bits(0..24));
        low.set_bits(56..64, ptr.get_bits(24..32));
        // limit (the `-1` in needed since the bound is inclusive)
        low.set_bits(0..16, (size_of::<TaskStateSegment>() - 1) as u64);
        // type (0b1001 = available 64-bit tss)
        low.set_bits(40..44, 0b1001);

        let mut high = 0;
        high.set_bits(0..32, ptr.get_bits(32..64));

        Descriptor::SystemSegment(low, high)
    }
}
<|endoftext|>