<commit_before>use std::collections::BTreeSet;

use value::Value;
use relation::{Relation, IndexSelect, ViewSelect};
use primitive::{Primitive, resolve_as_scalar};
use std::cmp::{min, max, Ordering};
use std::mem::replace;

#[derive(Clone, Debug)]
pub struct Table {
    pub insert: Option<IndexSelect>,
    pub remove: Option<IndexSelect>,
}

#[derive(Clone, Debug)]
pub struct Union {
    pub selects: Vec<IndexSelect>,
}

#[derive(Clone, Debug, Copy)]
pub enum Direction {
    Ascending,
    Descending,
}

#[derive(Clone, Debug)]
pub enum Input {
    Primitive{
        primitive: Primitive,
        input_bindings: Vec<(usize, usize)>,
    },
    View{
        input_ix: usize,
    },
}

#[derive(Clone, Debug)]
pub struct Source {
    pub input: Input,
    pub grouped_fields: Vec<usize>,
    pub sorted_fields: Vec<(usize, Direction)>,
    pub chunked: bool,
    pub constraint_bindings: Vec<(usize, usize)>,
    pub output_bindings: Vec<(usize, usize)>,
}

impl Source {
    fn prepare(&self, mut rows: Vec<Vec<Value>>) -> Vec<Vec<Value>> {
        // TODO compensate for inputs if primitive
        rows.sort_by(|a, b| {
            for &ix in self.grouped_fields.iter() {
                match a[ix].cmp(&b[ix]) {
                    Ordering::Greater => return Ordering::Greater,
                    Ordering::Less => return Ordering::Less,
                    Ordering::Equal => ()
                }
            }
            for &(ix, direction) in self.sorted_fields.iter() {
                match (a[ix].cmp(&b[ix]), direction) {
                    (Ordering::Greater, Direction::Ascending) => return Ordering::Greater,
                    (Ordering::Greater, Direction::Descending) => return Ordering::Less,
                    (Ordering::Less, Direction::Ascending) => return Ordering::Less,
                    (Ordering::Less, Direction::Descending) => return Ordering::Greater,
                    (Ordering::Equal, _) => ()
                }
            }
            return Ordering::Equal;
        });

        if self.chunked {
            // if the source is chunked, we need to turn all of the fields into columns for each grouping
            // we want the grouped fields to be scalars so that they can be joined against and every other
            // field should be a column of the values for that group
            let mut groups = Vec::new();
            let mut ordinal = 0;
            let grouped_fields = &self.grouped_fields;
            let row_len = rows.iter().next().map_or(0, |row| row.len());
            let ungrouped_fields = &(0..row_len).filter(|field_ix| !self.grouped_fields.contains(field_ix)).collect::<Vec<usize>>();
            let mut group: Vec<Value> = vec![Value::Column(vec![]); row_len]; // this creates a vector of the given length with all the values set to Value::Null
            for mut row in rows.into_iter() {
                // check if this is a new group
                let is_new_group = match self.grouped_fields.len() {
                    0 => false,
                    _ => !self.grouped_fields.iter().all(|&ix| row[ix] == group[ix])
                };
                if is_new_group {
                    // add the current ordinal to the group
                    group.push(Value::Float(ordinal as f64));
                    // push the old group
                    groups.push(group);
                    // create a new group with every field being a column
                    ordinal += 1;
                    group = vec![Value::Column(vec![]); row_len];
                    // any grouped field should have a scalar value for its index so that it can be joined on
                    // e.g. ["foo", [..], [..]] if you've grouped on index 0
                    for &grouped_field_ix in grouped_fields {
                        group[grouped_field_ix] = row[grouped_field_ix].clone();
                    }
                }
                // add the fields for this row
                for &field_ix in ungrouped_fields {
                    group[field_ix].as_column_mut().push(row[field_ix].clone());
                }
            }
            groups.push(group);
            // if there are grouped fields then the first group is the initial placeholder and
            // should be removed
            if self.grouped_fields.len() > 0 {
                groups.remove(0);
            }
            groups
        } else {
            // in the case where the source is not chunked all we need to worry about is setting the ordinal correctly
            // if there are no groups, then this is just normal rows and we number them monotonically
            // if there are groups, each new group resets the ordinal to 1
            let otherwise = vec![];
            let mut groups = Vec::new();
            let mut ordinal = 1;
            let mut prev_row = rows.iter().next().map_or(otherwise, |row| row.clone());
            for mut row in rows.into_iter() {
                let is_new_group = match self.grouped_fields.len() {
                    0 => false,
                    _ => !self.grouped_fields.iter().all(|&ix| row[ix] == prev_row[ix])
                };
                if is_new_group {
                    ordinal = 1;
                    prev_row = row.clone(); // @TODO: is there a way to do this without having to clone it?
                }
                row.push(Value::Float(ordinal as f64));
                groups.push(row);
                ordinal += 1;
            }
            groups
        }
    }
}

#[derive(Clone, Debug)]
pub struct Join {
    pub constants: Vec<Value>,
    pub sources: Vec<Source>,
    pub select: Vec<usize>,
}

#[derive(Clone, Debug)]
pub struct Reducer {
    pub primitive: Primitive,
    pub arguments: Vec<usize>,
}

#[derive(Clone, Debug)]
pub struct Aggregate {
    pub constants: Vec<Value>,
    pub outer: IndexSelect,
    pub inner: IndexSelect,
    pub directions: Vec<Direction>,
    pub limit_from: Option<usize>,
    pub limit_to: Option<usize>,
    pub reducers: Vec<Reducer>,
    pub selects_inner: bool,
    pub select: ViewSelect,
}

#[derive(Clone, Debug)]
pub enum View {
    Table(Table),
    Union(Union),
    Join(Join),
    Aggregate(Aggregate),
}

fn push_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for value in input.iter() {
        state.push(value);
    }
}

fn pop_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for _ in input.iter() {
        state.pop();
    }
}

// TODO this algorithm is incredibly naive and also clones excessively
fn join_step(join: &Join, ix: usize, inputs: &[Vec<Vec<Value>>], state: &mut Vec<Value>, index: &mut BTreeSet<Vec<Value>>) {
    if ix == join.sources.len() {
        index.insert(join.select.iter().map(|ix| state[*ix].clone()).collect());
    } else {
        let source = &join.sources[ix];
        match source.input {
            Input::View{..} => {
                for values in inputs[ix].iter() {
                    for &(field_ix, variable_ix) in source.output_bindings.iter() {
                        state[variable_ix] = values[field_ix].clone();
                    }
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        state[variable_ix] == values[field_ix]
                    ) {
                        join_step(join, ix+1, inputs, state, index);
                    }
                }
            }
            Input::Primitive{primitive, ref input_bindings} => {
                // values returned from primitives don't include inputs, so we will have to offset accesses by input_len
                let input_len = input_bindings.len();
                for mut values in primitive.eval_from_join(&input_bindings[..], &state[..]).into_iter() {
                    for &(field_ix, variable_ix) in source.output_bindings.iter() {
                        state[variable_ix] = replace(&mut values[field_ix - input_len], Value::Null);
                    }
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        state[variable_ix] == values[field_ix - input_len]
                    ) {
                        join_step(join, ix+1, inputs, state, index);
                    }
                }
            }
        }
    }
}

fn aggregate_step<'a>(aggregate: &Aggregate, input_sets: &'a [&[Vec<Value>]], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if input_sets.len() == 0 {
        index.insert(aggregate.select.select(&state[..]));
    } else {
        for values in input_sets[0].iter() {
            push_all(state, values);
            aggregate_step(aggregate, &input_sets[1..], state, index);
            pop_all(state, values);
        }
    }
}

fn compare_in_direction(xs: &[Value], ys: &[Value], directions: &[Direction]) -> Ordering {
    for ((x,y), direction) in xs.iter().zip(ys.iter()).zip(directions.iter()) {
        let cmp = match *direction {
            Direction::Ascending => x.cmp(y),
            Direction::Descending => y.cmp(x),
        };
        match cmp {
            Ordering::Greater => return Ordering::Greater,
            Ordering::Equal => (),
            Ordering::Less => return Ordering::Less,
        };
    }
    return Ordering::Equal;
}

impl View {
    pub fn run(&self, old_output: &Relation, upstream: &[&Relation]) -> Option<Relation> {
        let mut output = Relation::new(
            old_output.view.clone(),
            old_output.fields.clone(),
            old_output.names.clone()
            );
        match *self {
            View::Table(_) => None,
            View::Union(ref union) => {
                assert_eq!(union.selects.len(), upstream.len());
                for select in union.selects.iter() {
                    for values in select.select(&upstream[..]) {
                        output.index.insert(values);
                    }
                }
                Some(output)
            }
            View::Join(ref join) => {
                let mut state = join.constants.clone();
                let inputs = join.sources.iter().map(|source|
                    match source.input {
                        Input::Primitive{..} => {
                            vec![]
                        }
                        Input::View{input_ix, ..} => {
                            source.prepare(upstream[input_ix].index.iter().map(|row| row.clone()).collect())
                        }
                    }).collect::<Vec<_>>();
                join_step(join, 0, &inputs[..], &mut state, &mut output.index);
                Some(output)
            }
            View::Aggregate(ref aggregate) => {
                let mut outer = aggregate.outer.select(&upstream[..]);
                let mut inner = aggregate.inner.select(&upstream[..]);
                outer.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                outer.dedup();
                inner.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                let constants = &aggregate.constants[..];
                let mut group_start = 0;
                for outer_values in outer.into_iter() {
                    let mut group_end = group_start;
                    while (group_end < inner.len())
                    && (inner[group_end][0..outer_values.len()] == outer_values[..]) {
                        group_end += 1;
                    }
                    let (group, output_values) = {
                        let limit_from = match aggregate.limit_from {
                            None => group_start,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_to = match aggregate.limit_to {
                            None => group_end,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_from = min(max(limit_from, group_start), group_end);
                        let limit_to = min(max(limit_to, limit_from), group_end);
                        let group = &inner[limit_from..limit_to];
                        let output_values = aggregate.reducers.iter().map(|reducer| {
                            reducer.primitive.eval_from_aggregate(&reducer.arguments[..], constants, &outer_values[..], group)
                        }).collect::<Vec<_>>();
                        (group, output_values)
                    };
                    let mut output_sets = vec![];
                    let null = Value::Null;
                    let mut state = constants.iter().chain(outer_values.iter()).collect::<Vec<_>>();
                    if aggregate.selects_inner {
                        output_sets.push(group);
                    } else {
                        // nasty hack - fill null values for inner so that the ixes work out right
                        for _ in aggregate.inner.mapping.iter() {
                            state.push(&null);
                        }
                    }
                    for output in output_values.iter() {
                        output_sets.push(output);
                    }
                    aggregate_step(aggregate, &output_sets[..], &mut state, &mut output.index);
                    group_start = group_end;
                }
                Some(output)
            }
        }
    }
}<commit_msg>fix bug with ordinals on chunked sources<commit_after>use std::collections::BTreeSet;

use value::Value;
use relation::{Relation, IndexSelect, ViewSelect};
use primitive::{Primitive, resolve_as_scalar};
use std::cmp::{min, max, Ordering};
use std::mem::replace;

#[derive(Clone, Debug)]
pub struct Table {
    pub insert: Option<IndexSelect>,
    pub remove: Option<IndexSelect>,
}

#[derive(Clone, Debug)]
pub struct Union {
    pub selects: Vec<IndexSelect>,
}

#[derive(Clone, Debug, Copy)]
pub enum Direction {
    Ascending,
    Descending,
}

#[derive(Clone, Debug)]
pub enum Input {
    Primitive{
        primitive: Primitive,
        input_bindings: Vec<(usize, usize)>,
    },
    View{
        input_ix: usize,
    },
}

#[derive(Clone, Debug)]
pub struct Source {
    pub input: Input,
    pub grouped_fields: Vec<usize>,
    pub sorted_fields: Vec<(usize, Direction)>,
    pub chunked: bool,
    pub constraint_bindings: Vec<(usize, usize)>,
    pub output_bindings: Vec<(usize, usize)>,
}

impl Source {
    fn prepare(&self, mut rows: Vec<Vec<Value>>) -> Vec<Vec<Value>> {
        // TODO compensate for inputs if primitive
        rows.sort_by(|a, b| {
            for &ix in self.grouped_fields.iter() {
                match a[ix].cmp(&b[ix]) {
                    Ordering::Greater => return Ordering::Greater,
                    Ordering::Less => return Ordering::Less,
                    Ordering::Equal => ()
                }
            }
            for &(ix, direction) in self.sorted_fields.iter() {
                match (a[ix].cmp(&b[ix]), direction) {
                    (Ordering::Greater, Direction::Ascending) => return Ordering::Greater,
                    (Ordering::Greater, Direction::Descending) => return Ordering::Less,
                    (Ordering::Less, Direction::Ascending) => return Ordering::Less,
                    (Ordering::Less, Direction::Descending) => return Ordering::Greater,
                    (Ordering::Equal, _) => ()
                }
            }
            return Ordering::Equal;
        });

        if self.chunked {
            // if the source is chunked, we need to turn all of the fields into columns for each grouping
            // we want the grouped fields to be scalars so that they can be joined against and every other
            // field should be a column of the values for that group
            let mut groups = Vec::new();
            let mut ordinal = 0;
            let grouped_fields = &self.grouped_fields;
            let row_len = rows.iter().next().map_or(0, |row| row.len());
            let ungrouped_fields = &(0..row_len).filter(|field_ix| !self.grouped_fields.contains(field_ix)).collect::<Vec<usize>>();
            let mut group: Vec<Value> = vec![Value::Column(vec![]); row_len]; // this creates a vector of the given length with all the values set to Value::Null
            for mut row in rows.into_iter() {
                // check if this is a new group
                let is_new_group = match self.grouped_fields.len() {
                    0 => false,
                    _ => !self.grouped_fields.iter().all(|&ix| row[ix] == group[ix])
                };
                if is_new_group {
                    // add the current ordinal to the group
                    group.push(Value::Float(ordinal as f64));
                    // push the old group
                    groups.push(group);
                    // create a new group with every field being a column
                    ordinal += 1;
                    group = vec![Value::Column(vec![]); row_len];
                    // any grouped field should have a scalar value for its index so that it can be joined on
                    // e.g. ["foo", [..], [..]] if you've grouped on index 0
                    for &grouped_field_ix in grouped_fields {
                        group[grouped_field_ix] = row[grouped_field_ix].clone();
                    }
                }
                // add the fields for this row
                for &field_ix in ungrouped_fields {
                    group[field_ix].as_column_mut().push(row[field_ix].clone());
                }
            }
            group.push(Value::Float(ordinal as f64));
            groups.push(group);
            // if there are grouped fields then the first group is the initial placeholder and
            // should be removed
            if self.grouped_fields.len() > 0 {
                groups.remove(0);
            }
            groups
        } else {
            // in the case where the source is not chunked all we need to worry about is setting the ordinal correctly
            // if there are no groups, then this is just normal rows and we number them monotonically
            // if there are groups, each new group resets the ordinal to 1
            let otherwise = vec![];
            let mut groups = Vec::new();
            let mut ordinal = 1;
            let mut prev_row = rows.iter().next().map_or(otherwise, |row| row.clone());
            for mut row in rows.into_iter() {
                let is_new_group = match self.grouped_fields.len() {
                    0 => false,
                    _ => !self.grouped_fields.iter().all(|&ix| row[ix] == prev_row[ix])
                };
                if is_new_group {
                    ordinal = 1;
                    prev_row = row.clone(); // @TODO: is there a way to do this without having to clone it?
                }
                row.push(Value::Float(ordinal as f64));
                groups.push(row);
                ordinal += 1;
            }
            groups
        }
    }
}

#[derive(Clone, Debug)]
pub struct Join {
    pub constants: Vec<Value>,
    pub sources: Vec<Source>,
    pub select: Vec<usize>,
}

#[derive(Clone, Debug)]
pub struct Reducer {
    pub primitive: Primitive,
    pub arguments: Vec<usize>,
}

#[derive(Clone, Debug)]
pub struct Aggregate {
    pub constants: Vec<Value>,
    pub outer: IndexSelect,
    pub inner: IndexSelect,
    pub directions: Vec<Direction>,
    pub limit_from: Option<usize>,
    pub limit_to: Option<usize>,
    pub reducers: Vec<Reducer>,
    pub selects_inner: bool,
    pub select: ViewSelect,
}

#[derive(Clone, Debug)]
pub enum View {
    Table(Table),
    Union(Union),
    Join(Join),
    Aggregate(Aggregate),
}

fn push_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for value in input.iter() {
        state.push(value);
    }
}

fn pop_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for _ in input.iter() {
        state.pop();
    }
}

// TODO this algorithm is incredibly naive and also clones excessively
fn join_step(join: &Join, ix: usize, inputs: &[Vec<Vec<Value>>], state: &mut Vec<Value>, index: &mut BTreeSet<Vec<Value>>) {
    if ix == join.sources.len() {
        index.insert(join.select.iter().map(|ix| state[*ix].clone()).collect());
    } else {
        let source = &join.sources[ix];
        match source.input {
            Input::View{..} => {
                for values in inputs[ix].iter() {
                    for &(field_ix, variable_ix) in source.output_bindings.iter() {
                        state[variable_ix] = values[field_ix].clone();
                    }
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        state[variable_ix] == values[field_ix]
                    ) {
                        join_step(join, ix+1, inputs, state, index);
                    }
                }
            }
            Input::Primitive{primitive, ref input_bindings} => {
                // values returned from primitives don't include inputs, so we will have to offset accesses by input_len
                let input_len = input_bindings.len();
                for mut values in primitive.eval_from_join(&input_bindings[..], &state[..]).into_iter() {
                    for &(field_ix, variable_ix) in source.output_bindings.iter() {
                        state[variable_ix] = replace(&mut values[field_ix - input_len], Value::Null);
                    }
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        state[variable_ix] == values[field_ix - input_len]
                    ) {
                        join_step(join, ix+1, inputs, state, index);
                    }
                }
            }
        }
    }
}

fn aggregate_step<'a>(aggregate: &Aggregate, input_sets: &'a [&[Vec<Value>]], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if input_sets.len() == 0 {
        index.insert(aggregate.select.select(&state[..]));
    } else {
        for values in input_sets[0].iter() {
            push_all(state, values);
            aggregate_step(aggregate, &input_sets[1..], state, index);
            pop_all(state, values);
        }
    }
}

fn compare_in_direction(xs: &[Value], ys: &[Value], directions: &[Direction]) -> Ordering {
    for ((x,y), direction) in xs.iter().zip(ys.iter()).zip(directions.iter()) {
        let cmp = match *direction {
            Direction::Ascending => x.cmp(y),
            Direction::Descending => y.cmp(x),
        };
        match cmp {
            Ordering::Greater => return Ordering::Greater,
            Ordering::Equal => (),
            Ordering::Less => return Ordering::Less,
        };
    }
    return Ordering::Equal;
}

impl View {
    pub fn run(&self, old_output: &Relation, upstream: &[&Relation]) -> Option<Relation> {
        let mut output = Relation::new(
            old_output.view.clone(),
            old_output.fields.clone(),
            old_output.names.clone()
            );
        match *self {
            View::Table(_) => None,
            View::Union(ref union) => {
                assert_eq!(union.selects.len(), upstream.len());
                for select in union.selects.iter() {
                    for values in select.select(&upstream[..]) {
                        output.index.insert(values);
                    }
                }
                Some(output)
            }
            View::Join(ref join) => {
                let mut state = join.constants.clone();
                let inputs = join.sources.iter().map(|source|
                    match source.input {
                        Input::Primitive{..} => {
                            vec![]
                        }
                        Input::View{input_ix, ..} => {
                            source.prepare(upstream[input_ix].index.iter().map(|row| row.clone()).collect())
                        }
                    }).collect::<Vec<_>>();
                join_step(join, 0, &inputs[..], &mut state, &mut output.index);
                Some(output)
            }
            View::Aggregate(ref aggregate) => {
                let mut outer = aggregate.outer.select(&upstream[..]);
                let mut inner = aggregate.inner.select(&upstream[..]);
                outer.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                outer.dedup();
                inner.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                let constants = &aggregate.constants[..];
                let mut group_start = 0;
                for outer_values in outer.into_iter() {
                    let mut group_end = group_start;
                    while (group_end < inner.len())
                    && (inner[group_end][0..outer_values.len()] == outer_values[..]) {
                        group_end += 1;
                    }
                    let (group, output_values) = {
                        let limit_from = match aggregate.limit_from {
                            None => group_start,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_to = match aggregate.limit_to {
                            None => group_end,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_from = min(max(limit_from, group_start), group_end);
                        let limit_to = min(max(limit_to, limit_from), group_end);
                        let group = &inner[limit_from..limit_to];
                        let output_values = aggregate.reducers.iter().map(|reducer| {
                            reducer.primitive.eval_from_aggregate(&reducer.arguments[..], constants, &outer_values[..], group)
                        }).collect::<Vec<_>>();
                        (group, output_values)
                    };
                    let mut output_sets = vec![];
                    let null = Value::Null;
                    let mut state = constants.iter().chain(outer_values.iter()).collect::<Vec<_>>();
                    if aggregate.selects_inner {
                        output_sets.push(group);
                    } else {
                        // nasty hack - fill null values for inner so that the ixes work out right
                        for _ in aggregate.inner.mapping.iter() {
                            state.push(&null);
                        }
                    }
                    for output in output_values.iter() {
                        output_sets.push(output);
                    }
                    aggregate_step(aggregate, &output_sets[..], &mut state, &mut output.index);
                    group_start = group_end;
                }
                Some(output)
            }
        }
    }
}<|endoftext|>


<commit_before>
macro_rules! state_machine {
    (
        $mach_name:ident;
        $state_enum:ident;
        $event_enum:ident;
        [$($initial_state:tt)+];
        [$(
            [ $($state_name:tt)+ ] => [
                $( [ $($incoming_event:tt)+ ] [ $($new_state:tt)+ ] | $action:block ),+
            ]
        ),*]
    ) => (

        #[allow(dead_code)]
        #[derive(Debug)]
        pub struct $mach_name {
            current_state: $state_enum
        }

        impl $mach_name {
            pub fn new() -> $mach_name {
                $mach_name {
                    current_state: $initial_state
                }
            }

            pub fn process_event(&mut self, event: $event_enum) {
                match Some(&self.current_state) {
                    $(Some(&$state_name) => {
                        match Some(&event) {
                            $(Some(&$($incoming_event)+) => { {$action}; self.current_state = $new_state }),+
                            _ => {}
                        }
                    }),*
                    _ => {}
                }
            }
        }
    )

}
<commit_msg>Removed unused state machine macro<commit_after><|endoftext|>


<commit_before>use rustc_index::bit_set::HybridBitSet;
use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
use rustc_middle::mir::{self, BasicBlock, Local, Location};
use rustc_middle::ty::TyCtxt;

use crate::transform::{MirPass, MirSource};

/// This pass looks for MIR that always copies the same local into the return place and eliminates
/// the copy by renaming all uses of that local to `_0`.
///
/// This allows LLVM to perform an optimization similar to the named return value optimization
/// (NRVO) that is guaranteed in C++. This avoids a stack allocation and `memcpy` for the
/// relatively common pattern of allocating a buffer on the stack, mutating it, and returning it by
/// value like so:
///
/// ```rust
/// fn foo(init: fn(&mut [u8; 1024])) -> [u8; 1024] {
///     let mut buf = [0; 1024];
///     init(&mut buf);
///     buf
/// }
/// ```
///
/// For now, this pass is very simple and only capable of eliminating a single copy.
/// A more general version of copy propagation could yield even more benefits.
pub struct RenameReturnPlace;

impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut mir::Body<'tcx>) {
        if tcx.sess.opts.debugging_opts.mir_opt_level == 0 {
            return;
        }

        let returned_local = match local_eligible_for_nrvo(body) {
            Some(l) => l,
            None => {
                debug!("`{:?}` was ineligible for NRVO", src.def_id());
                return;
            }
        };

        debug!(
            "`{:?}` was eligible for NRVO, making {:?} the return place",
            src.def_id(),
            returned_local
        );

        RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body(body);

        // Clean up the `NOP`s we inserted for statements made useless by our renaming.
        for block_data in body.basic_blocks_mut() {
            block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
        }

        // Overwrite the debuginfo of `_0` with that of the renamed local.
        let (renamed_decl, ret_decl) =
            body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE);
        debug_assert_eq!(ret_decl.ty, renamed_decl.ty);
        ret_decl.clone_from(renamed_decl);
    }
}

/// MIR that is eligible for the NRVO must fulfill two conditions:
///   1. The return place must not be read prior to the `Return` terminator.
///   2. A simple assignment of a whole local to the return place (e.g., `_0 = _1`) must be the
///      only definition of the return place reaching the `Return` terminator.
///
/// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned
/// to the return place along all possible paths through the control-flow graph.
fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
    if IsReturnPlaceRead::run(body) {
        return None;
    }

    let mut copied_to_return_place = None;
    for block in body.basic_blocks().indices() {
        // Look for blocks with a `Return` terminator.
        if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) {
            continue;
        }

        // Look for an assignment of a single local to the return place prior to the `Return`.
        let returned_local = find_local_assigned_to_return_place(block, body)?;
        match body.local_kind(returned_local) {
            // FIXME: Can we do this for arguments as well?
            mir::LocalKind::Arg => return None,

            mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
            mir::LocalKind::Var | mir::LocalKind::Temp => {}
        }

        // If multiple different locals are copied to the return place. We can't pick a
        // single one to rename.
        if copied_to_return_place.map_or(false, |old| old != returned_local) {
            return None;
        }

        copied_to_return_place = Some(returned_local);
    }

    return copied_to_return_place;
}

fn find_local_assigned_to_return_place(
    start: BasicBlock,
    body: &mut mir::Body<'_>,
) -> Option<Local> {
    let mut block = start;
    let mut seen = HybridBitSet::new_empty(body.basic_blocks().len());

    // Iterate as long as `block` has exactly one predecessor that we have not yet visited.
    while seen.insert(block) {
        trace!("Looking for assignments to `_0` in {:?}", block);

        let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place);
        if local.is_some() {
            return local;
        }

        match body.predecessors()[block].as_slice() {
            &[pred] => block = pred,
            _ => return None,
        }
    }

    return None;
}

// If this statement is an assignment of an unprojected local to the return place,
// return that local.
fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> {
    if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
        if lhs.as_local() == Some(mir::RETURN_PLACE) {
            if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs {
                return rhs.as_local();
            }
        }
    }

    None
}

struct RenameToReturnPlace<'tcx> {
    to_rename: Local,
    tcx: TyCtxt<'tcx>,
}

/// Replaces all uses of `self.to_rename` with `_0`.
impl MutVisitor<'tcx> for RenameToReturnPlace<'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) {
        // Remove assignments of the local being replaced to the return place, since it is now the
        // return place:
        //     _0 = _1
        if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) {
            stmt.kind = mir::StatementKind::Nop;
            return;
        }

        // Remove storage annotations for the local being replaced:
        //     StorageLive(_1)
        if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) =
            stmt.kind
        {
            if local == self.to_rename {
                stmt.kind = mir::StatementKind::Nop;
                return;
            }
        }

        self.super_statement(stmt, loc)
    }

    fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }

    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
        assert_ne!(*l, mir::RETURN_PLACE);
        if *l == self.to_rename {
            *l = mir::RETURN_PLACE;
        }
    }
}

struct IsReturnPlaceRead(bool);

impl IsReturnPlaceRead {
    fn run(body: &mir::Body<'_>) -> bool {
        let mut vis = IsReturnPlaceRead(false);
        vis.visit_body(body);
        vis.0
    }
}

impl Visitor<'tcx> for IsReturnPlaceRead {
    fn visit_local(&mut self, &l: &Local, ctxt: PlaceContext, _: Location) {
        if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
            self.0 = true;
        }
    }

    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }
}
<commit_msg>Expand comment with possible improvements<commit_after>use rustc_index::bit_set::HybridBitSet;
use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
use rustc_middle::mir::{self, BasicBlock, Local, Location};
use rustc_middle::ty::TyCtxt;

use crate::transform::{MirPass, MirSource};

/// This pass looks for MIR that always copies the same local into the return place and eliminates
/// the copy by renaming all uses of that local to `_0`.
///
/// This allows LLVM to perform an optimization similar to the named return value optimization
/// (NRVO) that is guaranteed in C++. This avoids a stack allocation and `memcpy` for the
/// relatively common pattern of allocating a buffer on the stack, mutating it, and returning it by
/// value like so:
///
/// ```rust
/// fn foo(init: fn(&mut [u8; 1024])) -> [u8; 1024] {
///     let mut buf = [0; 1024];
///     init(&mut buf);
///     buf
/// }
/// ```
///
/// For now, this pass is very simple and only capable of eliminating a single copy. A more general
/// version of copy propagation, such as the one based on non-overlapping live ranges in [#47954] and
/// [#71003], could yield even more benefits.
///
/// [#47954]: https://github.com/rust-lang/rust/pull/47954
/// [#71003]: https://github.com/rust-lang/rust/pull/71003
pub struct RenameReturnPlace;

impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut mir::Body<'tcx>) {
        if tcx.sess.opts.debugging_opts.mir_opt_level == 0 {
            return;
        }

        let returned_local = match local_eligible_for_nrvo(body) {
            Some(l) => l,
            None => {
                debug!("`{:?}` was ineligible for NRVO", src.def_id());
                return;
            }
        };

        debug!(
            "`{:?}` was eligible for NRVO, making {:?} the return place",
            src.def_id(),
            returned_local
        );

        RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body(body);

        // Clean up the `NOP`s we inserted for statements made useless by our renaming.
        for block_data in body.basic_blocks_mut() {
            block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
        }

        // Overwrite the debuginfo of `_0` with that of the renamed local.
        let (renamed_decl, ret_decl) =
            body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE);
        debug_assert_eq!(ret_decl.ty, renamed_decl.ty);
        ret_decl.clone_from(renamed_decl);
    }
}

/// MIR that is eligible for the NRVO must fulfill two conditions:
///   1. The return place must not be read prior to the `Return` terminator.
///   2. A simple assignment of a whole local to the return place (e.g., `_0 = _1`) must be the
///      only definition of the return place reaching the `Return` terminator.
///
/// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned
/// to the return place along all possible paths through the control-flow graph.
fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
    if IsReturnPlaceRead::run(body) {
        return None;
    }

    let mut copied_to_return_place = None;
    for block in body.basic_blocks().indices() {
        // Look for blocks with a `Return` terminator.
        if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) {
            continue;
        }

        // Look for an assignment of a single local to the return place prior to the `Return`.
        let returned_local = find_local_assigned_to_return_place(block, body)?;
        match body.local_kind(returned_local) {
            // FIXME: Can we do this for arguments as well?
            mir::LocalKind::Arg => return None,

            mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
            mir::LocalKind::Var | mir::LocalKind::Temp => {}
        }

        // If multiple different locals are copied to the return place. We can't pick a
        // single one to rename.
        if copied_to_return_place.map_or(false, |old| old != returned_local) {
            return None;
        }

        copied_to_return_place = Some(returned_local);
    }

    return copied_to_return_place;
}

fn find_local_assigned_to_return_place(
    start: BasicBlock,
    body: &mut mir::Body<'_>,
) -> Option<Local> {
    let mut block = start;
    let mut seen = HybridBitSet::new_empty(body.basic_blocks().len());

    // Iterate as long as `block` has exactly one predecessor that we have not yet visited.
    while seen.insert(block) {
        trace!("Looking for assignments to `_0` in {:?}", block);

        let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place);
        if local.is_some() {
            return local;
        }

        match body.predecessors()[block].as_slice() {
            &[pred] => block = pred,
            _ => return None,
        }
    }

    return None;
}

// If this statement is an assignment of an unprojected local to the return place,
// return that local.
fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> {
    if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
        if lhs.as_local() == Some(mir::RETURN_PLACE) {
            if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs {
                return rhs.as_local();
            }
        }
    }

    None
}

struct RenameToReturnPlace<'tcx> {
    to_rename: Local,
    tcx: TyCtxt<'tcx>,
}

/// Replaces all uses of `self.to_rename` with `_0`.
impl MutVisitor<'tcx> for RenameToReturnPlace<'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) {
        // Remove assignments of the local being replaced to the return place, since it is now the
        // return place:
        //     _0 = _1
        if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) {
            stmt.kind = mir::StatementKind::Nop;
            return;
        }

        // Remove storage annotations for the local being replaced:
        //     StorageLive(_1)
        if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) =
            stmt.kind
        {
            if local == self.to_rename {
                stmt.kind = mir::StatementKind::Nop;
                return;
            }
        }

        self.super_statement(stmt, loc)
    }

    fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }

    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
        assert_ne!(*l, mir::RETURN_PLACE);
        if *l == self.to_rename {
            *l = mir::RETURN_PLACE;
        }
    }
}

struct IsReturnPlaceRead(bool);

impl IsReturnPlaceRead {
    fn run(body: &mir::Body<'_>) -> bool {
        let mut vis = IsReturnPlaceRead(false);
        vis.visit_body(body);
        vis.0
    }
}

impl Visitor<'tcx> for IsReturnPlaceRead {
    fn visit_local(&mut self, &l: &Local, ctxt: PlaceContext, _: Location) {
        if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
            self.0 = true;
        }
    }

    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) {
        // Ignore the implicit "use" of the return place in a `Return` statement.
        if let mir::TerminatorKind::Return = terminator.kind {
            return;
        }

        self.super_terminator(terminator, loc);
    }
}
<|endoftext|>


<commit_before>use cpu::common_defs::OpcodeExecInfo;
use cpu::CpuState;
use cpu::DecodeRegister;
use cpu::common_defs::address_mode::AddressMode;
use cpu::common_defs::opcode_class::OpcodeClass;
use memory::Memory;

macro_rules! set_zs {
    ($cpu_state:expr,$val:expr) => (set_z!($cpu_state,$val);set_s!($cpu_state,$val););
}
macro_rules! set_z {
    ($cpu_state:expr,$val:expr) => ($cpu_state.Z = $val == 0;);
}
macro_rules! set_s {
    ($cpu_state:expr,$val:expr) => ($cpu_state.S = $val >= 0x80;);
}
macro_rules! stack_push8 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write8(0x0100 as u16 | $cpu_state.sp as u16,$val)); $cpu_state.sp-=1;);
}
macro_rules! stack_push16 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write16(0x0100 as u16 | $cpu_state.sp as u16 -1,$val)); $cpu_state.sp-=2;);
}
macro_rules! stack_pull8 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.sp+=1;
            $mem.read8(0x0100 | $cpu_state.sp;)
        }
    }
}
macro_rules! stack_pull16 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.sp+=2;
            $mem.read16(0x0100 as u16 | $cpu_state.sp as u16 -1)
        }
    }
}

#[derive(Debug)]
pub enum ExecutionError {
  MemoryError(::memory::MemoryError),
  UnexpectedOpcode(String),
  UnexpectedAddressMode(String)
}
impl From<::memory::MemoryError> for ExecutionError {
    fn from(err: ::memory::MemoryError) -> ExecutionError {
        ExecutionError::MemoryError(err)
    }
}

pub struct CpuExecutor {
    op_table: Vec<OpcodeExecInfo>,
}

impl CpuExecutor {
	/// Construct a new CpuExecutor.
    pub fn new(opcodes: Vec<OpcodeExecInfo> ) -> CpuExecutor {
        CpuExecutor {
            op_table:  opcodes,
        }
    }

	/// Reset the given cpu_state.
    pub fn power_on(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        self.reset(cpu_state,mem);
    }
	/// Reset the given cpu_state.
    pub fn reset(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        cpu_state.pc = mem.read16(0xFFFC).unwrap();
        cpu_state.sp = 0xFD;
        cpu_state.pack_flags(0x24);
    }

    /// Execute a single instruction in the given memory and cpu_state context.
    pub fn step(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        try!(self.fetch_and_decode(cpu_state,mem));
        try!(self.execute(cpu_state,mem));
        Ok(())
    }

	/// Fetch the next instruction and perform address resolution.
    pub fn fetch_and_decode(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        cpu_state.instruction_register = mem.read8(cpu_state.pc).unwrap();
        cpu_state.decode_register = try!(self.decode(cpu_state,mem));
        cpu_state.pc = cpu_state.pc+1;
        Ok(())
    }

    /// Perform address resolution, returning the info in a DecodeRegister.
    fn decode(self: &CpuExecutor, cpu_state: &CpuState, mem: &Memory) -> Result<DecodeRegister,ExecutionError> {
        let mut dr = DecodeRegister {
            info : self.op_table[cpu_state.instruction_register as usize].clone(),
            ..Default::default()
        };
        match dr.info.address_mode {
            // no explicit addresses for the following modes
            AddressMode::Accumulator  => { },
            AddressMode::Implied      => { },

            // explicit addresses from here on out
            AddressMode::Absolute => {
                dr.addr_final  = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteX       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.x as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteY       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Immediate       => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Indirect        => {
                dr.addr_intermediate = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
            },
            AddressMode::IndexedIndirect => {
                dr.addr_intermediate = Some( (mem.read8(cpu_state.pc+1).unwrap() + cpu_state.x) as u16);
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::IndirectIndexed => {
                dr.addr_init         = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.addr_intermediate = Some(mem.read16(dr.addr_init.unwrap()).unwrap());
                dr.addr_final        = Some( ((dr.addr_intermediate.unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Relative        => {
                dr.value_final = Some(mem.read8(cpu_state.pc+1).unwrap());

                // NOTE: relative is from the end of the current instruction and relative
                // instructions are 2 bytes long, so we add 2 before adding in the specified offset
                //
                if dr.value_final.unwrap() < 0x80 { dr.addr_final  = Some(cpu_state.pc + 2 + dr.value_final.unwrap() as u16); }
                else                              { dr.addr_final  = Some(cpu_state.pc + 2 + dr.value_final.unwrap() as u16 - 0x100);}
            },
            AddressMode::ZeroPage        => {
                dr.addr_final  = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageX       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.x as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageY       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.y as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            _ => { return Err(ExecutionError::UnexpectedAddressMode(format!("unrecognized addressing mode '{:?}' while decoding instruction_register!",dr.info.address_mode))); }
        }

        return Ok(dr);
    }

    /// Perform the current instruction, returning the CpuState after execution.
    pub fn execute(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) -> Result<(),ExecutionError> {
    	// Figure out which opcode is being executed.
    	match cpu_state.decode_register.info.opcode_class {
    		OpcodeClass::ASL => {
				cpu_state.a = cpu_state.a << 1;
				cpu_state.C = (128 & cpu_state.a) > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::BCC => {
                if !cpu_state.C { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BCS => {
                if cpu_state.C { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BEQ => {
                if cpu_state.Z { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BIT => {
                cpu_state.V = (cpu_state.decode_register.value_final.unwrap() >> 6) & 1 > 0;
                set_z!(cpu_state,cpu_state.decode_register.value_final.unwrap() & cpu_state.a);
                set_s!(cpu_state,cpu_state.decode_register.value_final.unwrap());

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::BNE => {
                if !cpu_state.Z { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BPL => {
                if !cpu_state.S { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BVC => {
                if !cpu_state.V { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BVS => {
                if cpu_state.V { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::CLC => {
                cpu_state.C = false;
    		},
    		OpcodeClass::LDA => {
    			cpu_state.a = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.a);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LDX => {
    			cpu_state.x = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.x);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LDY => {
    			cpu_state.y = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.y);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LSR => {
    			cpu_state.a = cpu_state.a >> 1;
    			cpu_state.C = 1 & cpu_state.a > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::JMP => {
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::JSR => {
                stack_push16!(cpu_state,mem,cpu_state.pc+1);
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::NOP => {
            },
    		OpcodeClass::RTS => {
                cpu_state.pc = stack_pull16!(cpu_state,mem).unwrap() + 1;
            },
    		OpcodeClass::SEC => {
                cpu_state.C = true;
    		},
    		OpcodeClass::SED => {
                cpu_state.D = true;
    		},
    		OpcodeClass::SEI => {
                cpu_state.I = true;
    		},
    		OpcodeClass::STA => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.a));
                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::STX => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.x));
                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},

			_ => { return Err(ExecutionError::UnexpectedOpcode(format!("Unrecognised opcode class: {:?}", cpu_state.decode_register.info.opcode_class)));}

    	}
        Ok(())
    }
}

<commit_msg>implemented the Personal Home Page (PHP) opcode.  ahahahaha...<commit_after>use cpu::common_defs::OpcodeExecInfo;
use cpu::CpuState;
use cpu::DecodeRegister;
use cpu::common_defs::address_mode::AddressMode;
use cpu::common_defs::opcode_class::OpcodeClass;
use memory::Memory;

macro_rules! set_zs {
    ($cpu_state:expr,$val:expr) => (set_z!($cpu_state,$val);set_s!($cpu_state,$val););
}
macro_rules! set_z {
    ($cpu_state:expr,$val:expr) => ($cpu_state.Z = $val == 0;);
}
macro_rules! set_s {
    ($cpu_state:expr,$val:expr) => ($cpu_state.S = $val >= 0x80;);
}
macro_rules! stack_push8 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write8(0x0100 as u16 | $cpu_state.sp as u16,$val)); $cpu_state.sp-=1;);
}
macro_rules! stack_push16 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write16(0x0100 as u16 | $cpu_state.sp as u16 -1,$val)); $cpu_state.sp-=2;);
}
macro_rules! stack_pull8 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.sp+=1;
            $mem.read8(0x0100 | $cpu_state.sp;)
        }
    }
}
macro_rules! stack_pull16 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.sp+=2;
            $mem.read16(0x0100 as u16 | $cpu_state.sp as u16 -1)
        }
    }
}

#[derive(Debug)]
pub enum ExecutionError {
  MemoryError(::memory::MemoryError),
  UnexpectedOpcode(String),
  UnexpectedAddressMode(String)
}
impl From<::memory::MemoryError> for ExecutionError {
    fn from(err: ::memory::MemoryError) -> ExecutionError {
        ExecutionError::MemoryError(err)
    }
}

pub struct CpuExecutor {
    op_table: Vec<OpcodeExecInfo>,
}

impl CpuExecutor {
	/// Construct a new CpuExecutor.
    pub fn new(opcodes: Vec<OpcodeExecInfo> ) -> CpuExecutor {
        CpuExecutor {
            op_table:  opcodes,
        }
    }

	/// Reset the given cpu_state.
    pub fn power_on(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        self.reset(cpu_state,mem);
    }
	/// Reset the given cpu_state.
    pub fn reset(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        cpu_state.pc = mem.read16(0xFFFC).unwrap();
        cpu_state.sp = 0xFD;
        cpu_state.pack_flags(0x24);
    }

    /// Execute a single instruction in the given memory and cpu_state context.
    pub fn step(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        try!(self.fetch_and_decode(cpu_state,mem));
        try!(self.execute(cpu_state,mem));
        Ok(())
    }

	/// Fetch the next instruction and perform address resolution.
    pub fn fetch_and_decode(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        cpu_state.instruction_register = mem.read8(cpu_state.pc).unwrap();
        cpu_state.decode_register = try!(self.decode(cpu_state,mem));
        cpu_state.pc = cpu_state.pc+1;
        Ok(())
    }

    /// Perform address resolution, returning the info in a DecodeRegister.
    fn decode(self: &CpuExecutor, cpu_state: &CpuState, mem: &Memory) -> Result<DecodeRegister,ExecutionError> {
        let mut dr = DecodeRegister {
            info : self.op_table[cpu_state.instruction_register as usize].clone(),
            ..Default::default()
        };
        match dr.info.address_mode {
            // no explicit addresses for the following modes
            AddressMode::Accumulator  => { },
            AddressMode::Implied      => { },

            // explicit addresses from here on out
            AddressMode::Absolute => {
                dr.addr_final  = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteX       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.x as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteY       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Immediate       => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Indirect        => {
                dr.addr_intermediate = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
            },
            AddressMode::IndexedIndirect => {
                dr.addr_intermediate = Some( (mem.read8(cpu_state.pc+1).unwrap() + cpu_state.x) as u16);
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::IndirectIndexed => {
                dr.addr_init         = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.addr_intermediate = Some(mem.read16(dr.addr_init.unwrap()).unwrap());
                dr.addr_final        = Some( ((dr.addr_intermediate.unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Relative        => {
                dr.value_final = Some(mem.read8(cpu_state.pc+1).unwrap());

                // NOTE: relative is from the end of the current instruction and relative
                // instructions are 2 bytes long, so we add 2 before adding in the specified offset
                //
                if dr.value_final.unwrap() < 0x80 { dr.addr_final  = Some(cpu_state.pc + 2 + dr.value_final.unwrap() as u16); }
                else                              { dr.addr_final  = Some(cpu_state.pc + 2 + dr.value_final.unwrap() as u16 - 0x100);}
            },
            AddressMode::ZeroPage        => {
                dr.addr_final  = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageX       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.x as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageY       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.y as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            _ => { return Err(ExecutionError::UnexpectedAddressMode(format!("unrecognized addressing mode '{:?}' while decoding instruction_register!",dr.info.address_mode))); }
        }

        return Ok(dr);
    }

    /// Perform the current instruction, returning the CpuState after execution.
    pub fn execute(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) -> Result<(),ExecutionError> {
    	// Figure out which opcode is being executed.
    	match cpu_state.decode_register.info.opcode_class {
    		OpcodeClass::ASL => {
				cpu_state.a = cpu_state.a << 1;
				cpu_state.C = (128 & cpu_state.a) > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::BCC => {
                if !cpu_state.C { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BCS => {
                if cpu_state.C { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BEQ => {
                if cpu_state.Z { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BIT => {
                cpu_state.V = (cpu_state.decode_register.value_final.unwrap() >> 6) & 1 > 0;
                set_z!(cpu_state,cpu_state.decode_register.value_final.unwrap() & cpu_state.a);
                set_s!(cpu_state,cpu_state.decode_register.value_final.unwrap());

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::BNE => {
                if !cpu_state.Z { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BPL => {
                if !cpu_state.S { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BVC => {
                if !cpu_state.V { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else            { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::BVS => {
                if cpu_state.V { cpu_state.pc = cpu_state.decode_register.addr_final.unwrap(); }
                else           { cpu_state.pc += cpu_state.decode_register.info.len as u16-1;  }
    		},
    		OpcodeClass::CLC => {
                cpu_state.C = false;
    		},
    		OpcodeClass::LDA => {
    			cpu_state.a = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.a);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LDX => {
    			cpu_state.x = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.x);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LDY => {
    			cpu_state.y = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.y);

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::LSR => {
    			cpu_state.a = cpu_state.a >> 1;
    			cpu_state.C = 1 & cpu_state.a > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::JMP => {
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::JSR => {
                stack_push16!(cpu_state,mem,cpu_state.pc+1);
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::NOP => {
            },
    		OpcodeClass::PHP => {
                stack_push8!(cpu_state,mem,cpu_state.unpack_flags());
            },
    		OpcodeClass::RTS => {
                cpu_state.pc = stack_pull16!(cpu_state,mem).unwrap() + 1;
            },
    		OpcodeClass::SEC => {
                cpu_state.C = true;
    		},
    		OpcodeClass::SED => {
                cpu_state.D = true;
    		},
    		OpcodeClass::SEI => {
                cpu_state.I = true;
    		},
    		OpcodeClass::STA => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.a));
                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},
    		OpcodeClass::STX => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.x));
                cpu_state.pc += cpu_state.decode_register.info.len as u16-1;
    		},

			_ => { return Err(ExecutionError::UnexpectedOpcode(format!("Unrecognised opcode class: {:?}", cpu_state.decode_register.info.opcode_class)));}

    	}
        Ok(())
    }
}

<|endoftext|>


<commit_before>use cpu::common_defs::OpcodeExecInfo;
use cpu::CpuState;
use cpu::DecodeRegister;
use cpu::common_defs::address_mode::AddressMode;
use cpu::common_defs::opcode_class::OpcodeClass;
use memory::Memory;
use std::collections::HashMap;

macro_rules! set_zs {
    ($cpu_state:expr,$val:expr) => ($cpu_state.Z = $val == 0;$cpu_state.S = $val >= 0x80 );
}

macro_rules! stack_push8 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write8(0x0100 | $cpu_state.sp,$val)); $cpu_state.sp-=1;);
}
macro_rules! stack_push16 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write16(0x0100 as u16 | $cpu_state.sp as u16 -1,$val)); $cpu_state.sp-=2;);
}
macro_rules! stack_pull8 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.pc+=1;
            $mem.read8(0x0100 | $cpu_state.sp;)
        }
    }
}
macro_rules! stack_pull16 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.pc+=2;
            $mem.read16(0x0100 as u16 | $cpu_state.sp as u16 -1)
        }
    }
}

#[derive(Debug)]
pub enum ExecutionError {
  MemoryError(::memory::MemoryError),
  UnexpectedOpcode(String),
  UnexpectedAddressMode(String)
}
impl From<::memory::MemoryError> for ExecutionError {
    fn from(err: ::memory::MemoryError) -> ExecutionError {
        ExecutionError::MemoryError(err)
    }
}

pub struct CpuExecutor {
    op_table: HashMap<u8,OpcodeExecInfo>,
}

impl CpuExecutor {

	/// Construct a new CpuExecutor.
    pub fn new(opcodes: Vec<OpcodeExecInfo> ) -> CpuExecutor {
        // Construct the HashMap and copy it over to CpuExecutor's constructor.
        let mut op_table_temp: HashMap<u8, OpcodeExecInfo> = HashMap::new();
        for oc in opcodes {
            if op_table_temp.contains_key(&oc.opcode) {
                println!("Encountered duplicate opcode: {:#x}", oc.opcode);
                println!("Replacing previous opcode, which was {{ opcode: {:#x}, len: {}, cycles: {}, page_cycles: {} }}", oc.opcode, oc.len, oc.cycles, oc.page_cycles);
            }
            op_table_temp.insert(oc.opcode, oc);
        }

        CpuExecutor {
            op_table: op_table_temp,
        }
    }

	/// Reset the given cpu_state.
    pub fn power_on(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        self.reset(cpu_state,mem);
    }
	/// Reset the given cpu_state.
    pub fn reset(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        cpu_state.pc = mem.read16(0xFFFC).unwrap();
        cpu_state.sp = 0xFD;
        cpu_state.pack_flags(0x24);
    }

    /// Execute a single instruction in the given memory and cpu_state context.
    pub fn step(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        try!(self.fetch_and_decode(cpu_state,mem));
        try!(self.execute(cpu_state,mem));
        Ok(())
    }

	/// Fetch the next instruction and perform address resolution.
    pub fn fetch_and_decode(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        cpu_state.instruction_register = mem.read8(cpu_state.pc).unwrap();
        cpu_state.decode_register = try!(self.decode(cpu_state,mem));
        cpu_state.pc = cpu_state.pc+1;
        Ok(())
    }

    /// Perform address resolution, returning the info in a DecodeRegister.
    fn decode(self: &CpuExecutor, cpu_state: &CpuState, mem: &Memory) -> Result<DecodeRegister,ExecutionError> {
        let mut dr = DecodeRegister {
            info : self.op_table[&cpu_state.instruction_register].clone(),
            ..Default::default()
        };
        match dr.info.address_mode {
            // no explicit addresses for the following modes
            AddressMode::Accumulator  => { },
            AddressMode::Implied      => { },

            // explicit addresses from here on out
            AddressMode::Absolute => {
                dr.addr_final  = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteX       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.x as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteY       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Immediate       => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Indirect        => {
                dr.addr_intermediate = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
            },
            AddressMode::IndexedIndirect => {
                dr.addr_intermediate = Some( (mem.read8(cpu_state.pc+1).unwrap() + cpu_state.x) as u16);
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::IndirectIndexed => {
                dr.addr_init         = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.addr_intermediate = Some(mem.read16(dr.addr_init.unwrap()).unwrap());
                dr.addr_final        = Some( ((dr.addr_intermediate.unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Relative        => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPage        => {
                dr.addr_final  = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageX       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.x as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageY       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.y as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            _ => { return Err(ExecutionError::UnexpectedAddressMode(format!("unrecognized addressing mode '{:?}' while decoding instruction_register!",dr.info.address_mode))); }
        }

        return Ok(dr);
    }

    /// Perform the current instruction, returning the CpuState after execution.
    pub fn execute(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) -> Result<(),ExecutionError> {
    	// Figure out which opcode is being executed.
    	match cpu_state.decode_register.info.opcode_class {
    		OpcodeClass::ASL => {
				cpu_state.a = cpu_state.a << 1;
				cpu_state.C = (128 & cpu_state.a) > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDA => {
    			cpu_state.a = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.a);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDX => {
    			cpu_state.x = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.x);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDY => {
    			cpu_state.y = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.y);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LSR => {
    			cpu_state.a = cpu_state.a >> 1;
    			cpu_state.C = 1 & cpu_state.a > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::JMP => {
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::JSR => {
                stack_push16!(cpu_state,mem,cpu_state.pc-1);
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::NOP => {},
    		OpcodeClass::SEC => {
                cpu_state.C = true;
    		},
    		OpcodeClass::STX => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.decode_register.value_final.unwrap()));
                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},

			_ => { return Err(ExecutionError::UnexpectedOpcode(format!("Unrecognised opcode class: {:?}", cpu_state.decode_register.info.opcode_class)));}

    	}
        Ok(())
    }
}

<commit_msg>opcode lookup table is now just an array<commit_after>use cpu::common_defs::OpcodeExecInfo;
use cpu::CpuState;
use cpu::DecodeRegister;
use cpu::common_defs::address_mode::AddressMode;
use cpu::common_defs::opcode_class::OpcodeClass;
use memory::Memory;

macro_rules! set_zs {
    ($cpu_state:expr,$val:expr) => ($cpu_state.Z = $val == 0;$cpu_state.S = $val >= 0x80 );
}

macro_rules! stack_push8 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write8(0x0100 | $cpu_state.sp,$val)); $cpu_state.sp-=1;);
}
macro_rules! stack_push16 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write16(0x0100 as u16 | $cpu_state.sp as u16 -1,$val)); $cpu_state.sp-=2;);
}
macro_rules! stack_pull8 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.pc+=1;
            $mem.read8(0x0100 | $cpu_state.sp;)
        }
    }
}
macro_rules! stack_pull16 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.pc+=2;
            $mem.read16(0x0100 as u16 | $cpu_state.sp as u16 -1)
        }
    }
}

#[derive(Debug)]
pub enum ExecutionError {
  MemoryError(::memory::MemoryError),
  UnexpectedOpcode(String),
  UnexpectedAddressMode(String)
}
impl From<::memory::MemoryError> for ExecutionError {
    fn from(err: ::memory::MemoryError) -> ExecutionError {
        ExecutionError::MemoryError(err)
    }
}

pub struct CpuExecutor {
    op_table: Vec<OpcodeExecInfo>,
}

impl CpuExecutor {
	/// Construct a new CpuExecutor.
    pub fn new(opcodes: Vec<OpcodeExecInfo> ) -> CpuExecutor {
        CpuExecutor {
            op_table:  opcodes,
        }
    }

	/// Reset the given cpu_state.
    pub fn power_on(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        self.reset(cpu_state,mem);
    }
	/// Reset the given cpu_state.
    pub fn reset(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        cpu_state.pc = mem.read16(0xFFFC).unwrap();
        cpu_state.sp = 0xFD;
        cpu_state.pack_flags(0x24);
    }

    /// Execute a single instruction in the given memory and cpu_state context.
    pub fn step(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        try!(self.fetch_and_decode(cpu_state,mem));
        try!(self.execute(cpu_state,mem));
        Ok(())
    }

	/// Fetch the next instruction and perform address resolution.
    pub fn fetch_and_decode(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        cpu_state.instruction_register = mem.read8(cpu_state.pc).unwrap();
        cpu_state.decode_register = try!(self.decode(cpu_state,mem));
        cpu_state.pc = cpu_state.pc+1;
        Ok(())
    }

    /// Perform address resolution, returning the info in a DecodeRegister.
    fn decode(self: &CpuExecutor, cpu_state: &CpuState, mem: &Memory) -> Result<DecodeRegister,ExecutionError> {
        let mut dr = DecodeRegister {
            info : self.op_table[cpu_state.instruction_register as usize].clone(),
            ..Default::default()
        };
        match dr.info.address_mode {
            // no explicit addresses for the following modes
            AddressMode::Accumulator  => { },
            AddressMode::Implied      => { },

            // explicit addresses from here on out
            AddressMode::Absolute => {
                dr.addr_final  = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteX       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.x as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteY       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Immediate       => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Indirect        => {
                dr.addr_intermediate = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
            },
            AddressMode::IndexedIndirect => {
                dr.addr_intermediate = Some( (mem.read8(cpu_state.pc+1).unwrap() + cpu_state.x) as u16);
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::IndirectIndexed => {
                dr.addr_init         = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.addr_intermediate = Some(mem.read16(dr.addr_init.unwrap()).unwrap());
                dr.addr_final        = Some( ((dr.addr_intermediate.unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Relative        => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPage        => {
                dr.addr_final  = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageX       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.x as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageY       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.y as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            _ => { return Err(ExecutionError::UnexpectedAddressMode(format!("unrecognized addressing mode '{:?}' while decoding instruction_register!",dr.info.address_mode))); }
        }

        return Ok(dr);
    }

    /// Perform the current instruction, returning the CpuState after execution.
    pub fn execute(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) -> Result<(),ExecutionError> {
    	// Figure out which opcode is being executed.
    	match cpu_state.decode_register.info.opcode_class {
    		OpcodeClass::ASL => {
				cpu_state.a = cpu_state.a << 1;
				cpu_state.C = (128 & cpu_state.a) > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDA => {
    			cpu_state.a = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.a);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDX => {
    			cpu_state.x = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.x);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDY => {
    			cpu_state.y = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.y);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LSR => {
    			cpu_state.a = cpu_state.a >> 1;
    			cpu_state.C = 1 & cpu_state.a > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::JMP => {
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::JSR => {
                stack_push16!(cpu_state,mem,cpu_state.pc-1);
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::NOP => {},
    		OpcodeClass::SEC => {
                cpu_state.C = true;
    		},
    		OpcodeClass::STX => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.decode_register.value_final.unwrap()));
                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},

			_ => { return Err(ExecutionError::UnexpectedOpcode(format!("Unrecognised opcode class: {:?}", cpu_state.decode_register.info.opcode_class)));}

    	}
        Ok(())
    }
}

<|endoftext|>


<commit_before>use chrono::Utc;
use conduit_cookie::RequestSession;

use super::prelude::*;

use crate::middleware::log_request;
use crate::models::{ApiToken, User};
use crate::util::errors::{
    account_locked, forbidden, internal, AppError, AppResult, ChainError,
    InsecurelyGeneratedTokenRevoked,
};

#[derive(Debug)]
pub struct AuthenticatedUser {
    user: User,
    token_id: Option<i32>,
}

impl AuthenticatedUser {
    pub fn user_id(&self) -> i32 {
        self.user.id
    }

    pub fn api_token_id(&self) -> Option<i32> {
        self.token_id
    }

    pub fn user(self) -> User {
        self.user
    }
}

// The Origin header (https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Origin)
// is sent with CORS requests and POST requests, and indicates where the request comes from.
// We don't want to accept authenticated requests that originated from other sites, so this
// function returns an error if the Origin header doesn't match what we expect "this site" to
// be: https://crates.io in production, or http://localhost:port/ in development.
fn verify_origin(req: &dyn RequestExt) -> AppResult<()> {
    let headers = req.headers();
    let allowed_origins = &req.app().config.allowed_origins;

    let bad_origin = headers
        .get_all(header::ORIGIN)
        .iter()
        .filter_map(|value| value.to_str().ok())
        .find(|value| !allowed_origins.iter().any(|it| it == value));

    if let Some(bad_origin) = bad_origin {
        let error_message = format!(
            "only same-origin requests can be authenticated. got {:?}",
            bad_origin
        );
        return Err(internal(&error_message))
            .chain_error(|| Box::new(forbidden()) as Box<dyn AppError>);
    }
    Ok(())
}

fn authenticate_user(req: &dyn RequestExt) -> AppResult<AuthenticatedUser> {
    let conn = req.db_conn()?;

    let session = req.session();
    let user_id_from_session = session.get("user_id").and_then(|s| s.parse::<i32>().ok());

    let (user_id, token_id) = if let Some(id) = user_id_from_session {
        (id, None)
    } else {
        // Otherwise, look for an `Authorization` header on the request
        let maybe_authorization: Option<String> = {
            req.headers()
                .get(header::AUTHORIZATION)
                .and_then(|h| h.to_str().ok())
                .map(|h| h.to_string())
        };
        if let Some(header_value) = maybe_authorization {
            let (user_id, token_id) = ApiToken::find_by_api_token(&conn, &header_value)
                .map(|token| (token.user_id, Some(token.id)))
                .map_err(|e| {
                    if e.is::<InsecurelyGeneratedTokenRevoked>() {
                        e
                    } else {
                        e.chain(internal("invalid token")).chain(forbidden())
                    }
                })?;

            (user_id, token_id)
        } else {
            // Unable to authenticate the user
            return Err(internal("no cookie session or auth header found")).chain_error(forbidden);
        }
    };

    let user = User::find(&conn, user_id)
        .chain_error(|| internal("user_id from cookie or token not found in database"))?;

    Ok(AuthenticatedUser { user, token_id })
}

impl<'a> UserAuthenticationExt for dyn RequestExt + 'a {
    /// Obtain `AuthenticatedUser` for the request or return an `Forbidden` error
    fn authenticate(&mut self) -> AppResult<AuthenticatedUser> {
        verify_origin(self)?;

        let authenticated_user = authenticate_user(self)?;

        if let Some(reason) = &authenticated_user.user.account_lock_reason {
            let still_locked = if let Some(until) = authenticated_user.user.account_lock_until {
                until > Utc::now().naive_utc()
            } else {
                true
            };
            if still_locked {
                return Err(account_locked(
                    &reason,
                    authenticated_user.user.account_lock_until,
                ));
            }
        }

        log_request::add_custom_metadata(self, "uid", authenticated_user.user_id());
        if let Some(id) = authenticated_user.api_token_id() {
            log_request::add_custom_metadata(self, "tokenid", id);
        }

        Ok(authenticated_user)
    }
}
<commit_msg>controllers::util: Convert regular comment to doc comment<commit_after>use chrono::Utc;
use conduit_cookie::RequestSession;

use super::prelude::*;

use crate::middleware::log_request;
use crate::models::{ApiToken, User};
use crate::util::errors::{
    account_locked, forbidden, internal, AppError, AppResult, ChainError,
    InsecurelyGeneratedTokenRevoked,
};

#[derive(Debug)]
pub struct AuthenticatedUser {
    user: User,
    token_id: Option<i32>,
}

impl AuthenticatedUser {
    pub fn user_id(&self) -> i32 {
        self.user.id
    }

    pub fn api_token_id(&self) -> Option<i32> {
        self.token_id
    }

    pub fn user(self) -> User {
        self.user
    }
}

/// The Origin header (https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Origin)
/// is sent with CORS requests and POST requests, and indicates where the request comes from.
/// We don't want to accept authenticated requests that originated from other sites, so this
/// function returns an error if the Origin header doesn't match what we expect "this site" to
/// be: https://crates.io in production, or http://localhost:port/ in development.
fn verify_origin(req: &dyn RequestExt) -> AppResult<()> {
    let headers = req.headers();
    let allowed_origins = &req.app().config.allowed_origins;

    let bad_origin = headers
        .get_all(header::ORIGIN)
        .iter()
        .filter_map(|value| value.to_str().ok())
        .find(|value| !allowed_origins.iter().any(|it| it == value));

    if let Some(bad_origin) = bad_origin {
        let error_message = format!(
            "only same-origin requests can be authenticated. got {:?}",
            bad_origin
        );
        return Err(internal(&error_message))
            .chain_error(|| Box::new(forbidden()) as Box<dyn AppError>);
    }
    Ok(())
}

fn authenticate_user(req: &dyn RequestExt) -> AppResult<AuthenticatedUser> {
    let conn = req.db_conn()?;

    let session = req.session();
    let user_id_from_session = session.get("user_id").and_then(|s| s.parse::<i32>().ok());

    let (user_id, token_id) = if let Some(id) = user_id_from_session {
        (id, None)
    } else {
        // Otherwise, look for an `Authorization` header on the request
        let maybe_authorization: Option<String> = {
            req.headers()
                .get(header::AUTHORIZATION)
                .and_then(|h| h.to_str().ok())
                .map(|h| h.to_string())
        };
        if let Some(header_value) = maybe_authorization {
            let (user_id, token_id) = ApiToken::find_by_api_token(&conn, &header_value)
                .map(|token| (token.user_id, Some(token.id)))
                .map_err(|e| {
                    if e.is::<InsecurelyGeneratedTokenRevoked>() {
                        e
                    } else {
                        e.chain(internal("invalid token")).chain(forbidden())
                    }
                })?;

            (user_id, token_id)
        } else {
            // Unable to authenticate the user
            return Err(internal("no cookie session or auth header found")).chain_error(forbidden);
        }
    };

    let user = User::find(&conn, user_id)
        .chain_error(|| internal("user_id from cookie or token not found in database"))?;

    Ok(AuthenticatedUser { user, token_id })
}

impl<'a> UserAuthenticationExt for dyn RequestExt + 'a {
    /// Obtain `AuthenticatedUser` for the request or return an `Forbidden` error
    fn authenticate(&mut self) -> AppResult<AuthenticatedUser> {
        verify_origin(self)?;

        let authenticated_user = authenticate_user(self)?;

        if let Some(reason) = &authenticated_user.user.account_lock_reason {
            let still_locked = if let Some(until) = authenticated_user.user.account_lock_until {
                until > Utc::now().naive_utc()
            } else {
                true
            };
            if still_locked {
                return Err(account_locked(
                    &reason,
                    authenticated_user.user.account_lock_until,
                ));
            }
        }

        log_request::add_custom_metadata(self, "uid", authenticated_user.user_id());
        if let Some(id) = authenticated_user.api_token_id() {
            log_request::add_custom_metadata(self, "tokenid", id);
        }

        Ok(authenticated_user)
    }
}
<|endoftext|>