<commit_before>//! ODBC types those representation is compatible with the ODBC C API.
//!
//! This layer has not been created using automatic code generation. It is incomplete, i.e. it does
//! not contain every symbol or constant defined in the ODBC C headers. Symbols which are
//! deprecated since ODBC 3 have been left out intentionally. While some extra type safety has been
//! added by grouping some of C's `#define` constants into `enum`-types it mostly offers the same
//! power (all) and safety guarantess(none) as the wrapped C-API.
//! ODBC 4.0 is still under development by Microsoft, so these symbols are deactivated by default
//! in the cargo.toml

pub use self::{
    attributes::*, bulk_operation::*, c_data_type::*, fetch_orientation::*, info_type::*,
    nullable::*, param_type::*, sqlreturn::*, functions::*, interval::*
};
use std::os::raw::{c_int, c_void};

mod attributes;
mod bulk_operation;
mod c_data_type;
mod fetch_orientation;
mod functions;
mod info_type;
mod nullable;
mod param_type;
mod sqlreturn;
mod interval;

//These types can never be instantiated in Rust code.
pub enum Obj {}

pub enum Env {}

pub enum Dbc {}

pub enum Stmt {}

pub enum Desc {}

pub type Handle = *mut Obj;
pub type HEnv = *mut Env;
pub type HDesc = *mut Desc;

/// The connection handle references storage of all information about the connection to the data
/// source, including status, transaction state, and error information.
pub type HDbc = *mut Dbc;
pub type HStmt = *mut Stmt;

pub type SmallInt = i16;
pub type USmallInt = u16;
pub type Integer = i32;
pub type UInteger = u32;
pub type Pointer = *mut c_void;
pub type Char = u8;
pub type WChar = u16;

#[cfg(target_pointer_width = "64")]
pub type Len = i64;
#[cfg(target_pointer_width = "32")]
pub type Len = Integer;

#[cfg(target_pointer_width = "64")]
pub type ULen = u64;
#[cfg(target_pointer_width = "32")]
pub type ULen = UInteger;

pub type HWnd = Pointer;

pub type RetCode = i16;

// flags for null-terminated string
pub const NTS: SmallInt = -3;
pub const NTSL: Integer = -3;

/// Maximum message length
pub const MAX_MESSAGE_LENGTH: SmallInt = 512;
pub const SQLSTATE_SIZE: usize = 5;
pub const SQLSTATE_SIZEW: usize = 10;

// Special SQLGetData indicator values
pub const NULL_DATA: Len = -1;
pub const NO_TOTAL: Len = -4;
pub const SS_LENGTH_UNLIMITED: ULen = 0;

/// SQL Free Statement options
#[repr(u16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum FreeStmtOption {
    /// Closes the cursor associated with StatementHandle (if one was defined) and discards all
    /// pending results. The application can reopen this cursor later by executing a SELECT
    /// statement again with the same or different parameter values. If no cursor is open, this
    /// option has no effect for the application. `SQLCloseCursor` can also be called to close a
    /// cursor.
    Close = 0,
    // SQL_DROP = 1, is deprecated in favour of SQLFreeHandle
    /// Sets the `SQL_DESC_COUNT` field of the ARD to 0, releasing all column buffers bound by
    /// `SQLBindCol` for the given StatementHandle. This does not unbind the bookmark column; to do
    /// that, the `SQL_DESC_DATA_PTR` field of the ARD for the bookmark column is set to NULL.
    /// Notice that if this operation is performed on an explicitly allocated descriptor that is
    /// shared by more than one statement, the operation will affect the bindings of all statements
    /// that share the descriptor.
    Unbind = 2,
    /// Sets the `SQL_DESC_COUNT` field of the APD to 0, releasing all parameter buffers set by
    /// `SQLBindParameter` for the given StatementHandle. If this operation is performed on an
    /// explicitly allocated descriptor that is shared by more than one statement, this operation
    /// will affect the bindings of all the statements that share the descriptor.
    ResetParams = 3,
}

/// SQL Data Types
#[repr(i16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum SqlDataType {
    UnknownType = 0,
    // also called SQL_VARIANT_TYPE since odbc 4.0
    Char = 1,
    Numeric = 2,
    Decimal = 3,
    Integer = 4,
    Smallint = 5,
    Float = 6,
    Real = 7,
    Double = 8,
    Datetime = 9,
    Varchar = 12,
    #[cfg(feature = "odbc_version_4")]
    Udt = 17,
    #[cfg(feature = "odbc_version_4")]
    Row = 19,
    #[cfg(feature = "odbc_version_4")]
    Array = 50,
    #[cfg(feature = "odbc_version_4")]
    Multiset = 55,

    // one-parameter shortcuts for date/time data types
    Date = 91,
    Time = 92,
    Timestamp = 93,
    #[cfg(feature = "odbc_version_4")]
    TimeWithTimezone = 94,
    #[cfg(feature = "odbc_version_4")]
    TimestampWithTimezone = 95,

    //SQL extended datatypes:
    ExtTimeOrInterval = 10,
    ExtTimestamp = 11,
    ExtLongVarchar = -1,
    ExtBinary = -2,
    ExtVarBinary = -3,
    ExtLongVarBinary = -4,
    ExtBigInt = -5,
    ExtTinyInt = -6,
    ExtBit = -7,
    ExtWChar = -8,
    ExtWVarChar = -9,
    ExtWLongVarChar = -10,
    ExtGuid = -11,
    SsVariant = -150,
    SsUdt = -151,
    SsXml = -152,
    SsTable = -153,
    SsTime2 = -154,
    SsTimestampOffset = -155,
}

/// Represented in C headers as SQLSMALLINT
#[repr(i16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum HandleType {
    Env = 1,
    Dbc = 2,
    Stmt = 3,
    Desc = 4,
}

/// Options for `SQLDriverConnect`
#[repr(u16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DriverConnectOption {
    NoPrompt = 0,
    Complete = 1,
    Prompt = 2,
    CompleteRequired = 3,
}

#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum SqlAttributeStringLength {
    /// SQL_IS_POINTER
    Pointer = -4,
    /// SQL_IS_UINTEGER
    UInteger = -5,
    /// SQL_IS_INTEGER
    Integer = -6,
    /// SQL_IS_USMALLINT
    USmallInt = -7,
    /// SQL_IS_SMALLINT
    SmallInt = -8,
}

pub use self::SqlAttributeStringLength::*;

/// SQL_YEAR_MONTH_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct YearMonth {
    pub year: UInteger,
    pub month: UInteger,
}

/// SQL_DAY_SECOND_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct DaySecond {
    pub day: UInteger,
    pub hour: UInteger,
    pub minute: UInteger,
    pub second: UInteger,
    pub fraction: UInteger,
}

/// SQL_INTERVAL_UNION
#[repr(C)]
#[derive(Copy, Clone)]
pub union IntervalUnion {
    pub year_month: YearMonth,
    pub day_second: DaySecond,
}

/// SQL_INTERVAL_STRUCT
#[repr(C)]
#[derive(Clone, Copy)]
pub struct IntervalStruct {
    pub interval_type: c_int,
    pub interval_sign: SmallInt,
    pub interval_value: IntervalUnion,
}

/// SQL_DATE_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct Date {
    pub year: SmallInt,
    pub month: USmallInt,
    pub day: USmallInt,
}

/// SQL_TIME_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct Time {
    pub hour: USmallInt,
    pub minute: USmallInt,
    pub second: USmallInt,
}

/// SQL_TIMESTAMP_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct Timestamp {
    pub year: SmallInt,
    pub month: USmallInt,
    pub day: USmallInt,
    pub hour: USmallInt,
    pub minute: USmallInt,
    pub second: USmallInt,
    pub fraction: UInteger,
}

/// SQLGUID
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct Guid {
    pub d1: u32,
    pub d2: u16,
    pub d3: u16,
    pub d4: [u8; 8],
}

/// SQL_SS_TIME2_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct SsTime2 {
    pub hour: USmallInt,
    pub minute: USmallInt,
    pub second: USmallInt,
    pub fraction: UInteger,
}

/// SQL_SS_TIMESTAMPOFFSET_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct SsTimestampOffset {
    pub year: SmallInt,
    pub month: USmallInt,
    pub day: USmallInt,
    pub hour: USmallInt,
    pub minute: USmallInt,
    pub second: USmallInt,
    pub fraction: UInteger,
    pub timezone_hour: SmallInt,
    pub timezone_minute: SmallInt,
}

/// Statement attributes for `SQLSetStmtAttr`
#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum StatementAttribute {
    AsyncEnable = 4,
    ParamBindType = 18,
    ParamsetSize = 22,
    RowBindType = 5,
    RowArraySize = 27,
    RowsFetchedPtr = 26,
    AsyncStmtEvent = 29,
}

/// Connection attributes for `SQLSetConnectAttr`
#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ConnectionAttribute {
    AsyncEnable = 4,
    AccessMode = 101,
    AutoCommit = 102,
    LoginTimeout = 103,
    Trace = 104,
    TraceFile = 105,
    TranslateLib = 106,
    TranslateOption = 107,
    TxnIsolation = 108,
    CurrentCatalog = 109,
    OdbcCursors = 110,
    QuietMode = 111,
    PacketSize = 112,
    ConnectionTimeout = 113,
    DisconnectBehaviour = 114,
    AsyncDbcFunctionsEnable = 117,
    AsyncDbcEvent = 119,
    EnlistInDtc = 1207,
    EnlistInXa = 1208,
    ConnectionDead = 1209,
    AutoIpd = 10001,
    MetadataId = 10014,
}

/// `DiagIdentifier` for `SQLGetDiagField`
#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum HeaderDiagnosticIdentifier {
    /// SQL_DIAG_RETURNCODE
    ReturnCode = 1,
    /// SQL_DIAG_NUMBER
    Number = 2,
    /// SQL_DIAG_ROW_COUNT
    RowCount = 3,
    /// SQL_DIAG_SQLSTATE
    SqlState = 4,
    /// SQL_DIAG_NATIVE
    Native = 5,
    /// SQL_DIAG_MESSAGE_TEXT
    MessageText = 6,
    /// SQL_DIAG_DYNAMIC_FUNCTION
    DynamicFunction = 7,
    /// SQL_DIAG_CLASS_ORIGIN
    ClassOrigin = 8,
    /// SQL_DIAG_SUBCLASS_ORIGIN
    SubclassOrigin = 9,
    /// SQL_DIAG_CONNECTION_NAME
    ConnectionName = 10,
    /// SQL_DIAG_SERVER_NAME
    ServerName = 11,
    /// SQL_DIAG_DYNAMIC_FUNCTION_CODE
    DynamicFunctionCode = 12,
    /// SQL_DIAG_CURSOR_ROW_COUNT
    CursorRowCount = -1249,
    /// SQL_DIAG_ROW_NUMBER
    RowNumber = -1248,
    /// SQL_DIAG_COLUMN_NUMBER
    ColumnNumber = -1247,
}

#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum AsyncConnectionBehavior {
    /// SQL_ASYNC_DBC_ENABLE_ON
    On = 1,
    /// SQL_ASYNC_DBC_ENABLE_OFF = 0,
    Off = 0,
}

impl Default for AsyncConnectionBehavior {
    fn default() -> AsyncConnectionBehavior {
        AsyncConnectionBehavior::Off
    }
}

#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DynamicDiagnosticIdentifier {
    /// SQL_DIAG_ALTER_DOMAIN
    AlterDomain = 3,
    /// SQL_DIAG_ALTER_TABLE,
    AlterTable = 4,
    /// SQL_DIAG_CALL
    Call = 7,
    /// SQL_DIAG_CREATE_ASSERTION
    CreateAssertion = 6,
    /// SQL_DIAG_CREATE_CHARACTER_SET
    CreateCharacterSet = 8,
    /// SQL_DIAG_CREATE_COLLATION,
    CreateCollation = 10,
    /// SQL_DIAG_CREATE_DOMAIN
    CreateDomain = 23,
    /// SQL_DIAG_CREATE_INDEX
    CreateIndex = -1,
    /// SQL_DIAG_CREATE_SCHEMA
    CreateSchema = 64,
    /// SQL_DIAG_CREATE_TABLE
    CreateTable = 77,
    /// SQL_DIAG_CREATE_TRANSLATION
    CreateTranslation = 79,
    /// SQL_DIAG_CREATE_VIEW
    CreateView = 84,
    /// SQL_DIAG_DELETE_WHERE
    DeleteWhere = 19,
    /// SQL_DIAG_DROP_ASSERTION
    DropAssertion = 24,
    /// SQL_DIAG_DROP_CHARACTER_SET
    DropCharacterSet = 25,
    /// SQL_DIAG_DROP_COLLATION
    DropCollation = 26,
    /// SQL_DIAG_DROP_DOMAIN
    DropDomain = 27,
    /// SQL_DIAG_DROP_INDEX
    DropIndex = -2,
    /// SQL_DIAG_DROP_SCHEMA
    DropSchema = 31,
    /// SQL_DIAG_DROP_TABLE
    DropTable = 32,
    /// SQL_DIAG_DROP_TRANSLATION
    DropTranslation = 33,
    /// SQL_DIAG_DROP_VIEW
    DropView = 36,
    /// SQL_DIAG_DYNAMIC_DELETE_CURSOR
    DynamicDeleteCursor = 38,
    /// SQL_DIAG_DYNAMIC_UPDATE_CURSOR
    DynamicUpdateCursor = 81,
    /// SQL_DIAG_GRANT
    Grant = 48,
    /// SQL_DIAG_INSERT
    Insert = 50,
    /// SQL_DIAG_REVOKE
    Revoke = 59,
    // SQL_DIAG_SELECT_CURSOR
    SelectCursor = 85,
    /// SQL_DIAG_UNKNOWN_STATEMENT = 0,
    UnknownStatement = 0,
    /// SQL_DIAG_UPDATE_WHERE = 82,
    UpdateWhere = 82,
}

/// Completion types for `SQLEndTrans`
#[repr(i16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum CompletionType {
    Commit = 0,
    Rollback = 1,
}
<commit_msg>docstrings for some SqlType variants<commit_after>//! ODBC types those representation is compatible with the ODBC C API.
//!
//! This layer has not been created using automatic code generation. It is incomplete, i.e. it does
//! not contain every symbol or constant defined in the ODBC C headers. Symbols which are
//! deprecated since ODBC 3 have been left out intentionally. While some extra type safety has been
//! added by grouping some of C's `#define` constants into `enum`-types it mostly offers the same
//! power (all) and safety guarantess(none) as the wrapped C-API.
//! ODBC 4.0 is still under development by Microsoft, so these symbols are deactivated by default
//! in the cargo.toml

pub use self::{
    attributes::*, bulk_operation::*, c_data_type::*, fetch_orientation::*, info_type::*,
    nullable::*, param_type::*, sqlreturn::*, functions::*, interval::*
};
use std::os::raw::{c_int, c_void};

mod attributes;
mod bulk_operation;
mod c_data_type;
mod fetch_orientation;
mod functions;
mod info_type;
mod nullable;
mod param_type;
mod sqlreturn;
mod interval;

//These types can never be instantiated in Rust code.
pub enum Obj {}

pub enum Env {}

pub enum Dbc {}

pub enum Stmt {}

pub enum Desc {}

pub type Handle = *mut Obj;
pub type HEnv = *mut Env;
pub type HDesc = *mut Desc;

/// The connection handle references storage of all information about the connection to the data
/// source, including status, transaction state, and error information.
pub type HDbc = *mut Dbc;
pub type HStmt = *mut Stmt;

pub type SmallInt = i16;
pub type USmallInt = u16;
pub type Integer = i32;
pub type UInteger = u32;
pub type Pointer = *mut c_void;
pub type Char = u8;
pub type WChar = u16;

#[cfg(target_pointer_width = "64")]
pub type Len = i64;
#[cfg(target_pointer_width = "32")]
pub type Len = Integer;

#[cfg(target_pointer_width = "64")]
pub type ULen = u64;
#[cfg(target_pointer_width = "32")]
pub type ULen = UInteger;

pub type HWnd = Pointer;

pub type RetCode = i16;

// flags for null-terminated string
pub const NTS: SmallInt = -3;
pub const NTSL: Integer = -3;

/// Maximum message length
pub const MAX_MESSAGE_LENGTH: SmallInt = 512;
pub const SQLSTATE_SIZE: usize = 5;
pub const SQLSTATE_SIZEW: usize = 10;

// Special SQLGetData indicator values
pub const NULL_DATA: Len = -1;
pub const NO_TOTAL: Len = -4;
pub const SS_LENGTH_UNLIMITED: ULen = 0;

/// SQL Free Statement options
#[repr(u16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum FreeStmtOption {
    /// Closes the cursor associated with StatementHandle (if one was defined) and discards all
    /// pending results. The application can reopen this cursor later by executing a SELECT
    /// statement again with the same or different parameter values. If no cursor is open, this
    /// option has no effect for the application. `SQLCloseCursor` can also be called to close a
    /// cursor.
    Close = 0,
    // SQL_DROP = 1, is deprecated in favour of SQLFreeHandle
    /// Sets the `SQL_DESC_COUNT` field of the ARD to 0, releasing all column buffers bound by
    /// `SQLBindCol` for the given StatementHandle. This does not unbind the bookmark column; to do
    /// that, the `SQL_DESC_DATA_PTR` field of the ARD for the bookmark column is set to NULL.
    /// Notice that if this operation is performed on an explicitly allocated descriptor that is
    /// shared by more than one statement, the operation will affect the bindings of all statements
    /// that share the descriptor.
    Unbind = 2,
    /// Sets the `SQL_DESC_COUNT` field of the APD to 0, releasing all parameter buffers set by
    /// `SQLBindParameter` for the given StatementHandle. If this operation is performed on an
    /// explicitly allocated descriptor that is shared by more than one statement, this operation
    /// will affect the bindings of all the statements that share the descriptor.
    ResetParams = 3,
}

/// SQL Data Types
#[repr(i16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum SqlDataType {
    UnknownType = 0,
    // also called SQL_VARIANT_TYPE since odbc 4.0
    Char = 1,
    Numeric = 2,
    Decimal = 3,
    Integer = 4,
    Smallint = 5,
    Float = 6,
    Real = 7,
    /// Signed, approximate, numeric value with a binary precision 53 (zero or absolute value
    /// 10[-308] to 10[308]).
    Double = 8,
    Datetime = 9,
    Varchar = 12,
    #[cfg(feature = "odbc_version_4")]
    Udt = 17,
    #[cfg(feature = "odbc_version_4")]
    Row = 19,
    #[cfg(feature = "odbc_version_4")]
    Array = 50,
    #[cfg(feature = "odbc_version_4")]
    Multiset = 55,

    // one-parameter shortcuts for date/time data types
    Date = 91,
    Time = 92,
    /// Year, month, day, hour, minute, and second fields, with valid values as defined for the DATE
    /// and TIME data types.
    Timestamp = 93,
    #[cfg(feature = "odbc_version_4")]
    TimeWithTimezone = 94,
    #[cfg(feature = "odbc_version_4")]
    TimestampWithTimezone = 95,

    //SQL extended datatypes:
    ExtTimeOrInterval = 10,
    ExtTimestamp = 11,
    ExtLongVarchar = -1,
    ExtBinary = -2,
    ExtVarBinary = -3,
    ExtLongVarBinary = -4,
    ExtBigInt = -5,
    ExtTinyInt = -6,
    ExtBit = -7,
    ExtWChar = -8,
    ExtWVarChar = -9,
    ExtWLongVarChar = -10,
    ExtGuid = -11,
    SsVariant = -150,
    SsUdt = -151,
    SsXml = -152,
    SsTable = -153,
    SsTime2 = -154,
    SsTimestampOffset = -155,
}

/// Represented in C headers as SQLSMALLINT
#[repr(i16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum HandleType {
    Env = 1,
    Dbc = 2,
    Stmt = 3,
    Desc = 4,
}

/// Options for `SQLDriverConnect`
#[repr(u16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DriverConnectOption {
    NoPrompt = 0,
    Complete = 1,
    Prompt = 2,
    CompleteRequired = 3,
}

#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum SqlAttributeStringLength {
    /// SQL_IS_POINTER
    Pointer = -4,
    /// SQL_IS_UINTEGER
    UInteger = -5,
    /// SQL_IS_INTEGER
    Integer = -6,
    /// SQL_IS_USMALLINT
    USmallInt = -7,
    /// SQL_IS_SMALLINT
    SmallInt = -8,
}

pub use self::SqlAttributeStringLength::*;

/// SQL_YEAR_MONTH_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct YearMonth {
    pub year: UInteger,
    pub month: UInteger,
}

/// SQL_DAY_SECOND_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct DaySecond {
    pub day: UInteger,
    pub hour: UInteger,
    pub minute: UInteger,
    pub second: UInteger,
    pub fraction: UInteger,
}

/// SQL_INTERVAL_UNION
#[repr(C)]
#[derive(Copy, Clone)]
pub union IntervalUnion {
    pub year_month: YearMonth,
    pub day_second: DaySecond,
}

/// SQL_INTERVAL_STRUCT
#[repr(C)]
#[derive(Clone, Copy)]
pub struct IntervalStruct {
    pub interval_type: c_int,
    pub interval_sign: SmallInt,
    pub interval_value: IntervalUnion,
}

/// SQL_DATE_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct Date {
    pub year: SmallInt,
    pub month: USmallInt,
    pub day: USmallInt,
}

/// SQL_TIME_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct Time {
    pub hour: USmallInt,
    pub minute: USmallInt,
    pub second: USmallInt,
}

/// SQL_TIMESTAMP_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct Timestamp {
    pub year: SmallInt,
    pub month: USmallInt,
    pub day: USmallInt,
    pub hour: USmallInt,
    pub minute: USmallInt,
    pub second: USmallInt,
    pub fraction: UInteger,
}

/// SQLGUID
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct Guid {
    pub d1: u32,
    pub d2: u16,
    pub d3: u16,
    pub d4: [u8; 8],
}

/// SQL_SS_TIME2_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct SsTime2 {
    pub hour: USmallInt,
    pub minute: USmallInt,
    pub second: USmallInt,
    pub fraction: UInteger,
}

/// SQL_SS_TIMESTAMPOFFSET_STRUCT
#[repr(C)]
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Hash)]
pub struct SsTimestampOffset {
    pub year: SmallInt,
    pub month: USmallInt,
    pub day: USmallInt,
    pub hour: USmallInt,
    pub minute: USmallInt,
    pub second: USmallInt,
    pub fraction: UInteger,
    pub timezone_hour: SmallInt,
    pub timezone_minute: SmallInt,
}

/// Statement attributes for `SQLSetStmtAttr`
#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum StatementAttribute {
    AsyncEnable = 4,
    ParamBindType = 18,
    ParamsetSize = 22,
    RowBindType = 5,
    RowArraySize = 27,
    RowsFetchedPtr = 26,
    AsyncStmtEvent = 29,
}

/// Connection attributes for `SQLSetConnectAttr`
#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ConnectionAttribute {
    AsyncEnable = 4,
    AccessMode = 101,
    AutoCommit = 102,
    LoginTimeout = 103,
    Trace = 104,
    TraceFile = 105,
    TranslateLib = 106,
    TranslateOption = 107,
    TxnIsolation = 108,
    CurrentCatalog = 109,
    OdbcCursors = 110,
    QuietMode = 111,
    PacketSize = 112,
    ConnectionTimeout = 113,
    DisconnectBehaviour = 114,
    AsyncDbcFunctionsEnable = 117,
    AsyncDbcEvent = 119,
    EnlistInDtc = 1207,
    EnlistInXa = 1208,
    ConnectionDead = 1209,
    AutoIpd = 10001,
    MetadataId = 10014,
}

/// `DiagIdentifier` for `SQLGetDiagField`
#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum HeaderDiagnosticIdentifier {
    /// SQL_DIAG_RETURNCODE
    ReturnCode = 1,
    /// SQL_DIAG_NUMBER
    Number = 2,
    /// SQL_DIAG_ROW_COUNT
    RowCount = 3,
    /// SQL_DIAG_SQLSTATE
    SqlState = 4,
    /// SQL_DIAG_NATIVE
    Native = 5,
    /// SQL_DIAG_MESSAGE_TEXT
    MessageText = 6,
    /// SQL_DIAG_DYNAMIC_FUNCTION
    DynamicFunction = 7,
    /// SQL_DIAG_CLASS_ORIGIN
    ClassOrigin = 8,
    /// SQL_DIAG_SUBCLASS_ORIGIN
    SubclassOrigin = 9,
    /// SQL_DIAG_CONNECTION_NAME
    ConnectionName = 10,
    /// SQL_DIAG_SERVER_NAME
    ServerName = 11,
    /// SQL_DIAG_DYNAMIC_FUNCTION_CODE
    DynamicFunctionCode = 12,
    /// SQL_DIAG_CURSOR_ROW_COUNT
    CursorRowCount = -1249,
    /// SQL_DIAG_ROW_NUMBER
    RowNumber = -1248,
    /// SQL_DIAG_COLUMN_NUMBER
    ColumnNumber = -1247,
}

#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum AsyncConnectionBehavior {
    /// SQL_ASYNC_DBC_ENABLE_ON
    On = 1,
    /// SQL_ASYNC_DBC_ENABLE_OFF = 0,
    Off = 0,
}

impl Default for AsyncConnectionBehavior {
    fn default() -> AsyncConnectionBehavior {
        AsyncConnectionBehavior::Off
    }
}

#[repr(i32)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DynamicDiagnosticIdentifier {
    /// SQL_DIAG_ALTER_DOMAIN
    AlterDomain = 3,
    /// SQL_DIAG_ALTER_TABLE,
    AlterTable = 4,
    /// SQL_DIAG_CALL
    Call = 7,
    /// SQL_DIAG_CREATE_ASSERTION
    CreateAssertion = 6,
    /// SQL_DIAG_CREATE_CHARACTER_SET
    CreateCharacterSet = 8,
    /// SQL_DIAG_CREATE_COLLATION,
    CreateCollation = 10,
    /// SQL_DIAG_CREATE_DOMAIN
    CreateDomain = 23,
    /// SQL_DIAG_CREATE_INDEX
    CreateIndex = -1,
    /// SQL_DIAG_CREATE_SCHEMA
    CreateSchema = 64,
    /// SQL_DIAG_CREATE_TABLE
    CreateTable = 77,
    /// SQL_DIAG_CREATE_TRANSLATION
    CreateTranslation = 79,
    /// SQL_DIAG_CREATE_VIEW
    CreateView = 84,
    /// SQL_DIAG_DELETE_WHERE
    DeleteWhere = 19,
    /// SQL_DIAG_DROP_ASSERTION
    DropAssertion = 24,
    /// SQL_DIAG_DROP_CHARACTER_SET
    DropCharacterSet = 25,
    /// SQL_DIAG_DROP_COLLATION
    DropCollation = 26,
    /// SQL_DIAG_DROP_DOMAIN
    DropDomain = 27,
    /// SQL_DIAG_DROP_INDEX
    DropIndex = -2,
    /// SQL_DIAG_DROP_SCHEMA
    DropSchema = 31,
    /// SQL_DIAG_DROP_TABLE
    DropTable = 32,
    /// SQL_DIAG_DROP_TRANSLATION
    DropTranslation = 33,
    /// SQL_DIAG_DROP_VIEW
    DropView = 36,
    /// SQL_DIAG_DYNAMIC_DELETE_CURSOR
    DynamicDeleteCursor = 38,
    /// SQL_DIAG_DYNAMIC_UPDATE_CURSOR
    DynamicUpdateCursor = 81,
    /// SQL_DIAG_GRANT
    Grant = 48,
    /// SQL_DIAG_INSERT
    Insert = 50,
    /// SQL_DIAG_REVOKE
    Revoke = 59,
    // SQL_DIAG_SELECT_CURSOR
    SelectCursor = 85,
    /// SQL_DIAG_UNKNOWN_STATEMENT = 0,
    UnknownStatement = 0,
    /// SQL_DIAG_UPDATE_WHERE = 82,
    UpdateWhere = 82,
}

/// Completion types for `SQLEndTrans`
#[repr(i16)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum CompletionType {
    Commit = 0,
    Rollback = 1,
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 The 6502-rs Developers
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
// 3. Neither the names of the copyright holders nor the names of any
//    contributors may be used to endorse or promote products derived from this
//    software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

use log;

use std;

use address::{AddressDiff};
use instruction;
use instruction::{DecodedInstr};
use memory::Memory;
use registers::{ Registers, Status, StatusArgs };
use registers::{ PS_NEGATIVE, PS_OVERFLOW, PS_ZERO, PS_CARRY };

pub struct Machine {
    pub registers: Registers,
    pub memory:    Memory
}

impl Machine {
    pub fn new() -> Machine {
    	Machine{
    	    registers: Registers::new(),
    	    memory:    Memory::new()
    	}
    }

    pub fn reset(&mut self) {
    	*self = Machine::new();
    }

    pub fn fetch_next_and_decode(&mut self) -> Option<DecodedInstr> {
        let x: u8 = self.memory.get_byte(self.registers.program_counter);

        match instruction::OPCODES[x as uint] {
            Some((instr, am)) => {
                let extra_bytes = am.extra_bytes();
                let num_bytes = AddressDiff(1) + extra_bytes;

                let data_start = self.registers.program_counter
                               + AddressDiff(1);

                let slice = self.memory.get_slice(data_start, extra_bytes);
                let am_out = am.process(self, slice);

                // Increment program counter
                self.registers.program_counter =
                    self.registers.program_counter + num_bytes;

                Some((instr, am_out))
            }
            _ => None
        }
    }

    pub fn execute_instruction(&mut self, decoded_instr: DecodedInstr) {
        match decoded_instr {
            (instruction::ADC, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "add with carry immediate: {}", val);
                self.add_with_carry(val as i8);
            },
            (instruction::ADC, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr) as i8;
                log!(log::DEBUG, "add with carry. address: {}. value: {}",
                                 addr, val);
                self.add_with_carry(val);
            },

            (instruction::DEX, instruction::UseImplied) => {
                self.dec_x();
            }

            (instruction::LDA, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load A immediate: {}", val);
                self.load_accumulator(val as i8);
            },
            (instruction::LDA, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load A. address: {}. value: {}", addr, val);
                self.load_accumulator(val as i8);
            },

            (instruction::LDX, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load X immediate: {}", val);
                self.load_x_register(val as i8);
            },
            (instruction::LDX, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load X. address: {}. value: {}", addr, val);
                self.load_x_register(val as i8);
            },

            (instruction::LDY, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load Y immediate: {}", val);
                self.load_y_register(val as i8);
            },
            (instruction::LDY, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load Y. address: {}. value: {}", addr, val);
                self.load_y_register(val as i8);
            },

            (instruction::NOP, _) => {
                log!(log::DEBUG, "nop instr");
            },
            (_, _) => {
                log!(log::DEBUG, "attempting to execute unimplemented \
                                  instruction");
            },
        };
    }

    pub fn run(&mut self) {
        loop {
            if let Some(decoded_instr) = self.fetch_next_and_decode() {
                self.execute_instruction(decoded_instr);
            } else {
                break
            }
        }
    }

    fn load_register_with_flags(register: &mut i8,
                                status: &mut Status,
                                value: i8) {
        *register = value;

        let is_zero = value == 0;
        let is_negative = value < 0;

        status.set_with_mask(
            PS_ZERO | PS_NEGATIVE,
            Status::new(StatusArgs { zero: is_zero,
                                     negative: is_negative,
                                     ..StatusArgs::none() } ));
    }

    pub fn load_x_register(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.index_x,
                                          &mut self.registers.status,
                                          value);
    }

    pub fn load_y_register(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.index_y,
                                          &mut self.registers.status,
                                          value);
    }

    pub fn load_accumulator(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.accumulator,
                                          &mut self.registers.status,
                                          value);
    }

    // TODO akeeton: Implement binary-coded decimal.
    pub fn add_with_carry(&mut self, value: i8) {
        let a_before: i8 = self.registers.accumulator;
        let c_before: i8 = self.registers.status.get_carry();
        let a_after: i8 = a_before + c_before + value;

        debug_assert_eq!(a_after as u8, a_before as u8 + c_before as u8
                                        + value as u8);

        let did_carry = (a_after as u8) < (a_before as u8);

        let did_overflow   =
        	   (a_before < 0 && value < 0 && a_after >= 0)
        	|| (a_before > 0 && value > 0 && a_after <= 0);

        let mask = PS_CARRY | PS_OVERFLOW;

        self.registers.status.set_with_mask(mask,
            Status::new(StatusArgs { carry: did_carry,
                                     overflow: did_overflow,
                                     ..StatusArgs::none() } ));

        self.load_accumulator(a_after);

        log!(log::DEBUG, "accumulator: {}", self.registers.accumulator);
    }

    pub fn dec_x(&mut self) {
        let val = self.registers.index_x;
        self.load_x_register(val - 1);
    }
}

impl std::fmt::Show for Machine {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "Machine Dump:\n\nAccumulator: {}",
               self.registers.accumulator)
    }
}

#[test]
fn add_with_carry_test() {

    let mut machine = Machine::new();

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, 1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(-1);
    assert_eq!(machine.registers.accumulator, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    true);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, 2);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    let mut machine = Machine::new();

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, 127);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(-127);
    assert_eq!(machine.registers.accumulator, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),     true);
    assert_eq!(machine.registers.status.contains(PS_ZERO),      true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.registers.status.remove(PS_CARRY);
    machine.add_with_carry(-128);
    assert_eq!(machine.registers.accumulator, -128);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, -1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    let mut machine = Machine::new();

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, 127);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, -128);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW),  true);
}

#[test]
fn dec_x_test() {
    let mut machine = Machine::new();

    machine.dec_x();
    assert_eq!(machine.registers.index_x, -1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.dec_x();
    assert_eq!(machine.registers.index_x, -2);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.load_x_register(5);
    machine.dec_x();
    assert_eq!(machine.registers.index_x, 4);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.dec_x();
    machine.dec_x();
    machine.dec_x();
    machine.dec_x();

    assert_eq!(machine.registers.index_x, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.dec_x();
    assert_eq!(machine.registers.index_x, -1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);
}
<commit_msg>Implement the DEC instruction.<commit_after>// Copyright (C) 2014 The 6502-rs Developers
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
// 3. Neither the names of the copyright holders nor the names of any
//    contributors may be used to endorse or promote products derived from this
//    software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

use log;

use std;

use address::{Address, AddressDiff};
use instruction;
use instruction::{DecodedInstr};
use memory::Memory;
use registers::{ Registers, Status, StatusArgs };
use registers::{ PS_NEGATIVE, PS_OVERFLOW, PS_ZERO, PS_CARRY };

pub struct Machine {
    pub registers: Registers,
    pub memory:    Memory
}

impl Machine {
    pub fn new() -> Machine {
    	Machine{
    	    registers: Registers::new(),
    	    memory:    Memory::new()
    	}
    }

    pub fn reset(&mut self) {
    	*self = Machine::new();
    }

    pub fn fetch_next_and_decode(&mut self) -> Option<DecodedInstr> {
        let x: u8 = self.memory.get_byte(self.registers.program_counter);

        match instruction::OPCODES[x as uint] {
            Some((instr, am)) => {
                let extra_bytes = am.extra_bytes();
                let num_bytes = AddressDiff(1) + extra_bytes;

                let data_start = self.registers.program_counter
                               + AddressDiff(1);

                let slice = self.memory.get_slice(data_start, extra_bytes);
                let am_out = am.process(self, slice);

                // Increment program counter
                self.registers.program_counter =
                    self.registers.program_counter + num_bytes;

                Some((instr, am_out))
            }
            _ => None
        }
    }

    pub fn execute_instruction(&mut self, decoded_instr: DecodedInstr) {
        match decoded_instr {
            (instruction::ADC, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "add with carry immediate: {}", val);
                self.add_with_carry(val as i8);
            },
            (instruction::ADC, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr) as i8;
                log!(log::DEBUG, "add with carry. address: {}. value: {}",
                                 addr, val);
                self.add_with_carry(val);
            },

            (instruction::DEC, instruction::UseAddress(addr)) => {
                self.decrement_memory(addr)
            }

            (instruction::DEX, instruction::UseImplied) => {
                self.dec_x();
            }

            (instruction::LDA, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load A immediate: {}", val);
                self.load_accumulator(val as i8);
            },
            (instruction::LDA, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load A. address: {}. value: {}", addr, val);
                self.load_accumulator(val as i8);
            },

            (instruction::LDX, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load X immediate: {}", val);
                self.load_x_register(val as i8);
            },
            (instruction::LDX, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load X. address: {}. value: {}", addr, val);
                self.load_x_register(val as i8);
            },

            (instruction::LDY, instruction::UseImmediate(val)) => {
                log!(log::DEBUG, "load Y immediate: {}", val);
                self.load_y_register(val as i8);
            },
            (instruction::LDY, instruction::UseAddress(addr)) => {
                let val = self.memory.get_byte(addr);
                log!(log::DEBUG, "load Y. address: {}. value: {}", addr, val);
                self.load_y_register(val as i8);
            },

            (instruction::NOP, _) => {
                log!(log::DEBUG, "nop instr");
            },
            (_, _) => {
                log!(log::DEBUG, "attempting to execute unimplemented \
                                  instruction");
            },
        };
    }

    pub fn run(&mut self) {
        loop {
            if let Some(decoded_instr) = self.fetch_next_and_decode() {
                self.execute_instruction(decoded_instr);
            } else {
                break
            }
        }
    }

    fn load_register_with_flags(register: &mut i8,
                                status: &mut Status,
                                value: i8) {
        *register = value;

        let is_zero = value == 0;
        let is_negative = value < 0;

        status.set_with_mask(
            PS_ZERO | PS_NEGATIVE,
            Status::new(StatusArgs { zero: is_zero,
                                     negative: is_negative,
                                     ..StatusArgs::none() } ));
    }

    pub fn load_x_register(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.index_x,
                                          &mut self.registers.status,
                                          value);
    }

    pub fn load_y_register(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.index_y,
                                          &mut self.registers.status,
                                          value);
    }

    pub fn load_accumulator(&mut self, value: i8) {
        Machine::load_register_with_flags(&mut self.registers.accumulator,
                                          &mut self.registers.status,
                                          value);
    }

    // TODO akeeton: Implement binary-coded decimal.
    pub fn add_with_carry(&mut self, value: i8) {
        let a_before: i8 = self.registers.accumulator;
        let c_before: i8 = self.registers.status.get_carry();
        let a_after: i8 = a_before + c_before + value;

        debug_assert_eq!(a_after as u8, a_before as u8 + c_before as u8
                                        + value as u8);

        let did_carry = (a_after as u8) < (a_before as u8);

        let did_overflow   =
        	   (a_before < 0 && value < 0 && a_after >= 0)
        	|| (a_before > 0 && value > 0 && a_after <= 0);

        let mask = PS_CARRY | PS_OVERFLOW;

        self.registers.status.set_with_mask(mask,
            Status::new(StatusArgs { carry: did_carry,
                                     overflow: did_overflow,
                                     ..StatusArgs::none() } ));

        self.load_accumulator(a_after);

        log!(log::DEBUG, "accumulator: {}", self.registers.accumulator);
    }

    pub fn decrement_memory(&mut self, addr: Address) {
        let value_new = self.memory.get_byte(addr) - 1;

        self.memory.set_byte(addr, value_new);

        let is_negative = (value_new as i8) < 0;
        let is_zero     = value_new == 0;

        self.registers.status.set_with_mask(
            PS_NEGATIVE | PS_ZERO,
            Status::new(StatusArgs {
                negative: is_negative,
                zero:     is_zero,
                ..StatusArgs::none()
            })
        );
    }

    pub fn dec_x(&mut self) {
        let val = self.registers.index_x;
        self.load_x_register(val - 1);
    }
}

impl std::fmt::Show for Machine {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "Machine Dump:\n\nAccumulator: {}",
               self.registers.accumulator)
    }
}

#[test]
fn add_with_carry_test() {

    let mut machine = Machine::new();

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, 1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(-1);
    assert_eq!(machine.registers.accumulator, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    true);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, 2);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    let mut machine = Machine::new();

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, 127);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(-127);
    assert_eq!(machine.registers.accumulator, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),     true);
    assert_eq!(machine.registers.status.contains(PS_ZERO),      true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.registers.status.remove(PS_CARRY);
    machine.add_with_carry(-128);
    assert_eq!(machine.registers.accumulator, -128);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, -1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    let mut machine = Machine::new();

    machine.add_with_carry(127);
    assert_eq!(machine.registers.accumulator, 127);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.add_with_carry(1);
    assert_eq!(machine.registers.accumulator, -128);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE),  true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW),  true);
}

#[test]
fn decrement_memory_test() {
    let mut machine = Machine::new();
    let addr        = Address(0xA1B2);

    machine.memory.set_byte(addr, 5);

    machine.decrement_memory(addr);
    assert_eq!(machine.memory.get_byte(addr), 4);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);

    machine.decrement_memory(addr);
    assert_eq!(machine.memory.get_byte(addr), 3);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);

    machine.decrement_memory(addr);
    machine.decrement_memory(addr);
    machine.decrement_memory(addr);
    assert_eq!(machine.memory.get_byte(addr), 0);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);

    machine.decrement_memory(addr);
    assert_eq!(machine.memory.get_byte(addr) as i8, -1);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), true);
}

#[test]
fn dec_x_test() {
    let mut machine = Machine::new();

    machine.dec_x();
    assert_eq!(machine.registers.index_x, -1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.dec_x();
    assert_eq!(machine.registers.index_x, -2);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.load_x_register(5);
    machine.dec_x();
    assert_eq!(machine.registers.index_x, 4);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.dec_x();
    machine.dec_x();
    machine.dec_x();
    machine.dec_x();

    assert_eq!(machine.registers.index_x, 0);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     true);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), false);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);

    machine.dec_x();
    assert_eq!(machine.registers.index_x, -1);
    assert_eq!(machine.registers.status.contains(PS_CARRY),    false);
    assert_eq!(machine.registers.status.contains(PS_ZERO),     false);
    assert_eq!(machine.registers.status.contains(PS_NEGATIVE), true);
    assert_eq!(machine.registers.status.contains(PS_OVERFLOW), false);
}
<|endoftext|>


<commit_before>use std::vec::Vec;

pub struct Machine {
    memory: Vec<u32>,
    counter: u32,
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            memory: vec![0; 2097152],
            counter: 0,
        }
    }

    pub fn with_data(data: Vec<u32>) -> Machine {
        Machine {
            memory: data,
            counter: 0,
        }
    }

    pub fn loc(&self) -> u32 {
        self.counter
    }

    pub fn dump(&self) -> &Vec<u32> {
        &self.memory
    }

    pub fn load(&mut self, data: Vec<u32>) {
        self.memory = data;
    }

    pub fn exec(&mut self, opcode: u32, a: u32, b: u32) {
        match opcode {
            // PC <- M[A]
            1 => {
                self.counter = self.memory[a as usize];
            },

            // If M[B] = 0, then PC <- M[A]
            2 => {
                if 0 == self.memory[b as usize] {
                  self.counter = self.memory[a as usize];
                }
            },

            // M[A] <- PC
            3 => {
                self.memory[a as usize] = self.counter;
                self.counter += 4;
            },

            // M[A] <- M[B]
            4 => {
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[A] <- M[M[B]]
            5 => {
                let b = self.memory[b as usize];
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // Unknown opcode
            _ => {
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Machine;

    #[test]
    fn it_loads_data() {
        let mut m = Machine::new();
        assert_eq!(Some(&0), m.dump().first());

        m.load(vec![1]);
        assert_eq!(Some(&1), m.dump().first());
    }

    #[test]
    fn it_runs_opcode_1() {
        let mut m = Machine::with_data(vec![0, 2]);

        assert_eq!(0, m.loc());
        m.exec(1, 1, 0);
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_if() {
        let mut m = Machine::with_data(vec![0, 2]);

        assert_eq!(0, m.loc());
        m.exec(2, 1, 0);
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_else() {
        let mut m = Machine::with_data(vec![1, 2]);

        assert_eq!(0, m.loc());
        m.exec(2, 1, 0);
        assert_eq!(0, m.loc());
    }

    #[test]
    fn it_runs_opcode_3() {
        let mut m = Machine::with_data(vec![1]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(3, 0, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_4() {
        let mut m = Machine::with_data(vec![1, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(4, 0, 1);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_5() {
        let mut m = Machine::with_data(vec![1, 2, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(5, 0, 1);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }
}
<commit_msg>Test that opcode 6 works.<commit_after>use std::vec::Vec;

pub struct Machine {
    memory: Vec<u32>,
    counter: u32,
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            memory: vec![0; 2097152],
            counter: 0,
        }
    }

    pub fn with_data(data: Vec<u32>) -> Machine {
        Machine {
            memory: data,
            counter: 0,
        }
    }

    pub fn loc(&self) -> u32 {
        self.counter
    }

    pub fn dump(&self) -> &Vec<u32> {
        &self.memory
    }

    pub fn load(&mut self, data: Vec<u32>) {
        self.memory = data;
    }

    pub fn exec(&mut self, opcode: u32, a: u32, b: u32) {
        match opcode {
            // PC <- M[A]
            1 => {
                self.counter = self.memory[a as usize];
            },

            // If M[B] = 0, then PC <- M[A]
            2 => {
                if 0 == self.memory[b as usize] {
                  self.counter = self.memory[a as usize];
                }
            },

            // M[A] <- PC
            3 => {
                self.memory[a as usize] = self.counter;
                self.counter += 4;
            },

            // M[A] <- M[B]
            4 => {
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[A] <- M[M[B]]
            5 => {
                let b = self.memory[b as usize];
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[M[B]] <- M[A]
            6 => {
                let b = self.memory[b as usize];
                self.memory[b as usize] = self.memory[a as usize];
                self.counter += 4;
            },

            // Unknown opcode
            _ => {
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Machine;

    #[test]
    fn it_loads_data() {
        let mut m = Machine::new();
        assert_eq!(Some(&0), m.dump().first());

        m.load(vec![1]);
        assert_eq!(Some(&1), m.dump().first());
    }

    #[test]
    fn it_runs_opcode_1() {
        let mut m = Machine::with_data(vec![0, 2]);

        assert_eq!(0, m.loc());
        m.exec(1, 1, 0);
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_if() {
        let mut m = Machine::with_data(vec![0, 2]);

        assert_eq!(0, m.loc());
        m.exec(2, 1, 0);
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_else() {
        let mut m = Machine::with_data(vec![1, 2]);

        assert_eq!(0, m.loc());
        m.exec(2, 1, 0);
        assert_eq!(0, m.loc());
    }

    #[test]
    fn it_runs_opcode_3() {
        let mut m = Machine::with_data(vec![1]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(3, 0, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_4() {
        let mut m = Machine::with_data(vec![1, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(4, 0, 1);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_5() {
        let mut m = Machine::with_data(vec![1, 2, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(5, 0, 1);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_6() {
        let mut m = Machine::with_data(vec![1, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(6, 1, 1);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }
}
<|endoftext|>


<commit_before>extern crate hyper;
extern crate pbr;

use std::sync::{Arc, Mutex};

pub mod client;
pub mod contentlength;
pub mod download;
pub mod http_version;
pub mod response;
pub mod write;

/// Represents a number of bytes, as `u64`.
type Bytes = u64;
/// Represents a 'chunk', which is just a piece of bytes.
type Chunk = Vec<u8>;
/// Represents a list of chunks
pub type Chunks = Vec<Chunk>;
/// Represents a shared mutable reference of chunks
pub type SChunks = Arc<Mutex<Chunks>>;
<commit_msg>Move private Bytes type to public<commit_after>extern crate hyper;
extern crate pbr;

use std::sync::{Arc, Mutex};

pub mod client;
pub mod contentlength;
pub mod download;
pub mod http_version;
pub mod response;
pub mod write;

/// Represents a number of bytes, as `u64`.
pub type Bytes = u64;
/// Represents a 'chunk', which is just a piece of bytes.
type Chunk = Vec<u8>;
/// Represents a list of chunks
pub type Chunks = Vec<Chunk>;
/// Represents a shared mutable reference of chunks
pub type SChunks = Arc<Mutex<Chunks>>;
<|endoftext|>


<commit_before>// Copyright 2015 Keegan McAllister.
// Copyright 2016 Corey Farwell.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// See `LICENSE` in this repository.

use std::io::{self, Read};
use std::{panic, process};

/// Utility that reads a `Vec` of bytes from standard input (stdin)
/// and passes it to `closure`. All panics that occur within
/// `closure` will be treated as aborts. This is done so that
/// AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Image;
/// # impl Image {
/// #     fn parse(_: &[u8]) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_bytes(|read| {
///         Image::parse(&read)
///     })
/// }
/// ```
#[deprecated(since = "0.3.3",
note="This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead.")]
pub fn read_stdio_bytes<F>(closure: F)
where
    F: Fn(Vec<u8>) + panic::RefUnwindSafe,
{
    let mut input = vec![];
    let result = io::stdin().read_to_end(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

/// Utility that reads a `String` from standard input (stdin) and
/// passes it to `closure`. If a `String` cannot be constructed from
/// the data provided by standard input, `closure` will _not_ be
/// called. All panics that occur within `closure` will be treated as
/// aborts. This is done so that AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Url;
/// # impl Url {
/// #     fn parse(_: &str) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_string(|string| {
///         Url::parse(&string)
///     })
/// }
/// ```
#[deprecated(since = "0.3.3",
note="This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead.")]
pub fn read_stdio_string<F>(closure: F)
where
    F: Fn(String) + panic::RefUnwindSafe,
{
    let mut input = String::new();
    let result = io::stdin().read_to_string(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

// those functions are provided by the afl-llvm-rt static library
extern "C" {
    fn __afl_persistent_loop(counter: usize) -> isize;
    fn __afl_manual_init();
}

/// Fuzz a closure by passing it a `&[u8]`
///
/// This slice contains a "random" quantity of "random" data.
///
/// ```rust,ignore
/// # extern crate afl;
/// # use afl::fuzz;
/// # fn main() {
/// fuzz(|data|{
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
pub fn fuzz<F>(closure: F) where F: Fn(&[u8]) + std::panic::RefUnwindSafe {
    // this marker strings needs to be in the produced executable for
    // afl-fuzz to detect `persistent mode` and `defered mode`
    static PERSIST_MARKER: &'static str = "##SIG_AFL_PERSISTENT##\0";
    static DEFERED_MARKER: &'static str = "##SIG_AFL_DEFER_FORKSRV##\0";

    // we now need a fake instruction to prevent the compiler from optimizing out
    // those marker strings
    unsafe{std::ptr::read_volatile(&PERSIST_MARKER)}; // hack used in https://github.com/bluss/bencher for black_box()
    unsafe{std::ptr::read_volatile(&DEFERED_MARKER)};
    // unsafe { asm!("" : : "r"(&PERSIST_MARKER)) }; // hack used in nightly's back_box(), requires feature asm
    // unsafe { asm!("" : : "r"(&DEFERED_MARKER)) };

    // sets panic hook to abort
    std::panic::set_hook(Box::new(|_| {
        std::process::abort();
    }));

    let mut input = vec![];

    // initialize forkserver there
    unsafe{__afl_manual_init()};

    while unsafe{__afl_persistent_loop(1000)} != 0 {
        // get buffer from AFL through stdin
        let result = io::stdin().read_to_end(&mut input);
        if result.is_err() {
            return;
        }

        // We still catch unwinding panics just in case the fuzzed code modifies
        // the panic hook.
        // If so, the fuzzer will be unable to tell different bugs appart and you will
        // only be able to find one bug at a time before fixing it to then find a new one.
        let did_panic = std::panic::catch_unwind(|| {
            closure(&input);
        }).is_err();

        if did_panic {
            // hopefully the custom panic hook will be called before and abort the
            // process before the stack frames are unwinded.
            std::process::abort();
        }
        input.clear();
    }
}

/// Fuzz a closure-like block of code by passing it an object of arbitrary type.
///
/// You can choose the type of the argument using the syntax as in the example below.
/// Please check out the `arbitrary` crate to see which types are available.
///
/// For performance reasons, it is recommended that you use the native type `&[u8]` when possible.
///
/// ```rust,ignore
/// # #[macro_use] extern crate afl;
/// # fn main() {
/// fuzz!(|data: &[u8]| {
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
#[macro_export]
macro_rules! fuzz {
    (|$buf:ident| $body:block) => {
        afl::fuzz(|$buf| $body);
    };
    (|$buf:ident: &[u8]| $body:block) => {
        afl::fuzz(|$buf| $body);
    };
    (|$buf:ident: $dty: ty| $body:block) => {
        afl::fuzz(|$buf| {
            let $buf: $dty = {
                use arbitrary::{Arbitrary, RingBuffer};
                if let Ok(d) = RingBuffer::new($buf, $buf.len()).and_then(|mut b|{
                        Arbitrary::arbitrary(&mut b).map_err(|_| "")
                    }) {
                    d
                } else {
                    return
                }
            };

            $body
        });
    };
}

#[cfg(test)]
mod test {
    /*
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::thread;
    use std::time;

    extern crate libc;
    extern crate tempdir;

    fn target_path() -> PathBuf {
        if PathBuf::from("../target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("../target/debug/")
        } else if PathBuf::from("target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("target/debug/")
        } else {
            panic!("Could not find cargo-afl-fuzz!");
        }
    }

    #[test]
    fn test_cargo_afl_fuzz() {
        let temp_dir = tempdir::TempDir::new("aflrs").expect("Could not create temporary directory");
        let temp_dir_path = temp_dir.path();
        let mut child = Command::new(target_path().join("cargo-afl-fuzz"))
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .arg("-i")
            .arg(".")
            .arg("-o")
            .arg(temp_dir_path)
            .arg(target_path().join("examples").join("hello"))
            .spawn()
            .expect("Could not run cargo-afl-fuzz");
        thread::sleep(time::Duration::from_secs(7));
        child.kill().unwrap();
        assert!(temp_dir_path.join("fuzzer_stats").is_file());
    }
    */
}
<commit_msg>Update lib.rs<commit_after>// Copyright 2015 Keegan McAllister.
// Copyright 2016 Corey Farwell.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// See `LICENSE` in this repository.

use std::io::{self, Read};
use std::{panic, process};

/// Utility that reads a `Vec` of bytes from standard input (stdin)
/// and passes it to `closure`. All panics that occur within
/// `closure` will be treated as aborts. This is done so that
/// AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Image;
/// # impl Image {
/// #     fn parse(_: &[u8]) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_bytes(|read| {
///         Image::parse(&read)
///     })
/// }
/// ```
#[deprecated(since = "0.3.3",
note="This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead.")]
pub fn read_stdio_bytes<F>(closure: F)
where
    F: Fn(Vec<u8>) + panic::RefUnwindSafe,
{
    let mut input = vec![];
    let result = io::stdin().read_to_end(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

/// Utility that reads a `String` from standard input (stdin) and
/// passes it to `closure`. If a `String` cannot be constructed from
/// the data provided by standard input, `closure` will _not_ be
/// called. All panics that occur within `closure` will be treated as
/// aborts. This is done so that AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Url;
/// # impl Url {
/// #     fn parse(_: &str) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_string(|string| {
///         Url::parse(&string)
///     })
/// }
/// ```
#[deprecated(since = "0.3.3",
note="This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead.")]
pub fn read_stdio_string<F>(closure: F)
where
    F: Fn(String) + panic::RefUnwindSafe,
{
    let mut input = String::new();
    let result = io::stdin().read_to_string(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

// those functions are provided by the afl-llvm-rt static library
extern "C" {
    fn __afl_persistent_loop(counter: usize) -> isize;
    fn __afl_manual_init();
}

/// Fuzz a closure by passing it a `&[u8]`
///
/// This slice contains a "random" quantity of "random" data.
///
/// ```rust,no_run
/// # extern crate afl;
/// # use afl::fuzz;
/// # fn main() {
/// fuzz(|data|{
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
pub fn fuzz<F>(closure: F) where F: Fn(&[u8]) + std::panic::RefUnwindSafe {
    // this marker strings needs to be in the produced executable for
    // afl-fuzz to detect `persistent mode` and `defered mode`
    static PERSIST_MARKER: &'static str = "##SIG_AFL_PERSISTENT##\0";
    static DEFERED_MARKER: &'static str = "##SIG_AFL_DEFER_FORKSRV##\0";

    // we now need a fake instruction to prevent the compiler from optimizing out
    // those marker strings
    unsafe{std::ptr::read_volatile(&PERSIST_MARKER)}; // hack used in https://github.com/bluss/bencher for black_box()
    unsafe{std::ptr::read_volatile(&DEFERED_MARKER)};
    // unsafe { asm!("" : : "r"(&PERSIST_MARKER)) }; // hack used in nightly's back_box(), requires feature asm
    // unsafe { asm!("" : : "r"(&DEFERED_MARKER)) };

    // sets panic hook to abort
    std::panic::set_hook(Box::new(|_| {
        std::process::abort();
    }));

    let mut input = vec![];

    // initialize forkserver there
    unsafe{__afl_manual_init()};

    while unsafe{__afl_persistent_loop(1000)} != 0 {
        // get buffer from AFL through stdin
        let result = io::stdin().read_to_end(&mut input);
        if result.is_err() {
            return;
        }

        // We still catch unwinding panics just in case the fuzzed code modifies
        // the panic hook.
        // If so, the fuzzer will be unable to tell different bugs appart and you will
        // only be able to find one bug at a time before fixing it to then find a new one.
        let did_panic = std::panic::catch_unwind(|| {
            closure(&input);
        }).is_err();

        if did_panic {
            // hopefully the custom panic hook will be called before and abort the
            // process before the stack frames are unwinded.
            std::process::abort();
        }
        input.clear();
    }
}

/// Fuzz a closure-like block of code by passing it an object of arbitrary type.
///
/// You can choose the type of the argument using the syntax as in the example below.
/// Please check out the `arbitrary` crate to see which types are available.
///
/// For performance reasons, it is recommended that you use the native type `&[u8]` when possible.
///
/// ```rust,no_run
/// # #[macro_use] extern crate afl;
/// # fn main() {
/// fuzz!(|data: &[u8]| {
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
#[macro_export]
macro_rules! fuzz {
    (|$buf:ident| $body:block) => {
        afl::fuzz(|$buf| $body);
    };
    (|$buf:ident: &[u8]| $body:block) => {
        afl::fuzz(|$buf| $body);
    };
    (|$buf:ident: $dty: ty| $body:block) => {
        afl::fuzz(|$buf| {
            let $buf: $dty = {
                use arbitrary::{Arbitrary, RingBuffer};
                if let Ok(d) = RingBuffer::new($buf, $buf.len()).and_then(|mut b|{
                        Arbitrary::arbitrary(&mut b).map_err(|_| "")
                    }) {
                    d
                } else {
                    return
                }
            };

            $body
        });
    };
}

#[cfg(test)]
mod test {
    /*
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::thread;
    use std::time;

    extern crate libc;
    extern crate tempdir;

    fn target_path() -> PathBuf {
        if PathBuf::from("../target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("../target/debug/")
        } else if PathBuf::from("target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("target/debug/")
        } else {
            panic!("Could not find cargo-afl-fuzz!");
        }
    }

    #[test]
    fn test_cargo_afl_fuzz() {
        let temp_dir = tempdir::TempDir::new("aflrs").expect("Could not create temporary directory");
        let temp_dir_path = temp_dir.path();
        let mut child = Command::new(target_path().join("cargo-afl-fuzz"))
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .arg("-i")
            .arg(".")
            .arg("-o")
            .arg(temp_dir_path)
            .arg(target_path().join("examples").join("hello"))
            .spawn()
            .expect("Could not run cargo-afl-fuzz");
        thread::sleep(time::Duration::from_secs(7));
        child.kill().unwrap();
        assert!(temp_dir_path.join("fuzzer_stats").is_file());
    }
    */
}
<|endoftext|>


<commit_before>/*!
Simple persistent generic HashMap/Key-value store, using file locking to limit writing between threads.

This is in a beta state at the moment.

Basic usage:

```
extern crate typedb;

use typedb::{ KV, Value };

fn main() {
    let mut test_store = KV::<String, Value>::new("./db.cab").unwrap();

    let _ = test_store.insert("key".to_string(), Value::String("value".to_string()));
    println!("{:?}", test_store.get("key".to_string()));
    let _ = test_store.remove("key".to_string());
}
```

Usage with user defined Key and Value types:

```
#[macro_use]
extern crate typedb;
#[macro_use]
extern crate serde_derive;

use typedb::KV;

key!(MyKey:
    String(String),
    Int(i32),
);

value!(MyValue:
    String(String),
    Int(i32),
);

fn main() {
    let mut test_store = KV::<MyKey, MyValue>::new("./db.cab").unwrap();

    let _ = test_store.insert(MyKey::Int(1i32), MyValue::String("value".to_string()));
    println!("{:?}", test_store.get(MyKey::Int(1i32)));
    let _ = test_store.remove(MyKey::Int(1i32));
}
```

*/

#![deny(missing_docs)]

extern crate bincode;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate log;
extern crate fs2;

/// Macros for simplifying custom key and value types definition
pub mod macros;

use std::collections::HashMap;
use std::fs::{ File, OpenOptions };
use std::io::prelude::*;
use std::io::SeekFrom;
use std::hash::Hash;

use bincode::{serialize, deserialize, Infinite};
use serde::ser::Serialize;
use serde::de::Deserialize;

use fs2::FileExt;

/// The maximum number of retries the cab will make
const MAX_RETRIES:i32 = 10;

/// Definition of a macro for retrying an operation
macro_rules! retry {
    ($i:ident, $b:block) => (
        for $i in 0..MAX_RETRIES {
            if $i != 0 {
                std::thread::park_timeout(std::time::Duration::new(0u64, 1000u32));
            }

            $b
        }
    )
}

/// Macro for what to do on last retry
macro_rules! last_retry {
    ($i:ident, $b:stmt) => (
        if $i >= MAX_RETRIES - 1 {
            $b;
        }
    )
}

/// A default value type to use with KV
#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
pub enum Value {
    /// Cab default Value type for strings
    String(String),
    /// Cab default Value type for intergers
    Int(i32),
    /// Cab default Value type for floats
    Float(f32),
    /// Cab default Value type for a sub map
    Map(HashMap<String, Value>),
    /// Cab default Value type for list
    List(Vec<Value>),
}

/// Type alias for results from KV
type KVResult = Result<bool, KVError>;

/// Errors that KV might have
#[derive(Debug, PartialEq)]
pub enum KVError {
    /// Could not deserialize the cab from disk
    CouldntDecode,
    /// Couldn't serialize the hashmap for writing to disk
    CouldntEncode,
    /// Could not write to the cab on disk
    CouldntWrite,
    /// Failed to read the cab from disk
    FailedToRead,
    /// Could not flush the data to disk
    CouldntFlush,
    /// Could not open the file from disk
    CouldntOpen,
    /// Could not load the file from disk
    CouldntLoad,
    /// Failed to write lock the cab
    CouldntWriteLock,
    /// Failed to read lock the cab
    CouldntReadLock,
    /// Failed to unlock the cab
    CouldntUnlock,
    /// Failed to set SeekFrom::Start
    CouldntSeekToStart,
}

/// The type that represents the key-value store
pub struct KV<K,V> {
    cab: HashMap<K, V>,
    file: File,
}

impl<K: Clone + Serialize + Deserialize + Eq + Hash, V: Clone + Serialize + Deserialize> KV<K,V> {
    /// Creates a new instance of the KV store
    pub fn new(p:&'static str) -> Result<KV<K,V>, KVError> {
        // create the KV instance
        let mut store = KV {
            cab: HashMap::new(),
            file: KV::<K, V>::retry_get_file(p).unwrap(),
        };

        match store.load_from_persist(false) {
            Ok(f) => trace!("{}", f),
            Err(e) => {
                return Err(e);
            },
        };

        Ok(store)
    }

    /// Retry to get a reference to the cab file at path p and create if doesn't exist
    fn retry_get_file(p:&'static str) -> Result<File, KVError> {
        retry!(i, {
            match OpenOptions::new().read(true).write(true).create(true).open(p) {
                Ok(f) => { return Ok(f); },
                Err(e) => {
                    error!("{}", e);
                    last_retry!(i, return Err(KVError::CouldntOpen));
                    continue;
                }
            };
        });

        Err(KVError::CouldntOpen)
    }

    /// Inserta a key, value pair into the key-value store
    pub fn insert(&mut self, key: K, value: V) -> KVResult {
        // check that can write to the cab
        if let Err(e) = self.file.lock_exclusive() {
            error!("{}", e);
            return Err(KVError::CouldntWriteLock);
        }
        // make sure mem version up to date
        self.load_from_persist(true)?;
        // insert into the HashMap
        self.cab.insert(key, value);
        // persist
        self.write_to_persist()
    }

    /// Get the value from a key
    pub fn get(&mut self, key: K) -> Result<Option<V>, KVError> {
        // make sure mem version up to date
        self.load_from_persist(false)?;
        // get the value from the cab
        match self.cab.get(&key) {
            Some(v) => Ok(Some((*v).clone())),
            None => Ok(None),
        }
    }

    /// Removes a key and associated value from the key-value Store
    pub fn remove(&mut self, key: K) -> KVResult {
        // check that can write to the cab
        if let Err(e) = self.file.lock_exclusive() {
            error!("{}", e);
            return Err(KVError::CouldntWriteLock);
        }
        // make sure mem version up to date
        self.load_from_persist(true)?;
        // remove from the HashMap
        self.cab.remove(&key);
        // persist
        self.write_to_persist()
    }

    /// get all the keys contained in the KV Store
    pub fn keys(&mut self) -> Result<Vec<K>, KVError> {
        // make sure mem version up to date
        self.load_from_persist(false)?;
        // create a vec from the cabs keys
        Ok(self.cab.keys().map(|k| k.clone()).collect())
    }

    /// Writes the key-value Store to file
    fn write_to_persist(&mut self) -> KVResult {
        // attempt to write to the cab
        retry!(i, {
            // check that can write to the cab
            if let Err(e) = self.file.lock_exclusive() {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntWriteLock));
                continue;
            }

            // serialize the cab as a u8 vec
            let byte_vec: Vec<u8> = match serialize(&mut self.cab, Infinite) {
                Ok(bv) => bv,
                Err(e) => {
                    error!("serialize: {}", e);
                    return Err(KVError::CouldntEncode);
                },
            };

            // seek to the start
            if let Err(e) = self.file.seek(SeekFrom::Start(0)) {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntSeekToStart));
                continue;
            }

            // write the bytes to it
            if let Err(e) = self.file.write_all(byte_vec.as_slice()) {
                error!("file.write_all/retry: {}", e);
                last_retry!(i, return Err(KVError::CouldntWrite));
                continue;
            }

            // flush to disk
            if let Err(e) = self.file.flush() {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntFlush));
                continue;
            }

            // unlock the file after write has completed
            if let Err(e) = self.file.unlock() {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntUnlock));
                continue;
            }
            return Ok(true);
        });

        Err(KVError::CouldntWrite)
    }

    /// Loads key-value store from file
    fn load_from_persist(&mut self, already_locked:bool) -> KVResult {
        retry!(i, {
            // byte vec to read into
            let mut byte_vec = Vec::new();

            // wait/lock the cab and read the bytes
            if !already_locked {
                if let Err(e) = self.file.lock_shared() {
                    error!("{}", e);
                    return Err(KVError::CouldntReadLock);
                }
            }

            // seek to the start
            if let Err(e) = self.file.seek(SeekFrom::Start(0)) {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntSeekToStart));
                continue;
            }

            // read the file into the buffer
            match self.file.read_to_end(&mut byte_vec) {
                Ok(count) => {
                    // don't attempt to deserialize as empty
                    if count == 0 {
                        if !already_locked {
                            if let Err(e) = self.file.unlock() {
                                error!("{}", e);
                                last_retry!(i, return Err(KVError::CouldntUnlock));
                            }
                        }
                        return Ok(true);
                    }
                },
                Err(e) => {
                    error!("{}", e);
                    return Err(KVError::FailedToRead);
                },
            }

            // deserialize u8 vec back into HashMap
            match deserialize(byte_vec.as_slice()) {
                Ok(f) => {
                    // assign read HashMap back to self
                    self.cab = f;
                    if !already_locked {
                        if let Err(e) = self.file.unlock() {
                            error!("{}", e);
                            last_retry!(i, return Err(KVError::CouldntUnlock));
                        }
                    }
                    return Ok(true);
                },
                Err(e) => {
                    error!("{}", e);
                    last_retry!(i, return Err(KVError::CouldntDecode));
                    continue;
                },
            };
        });

        Err(KVError::CouldntLoad)
    }
}
<commit_msg>Made KVError enum consistent<commit_after>/*!
Simple persistent generic HashMap/Key-value store, using file locking to limit writing between threads.

This is in a beta state at the moment.

Basic usage:

```
extern crate typedb;

use typedb::{ KV, Value };

fn main() {
    let mut test_store = KV::<String, Value>::new("./db.cab").unwrap();

    let _ = test_store.insert("key".to_string(), Value::String("value".to_string()));
    println!("{:?}", test_store.get("key".to_string()));
    let _ = test_store.remove("key".to_string());
}
```

Usage with user defined Key and Value types:

```
#[macro_use]
extern crate typedb;
#[macro_use]
extern crate serde_derive;

use typedb::KV;

key!(MyKey:
    String(String),
    Int(i32),
);

value!(MyValue:
    String(String),
    Int(i32),
);

fn main() {
    let mut test_store = KV::<MyKey, MyValue>::new("./db.cab").unwrap();

    let _ = test_store.insert(MyKey::Int(1i32), MyValue::String("value".to_string()));
    println!("{:?}", test_store.get(MyKey::Int(1i32)));
    let _ = test_store.remove(MyKey::Int(1i32));
}
```

*/

#![deny(missing_docs)]

extern crate bincode;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate log;
extern crate fs2;

/// Macros for simplifying custom key and value types definition
pub mod macros;

use std::collections::HashMap;
use std::fs::{ File, OpenOptions };
use std::io::prelude::*;
use std::io::SeekFrom;
use std::hash::Hash;

use bincode::{serialize, deserialize, Infinite};
use serde::ser::Serialize;
use serde::de::Deserialize;

use fs2::FileExt;

/// The maximum number of retries the cab will make
const MAX_RETRIES:i32 = 10;

/// Definition of a macro for retrying an operation
macro_rules! retry {
    ($i:ident, $b:block) => (
        for $i in 0..MAX_RETRIES {
            if $i != 0 {
                std::thread::park_timeout(std::time::Duration::new(0u64, 1000u32));
            }

            $b
        }
    )
}

/// Macro for what to do on last retry
macro_rules! last_retry {
    ($i:ident, $b:stmt) => (
        if $i >= MAX_RETRIES - 1 {
            $b;
        }
    )
}

/// A default value type to use with KV
#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
pub enum Value {
    /// Cab default Value type for strings
    String(String),
    /// Cab default Value type for intergers
    Int(i32),
    /// Cab default Value type for floats
    Float(f32),
    /// Cab default Value type for a sub map
    Map(HashMap<String, Value>),
    /// Cab default Value type for list
    List(Vec<Value>),
}

/// Type alias for results from KV
type KVResult = Result<bool, KVError>;

/// Errors that KV might have
#[derive(Debug, PartialEq)]
pub enum KVError {
    /// Could not deserialize the cab from disk
    CouldntDecode,
    /// Couldn't serialize the hashmap for writing to disk
    CouldntEncode,
    /// Could not write to the cab on disk
    CouldntWrite,
    /// Failed to read the cab from disk
    CouldntRead,
    /// Could not flush the data to disk
    CouldntFlush,
    /// Could not open the file from disk
    CouldntOpen,
    /// Could not load the file from disk
    CouldntLoad,
    /// Failed to write lock the cab
    CouldntWriteLock,
    /// Failed to read lock the cab
    CouldntReadLock,
    /// Failed to unlock the cab
    CouldntUnlock,
    /// Failed to set SeekFrom::Start
    CouldntSeekToStart,
}

/// The type that represents the key-value store
pub struct KV<K,V> {
    cab: HashMap<K, V>,
    file: File,
}

impl<K: Clone + Serialize + Deserialize + Eq + Hash, V: Clone + Serialize + Deserialize> KV<K,V> {
    /// Creates a new instance of the KV store
    pub fn new(p:&'static str) -> Result<KV<K,V>, KVError> {
        // create the KV instance
        let mut store = KV {
            cab: HashMap::new(),
            file: KV::<K, V>::retry_get_file(p).unwrap(),
        };

        match store.load_from_persist(false) {
            Ok(f) => trace!("{}", f),
            Err(e) => {
                return Err(e);
            },
        };

        Ok(store)
    }

    /// Retry to get a reference to the cab file at path p and create if doesn't exist
    fn retry_get_file(p:&'static str) -> Result<File, KVError> {
        retry!(i, {
            match OpenOptions::new().read(true).write(true).create(true).open(p) {
                Ok(f) => { return Ok(f); },
                Err(e) => {
                    error!("{}", e);
                    last_retry!(i, return Err(KVError::CouldntOpen));
                    continue;
                }
            };
        });

        Err(KVError::CouldntOpen)
    }

    /// Inserta a key, value pair into the key-value store
    pub fn insert(&mut self, key: K, value: V) -> KVResult {
        // check that can write to the cab
        if let Err(e) = self.file.lock_exclusive() {
            error!("{}", e);
            return Err(KVError::CouldntWriteLock);
        }
        // make sure mem version up to date
        self.load_from_persist(true)?;
        // insert into the HashMap
        self.cab.insert(key, value);
        // persist
        self.write_to_persist()
    }

    /// Get the value from a key
    pub fn get(&mut self, key: K) -> Result<Option<V>, KVError> {
        // make sure mem version up to date
        self.load_from_persist(false)?;
        // get the value from the cab
        match self.cab.get(&key) {
            Some(v) => Ok(Some((*v).clone())),
            None => Ok(None),
        }
    }

    /// Removes a key and associated value from the key-value Store
    pub fn remove(&mut self, key: K) -> KVResult {
        // check that can write to the cab
        if let Err(e) = self.file.lock_exclusive() {
            error!("{}", e);
            return Err(KVError::CouldntWriteLock);
        }
        // make sure mem version up to date
        self.load_from_persist(true)?;
        // remove from the HashMap
        self.cab.remove(&key);
        // persist
        self.write_to_persist()
    }

    /// get all the keys contained in the KV Store
    pub fn keys(&mut self) -> Result<Vec<K>, KVError> {
        // make sure mem version up to date
        self.load_from_persist(false)?;
        // create a vec from the cabs keys
        Ok(self.cab.keys().map(|k| k.clone()).collect())
    }

    /// Writes the key-value Store to file
    fn write_to_persist(&mut self) -> KVResult {
        // attempt to write to the cab
        retry!(i, {
            // check that can write to the cab
            if let Err(e) = self.file.lock_exclusive() {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntWriteLock));
                continue;
            }

            // serialize the cab as a u8 vec
            let byte_vec: Vec<u8> = match serialize(&mut self.cab, Infinite) {
                Ok(bv) => bv,
                Err(e) => {
                    error!("serialize: {}", e);
                    return Err(KVError::CouldntEncode);
                },
            };

            // seek to the start
            if let Err(e) = self.file.seek(SeekFrom::Start(0)) {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntSeekToStart));
                continue;
            }

            // write the bytes to it
            if let Err(e) = self.file.write_all(byte_vec.as_slice()) {
                error!("file.write_all/retry: {}", e);
                last_retry!(i, return Err(KVError::CouldntWrite));
                continue;
            }

            // flush to disk
            if let Err(e) = self.file.flush() {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntFlush));
                continue;
            }

            // unlock the file after write has completed
            if let Err(e) = self.file.unlock() {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntUnlock));
                continue;
            }
            return Ok(true);
        });

        Err(KVError::CouldntWrite)
    }

    /// Loads key-value store from file
    fn load_from_persist(&mut self, already_locked:bool) -> KVResult {
        retry!(i, {
            // byte vec to read into
            let mut byte_vec = Vec::new();

            // wait/lock the cab and read the bytes
            if !already_locked {
                if let Err(e) = self.file.lock_shared() {
                    error!("{}", e);
                    return Err(KVError::CouldntReadLock);
                }
            }

            // seek to the start
            if let Err(e) = self.file.seek(SeekFrom::Start(0)) {
                error!("{}", e);
                last_retry!(i, return Err(KVError::CouldntSeekToStart));
                continue;
            }

            // read the file into the buffer
            match self.file.read_to_end(&mut byte_vec) {
                Ok(count) => {
                    // don't attempt to deserialize as empty
                    if count == 0 {
                        if !already_locked {
                            if let Err(e) = self.file.unlock() {
                                error!("{}", e);
                                last_retry!(i, return Err(KVError::CouldntUnlock));
                            }
                        }
                        return Ok(true);
                    }
                },
                Err(e) => {
                    error!("{}", e);
                    return Err(KVError::CouldntRead);
                },
            }

            // deserialize u8 vec back into HashMap
            match deserialize(byte_vec.as_slice()) {
                Ok(f) => {
                    // assign read HashMap back to self
                    self.cab = f;
                    if !already_locked {
                        if let Err(e) = self.file.unlock() {
                            error!("{}", e);
                            last_retry!(i, return Err(KVError::CouldntUnlock));
                        }
                    }
                    return Ok(true);
                },
                Err(e) => {
                    error!("{}", e);
                    last_retry!(i, return Err(KVError::CouldntDecode));
                    continue;
                },
            };
        });

        Err(KVError::CouldntLoad)
    }
}
<|endoftext|>