<commit_before>use crate::common::Result;
use crate::config::Config;
use crate::state::{State, StateUpdate};
use async_std::io;
use async_std::prelude::*;
use async_std::task;
use std::fmt::{self, Write};
use termion::{clear, cursor, style};
use unicode_truncate::UnicodeTruncateStr;

const ALTERNATE_SCREEN: &'static str = csi!("?1049h");
const MAIN_SCREEN: &'static str = csi!("?1049l");

trait Component {
    fn render(&mut self, state: &State) -> Result<()>;
}

#[derive(Debug, Clone, Default)]
struct Prompt {
    symbol: String,
    query: String,
}

// TODO: From Config
impl Prompt {
    fn new(_config: &Config) -> Self {
        let symbol = "$".into();
        let query = "".into();

        Self { symbol, query }
    }
}

impl Component for Prompt {
    fn render(&mut self, state: &State) -> Result<()> {
        self.query = state.query();

        Ok(())
    }
}

impl fmt::Display for Prompt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} {}", self.symbol, self.query)
    }
}

#[derive(Debug, Clone, Default)]
struct Meter {
    current: usize,
    total: usize,
}

// TODO: From Config
impl Meter {
    fn new(_config: &Config) -> Self {
        let current = 0;
        let total = 0;

        Self { current, total }
    }
}

impl Component for Meter {
    fn render(&mut self, state: &State) -> Result<()> {
        self.current = state.matches().len();
        self.total = state.pool_len();

        Ok(())
    }
}

impl fmt::Display for Meter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "  {}/{}", self.current, self.total)
    }
}

#[derive(Debug, Clone)]
enum Mode {
    Full,
    Inline(usize),
}

impl Mode {
    pub fn setup(&self) -> Option<String> {
        let setup = match self {
            Self::Full => format!("{}{}", ALTERNATE_SCREEN, cursor::Goto(1, 1)),
            Self::Inline(height) => {
                let room = std::iter::repeat("\n")
                    .take(*height)
                    .collect::<Vec<&str>>()
                    .join("");

                let up = *height as u16;

                format!("{}{}\r", room, cursor::Up(up))
            }
        };

        Some(setup)
    }

    pub fn teardown(&self) -> Option<String> {
        let teardown = match self {
            Self::Full => MAIN_SCREEN.to_string(),
            Self::Inline(_) => format!("{}{}{}", clear::CurrentLine, clear::AfterCursor, "\r"),
        };

        Some(teardown)
    }
}

#[derive(Debug, Clone)]
pub struct Layout<W: io::Write + Send + Unpin + 'static> {
    size: (usize, usize),
    offset: usize,
    mode: Mode,
    writer: W,

    prompt: Prompt,
    meter: Meter,
}

impl<W: io::Write + Send + Unpin + 'static> Layout<W> {
    pub async fn new(config: &Config, writer: W) -> Result<Self> {
        let size = config.screen.size;
        let offset = 0;
        let mode = if config.screen.full {
            Mode::Full
        } else {
            let (_, height) = size;
            Mode::Inline(height)
        };

        let prompt = Prompt::new(config);
        let meter = Meter::new(config);

        let mut layout = Self {
            size,
            offset,
            mode,
            writer,
            prompt,
            meter,
        };

        if let Some(setup) = layout.mode.setup() {
            layout.write(&setup).await?;
        }

        Ok(layout)
    }

    pub async fn write<D: std::fmt::Display>(&mut self, display: &D) -> Result<()> {
        let rendered = format!("{}", display,);
        self.writer.write_all(rendered.as_bytes()).await?;
        self.writer.flush().await?;

        Ok(())
    }

    pub async fn render(&mut self, state: &State) -> Result<()> {
        let d = self.draw(state)?;

        self.write(&d).await?;

        Ok(())
    }

    pub fn draw(&mut self, state: &State) -> Result<String> {
        let mut display = String::new();

        match state.last_update() {
            StateUpdate::Query => {
                let prompt = self.draw_prompt(&state)?;
                write!(&mut display, "{}", prompt)?;
            }
            _ => {
                let list = self.draw_list(&state)?;
                let prompt = self.draw_prompt(&state)?;
                write!(&mut display, "{}{}", list, prompt)?;
            }
        }

        Ok(display)
    }

    fn draw_prompt(&mut self, state: &State) -> Result<String> {
        self.prompt.render(state)?;
        let prompt = format!("{}\r{}", clear::CurrentLine, self.prompt);

        Ok(prompt)
    }

    fn draw_list(&mut self, state: &State) -> Result<String> {
        let mut display = String::new();

        self.meter.render(state)?;
        let counter = format!(
            "{}  {}",
            clear::CurrentLine,
            self.meter
        );

        let invert = format!("{}", style::Invert);
        let no_invert = format!("{}", style::NoInvert);

        let (width, _) = self.size;
        let line_len = width - 2;
        let (offset, lines) = self.scroll(&state);
        let mut list: Vec<String> = state
            .matches()
            .iter()
            .cloned()
            .enumerate()
            .skip(offset)
            .take(lines)
            .map(|(idx, c)| (idx, c.text))
            .map(|(index, candidate)| {
                let (truncated, _) = candidate.unicode_truncate(line_len);
                if index == state.selection_idx() {
                    format!(
                        "{}{}> {}{}",
                        clear::CurrentLine,
                        invert,
                        truncated,
                        no_invert
                    )
                } else {
                    format!("{}  {}", clear::CurrentLine, truncated)
                }
            })
            .collect();

        // The counter is another element of the list
        list.insert(0, counter);

        write!(
            &mut display,
            "{}\r{}{}{}",
            cursor::Down(1),
            list.join("\n"),
            clear::AfterCursor,
            // We always need to reprint the prompt after
            // going up to set the cursor in the last
            // position
            cursor::Up(list.len() as u16),
        )?;

        Ok(display)
    }

    fn scroll(&mut self, state: &State) -> (usize, usize) {
        let (_, height) = self.size;
        let lines_len = height - 2;

        let selection = state.selection_idx();

        let top_position = self.offset;
        let last_position = (lines_len + self.offset) - 1;

        if selection > last_position {
            self.offset += selection - last_position;
        } else if selection < top_position {
            self.offset -= top_position - selection;
        }

        (self.offset, lines_len)
    }
}

impl<W: io::Write + Send + Unpin + 'static> Drop for Layout<W> {
    fn drop(&mut self) {
        task::block_on(async {
            if let Some(teardown) = self.mode.teardown() {
                self.write(&teardown)
                    .await
                    .expect("Error writing to output");
            }
        });
    }
}
<commit_msg>Basic List UI component<commit_after>use crate::common::Result;
use crate::config::Config;
use crate::state::{State, StateUpdate};
use async_std::io;
use async_std::prelude::*;
use async_std::task;
use std::fmt::{self, Write};
use termion::{clear, cursor, style};
use unicode_truncate::UnicodeTruncateStr;

const ALTERNATE_SCREEN: &'static str = csi!("?1049h");
const MAIN_SCREEN: &'static str = csi!("?1049l");

trait Component {
    fn render(&mut self, state: &State) -> Result<()>;
}

#[derive(Debug, Clone, Default)]
struct Prompt {
    symbol: String,
    query: String,
}

// TODO: From Config
impl Prompt {
    fn new(_config: &Config) -> Self {
        let symbol = "$".into();
        let query = "".into();

        Self { symbol, query }
    }
}

impl Component for Prompt {
    fn render(&mut self, state: &State) -> Result<()> {
        self.query = state.query();

        Ok(())
    }
}

impl fmt::Display for Prompt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} {}", self.symbol, self.query)
    }
}

#[derive(Debug, Clone, Default)]
struct Meter {
    current: usize,
    total: usize,
}

// TODO: From Config
impl Meter {
    fn new(_config: &Config) -> Self {
        let current = 0;
        let total = 0;

        Self { current, total }
    }
}

impl Component for Meter {
    fn render(&mut self, state: &State) -> Result<()> {
        self.current = state.matches().len();
        self.total = state.pool_len();

        Ok(())
    }
}

impl fmt::Display for Meter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "  {}/{}", self.current, self.total)
    }
}

#[derive(Debug, Clone, Default)]
struct List {
    size: (usize, usize),
    offset: usize,
    items: Vec<String>,
}

// TODO: From Config
impl List {
    fn new(config: &Config) -> Self {
        let size = config.screen.size;
        let offset = 0;
        let items = vec![];

        Self {
            size,
            offset,
            items,
        }
    }

    pub fn len(&self) -> usize {
        self.items.len()
    }

    fn scroll(&mut self, state: &State) -> (usize, usize) {
        let (_, height) = self.size;
        let len = height - 2;

        let selection = state.selection_idx();

        let top_position = self.offset;
        let last_position = (len + self.offset) - 1;

        // cycle through the list
        if selection > last_position {
            self.offset += selection - last_position;
        } else if selection < top_position {
            self.offset -= top_position - selection;
        }

        (self.offset, len)
    }
}

impl Component for List {
    fn render(&mut self, state: &State) -> Result<()> {
        let invert = format!("{}", style::Invert);
        let no_invert = format!("{}", style::NoInvert);

        let (width, _) = self.size;
        let line_len = width - 2;
        let (offset, lines) = self.scroll(state);

        self.items = state
            .matches()
            .iter()
            .cloned()
            .enumerate()
            .skip(offset)
            .take(lines)
            .map(|(idx, c)| (idx, c.text))
            .map(|(index, candidate)| {
                let (truncated, _) = candidate.unicode_truncate(line_len);
                if index == state.selection_idx() {
                    format!(
                        "{}{}> {}{}",
                        clear::CurrentLine,
                        invert,
                        truncated,
                        no_invert
                    )
                } else {
                    format!("{}  {}", clear::CurrentLine, truncated)
                }
            })
            .collect();

        Ok(())
    }
}

impl fmt::Display for List {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.items.join("\n"))
    }
}

#[derive(Debug, Clone)]
enum Mode {
    Full,
    Inline(usize),
}

impl Mode {
    pub fn setup(&self) -> Option<String> {
        let setup = match self {
            Self::Full => format!("{}{}", ALTERNATE_SCREEN, cursor::Goto(1, 1)),
            Self::Inline(height) => {
                let room = std::iter::repeat("\n")
                    .take(*height)
                    .collect::<Vec<&str>>()
                    .join("");

                let up = *height as u16;

                format!("{}{}\r", room, cursor::Up(up))
            }
        };

        Some(setup)
    }

    pub fn teardown(&self) -> Option<String> {
        let teardown = match self {
            Self::Full => MAIN_SCREEN.to_string(),
            Self::Inline(_) => format!("{}{}{}", clear::CurrentLine, clear::AfterCursor, "\r"),
        };

        Some(teardown)
    }
}

#[derive(Debug, Clone)]
pub struct Layout<W: io::Write + Send + Unpin + 'static> {
    mode: Mode,
    writer: W,
    prompt: Prompt,
    meter: Meter,
    list: List,
}

impl<W: io::Write + Send + Unpin + 'static> Layout<W> {
    pub async fn new(config: &Config, writer: W) -> Result<Self> {
        let size = config.screen.size;
        let mode = if config.screen.full {
            Mode::Full
        } else {
            let (_, height) = size;
            Mode::Inline(height)
        };

        let prompt = Prompt::new(config);
        let meter = Meter::new(config);
        let list = List::new(config);

        let mut layout = Self {
            mode,
            writer,
            prompt,
            meter,
            list,
        };

        if let Some(setup) = layout.mode.setup() {
            layout.write(&setup).await?;
        }

        Ok(layout)
    }

    pub async fn write<D: std::fmt::Display>(&mut self, display: &D) -> Result<()> {
        let rendered = format!("{}", display,);
        self.writer.write_all(rendered.as_bytes()).await?;
        self.writer.flush().await?;

        Ok(())
    }

    pub async fn render(&mut self, state: &State) -> Result<()> {
        let d = self.draw(state)?;

        self.write(&d).await?;

        Ok(())
    }

    pub fn draw(&mut self, state: &State) -> Result<String> {
        let mut display = String::new();

        match state.last_update() {
            StateUpdate::Query => {
                let prompt = self.draw_prompt(&state)?;
                write!(&mut display, "{}", prompt)?;
            }
            _ => {
                let list = self.draw_list(&state)?;
                let prompt = self.draw_prompt(&state)?;
                write!(&mut display, "{}{}", list, prompt)?;
            }
        }

        Ok(display)
    }

    fn draw_prompt(&mut self, state: &State) -> Result<String> {
        self.prompt.render(state)?;
        let prompt = format!("{}\r{}", clear::CurrentLine, self.prompt);

        Ok(prompt)
    }

    fn draw_list(&mut self, state: &State) -> Result<String> {
        let mut display = String::new();

        self.meter.render(state)?;
        self.list.render(state)?;

        let list_len = self.list.len() as u16;

        // Only add a new line if we are going to print
        // items
        let meter_separator = if list_len == 0 { "" } else { "\n" };

        write!(
            &mut display,
            "{}{}\r{}{}{}{}{}",
            cursor::Down(1),
            clear::CurrentLine,
            self.meter,
            meter_separator,
            self.list,
            clear::AfterCursor,
            // We always need to reprint the prompt after
            // going up to set the cursor in the last
            // position
            cursor::Up(list_len + 1),
        )?;

        Ok(display)
    }
}

impl<W: io::Write + Send + Unpin + 'static> Drop for Layout<W> {
    fn drop(&mut self) {
        task::block_on(async {
            if let Some(teardown) = self.mode.teardown() {
                self.write(&teardown)
                    .await
                    .expect("Error writing to output");
            }
        });
    }
}
<|endoftext|>


<commit_before>extern crate rustbox;

use self::rustbox::{Color, Style, RustBox};

use super::board::{Board, HEIGHT, WIDTH};
use super::tetromino::{Tetromino, TetrominoType};
use super::window::Window;

// Default scaling factor for the board
const SCALE: usize = 2;

// Default values for styling terminal output
const DEFAULT_STYLE: Style = rustbox::RB_NORMAL;
const DEFAULT_FG: Color = Color::White;
const DEFAULT_BG: Color = Color::Black;

// Default character and color pairs
const I: (char, Color) = ('@', Color::Cyan);
const J: (char, Color) = ('#', Color::Blue);
const L: (char, Color) = ('&', Color::White);
const O: (char, Color) = ('%', Color::Yellow);
const S: (char, Color) = ('$', Color::Green);
const T: (char, Color) = ('*', Color::Magenta);
const Z: (char, Color) = ('O', Color::Red);

/// A collection of Window structs representing the user interface
pub struct Ui<'a> {
    board: Window<'a>,
    score: Window<'a>,
    level: Window<'a>,
    lines: Window<'a>,
    next_piece: Window<'a>,
}

impl<'a> Ui<'a> {

    /// Initializes a new Ui struct
    pub fn new(rb: &'a RustBox) -> Self {
        Ui {
            board: Window::new(0, 0, (11 * SCALE) - 1, 21, rb),
            score: Window::new(12 * SCALE, 0, 11, 2, rb),
            level: Window::new(12 * SCALE, 4, 11, 2, rb),
            lines: Window::new(12 * SCALE, 8, 11, 2, rb),
            next_piece: Window::new(12 * SCALE, 12, 11, 6, rb)
        }
    }

    /// Setup the default elements of the user interface
    pub fn setup(&self) {

        // Print out the borders for the player stats
        self.board.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
        self.score.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
        self.level.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
        self.lines.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
        self.next_piece.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
    }

    /// Print the state of the board
    pub fn print_board(&self, board: &Board) {

        // Start at 2 because only 20 of the board's rows should be displayed
        for y in 2..HEIGHT {
            for x in 0..WIDTH {
                match board.field()[y][x] {

                    // When printing the board, offset x and y to compensate
                    // for the Window's borders and showing only 20 rows
                    Some(ref mino) =>  {
                        let (rune, color) = self.get_style(mino);
                        self.board.print_char((x * SCALE) + 1, y - 1, DEFAULT_STYLE, color, DEFAULT_BG, rune);
                        self.board.print_char((x * SCALE) + 2, y - 1, DEFAULT_STYLE, color, DEFAULT_BG, rune);
                    }

                    None => {
                        self.board.print_char((x * SCALE) + 1, y - 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, ' ');
                        self.board.print_char((x * SCALE) + 2, y - 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, '.');
                    },
                }
            }
        }
    }

    /// Get the print style for a Mino
    fn get_style(&self, mino: &TetrominoType) -> (char, Color) {
        match *mino {
            TetrominoType::I => I,
            TetrominoType::J => J,
            TetrominoType::L => L,
            TetrominoType::O => O,
            TetrominoType::S => S,
            TetrominoType::T => T,
            TetrominoType::Z => Z,
        }
    }

    /// Print the next Tetromino
    pub fn print_next(&self, tetromino: Tetromino) {
        for i in 1..self.next_piece.w {
            for j in 1..self.next_piece.h {
                self.next_piece.print_char(i, j, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, ' ');
            }
        }

        for &mino in tetromino.minos().iter() {
            let (rune, color) = self.get_style(&tetromino.tetromino_type());
            self.next_piece.print_char((mino.x + 4) as usize, (mino.y + 2) as usize, DEFAULT_STYLE, color, DEFAULT_BG, rune);
        }
    }

    /// Print the player's score
    pub fn print_score(&self, score: usize) {
        self.score.print(1, 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, &format!("{:}", score));
    }

    /// Print the difficulty level
    pub fn print_level(&self, level: usize) {
        self.level.print(1, 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, &format!("{:}", level));
    }

    /// Print the number of lines cleared
    pub fn print_lines(&self, lines: usize) {
        self.lines.print(1, 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, &format!("{:}", lines));
    }
}<commit_msg>Use '■' as the only ASCII character for Minos<commit_after>extern crate rustbox;

use self::rustbox::{Color, Style, RustBox};

use super::board::{Board, HEIGHT, WIDTH};
use super::tetromino::{Tetromino, TetrominoType};
use super::window::Window;

// Default scaling factor for the board
const SCALE: usize = 2;

// Default values for styling terminal output
const DEFAULT_STYLE: Style = rustbox::RB_NORMAL;
const DEFAULT_FG: Color = Color::White;
const DEFAULT_BG: Color = Color::Black;

// Default colors for Tetrominos
const I: Color = Color::Cyan;
const J: Color = Color::Blue;
const L: Color = Color::White;
const O: Color = Color::Yellow;
const S: Color = Color::Green;
const T: Color = Color::Magenta;
const Z: Color = Color::Red;

/// A collection of Window structs representing the user interface
pub struct Ui<'a> {
    board: Window<'a>,
    score: Window<'a>,
    level: Window<'a>,
    lines: Window<'a>,
    next_piece: Window<'a>,
}

impl<'a> Ui<'a> {

    /// Initializes a new Ui struct
    pub fn new(rb: &'a RustBox) -> Self {
        Ui {
            board: Window::new(0, 0, (11 * SCALE) - 1, 21, rb),
            score: Window::new(12 * SCALE, 0, 11, 2, rb),
            level: Window::new(12 * SCALE, 4, 11, 2, rb),
            lines: Window::new(12 * SCALE, 8, 11, 2, rb),
            next_piece: Window::new(12 * SCALE, 12, 11, 6, rb)
        }
    }

    /// Setup the default elements of the user interface
    pub fn setup(&self) {

        // Print out the borders for the player stats
        self.board.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
        self.score.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
        self.level.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
        self.lines.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
        self.next_piece.print_borders(DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG);
    }

    /// Print the state of the board
    pub fn print_board(&self, board: &Board) {

        // Start at 2 because only 20 of the board's rows should be displayed
        for y in 2..HEIGHT {
            for x in 0..WIDTH {
                match board.field()[y][x] {

                    // When printing the board, offset x and y to compensate
                    // for the Window's borders and showing only 20 rows
                    Some(ref mino) =>  {
                        let color = self.get_tetromino_color(mino);
                        self.board.print_char((x * SCALE) + 1, y - 1, DEFAULT_STYLE, color, DEFAULT_BG, '■');
                        self.board.print_char((x * SCALE) + 2, y - 1, DEFAULT_STYLE, color, DEFAULT_BG, '■');
                    }

                    None => {
                        self.board.print_char((x * SCALE) + 1, y - 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, ' ');
                        self.board.print_char((x * SCALE) + 2, y - 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, '.');
                    },
                }
            }
        }
    }

    /// Get the color associated with a TetrominoType
    fn get_tetromino_color(&self, tetromino_type: &TetrominoType) -> Color {
        match tetromino_type {
            &TetrominoType::I => I,
            &TetrominoType::J => J,
            &TetrominoType::L => L,
            &TetrominoType::O => O,
            &TetrominoType::S => S,
            &TetrominoType::T => T,
            &TetrominoType::Z => Z,
        }
    }

    /// Print the next Tetromino
    pub fn print_next(&self, tetromino: Tetromino) {
        for i in 1..self.next_piece.w {
            for j in 1..self.next_piece.h {
                self.next_piece.print_char(i, j, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, ' ');
            }
        }

        for &mino in tetromino.minos().iter() {
            let color = self.get_tetromino_color(&tetromino.tetromino_type());
            self.next_piece.print_char((mino.x + 4) as usize, (mino.y + 2) as usize, DEFAULT_STYLE, color, DEFAULT_BG, '■');
        }
    }

    /// Print the player's score
    pub fn print_score(&self, score: usize) {
        self.score.print(1, 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, &format!("{:}", score));
    }

    /// Print the difficulty level
    pub fn print_level(&self, level: usize) {
        self.level.print(1, 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, &format!("{:}", level));
    }

    /// Print the number of lines cleared
    pub fn print_lines(&self, lines: usize) {
        self.lines.print(1, 1, DEFAULT_STYLE, DEFAULT_FG, DEFAULT_BG, &format!("{:}", lines));
    }
}<|endoftext|>


<commit_before>use std::io;
use std::io::Read;
use command::Command;

type Cell = u8;

pub struct VM {
    pub cells: [Cell; 30_000],
    pub data_pointer: usize,
}

impl VM {
    pub fn new() -> VM {
        VM { cells: [0; 30_000], data_pointer: 0 }
    }

    pub fn output(&self) {
        print!("{}", self.cells[self.data_pointer] as char);
    }

    pub fn input(&mut self) {
        match io::stdin().bytes().next() {
            Some(res) => match res {
                Ok(value) => self.cells[self.data_pointer] = value,
                Err(error) => panic!(error),
            },
            None => {}
        }
        println!("");
    }

    pub fn increment(&mut self) {
        self.cells[self.data_pointer] = self.cells[self.data_pointer].checked_add(1).unwrap_or(0);
    }

    pub fn decrement(&mut self) {
        self.cells[self.data_pointer] = self.cells[self.data_pointer]
            .checked_sub(1)
            .unwrap_or(Cell::max_value());
    }

    pub fn right_shift(&mut self) {
        self.data_pointer += 1;
    }

    pub fn left_shift(&mut self) {
        self.data_pointer -= 1;
    }

    pub fn apply(&mut self, command: Command) {
        match command {
            Command::Output => self.output(),
            Command::Input => self.input(),
            Command::Increment => self.increment(),
            Command::Decrement => self.decrement(),
            Command::RightShift => self.right_shift(),
            Command::LeftShift => self.left_shift(),
            _ => {},
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_new() {
        let vm = VM::new();
        assert!(vm.cells[0] == 0);
        assert!(vm.cells[30_000 - 1] == 0);
        assert!(vm.cells[0] == 0);
    }

    #[test]
    fn test_increment() {
        let mut vm = VM::new();
        vm.increment();
        assert!(vm.cells[0] == 1);
    }

    #[test]
    fn test_decrement() {
        const STARTING_VALUE: u8 = 2;

        let mut vm = VM::new();
        vm.cells[0] = STARTING_VALUE;
        vm.decrement();
        assert!(vm.cells[0] == STARTING_VALUE - 1);
    }

    #[test]
    fn test_right_shift() {
        let mut vm = VM::new();
        vm.right_shift();
        assert!(vm.data_pointer == 1);
    }
}
<commit_msg>Add unit test for VM.left_shift()<commit_after>use std::io;
use std::io::Read;
use command::Command;

type Cell = u8;

pub struct VM {
    pub cells: [Cell; 30_000],
    pub data_pointer: usize,
}

impl VM {
    pub fn new() -> VM {
        VM { cells: [0; 30_000], data_pointer: 0 }
    }

    pub fn output(&self) {
        print!("{}", self.cells[self.data_pointer] as char);
    }

    pub fn input(&mut self) {
        match io::stdin().bytes().next() {
            Some(res) => match res {
                Ok(value) => self.cells[self.data_pointer] = value,
                Err(error) => panic!(error),
            },
            None => {}
        }
        println!("");
    }

    pub fn increment(&mut self) {
        self.cells[self.data_pointer] = self.cells[self.data_pointer].checked_add(1).unwrap_or(0);
    }

    pub fn decrement(&mut self) {
        self.cells[self.data_pointer] = self.cells[self.data_pointer]
            .checked_sub(1)
            .unwrap_or(Cell::max_value());
    }

    pub fn right_shift(&mut self) {
        self.data_pointer += 1;
    }

    pub fn left_shift(&mut self) {
        self.data_pointer -= 1;
    }

    pub fn apply(&mut self, command: Command) {
        match command {
            Command::Output => self.output(),
            Command::Input => self.input(),
            Command::Increment => self.increment(),
            Command::Decrement => self.decrement(),
            Command::RightShift => self.right_shift(),
            Command::LeftShift => self.left_shift(),
            _ => {},
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_new() {
        let vm = VM::new();
        assert!(vm.cells[0] == 0);
        assert!(vm.cells[30_000 - 1] == 0);
        assert!(vm.cells[0] == 0);
    }

    #[test]
    fn test_increment() {
        let mut vm = VM::new();
        vm.increment();
        assert!(vm.cells[0] == 1);
    }

    #[test]
    fn test_decrement() {
        const STARTING_VALUE: u8 = 2;

        let mut vm = VM::new();
        vm.cells[0] = STARTING_VALUE;
        vm.decrement();
        assert!(vm.cells[0] == STARTING_VALUE - 1);
    }

    #[test]
    fn test_right_shift() {
        let mut vm = VM::new();
        vm.right_shift();
        assert!(vm.data_pointer == 1);
    }

    #[test]
    fn test_left_shift() {
        const STARTING_POSITION: usize = 2;

        let mut vm = VM::new();
        vm.data_pointer  = STARTING_POSITION;
        vm.left_shift();
        assert!(vm.data_pointer == 1);
    }
}
<|endoftext|>


<commit_before>use CL::*;
use hl::*;
use error::check;

// These are basically types that can be safely memcpyed.
trait VectorType {}

impl int: VectorType;
impl i32: VectorType;
impl u32: VectorType;
impl float: VectorType;
impl f64: VectorType;
impl f32: VectorType;

pub struct Vector<T: VectorType> {
    cl_buffer: cl_mem,
    size: uint,
    context: @ComputeContext,

    drop {
        clReleaseMemObject(self.cl_buffer);
    }
}

pub impl<T: VectorType> Vector<T> {
    static pub fn from_vec(ctx: @ComputeContext, v: &[const T]) -> Vector<T> {
        do vec::as_const_buf(v) |p, len| {
            let mut status = 0;
            let byte_size = len * sys::size_of::<T>() as libc::size_t;

            let buf = clCreateBuffer(ctx.ctx.ctx,
                                     CL_MEM_READ_WRITE | CL_MEM_COPY_HOST_PTR,
                                     byte_size,
                                     p as *libc::c_void,
                                     ptr::addr_of(&status));
            check(status, "Could not allocate buffer");

            //let status = clEnqueueWriteBuffer(
            //    ctx.q, buf, CL_TRUE, 0, byte_size, p as *libc::c_void,
            //    0, ptr::null(), ptr::null());

            Vector {
                cl_buffer: buf,
                size: len,
                context: ctx,
            }
        }
    }

    pub fn to_vec(self) -> ~[T] unsafe {
        let mut result = ~[];
        vec::reserve(&mut result, self.size);
        vec::raw::set_len(&mut result, self.size);
        do vec::as_imm_buf(result) |p, len| {
            clEnqueueReadBuffer(
                self.context.q.cqueue, self.cl_buffer, CL_TRUE, 0,
                len * sys::size_of::<T>() as libc::size_t,
                p as *libc::c_void, 0, ptr::null(), ptr::null());

        }
        move result
    }
}

impl<T: VectorType> Vector<T>: ::hl::KernelArg {
    pure fn get_value(&self) -> (libc::size_t, *libc::c_void) {
        (sys::size_of::<cl_mem>() as libc::size_t,
         ptr::addr_of(&self.cl_buffer) as *libc::c_void)
    }
}

#[cfg(test)]
mod test {
    macro_rules! expect (
        ($test: expr, $expected: expr) => ({
            let test = $test;
            let expected = $expected;
            if test != expected {
                fail fmt!("Test failure in %s: expected %?, got %?",
                          stringify!($test),
                          expected, test)
            }
        })
    )

    #[test]
    fn gpu_vector() {
        let ctx = create_compute_context();

        let x = ~[1, 2, 3, 4, 5];
        let gx = Vector::from_vec(ctx, x);
        let y = gx.to_vec();
        expect!(x, y);
    }
}
<commit_msg>Making VectorType public<commit_after>use CL::*;
use hl::*;
use error::check;

// These are basically types that can be safely memcpyed.
pub trait VectorType {}

impl int: VectorType;
impl i32: VectorType;
impl u32: VectorType;
impl float: VectorType;
impl f64: VectorType;
impl f32: VectorType;

pub struct Vector<T: VectorType> {
    cl_buffer: cl_mem,
    size: uint,
    context: @ComputeContext,

    drop {
        clReleaseMemObject(self.cl_buffer);
    }
}

pub impl<T: VectorType> Vector<T> {
    static pub fn from_vec(ctx: @ComputeContext, v: &[const T]) -> Vector<T> {
        do vec::as_const_buf(v) |p, len| {
            let mut status = 0;
            let byte_size = len * sys::size_of::<T>() as libc::size_t;

            let buf = clCreateBuffer(ctx.ctx.ctx,
                                     CL_MEM_READ_WRITE | CL_MEM_COPY_HOST_PTR,
                                     byte_size,
                                     p as *libc::c_void,
                                     ptr::addr_of(&status));
            check(status, "Could not allocate buffer");

            //let status = clEnqueueWriteBuffer(
            //    ctx.q, buf, CL_TRUE, 0, byte_size, p as *libc::c_void,
            //    0, ptr::null(), ptr::null());

            Vector {
                cl_buffer: buf,
                size: len,
                context: ctx,
            }
        }
    }

    pub fn to_vec(self) -> ~[T] unsafe {
        let mut result = ~[];
        vec::reserve(&mut result, self.size);
        vec::raw::set_len(&mut result, self.size);
        do vec::as_imm_buf(result) |p, len| {
            clEnqueueReadBuffer(
                self.context.q.cqueue, self.cl_buffer, CL_TRUE, 0,
                len * sys::size_of::<T>() as libc::size_t,
                p as *libc::c_void, 0, ptr::null(), ptr::null());

        }
        move result
    }
}

impl<T: VectorType> Vector<T>: ::hl::KernelArg {
    pure fn get_value(&self) -> (libc::size_t, *libc::c_void) {
        (sys::size_of::<cl_mem>() as libc::size_t,
         ptr::addr_of(&self.cl_buffer) as *libc::c_void)
    }
}

#[cfg(test)]
mod test {
    macro_rules! expect (
        ($test: expr, $expected: expr) => ({
            let test = $test;
            let expected = $expected;
            if test != expected {
                fail fmt!("Test failure in %s: expected %?, got %?",
                          stringify!($test),
                          expected, test)
            }
        })
    )

    #[test]
    fn gpu_vector() {
        let ctx = create_compute_context();

        let x = ~[1, 2, 3, 4, 5];
        let gx = Vector::from_vec(ctx, x);
        let y = gx.to_vec();
        expect!(x, y);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
pub mod reply_markup;

use serde::de::Deserialize;
use serde::ser::{Serialize, Serializer};

use types::*;

pub use self::reply_markup::*;

pub trait Request: Serialize {
    type Response: Deserialize;

    fn name(&self) -> &'static str;
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum ChatId<'a> {
    ById(Integer),
    ByUsername(&'a str),
}

impl<'a> ChatId<'a> {
    pub fn from_id(id: Integer) -> Self {
        ChatId::ById(id)
    }

    pub fn from_username(username: &'a str) -> Self {
        ChatId::ByUsername(username)
    }
}

impl<'a> Serialize for ChatId<'a> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer {
        match *self {
            ChatId::ById(id) => serializer.serialize_i64(id),
            ChatId::ByUsername(username) => serializer.serialize_str(username),
        }
    }
}

impl<'a> From<Integer> for ChatId<'a> {
    fn from(value: Integer) -> ChatId<'a> {
        ChatId::from_id(value)
    }
}

impl<'a, 'b> From<&'b Chat> for ChatId<'a> {
    fn from(value: &'b Chat) -> ChatId<'a> {
        ChatId::from_id(value.id())
    }
}

impl<'a, 'b> From<&'b ForwardFrom> for ChatId<'a> {
    fn from(value: &'b ForwardFrom) -> ChatId<'a> {
        let id = match *value {
            ForwardFrom::User {ref user, ..} => user.id,
            ForwardFrom::Channel {ref channel, ..} => channel.id,
        };
        ChatId::from_id(id)
    }
}

macro_rules! from_chat_type {
    ($chat: ident) => {
        impl<'a, 'b> From<&'b $chat> for ChatId<'a> {
            fn from(value: &'b $chat) -> ChatId<'a> {
                ChatId::from_id(value.id)
            }
        }
    };
}

from_chat_type!(User);
from_chat_type!(Group);
from_chat_type!(Supergroup);
from_chat_type!(Channel);

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Serialize)]
pub enum ParseMode {
    Markdown,
    #[serde(rename = "HTML")]
    Html
}

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
pub struct MessageId(pub Integer);

impl Serialize for MessageId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer {
        serializer.serialize_i64(self.0)
    }
}

impl From<Integer> for MessageId {
    fn from(value: Integer) -> Self {
        MessageId(value)
    }
}

impl<'a> From<&'a Message> for MessageId {
    fn from(value: &Message) -> Self {
        MessageId(value.id)
    }
}
<commit_msg>feat: add UserId struct<commit_after>#[macro_use]
pub mod reply_markup;

use serde::de::Deserialize;
use serde::ser::{Serialize, Serializer};

use types::*;

pub use self::reply_markup::*;

pub trait Request: Serialize {
    type Response: Deserialize;

    fn name(&self) -> &'static str;
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum ChatId<'a> {
    ById(Integer),
    ByUsername(&'a str),
}

impl<'a> ChatId<'a> {
    pub fn from_id(id: Integer) -> Self {
        ChatId::ById(id)
    }

    pub fn from_username(username: &'a str) -> Self {
        ChatId::ByUsername(username)
    }
}

impl<'a> Serialize for ChatId<'a> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer {
        match *self {
            ChatId::ById(id) => serializer.serialize_i64(id),
            ChatId::ByUsername(username) => serializer.serialize_str(username),
        }
    }
}

impl<'a> From<Integer> for ChatId<'a> {
    fn from(value: Integer) -> ChatId<'a> {
        ChatId::from_id(value)
    }
}

impl<'a, 'b> From<&'b Chat> for ChatId<'a> {
    fn from(value: &'b Chat) -> ChatId<'a> {
        ChatId::from_id(value.id())
    }
}

impl<'a, 'b> From<&'b ForwardFrom> for ChatId<'a> {
    fn from(value: &'b ForwardFrom) -> ChatId<'a> {
        let id = match *value {
            ForwardFrom::User {ref user, ..} => user.id,
            ForwardFrom::Channel {ref channel, ..} => channel.id,
        };
        ChatId::from_id(id)
    }
}

macro_rules! from_chat_type {
    ($chat: ident) => {
        impl<'a, 'b> From<&'b $chat> for ChatId<'a> {
            fn from(value: &'b $chat) -> ChatId<'a> {
                ChatId::from_id(value.id)
            }
        }
    };
}

from_chat_type!(User);
from_chat_type!(Group);
from_chat_type!(Supergroup);
from_chat_type!(Channel);

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct UserId {
    inner: Integer
}

impl<'a> Serialize for UserId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer {
        serializer.serialize_i64(self.inner)
    }
}

impl From<Integer> for UserId {
    fn from(value: Integer) -> UserId {
        UserId { inner: value }
    }
}

impl<'b> From<&'b User> for UserId {
    fn from(value: &User) -> UserId {
        UserId { inner: value.id }
    }
}

impl<'b> From<&'b ChatMember> for UserId {
    fn from(value: &ChatMember) -> UserId {
        UserId { inner: value.user.id }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Serialize)]
pub enum ParseMode {
    Markdown,
    #[serde(rename = "HTML")]
    Html
}

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
pub struct MessageId(pub Integer);

impl Serialize for MessageId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer {
        serializer.serialize_i64(self.0)
    }
}

impl From<Integer> for MessageId {
    fn from(value: Integer) -> Self {
        MessageId(value)
    }
}

impl<'a> From<&'a Message> for MessageId {
    fn from(value: &Message) -> Self {
        MessageId(value.id)
    }
}
<|endoftext|>


<commit_before>use std::collections::{BTreeMap, VecDeque};

use anyhow::Result;
use log::debug;
use smallvec::smallvec;

use crate::{
    analysis::{
        cfg::{
            flow,
            flow::{Flow, Flows},
            BasicBlock, CFG,
        },
        dis,
    },
    aspace::AddressSpace,
    module::Module,
    VA,
};

struct InstructionDescriptor {
    length:     u64,
    successors: Flows,
}

fn fallthrough_flows<'a>(flows: &'a Flows) -> Box<dyn Iterator<Item = &'a Flow> + 'a> {
    Box::new(flows.iter().filter(|flow| matches!(flow, Flow::Fallthrough(_))))
}

fn non_fallthrough_flows<'a>(flows: &'a Flows) -> Box<dyn Iterator<Item = &'a Flow> + 'a> {
    Box::new(flows.iter().filter(|flow| !matches!(flow, Flow::Fallthrough(_))))
}

fn empty<'a, T>(mut i: Box<dyn Iterator<Item = T> + 'a>) -> bool {
    i.next().is_none()
}

fn read_insn_descriptors(module: &Module, va: VA) -> Result<BTreeMap<VA, InstructionDescriptor>> {
    let decoder = dis::get_disassembler(module)?;
    let mut insn_buf = [0u8; 16];

    let mut queue: VecDeque<VA> = Default::default();
    queue.push_back(va);

    let mut insns: BTreeMap<VA, InstructionDescriptor> = Default::default();

    loop {
        let va = match queue.pop_back() {
            None => break,
            Some(va) => va,
        };

        if insns.contains_key(&va) {
            continue;
        }

        // TODO: optimize here by re-using buffers.
        if module.address_space.read_into(va, &mut insn_buf).is_ok() {
            if let Ok(Some(insn)) = decoder.decode(&insn_buf) {
                let successors: Flows = flow::get_insn_flow(module, va, &insn)?
                    // remove CALL instructions for cfg reconstruction.
                    .into_iter()
                    .filter(|succ| !matches!(succ, Flow::Call(_)))
                    .collect();

                for target in successors.iter() {
                    queue.push_back(target.va());
                }

                let desc = InstructionDescriptor {
                    length: insn.length as u64,
                    successors,
                };

                insns.insert(va, desc);
            }
        }
    }

    Ok(insns)
}

/// compute successors for an instruction (specified by VA).
/// this function should not fail.
fn compute_successors(insns: &BTreeMap<VA, InstructionDescriptor>) -> BTreeMap<VA, Flows> {
    let mut successors: BTreeMap<VA, Flows> = Default::default();

    for &va in insns.keys() {
        successors.insert(va, smallvec![]);
    }

    for (&va, desc) in insns.iter() {
        successors
            .entry(va)
            .and_modify(|l: &mut Flows| l.extend(desc.successors.clone()));
    }

    successors
}

/// compute predecessors for an instruction (specified by VA).
/// this function should not fail.
fn compute_predecessors(insns: &BTreeMap<VA, InstructionDescriptor>) -> BTreeMap<VA, Flows> {
    let mut predecessors: BTreeMap<VA, Flows> = Default::default();

    for &va in insns.keys() {
        predecessors.insert(va, smallvec![]);
    }

    for (&va, desc) in insns.iter() {
        for succ in desc.successors.iter() {
            let flow = succ.swap(va);
            predecessors.entry(succ.va()).and_modify(|l: &mut Flows| l.push(flow));
        }
    }

    predecessors
}

/// this function should not fail.
fn compute_basic_blocks(
    insns: &BTreeMap<VA, InstructionDescriptor>,
    predecessors: &BTreeMap<VA, Flows>,
    successors: &BTreeMap<VA, Flows>,
) -> BTreeMap<VA, BasicBlock> {
    // find all the basic block start addresses.
    //
    // scan through all instructions, looking for:
    //  1. instruction with nothing before it, or
    //  2. instruction with a non-fallthrough flow to it (e.g. jmp target), or
    //  3. the prior instruction also branched elsewhere
    let starts: Vec<VA> = insns
        .keys()
        .filter(|&va| {
            let preds = &predecessors[va];

            // its a root, because nothing flows here.
            if preds.is_empty() {
                return true;
            }

            // its a bb start, because there's a branch to here.
            if !empty(non_fallthrough_flows(preds)) {
                return true;
            }

            // its a bb start, because the instruction that fallthrough here
            // also branched somewhere else.
            for pred in fallthrough_flows(preds) {
                if !empty(non_fallthrough_flows(&successors[&pred.va()])) {
                    return true;
                }
            }

            false
        })
        .cloned()
        .collect();

    let mut basic_blocks: BTreeMap<VA, BasicBlock> = Default::default();
    let mut basic_blocks_by_last_insn: BTreeMap<VA, VA> = Default::default();

    // compute the basic block instances.
    //
    // for each basic block start,
    // scan forward to find the end of the block.
    //
    // set the bb successors.
    // leave the predecessors for a subsequent pass.
    // this is because we need the basic blocks indexed by final address.
    for &start in starts.iter() {
        let mut va = start;
        let mut insn = &insns[&va];

        let mut bb = BasicBlock {
            address:      va,
            length:       0,
            predecessors: Default::default(),
            successors:   Default::default(),
        };

        loop {
            let flows = &successors[&va];

            // there is no fallthrough here, so end of bb.
            // for example: ret has no fallthrough, and is end of basic block.
            if empty(fallthrough_flows(flows)) {
                break;
            }

            // there is a non-fallthrough flow, so end of bb.
            // for example: jnz has a non-fallthrough flow from it.
            if !empty(non_fallthrough_flows(flows)) {
                break;
            }

            // now we need to check the next instruction.
            // if its the start of a basic block,
            // then the current basic block must end.

            let next_va = va + insn.length;

            // there is not a subsequent instruction, so end of bb.
            // this might be considered an analysis error.
            // but, we don't have a better framework for this right now.
            // options include:
            //  1. pretend the current BB is ok, or
            //  2. fail the whole analysis
            // for right now, we're doing (1).
            if !predecessors.contains_key(&next_va) {
                log::warn!("cfg: {:#x}: no subsequent instruction at {:#x}", va, next_va);
                break;
            }

            // the next instruction has other flows to it, so its a new bb.
            // the next instruction is not part of this bb.
            // for example, the target of a fallthrough AND a jump from elsewhere.
            if !empty(non_fallthrough_flows(&predecessors[&next_va])) {
                break;
            }

            bb.length += insn.length;

            va = next_va;
            insn = &insns[&next_va];
        }
        // insn is the last instruction of the current basic block.
        // va is the address of the last instruction of the current basic block.

        bb.length += insn.length;
        bb.successors = successors.get(&va).unwrap_or(&smallvec![]).clone();

        basic_blocks.insert(start, bb);
        basic_blocks_by_last_insn.insert(va, start);
    }

    for (va, bb) in basic_blocks.iter_mut() {
        bb.predecessors.extend(
            predecessors
                .get(va)
                .unwrap_or(&smallvec![])
                .iter()
                .filter(|pred| basic_blocks_by_last_insn.contains_key(&pred.va()))
                .map(|pred| pred.swap(basic_blocks_by_last_insn[&pred.va()])),
        )
    }

    basic_blocks
}

pub fn build_cfg(module: &Module, va: VA) -> Result<CFG> {
    debug!("cfg: {:#x}", va);

    let insns = read_insn_descriptors(module, va)?;
    debug!("cfg: {:#x}: {} instructions", va, insns.len());

    let successors = compute_successors(&insns);
    let predecessors = compute_predecessors(&insns);

    let bbs = compute_basic_blocks(&insns, &predecessors, &successors);
    debug!("cfg: {:#x}: {} basic blocks", va, bbs.len());

    Ok(CFG { basic_blocks: bbs })
}

#[cfg(test)]
mod tests {
    use crate::{analysis::cfg::local::*, rsrc::*};
    use anyhow::Result;

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        // this function has:
        //   - api call at 0x1800527F8
        //   - cmovns at 0x180052841
        //   - conditional branches at 0x180052829
        let cfg = build_cfg(&pe.module, 0x1800527B0)?;
        assert_eq!(cfg.basic_blocks.len(), 4);

        Ok(())
    }
}
<commit_msg>core: dis: insn length is u8<commit_after>use std::collections::{BTreeMap, VecDeque};

use anyhow::Result;
use log::debug;
use smallvec::smallvec;

use crate::{
    analysis::{
        cfg::{
            flow,
            flow::{Flow, Flows},
            BasicBlock, CFG,
        },
        dis,
    },
    aspace::AddressSpace,
    module::Module,
    VA,
};

struct InstructionDescriptor {
    length:     u8,
    successors: Flows,
}

fn fallthrough_flows<'a>(flows: &'a Flows) -> Box<dyn Iterator<Item = &'a Flow> + 'a> {
    Box::new(flows.iter().filter(|flow| matches!(flow, Flow::Fallthrough(_))))
}

fn non_fallthrough_flows<'a>(flows: &'a Flows) -> Box<dyn Iterator<Item = &'a Flow> + 'a> {
    Box::new(flows.iter().filter(|flow| !matches!(flow, Flow::Fallthrough(_))))
}

fn empty<'a, T>(mut i: Box<dyn Iterator<Item = T> + 'a>) -> bool {
    i.next().is_none()
}

fn read_insn_descriptors(module: &Module, va: VA) -> Result<BTreeMap<VA, InstructionDescriptor>> {
    let decoder = dis::get_disassembler(module)?;
    let mut insn_buf = [0u8; 16];

    let mut queue: VecDeque<VA> = Default::default();
    queue.push_back(va);

    let mut insns: BTreeMap<VA, InstructionDescriptor> = Default::default();

    loop {
        let va = match queue.pop_back() {
            None => break,
            Some(va) => va,
        };

        if insns.contains_key(&va) {
            continue;
        }

        // TODO: optimize here by re-using buffers.
        if module.address_space.read_into(va, &mut insn_buf).is_ok() {
            if let Ok(Some(insn)) = decoder.decode(&insn_buf) {
                let successors: Flows = flow::get_insn_flow(module, va, &insn)?
                    // remove CALL instructions for cfg reconstruction.
                    .into_iter()
                    .filter(|succ| !matches!(succ, Flow::Call(_)))
                    .collect();

                for target in successors.iter() {
                    queue.push_back(target.va());
                }

                let desc = InstructionDescriptor {
                    length: insn.length,
                    successors,
                };

                insns.insert(va, desc);
            }
        }
    }

    Ok(insns)
}

/// compute successors for an instruction (specified by VA).
/// this function should not fail.
fn compute_successors(insns: &BTreeMap<VA, InstructionDescriptor>) -> BTreeMap<VA, Flows> {
    let mut successors: BTreeMap<VA, Flows> = Default::default();

    for &va in insns.keys() {
        successors.insert(va, smallvec![]);
    }

    for (&va, desc) in insns.iter() {
        successors
            .entry(va)
            .and_modify(|l: &mut Flows| l.extend(desc.successors.clone()));
    }

    successors
}

/// compute predecessors for an instruction (specified by VA).
/// this function should not fail.
fn compute_predecessors(insns: &BTreeMap<VA, InstructionDescriptor>) -> BTreeMap<VA, Flows> {
    let mut predecessors: BTreeMap<VA, Flows> = Default::default();

    for &va in insns.keys() {
        predecessors.insert(va, smallvec![]);
    }

    for (&va, desc) in insns.iter() {
        for succ in desc.successors.iter() {
            let flow = succ.swap(va);
            predecessors.entry(succ.va()).and_modify(|l: &mut Flows| l.push(flow));
        }
    }

    predecessors
}

/// this function should not fail.
fn compute_basic_blocks(
    insns: &BTreeMap<VA, InstructionDescriptor>,
    predecessors: &BTreeMap<VA, Flows>,
    successors: &BTreeMap<VA, Flows>,
) -> BTreeMap<VA, BasicBlock> {
    // find all the basic block start addresses.
    //
    // scan through all instructions, looking for:
    //  1. instruction with nothing before it, or
    //  2. instruction with a non-fallthrough flow to it (e.g. jmp target), or
    //  3. the prior instruction also branched elsewhere
    let starts: Vec<VA> = insns
        .keys()
        .filter(|&va| {
            let preds = &predecessors[va];

            // its a root, because nothing flows here.
            if preds.is_empty() {
                return true;
            }

            // its a bb start, because there's a branch to here.
            if !empty(non_fallthrough_flows(preds)) {
                return true;
            }

            // its a bb start, because the instruction that fallthrough here
            // also branched somewhere else.
            for pred in fallthrough_flows(preds) {
                if !empty(non_fallthrough_flows(&successors[&pred.va()])) {
                    return true;
                }
            }

            false
        })
        .cloned()
        .collect();

    let mut basic_blocks: BTreeMap<VA, BasicBlock> = Default::default();
    let mut basic_blocks_by_last_insn: BTreeMap<VA, VA> = Default::default();

    // compute the basic block instances.
    //
    // for each basic block start,
    // scan forward to find the end of the block.
    //
    // set the bb successors.
    // leave the predecessors for a subsequent pass.
    // this is because we need the basic blocks indexed by final address.
    for &start in starts.iter() {
        let mut va = start;
        let mut insn = &insns[&va];

        let mut bb = BasicBlock {
            address:      va,
            length:       0,
            predecessors: Default::default(),
            successors:   Default::default(),
        };

        loop {
            let flows = &successors[&va];

            // there is no fallthrough here, so end of bb.
            // for example: ret has no fallthrough, and is end of basic block.
            if empty(fallthrough_flows(flows)) {
                break;
            }

            // there is a non-fallthrough flow, so end of bb.
            // for example: jnz has a non-fallthrough flow from it.
            if !empty(non_fallthrough_flows(flows)) {
                break;
            }

            // now we need to check the next instruction.
            // if its the start of a basic block,
            // then the current basic block must end.

            let next_va = va + (insn.length as u64);

            // there is not a subsequent instruction, so end of bb.
            // this might be considered an analysis error.
            // but, we don't have a better framework for this right now.
            // options include:
            //  1. pretend the current BB is ok, or
            //  2. fail the whole analysis
            // for right now, we're doing (1).
            if !predecessors.contains_key(&next_va) {
                log::warn!("cfg: {:#x}: no subsequent instruction at {:#x}", va, next_va);
                break;
            }

            // the next instruction has other flows to it, so its a new bb.
            // the next instruction is not part of this bb.
            // for example, the target of a fallthrough AND a jump from elsewhere.
            if !empty(non_fallthrough_flows(&predecessors[&next_va])) {
                break;
            }

            bb.length += insn.length as u64;

            va = next_va;
            insn = &insns[&next_va];
        }
        // insn is the last instruction of the current basic block.
        // va is the address of the last instruction of the current basic block.

        bb.length += insn.length as u64;
        bb.successors = successors.get(&va).unwrap_or(&smallvec![]).clone();

        basic_blocks.insert(start, bb);
        basic_blocks_by_last_insn.insert(va, start);
    }

    for (va, bb) in basic_blocks.iter_mut() {
        bb.predecessors.extend(
            predecessors
                .get(va)
                .unwrap_or(&smallvec![])
                .iter()
                .filter(|pred| basic_blocks_by_last_insn.contains_key(&pred.va()))
                .map(|pred| pred.swap(basic_blocks_by_last_insn[&pred.va()])),
        )
    }

    basic_blocks
}

pub fn build_cfg(module: &Module, va: VA) -> Result<CFG> {
    debug!("cfg: {:#x}", va);

    let insns = read_insn_descriptors(module, va)?;
    debug!("cfg: {:#x}: {} instructions", va, insns.len());

    let successors = compute_successors(&insns);
    let predecessors = compute_predecessors(&insns);

    let bbs = compute_basic_blocks(&insns, &predecessors, &successors);
    debug!("cfg: {:#x}: {} basic blocks", va, bbs.len());

    Ok(CFG { basic_blocks: bbs })
}

#[cfg(test)]
mod tests {
    use crate::{analysis::cfg::local::*, rsrc::*};
    use anyhow::Result;

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        // this function has:
        //   - api call at 0x1800527F8
        //   - cmovns at 0x180052841
        //   - conditional branches at 0x180052829
        let cfg = build_cfg(&pe.module, 0x1800527B0)?;
        assert_eq!(cfg.basic_blocks.len(), 4);

        Ok(())
    }
}
<|endoftext|>


<commit_before>//! HIR for references to types. Paths in these are not yet resolved. They can
//! be directly created from an ast::TypeRef, without further queries.

use hir_expand::{name::Name, AstId, InFile};
use std::convert::TryInto;
use syntax::ast;

use crate::{body::LowerCtx, path::Path};

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub enum Mutability {
    Shared,
    Mut,
}

impl Mutability {
    pub fn from_mutable(mutable: bool) -> Mutability {
        if mutable {
            Mutability::Mut
        } else {
            Mutability::Shared
        }
    }

    pub fn as_keyword_for_ref(self) -> &'static str {
        match self {
            Mutability::Shared => "",
            Mutability::Mut => "mut ",
        }
    }

    pub fn as_keyword_for_ptr(self) -> &'static str {
        match self {
            Mutability::Shared => "const ",
            Mutability::Mut => "mut ",
        }
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub enum Rawness {
    RawPtr,
    Ref,
}

impl Rawness {
    pub fn from_raw(is_raw: bool) -> Rawness {
        if is_raw {
            Rawness::RawPtr
        } else {
            Rawness::Ref
        }
    }
}

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct TraitRef {
    pub path: Path,
}

impl TraitRef {
    /// Converts an `ast::PathType` to a `hir::TraitRef`.
    pub(crate) fn from_ast(ctx: &LowerCtx, node: ast::Type) -> Option<Self> {
        // FIXME: Use `Path::from_src`
        match node {
            ast::Type::PathType(path) => {
                path.path().and_then(|it| ctx.lower_path(it)).map(|path| TraitRef { path })
            }
            _ => None,
        }
    }
}

/// Compare ty::Ty
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum TypeRef {
    Never,
    Placeholder,
    Tuple(Vec<TypeRef>),
    Path(Path),
    RawPtr(Box<TypeRef>, Mutability),
    Reference(Box<TypeRef>, Option<LifetimeRef>, Mutability),
    // FIXME: for full const generics, the latter element (length) here is going to have to be an
    // expression that is further lowered later in hir_ty.
    Array(Box<TypeRef>, ConstScalar),
    Slice(Box<TypeRef>),
    /// A fn pointer. Last element of the vector is the return type.
    Fn(Vec<TypeRef>, bool /*varargs*/),
    // For
    ImplTrait(Vec<TypeBound>),
    DynTrait(Vec<TypeBound>),
    Macro(AstId<ast::MacroCall>),
    Error,
}

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct LifetimeRef {
    pub name: Name,
}

impl LifetimeRef {
    pub(crate) fn new_name(name: Name) -> Self {
        LifetimeRef { name }
    }

    pub(crate) fn new(lifetime: &ast::Lifetime) -> Self {
        LifetimeRef { name: Name::new_lifetime(lifetime) }
    }

    pub fn missing() -> LifetimeRef {
        LifetimeRef { name: Name::missing() }
    }
}

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum TypeBound {
    Path(Path),
    // ForLifetime(Vec<LifetimeRef>, Path), FIXME ForLifetime
    Lifetime(LifetimeRef),
    Error,
}

impl TypeRef {
    /// Converts an `ast::TypeRef` to a `hir::TypeRef`.
    pub fn from_ast(ctx: &LowerCtx, node: ast::Type) -> Self {
        match node {
            ast::Type::ParenType(inner) => TypeRef::from_ast_opt(&ctx, inner.ty()),
            ast::Type::TupleType(inner) => {
                TypeRef::Tuple(inner.fields().map(|it| TypeRef::from_ast(ctx, it)).collect())
            }
            ast::Type::NeverType(..) => TypeRef::Never,
            ast::Type::PathType(inner) => {
                // FIXME: Use `Path::from_src`
                inner
                    .path()
                    .and_then(|it| ctx.lower_path(it))
                    .map(TypeRef::Path)
                    .unwrap_or(TypeRef::Error)
            }
            ast::Type::PtrType(inner) => {
                let inner_ty = TypeRef::from_ast_opt(&ctx, inner.ty());
                let mutability = Mutability::from_mutable(inner.mut_token().is_some());
                TypeRef::RawPtr(Box::new(inner_ty), mutability)
            }
            ast::Type::ArrayType(inner) => {
                // FIXME: This is a hack. We should probably reuse the machinery of
                // `hir_def::body::lower` to lower this into an `Expr` and then evaluate it at the
                // `hir_ty` level, which would allow knowing the type of:
                // let v: [u8; 2 + 2] = [0u8; 4];
                let len = inner
                    .expr()
                    .map(ConstScalar::usize_from_literal_expr)
                    .unwrap_or(ConstScalar::Unknown);

                TypeRef::Array(Box::new(TypeRef::from_ast_opt(&ctx, inner.ty())), len)
            }
            ast::Type::SliceType(inner) => {
                TypeRef::Slice(Box::new(TypeRef::from_ast_opt(&ctx, inner.ty())))
            }
            ast::Type::RefType(inner) => {
                let inner_ty = TypeRef::from_ast_opt(&ctx, inner.ty());
                let lifetime = inner.lifetime().map(|lt| LifetimeRef::new(&lt));
                let mutability = Mutability::from_mutable(inner.mut_token().is_some());
                TypeRef::Reference(Box::new(inner_ty), lifetime, mutability)
            }
            ast::Type::InferType(_inner) => TypeRef::Placeholder,
            ast::Type::FnPtrType(inner) => {
                let ret_ty = inner
                    .ret_type()
                    .and_then(|rt| rt.ty())
                    .map(|it| TypeRef::from_ast(ctx, it))
                    .unwrap_or_else(|| TypeRef::Tuple(Vec::new()));
                let mut is_varargs = false;
                let mut params = if let Some(pl) = inner.param_list() {
                    if let Some(param) = pl.params().last() {
                        is_varargs = param.dotdotdot_token().is_some();
                    }

                    pl.params().map(|p| p.ty()).map(|it| TypeRef::from_ast_opt(&ctx, it)).collect()
                } else {
                    Vec::new()
                };
                params.push(ret_ty);
                TypeRef::Fn(params, is_varargs)
            }
            // for types are close enough for our purposes to the inner type for now...
            ast::Type::ForType(inner) => TypeRef::from_ast_opt(&ctx, inner.ty()),
            ast::Type::ImplTraitType(inner) => {
                TypeRef::ImplTrait(type_bounds_from_ast(ctx, inner.type_bound_list()))
            }
            ast::Type::DynTraitType(inner) => {
                TypeRef::DynTrait(type_bounds_from_ast(ctx, inner.type_bound_list()))
            }
            ast::Type::MacroType(mt) => match mt.macro_call() {
                Some(mc) => ctx
                    .ast_id(&mc)
                    .map(|mc| TypeRef::Macro(InFile::new(ctx.file_id(), mc)))
                    .unwrap_or(TypeRef::Error),
                None => TypeRef::Error,
            },
        }
    }

    pub(crate) fn from_ast_opt(ctx: &LowerCtx, node: Option<ast::Type>) -> Self {
        if let Some(node) = node {
            TypeRef::from_ast(ctx, node)
        } else {
            TypeRef::Error
        }
    }

    pub(crate) fn unit() -> TypeRef {
        TypeRef::Tuple(Vec::new())
    }

    pub fn walk(&self, f: &mut impl FnMut(&TypeRef)) {
        go(self, f);

        fn go(type_ref: &TypeRef, f: &mut impl FnMut(&TypeRef)) {
            f(type_ref);
            match type_ref {
                TypeRef::Fn(types, _) | TypeRef::Tuple(types) => {
                    types.iter().for_each(|t| go(t, f))
                }
                TypeRef::RawPtr(type_ref, _)
                | TypeRef::Reference(type_ref, ..)
                | TypeRef::Array(type_ref, _)
                | TypeRef::Slice(type_ref) => go(&type_ref, f),
                TypeRef::ImplTrait(bounds) | TypeRef::DynTrait(bounds) => {
                    for bound in bounds {
                        match bound {
                            TypeBound::Path(path) => go_path(path, f),
                            TypeBound::Lifetime(_) | TypeBound::Error => (),
                        }
                    }
                }
                TypeRef::Path(path) => go_path(path, f),
                TypeRef::Never | TypeRef::Placeholder | TypeRef::Macro(_) | TypeRef::Error => {}
            };
        }

        fn go_path(path: &Path, f: &mut impl FnMut(&TypeRef)) {
            if let Some(type_ref) = path.type_anchor() {
                go(type_ref, f);
            }
            for segment in path.segments().iter() {
                if let Some(args_and_bindings) = segment.args_and_bindings {
                    for arg in &args_and_bindings.args {
                        match arg {
                            crate::path::GenericArg::Type(type_ref) => {
                                go(type_ref, f);
                            }
                            crate::path::GenericArg::Lifetime(_) => {}
                        }
                    }
                    for binding in &args_and_bindings.bindings {
                        if let Some(type_ref) = &binding.type_ref {
                            go(type_ref, f);
                        }
                        for bound in &binding.bounds {
                            match bound {
                                TypeBound::Path(path) => go_path(path, f),
                                TypeBound::Lifetime(_) | TypeBound::Error => (),
                            }
                        }
                    }
                }
            }
        }
    }
}

pub(crate) fn type_bounds_from_ast(
    lower_ctx: &LowerCtx,
    type_bounds_opt: Option<ast::TypeBoundList>,
) -> Vec<TypeBound> {
    if let Some(type_bounds) = type_bounds_opt {
        type_bounds.bounds().map(|it| TypeBound::from_ast(lower_ctx, it)).collect()
    } else {
        vec![]
    }
}

impl TypeBound {
    pub(crate) fn from_ast(ctx: &LowerCtx, node: ast::TypeBound) -> Self {
        match node.kind() {
            ast::TypeBoundKind::PathType(path_type) => {
                let path = match path_type.path() {
                    Some(p) => p,
                    None => return TypeBound::Error,
                };

                let path = match ctx.lower_path(path) {
                    Some(p) => p,
                    None => return TypeBound::Error,
                };
                TypeBound::Path(path)
            }
            ast::TypeBoundKind::ForType(_) => TypeBound::Error, // FIXME ForType
            ast::TypeBoundKind::Lifetime(lifetime) => {
                TypeBound::Lifetime(LifetimeRef::new(&lifetime))
            }
        }
    }

    pub fn as_path(&self) -> Option<&Path> {
        match self {
            TypeBound::Path(p) => Some(p),
            _ => None,
        }
    }
}

/// A concrete constant value
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ConstScalar {
    // for now, we only support the trivial case of constant evaluating the length of an array
    // Note that this is u64 because the target usize may be bigger than our usize
    Usize(u64),

    /// Case of an unknown value that rustc might know but we don't
    // FIXME: this is a hack to get around chalk not being able to represent unevaluatable
    // constants
    // https://github.com/rust-analyzer/rust-analyzer/pull/8813#issuecomment-840679177
    // https://rust-lang.zulipchat.com/#narrow/stream/144729-wg-traits/topic/Handling.20non.20evaluatable.20constants'.20equality/near/238386348
    Unknown,
}

impl std::fmt::Display for ConstScalar {
    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        match self {
            ConstScalar::Usize(us) => write!(fmt, "{}", us),
            ConstScalar::Unknown => write!(fmt, "_"),
        }
    }
}

impl ConstScalar {
    /// Gets a target usize out of the ConstScalar
    pub fn as_usize(&self) -> Option<u64> {
        match self {
            &ConstScalar::Usize(us) => Some(us),
            _ => None,
        }
    }

    // FIXME: as per the comments on `TypeRef::Array`, this evaluation should not happen at this
    // parse stage.
    fn usize_from_literal_expr(expr: ast::Expr) -> ConstScalar {
        match expr {
            ast::Expr::Literal(lit) => {
                let lkind = lit.kind();
                match lkind {
                    ast::LiteralKind::IntNumber(num)
                        if num.suffix() == None || num.suffix() == Some("usize") =>
                    {
                        num.value().and_then(|v| v.try_into().ok())
                    }
                    _ => None,
                }
            }
            _ => None,
        }
        .map(ConstScalar::Usize)
        .unwrap_or(ConstScalar::Unknown)
    }
}
<commit_msg>Explain why nested `TypeRef` aren't interned<commit_after>//! HIR for references to types. Paths in these are not yet resolved. They can
//! be directly created from an ast::TypeRef, without further queries.

use hir_expand::{name::Name, AstId, InFile};
use std::convert::TryInto;
use syntax::ast;

use crate::{body::LowerCtx, path::Path};

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub enum Mutability {
    Shared,
    Mut,
}

impl Mutability {
    pub fn from_mutable(mutable: bool) -> Mutability {
        if mutable {
            Mutability::Mut
        } else {
            Mutability::Shared
        }
    }

    pub fn as_keyword_for_ref(self) -> &'static str {
        match self {
            Mutability::Shared => "",
            Mutability::Mut => "mut ",
        }
    }

    pub fn as_keyword_for_ptr(self) -> &'static str {
        match self {
            Mutability::Shared => "const ",
            Mutability::Mut => "mut ",
        }
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub enum Rawness {
    RawPtr,
    Ref,
}

impl Rawness {
    pub fn from_raw(is_raw: bool) -> Rawness {
        if is_raw {
            Rawness::RawPtr
        } else {
            Rawness::Ref
        }
    }
}

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct TraitRef {
    pub path: Path,
}

impl TraitRef {
    /// Converts an `ast::PathType` to a `hir::TraitRef`.
    pub(crate) fn from_ast(ctx: &LowerCtx, node: ast::Type) -> Option<Self> {
        // FIXME: Use `Path::from_src`
        match node {
            ast::Type::PathType(path) => {
                path.path().and_then(|it| ctx.lower_path(it)).map(|path| TraitRef { path })
            }
            _ => None,
        }
    }
}

/// Compare ty::Ty
///
/// Note: Most users of `TypeRef` that end up in the salsa database intern it using
/// `Interned<TypeRef>` to save space. But notably, nested `TypeRef`s are not interned, since that
/// does not seem to save any noticeable amount of memory.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum TypeRef {
    Never,
    Placeholder,
    Tuple(Vec<TypeRef>),
    Path(Path),
    RawPtr(Box<TypeRef>, Mutability),
    Reference(Box<TypeRef>, Option<LifetimeRef>, Mutability),
    // FIXME: for full const generics, the latter element (length) here is going to have to be an
    // expression that is further lowered later in hir_ty.
    Array(Box<TypeRef>, ConstScalar),
    Slice(Box<TypeRef>),
    /// A fn pointer. Last element of the vector is the return type.
    Fn(Vec<TypeRef>, bool /*varargs*/),
    // For
    ImplTrait(Vec<TypeBound>),
    DynTrait(Vec<TypeBound>),
    Macro(AstId<ast::MacroCall>),
    Error,
}

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct LifetimeRef {
    pub name: Name,
}

impl LifetimeRef {
    pub(crate) fn new_name(name: Name) -> Self {
        LifetimeRef { name }
    }

    pub(crate) fn new(lifetime: &ast::Lifetime) -> Self {
        LifetimeRef { name: Name::new_lifetime(lifetime) }
    }

    pub fn missing() -> LifetimeRef {
        LifetimeRef { name: Name::missing() }
    }
}

#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum TypeBound {
    Path(Path),
    // ForLifetime(Vec<LifetimeRef>, Path), FIXME ForLifetime
    Lifetime(LifetimeRef),
    Error,
}

impl TypeRef {
    /// Converts an `ast::TypeRef` to a `hir::TypeRef`.
    pub fn from_ast(ctx: &LowerCtx, node: ast::Type) -> Self {
        match node {
            ast::Type::ParenType(inner) => TypeRef::from_ast_opt(&ctx, inner.ty()),
            ast::Type::TupleType(inner) => {
                TypeRef::Tuple(inner.fields().map(|it| TypeRef::from_ast(ctx, it)).collect())
            }
            ast::Type::NeverType(..) => TypeRef::Never,
            ast::Type::PathType(inner) => {
                // FIXME: Use `Path::from_src`
                inner
                    .path()
                    .and_then(|it| ctx.lower_path(it))
                    .map(TypeRef::Path)
                    .unwrap_or(TypeRef::Error)
            }
            ast::Type::PtrType(inner) => {
                let inner_ty = TypeRef::from_ast_opt(&ctx, inner.ty());
                let mutability = Mutability::from_mutable(inner.mut_token().is_some());
                TypeRef::RawPtr(Box::new(inner_ty), mutability)
            }
            ast::Type::ArrayType(inner) => {
                // FIXME: This is a hack. We should probably reuse the machinery of
                // `hir_def::body::lower` to lower this into an `Expr` and then evaluate it at the
                // `hir_ty` level, which would allow knowing the type of:
                // let v: [u8; 2 + 2] = [0u8; 4];
                let len = inner
                    .expr()
                    .map(ConstScalar::usize_from_literal_expr)
                    .unwrap_or(ConstScalar::Unknown);

                TypeRef::Array(Box::new(TypeRef::from_ast_opt(&ctx, inner.ty())), len)
            }
            ast::Type::SliceType(inner) => {
                TypeRef::Slice(Box::new(TypeRef::from_ast_opt(&ctx, inner.ty())))
            }
            ast::Type::RefType(inner) => {
                let inner_ty = TypeRef::from_ast_opt(&ctx, inner.ty());
                let lifetime = inner.lifetime().map(|lt| LifetimeRef::new(&lt));
                let mutability = Mutability::from_mutable(inner.mut_token().is_some());
                TypeRef::Reference(Box::new(inner_ty), lifetime, mutability)
            }
            ast::Type::InferType(_inner) => TypeRef::Placeholder,
            ast::Type::FnPtrType(inner) => {
                let ret_ty = inner
                    .ret_type()
                    .and_then(|rt| rt.ty())
                    .map(|it| TypeRef::from_ast(ctx, it))
                    .unwrap_or_else(|| TypeRef::Tuple(Vec::new()));
                let mut is_varargs = false;
                let mut params = if let Some(pl) = inner.param_list() {
                    if let Some(param) = pl.params().last() {
                        is_varargs = param.dotdotdot_token().is_some();
                    }

                    pl.params().map(|p| p.ty()).map(|it| TypeRef::from_ast_opt(&ctx, it)).collect()
                } else {
                    Vec::new()
                };
                params.push(ret_ty);
                TypeRef::Fn(params, is_varargs)
            }
            // for types are close enough for our purposes to the inner type for now...
            ast::Type::ForType(inner) => TypeRef::from_ast_opt(&ctx, inner.ty()),
            ast::Type::ImplTraitType(inner) => {
                TypeRef::ImplTrait(type_bounds_from_ast(ctx, inner.type_bound_list()))
            }
            ast::Type::DynTraitType(inner) => {
                TypeRef::DynTrait(type_bounds_from_ast(ctx, inner.type_bound_list()))
            }
            ast::Type::MacroType(mt) => match mt.macro_call() {
                Some(mc) => ctx
                    .ast_id(&mc)
                    .map(|mc| TypeRef::Macro(InFile::new(ctx.file_id(), mc)))
                    .unwrap_or(TypeRef::Error),
                None => TypeRef::Error,
            },
        }
    }

    pub(crate) fn from_ast_opt(ctx: &LowerCtx, node: Option<ast::Type>) -> Self {
        if let Some(node) = node {
            TypeRef::from_ast(ctx, node)
        } else {
            TypeRef::Error
        }
    }

    pub(crate) fn unit() -> TypeRef {
        TypeRef::Tuple(Vec::new())
    }

    pub fn walk(&self, f: &mut impl FnMut(&TypeRef)) {
        go(self, f);

        fn go(type_ref: &TypeRef, f: &mut impl FnMut(&TypeRef)) {
            f(type_ref);
            match type_ref {
                TypeRef::Fn(types, _) | TypeRef::Tuple(types) => {
                    types.iter().for_each(|t| go(t, f))
                }
                TypeRef::RawPtr(type_ref, _)
                | TypeRef::Reference(type_ref, ..)
                | TypeRef::Array(type_ref, _)
                | TypeRef::Slice(type_ref) => go(&type_ref, f),
                TypeRef::ImplTrait(bounds) | TypeRef::DynTrait(bounds) => {
                    for bound in bounds {
                        match bound {
                            TypeBound::Path(path) => go_path(path, f),
                            TypeBound::Lifetime(_) | TypeBound::Error => (),
                        }
                    }
                }
                TypeRef::Path(path) => go_path(path, f),
                TypeRef::Never | TypeRef::Placeholder | TypeRef::Macro(_) | TypeRef::Error => {}
            };
        }

        fn go_path(path: &Path, f: &mut impl FnMut(&TypeRef)) {
            if let Some(type_ref) = path.type_anchor() {
                go(type_ref, f);
            }
            for segment in path.segments().iter() {
                if let Some(args_and_bindings) = segment.args_and_bindings {
                    for arg in &args_and_bindings.args {
                        match arg {
                            crate::path::GenericArg::Type(type_ref) => {
                                go(type_ref, f);
                            }
                            crate::path::GenericArg::Lifetime(_) => {}
                        }
                    }
                    for binding in &args_and_bindings.bindings {
                        if let Some(type_ref) = &binding.type_ref {
                            go(type_ref, f);
                        }
                        for bound in &binding.bounds {
                            match bound {
                                TypeBound::Path(path) => go_path(path, f),
                                TypeBound::Lifetime(_) | TypeBound::Error => (),
                            }
                        }
                    }
                }
            }
        }
    }
}

pub(crate) fn type_bounds_from_ast(
    lower_ctx: &LowerCtx,
    type_bounds_opt: Option<ast::TypeBoundList>,
) -> Vec<TypeBound> {
    if let Some(type_bounds) = type_bounds_opt {
        type_bounds.bounds().map(|it| TypeBound::from_ast(lower_ctx, it)).collect()
    } else {
        vec![]
    }
}

impl TypeBound {
    pub(crate) fn from_ast(ctx: &LowerCtx, node: ast::TypeBound) -> Self {
        match node.kind() {
            ast::TypeBoundKind::PathType(path_type) => {
                let path = match path_type.path() {
                    Some(p) => p,
                    None => return TypeBound::Error,
                };

                let path = match ctx.lower_path(path) {
                    Some(p) => p,
                    None => return TypeBound::Error,
                };
                TypeBound::Path(path)
            }
            ast::TypeBoundKind::ForType(_) => TypeBound::Error, // FIXME ForType
            ast::TypeBoundKind::Lifetime(lifetime) => {
                TypeBound::Lifetime(LifetimeRef::new(&lifetime))
            }
        }
    }

    pub fn as_path(&self) -> Option<&Path> {
        match self {
            TypeBound::Path(p) => Some(p),
            _ => None,
        }
    }
}

/// A concrete constant value
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ConstScalar {
    // for now, we only support the trivial case of constant evaluating the length of an array
    // Note that this is u64 because the target usize may be bigger than our usize
    Usize(u64),

    /// Case of an unknown value that rustc might know but we don't
    // FIXME: this is a hack to get around chalk not being able to represent unevaluatable
    // constants
    // https://github.com/rust-analyzer/rust-analyzer/pull/8813#issuecomment-840679177
    // https://rust-lang.zulipchat.com/#narrow/stream/144729-wg-traits/topic/Handling.20non.20evaluatable.20constants'.20equality/near/238386348
    Unknown,
}

impl std::fmt::Display for ConstScalar {
    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        match self {
            ConstScalar::Usize(us) => write!(fmt, "{}", us),
            ConstScalar::Unknown => write!(fmt, "_"),
        }
    }
}

impl ConstScalar {
    /// Gets a target usize out of the ConstScalar
    pub fn as_usize(&self) -> Option<u64> {
        match self {
            &ConstScalar::Usize(us) => Some(us),
            _ => None,
        }
    }

    // FIXME: as per the comments on `TypeRef::Array`, this evaluation should not happen at this
    // parse stage.
    fn usize_from_literal_expr(expr: ast::Expr) -> ConstScalar {
        match expr {
            ast::Expr::Literal(lit) => {
                let lkind = lit.kind();
                match lkind {
                    ast::LiteralKind::IntNumber(num)
                        if num.suffix() == None || num.suffix() == Some("usize") =>
                    {
                        num.value().and_then(|v| v.try_into().ok())
                    }
                    _ => None,
                }
            }
            _ => None,
        }
        .map(ConstScalar::Usize)
        .unwrap_or(ConstScalar::Unknown)
    }
}
<|endoftext|>