rox

#[macro_use]
extern crate rustler;
#[macro_use]
extern crate lazy_static;
extern crate rocksdb;

use std::io::Write;
use std::ops::Deref;
use std::path::Path;
use std::sync::{Arc, RwLock};

use rustler::resource::ResourceArc;

use rustler::{Decoder, Encoder, Env, Error, NifResult, Term};

use rocksdb::{
    ColumnFamily, DBCompressionType, DBIterator, Direction, IteratorMode, Options, Snapshot,
    WriteBatch, WriteOptions, DB,
};
use rustler::dynamic::TermType;
use rustler::types::atom::Atom;
use rustler::types::binary::{Binary, OwnedBinary};
use rustler::types::list::ListIterator;

mod atoms {
    rustler_atoms! {
        atom ok;
        atom error;
        atom not_found;

        // Batch Operation Atoms
        atom put;
        atom put_cf;
        atom delete;
        atom delete_cf;

        // Compression Type Atoms
        atom snappy;
        atom zlib;
        atom bzip2;
        atom lz4;
        atom lz4h;
        atom none;

        // Iterator Atoms
        atom forward;
        atom reverse;
        atom start;
        atom end;
        atom from;
        atom done;


        // Block Based Table Option atoms
        // atom no_block_cache;
        // atom block_size;
        // atom block_cache_size;
        // atom bloom_filter_policy;
        // atom format_version;
        // atom skip_table_builder_flush;
        // atom cache_index_and_filter_blocks;

        // CF Options Related atoms
        // atom block_cache_size_mb_for_point_lookup;
        // atom memtable_memory_budget;
        // atom write_buffer_size;
        // atom max_write_buffer_number;
        // atom min_write_buffer_number_to_merge;
        // atom compression;
        // atom num_levels;
        // atom level0_file_num_compaction_trigger;
        // atom level0_slowdown_writes_trigger;
        // atom level0_stop_writes_trigger;
        // atom max_mem_compaction_level;
        // atom target_file_size_base;
        // atom target_file_size_multiplier;
        // atom max_bytes_for_level_base;
        // atom max_bytes_for_level_multiplier;
        // atom expand_compaction_factor;
        // atom source_compaction_factor;
        // atom max_grandparent_overlap_factor;
        // atom soft_rate_limit;
        // atom hard_rate_limit;
        // atom arena_block_size;
        // atom disable_auto_compaction;
        // atom purge_redundant_kvs_while_flush;
        // atom compaction_style;
        // atom verify_checksums_in_compaction;
        // atom filter_deletes;
        // atom max_sequential_kip_in_iterations;
        // atom inplace_update_support;
        // atom inplace_update_num_locks;
        // atom table_factory_block_cache_size;
        // atom in_memory_mode;
        // atom block_based_table_options;

        // DB Options
        atom total_threads;
        atom optimize_level_type_compaction_memtable_memory_budget;
        atom create_if_missing;
        atom max_open_files;
        atom compression_type;
        atom use_fsync;
        atom bytes_per_sync;
        atom allow_os_buffer;
        atom table_cache_num_shard_bits;
        atom min_write_buffer_number;
        atom max_write_buffer_number;
        atom write_buffer_size;
        atom max_bytes_for_level_base;
        atom max_bytes_for_level_multiplier;
        atom max_manifest_file_size;
        atom target_file_size_base;
        atom min_write_buffer_number_to_merge;
        atom level_zero_file_num_compaction_trigger;
        atom level_zero_slowdown_writes_trigger;
        atom level_zero_stop_writes_trigger;
        // atom compaction_style;
        atom max_background_compactions;
        atom max_background_flushes;
        atom disable_auto_compactions;
        atom report_bg_io_stats;
        atom num_levels;
        // atom max_total_wal_size;
        // atom use_fsync;
        // atom db_paths;
        // atom db_log_dir;
        // atom wal_dir;
        // atom delete_obsolete_files_period_micros;
        // atom max_background_compactions;
        // atom max_background_flushes;
        // atom max_log_file_size;
        // atom log_file_time_to_roll;
        // atom keep_log_file_num;
        // atom max_manifest_file_size;
        // atom table_cache_numshardbits;
        // atom wal_ttl_seconds;
        // atom wal_size_limit_mb;
        // atom manifest_preallocation_size;
        // atom allow_os_buffer;
        // atom allow_mmap_reads;
        // atom allow_mmap_writes;
        // atom is_fd_close_on_exec;
        // atom skip_log_error_on_recovery;
        // atom stats_dump_period_sec;
        // atom advise_random_on_open;
        // atom access_hint;
        // atom compaction_readahead_size;
        // atom use_adaptive_mutex;
        // atom bytes_per_sync;
        // atom skip_stats_update_on_db_open;
        // atom wal_recovery_mode;
        atom use_direct_io_for_flush_and_compaction;

        // Read Options
        // atom fill_cache;
        // atom iterate_upper_bound;
        // atom snapshot;

        // Write Options
        atom sync;
        atom disable_wal;
    }
}

struct SnapshotWrapper {
    pub snapshot: Option<Snapshot<'static>>,

    #[allow(dead_code)]
    db: Arc<RwLock<DB>>,
}

unsafe impl Sync for SnapshotWrapper {}
unsafe impl Send for SnapshotWrapper {}

unsafe fn shorten_snapshot<'a>(s: Snapshot<'static>) -> Snapshot<'a> {
    // More dragons: After working mad science to lengthen the lifetime of the snapshot struct when
    // creating it, we need to *shorten* its lifetime when its surrounding handle is being dropped,
    // to avoid memory leaks
    std::mem::transmute(s)
}

impl Deref for SnapshotWrapper {
    type Target = Snapshot<'static>;

    fn deref(&self) -> &Self::Target {
        self.snapshot
            .as_ref()
            .expect("Invariant violation: Attempting to use a freed snapshot handle. Good luck.")
    }
}

impl Drop for SnapshotWrapper {
    fn drop(&mut self) {
        let snapshot = std::mem::replace(&mut self.snapshot, None)
            .expect("Invariant violation: Dropping a snapshot handle that's already been dropped.");
        unsafe { shorten_snapshot(snapshot) };
    }
}

struct SnapshotHandle {
    pub snapshot: Arc<SnapshotWrapper>,
}

unsafe impl Sync for SnapshotHandle {}
unsafe impl Send for SnapshotHandle {}

impl SnapshotHandle {
    fn new<'a>(db: &DBHandle) -> Self {
        let read_db = db.read().unwrap();
        let snapshot = read_db.snapshot();

        // here be dragons!
        // The rocksdb snapshot type is (rightfully) parametrized to have the
        // same lifetime as its owning DB - this doesn't work for us, however, as we're not using
        // Rust lifetimes at *all* to control the lifetime of our values - instead, we rely
        // entirely on reference counting and the Erlang garbage collector to free variables.
        // Extending this lifetime to static essentially tells Rust to let it live forever, and
        // it'll be dropped on its own by the GC handling inside Rustler.
        let eternal_snapshot: Snapshot<'static> = unsafe { std::mem::transmute(snapshot) };

        let wrapper = SnapshotWrapper {
            snapshot: Some(eternal_snapshot),
            db: db.deref().clone(),
        };

        SnapshotHandle {
            snapshot: Arc::new(wrapper),
        }
    }
}

impl Deref for SnapshotHandle {
    type Target = SnapshotWrapper;

    fn deref(&self) -> &Self::Target {
        self.snapshot.deref()
    }
}

struct DBHandle {
    pub db: Arc<RwLock<DB>>,
}

impl Deref for DBHandle {
    type Target = Arc<RwLock<DB>>;

    fn deref(&self) -> &Self::Target {
        &self.db
    }
}

///
/// Struct representing a held reference to a database for use in refcount pools.
///
/// CFHandle, IteratorHandle, etc need to hold onto these references to prevent the parent database
/// from being dropped before the dependent child objects are
///
enum DatabaseRef {
    DB(Arc<RwLock<DB>>),
    Snapshot(Arc<SnapshotWrapper>),
}

impl<'a> From<&'a Arc<RwLock<DB>>> for DatabaseRef {
    fn from(db: &Arc<RwLock<DB>>) -> Self {
        DatabaseRef::DB(db.clone())
    }
}

impl<'a> From<&'a Arc<SnapshotWrapper>> for DatabaseRef {
    fn from(snapshot: &Arc<SnapshotWrapper>) -> Self {
        DatabaseRef::Snapshot(snapshot.clone())
    }
}

struct CFHandle {
    pub cf: ColumnFamily,
    #[allow(dead_code)]
    db: DatabaseRef,
}

unsafe impl Sync for CFHandle {}
unsafe impl Send for CFHandle {}

struct IteratorHandle {
    pub iter: RwLock<DBIterator>,
    #[allow(dead_code)]
    db: DatabaseRef,
}

unsafe impl Sync for IteratorHandle {}
unsafe impl Send for IteratorHandle {}

struct CompressionType {
    pub raw: DBCompressionType,
}
impl<'a> Decoder<'a> for CompressionType {
    fn decode(term: Term<'a>) -> NifResult<Self> {
        if atoms::none() == term {
            Ok(CompressionType {
                raw: DBCompressionType::None,
            })
        } else if atoms::snappy() == term {
            Ok(CompressionType {
                raw: DBCompressionType::Snappy,
            })
        } else if atoms::zlib() == term {
            Ok(CompressionType {
                raw: DBCompressionType::Zlib,
            })
        } else if atoms::bzip2() == term {
            Ok(CompressionType {
                raw: DBCompressionType::Bz2,
            })
        } else if atoms::lz4() == term {
            Ok(CompressionType {
                raw: DBCompressionType::Lz4,
            })
        } else if atoms::lz4h() == term {
            Ok(CompressionType {
                raw: DBCompressionType::Lz4hc,
            })
        } else {
            Err(Error::BadArg)
        }
    }
}

impl Into<DBCompressionType> for CompressionType {
    fn into(self) -> DBCompressionType {
        self.raw
    }
}

enum BatchOperation<'a> {
    Put(&'a [u8], &'a [u8]),
    PutCf(ColumnFamily, &'a [u8], &'a [u8]),
    Delete(&'a [u8]),
    DeleteCf(ColumnFamily, &'a [u8]),
}

impl<'a> Decoder<'a> for BatchOperation<'a> {
    fn decode(term: Term<'a>) -> NifResult<Self> {
        let (operation, details): (Term, Term) = term.decode()?;
        if atoms::put() == operation {
            let (key, val): (Binary, Binary) = details.decode()?;
            Ok(BatchOperation::Put(key.as_slice(), val.as_slice()))
        } else if atoms::put_cf() == operation {
            let (cf, key, val): (ResourceArc<CFHandle>, Binary, Binary) = details.decode()?;
            Ok(BatchOperation::PutCf(cf.cf, key.as_slice(), val.as_slice()))
        } else if atoms::delete() == operation {
            let key: Binary = details.decode()?;
            Ok(BatchOperation::Delete(key.as_slice()))
        } else if atoms::delete_cf() == operation {
            let (cf, key): (ResourceArc<CFHandle>, Binary) = details.decode()?;
            Ok(BatchOperation::DeleteCf(cf.cf, key.as_slice()))
        } else {
            Err(Error::BadArg)
        }
    }
}

macro_rules! handle_error {
    ($env:expr, $e:expr) => {
        match $e {
            Ok(inner) => inner,
            Err(err) => return Ok((atoms::error(), err.to_string().encode($env)).encode($env)),
        }
    };
}

fn decode_write_options<'a>(env: Env<'a>, arg: Term<'a>) -> NifResult<WriteOptions> {
    let mut opts = WriteOptions::new();

    if let Ok(sync) = arg.map_get(atoms::sync().to_term(env)) {
        opts.set_sync(sync.decode()?);
    }

    if let Ok(disable_wal) = arg.map_get(atoms::disable_wal().to_term(env)) {
        opts.disable_wal(disable_wal.decode()?);
    }

    Ok(opts)
}

// fn decode_read_options<'a>(env: Env<'a>, arg: Term<'a>) -> NifResult<ReadOptions> {
//     let mut opts = ReadOptions::default();

//     if let Ok(fill_cache) = arg.map_get(atoms::fill_cache().to_term(env)) {
//         opts.fill_cache(fill_cache.decode()?);
//     }

//     // TODO implement handling snapshots
//     if let Ok(iter_upper) = arg.map_get(atoms::iterate_upper_bound().to_term(env)) {
//         let key: &str = iter_upper.decode()?;
//         opts.set_iterate_upper_bound(key);
//     }

//     Ok(opts)
// }

fn decode_db_options<'a>(env: Env<'a>, arg: Term<'a>) -> NifResult<Options> {
    let mut opts = Options::default();

    if let Ok(count) = arg.map_get(atoms::total_threads().to_term(env)) {
        opts.increase_parallelism(count.decode()?);
    }

    if let Ok(memtable_budget) =
        arg.map_get(atoms::optimize_level_type_compaction_memtable_memory_budget().to_term(env))
    {
        let i_size: u64 = memtable_budget.decode()?;
        opts.optimize_level_style_compaction(i_size as usize);
    }

    if let Ok(create_if_missing) = arg.map_get(atoms::create_if_missing().to_term(env)) {
        opts.create_if_missing(create_if_missing.decode()?);
    }

    if let Ok(compression_type_opt) = arg.map_get(atoms::compression_type().to_term(env)) {
        let compression_type: CompressionType = compression_type_opt.decode()?;
        opts.set_compression_type(compression_type.into());
    }

    // TODO: Set Compression Type Per Level

    if let Ok(max_open_files) = arg.map_get(atoms::max_open_files().to_term(env)) {
        opts.set_max_open_files(max_open_files.decode()?);
    }

    if let Ok(use_fsync) = arg.map_get(atoms::use_fsync().to_term(env)) {
        opts.set_use_fsync(use_fsync.decode()?);
    }

    if let Ok(bytes_per_sync) = arg.map_get(atoms::bytes_per_sync().to_term(env)) {
        opts.set_bytes_per_sync(bytes_per_sync.decode()?);
    }

    if let Ok(allow_os_buffer) = arg.map_get(atoms::allow_os_buffer().to_term(env)) {
        #[allow(deprecated)]
        opts.set_allow_os_buffer(allow_os_buffer.decode()?);
    }

    if let Ok(nbits) = arg.map_get(atoms::table_cache_num_shard_bits().to_term(env)) {
        opts.set_table_cache_num_shard_bits(nbits.decode()?);
    }

    if let Ok(nbuf) = arg.map_get(atoms::min_write_buffer_number().to_term(env)) {
        opts.set_min_write_buffer_number(nbuf.decode()?);
    }

    if let Ok(nbuf) = arg.map_get(atoms::max_write_buffer_number().to_term(env)) {
        opts.set_max_write_buffer_number(nbuf.decode()?);
    }

    if let Ok(size) = arg.map_get(atoms::write_buffer_size().to_term(env)) {
        let i_size: u64 = size.decode()?;
        opts.set_write_buffer_size(i_size as usize);
    }

    if let Ok(max_bytes) = arg.map_get(atoms::max_bytes_for_level_base().to_term(env)) {
        opts.set_max_bytes_for_level_base(max_bytes.decode()?);
    }

    if let Ok(multiplier) = arg.map_get(atoms::max_bytes_for_level_multiplier().to_term(env)) {
        opts.set_max_bytes_for_level_multiplier(multiplier.decode()?);
    }

    if let Ok(max_size) = arg.map_get(atoms::max_manifest_file_size().to_term(env)) {
        let i_size: u64 = max_size.decode()?;
        opts.set_max_manifest_file_size(i_size as usize);
    }

    if let Ok(target_size) = arg.map_get(atoms::target_file_size_base().to_term(env)) {
        opts.set_target_file_size_base(target_size.decode()?);
    }

    if let Ok(to_merge) = arg.map_get(atoms::min_write_buffer_number_to_merge().to_term(env)) {
        opts.set_min_write_buffer_number_to_merge(to_merge.decode()?);
    }

    if let Ok(n) = arg.map_get(atoms::level_zero_file_num_compaction_trigger().to_term(env)) {
        opts.set_level_zero_file_num_compaction_trigger(n.decode()?);
    }

    if let Ok(n) = arg.map_get(atoms::level_zero_slowdown_writes_trigger().to_term(env)) {
        opts.set_level_zero_slowdown_writes_trigger(n.decode()?);
    }

    if let Ok(n) = arg.map_get(atoms::level_zero_stop_writes_trigger().to_term(env)) {
        opts.set_level_zero_stop_writes_trigger(n.decode()?);
    }

    // Todo set compaction style

    if let Ok(n) = arg.map_get(atoms::max_background_compactions().to_term(env)) {
        opts.set_max_background_compactions(n.decode()?);
    }

    if let Ok(n) = arg.map_get(atoms::max_background_flushes().to_term(env)) {
        opts.set_max_background_flushes(n.decode()?);
    }

    if let Ok(disable) = arg.map_get(atoms::disable_auto_compactions().to_term(env)) {
        opts.set_disable_auto_compactions(disable.decode()?);
    }

    if let Ok(bg_io_stats) = arg.map_get(atoms::report_bg_io_stats().to_term(env)) {
        opts.set_report_bg_io_stats(bg_io_stats.decode()?);
    }

    // Todo: set WAL Recovery Mode

    if let Ok(num_levels) = arg.map_get(atoms::num_levels().to_term(env)) {
        opts.set_num_levels(num_levels.decode()?);
    }

    if let Ok(enabled) = arg.map_get(atoms::use_direct_io_for_flush_and_compaction().to_term(env)) {
        opts.set_use_direct_io_for_flush_and_compaction(enabled.decode()?);
    }

    Ok(opts)
}

fn decode_iterator_mode<'a>(arg: Term<'a>) -> NifResult<IteratorMode<'a>> {
    match arg.get_type() {
        TermType::Atom => {
            let atom: Atom = arg.decode()?;

            if atom == atoms::start() {
                Ok(IteratorMode::Start)
            } else if atom == atoms::end() {
                Ok(IteratorMode::End)
            } else {
                Err(Error::BadArg)
            }
        }

        TermType::Tuple => {
            let (atom_from, key, atom_dir): (Atom, Binary, Atom) = arg.decode()?;

            if atom_from == atoms::from()
                && (atom_dir == atoms::forward() || atom_dir == atoms::reverse())
            {
                let dir: Direction = if atom_dir == atoms::forward() {
                    Direction::Forward
                } else {
                    Direction::Reverse
                };
                Ok(IteratorMode::From(key.as_slice(), dir))
            } else {
                Err(Error::BadArg)
            }
        }

        _ => Err(Error::BadArg),
    }
}

fn open<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let path: &Path = Path::new(args[0].decode()?);

    let db_opts = if args[1].map_size()? > 0 {
        decode_db_options(env, args[1])?
    } else {
        Options::default()
    };

    let cf: Vec<&str> = if args[2].list_length()? == 0 {
        vec![]
    } else {
        let iter: ListIterator = args[2].decode()?;
        let result: Vec<&str> = try!(iter.map(|x| x.decode()).collect::<NifResult<Vec<&str>>>());

        result
    };

    let db: DB = handle_error!(env, DB::open_cf(&db_opts, path, &cf));

    let resp = (
        atoms::ok(),
        ResourceArc::new(DBHandle {
            db: Arc::new(RwLock::new(db)),
        }),
    ).encode(env);

    Ok(resp)
}

fn create_snapshot<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let db_arc: ResourceArc<DBHandle> = args[0].decode()?;
    let db_handle = db_arc.deref();

    let resp = (
        atoms::ok(),
        ResourceArc::new(SnapshotHandle::new(db_handle)),
    ).encode(env);

    Ok(resp)
}

fn count<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let db_arc: ResourceArc<DBHandle> = args[0].decode()?;
    let db_handle = db_arc.deref();

    let iterator = db_handle.db.read().unwrap().iterator(IteratorMode::Start);

    let count = iterator.count();

    Ok((count as u64).encode(env))
}

fn count_cf<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let db_arc: ResourceArc<DBHandle> = args[0].decode()?;
    let db_handle = db_arc.deref();

    let cf_arc: ResourceArc<CFHandle> = args[1].decode()?;
    let cf = cf_arc.deref().cf;

    let iterator = handle_error!(
        env,
        db_handle
            .db
            .read()
            .unwrap()
            .iterator_cf(cf, IteratorMode::Start)
    );

    let count = iterator.count();

    Ok((count as u64).encode(env))
}

fn create_cf<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let db_arc: ResourceArc<DBHandle> = args[0].decode()?;
    let mut db = db_arc.db.write().unwrap();

    let name: &str = args[1].decode()?;
    let has_db_opts = args[2].map_size()? > 0;
    let opts = if has_db_opts {
        decode_db_options(env, args[2])?
    } else {
        Options::default()
    };

    let cf = handle_error!(env, db.create_cf(name, &opts));

    let resp = (
        atoms::ok(),
        ResourceArc::new(CFHandle {
            cf: cf,
            db: DatabaseRef::from(&db_arc.db),
        }),
    ).encode(env);

    Ok(resp)
}

fn list_cf<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let path: &Path = Path::new(args[0].decode()?);

    let db_opts = if args[1].map_size()? > 0 {
        decode_db_options(env, args[1])?
    } else {
        Options::default()
    };

    let paths: Vec<String> = handle_error!(env, DB::list_cf(&db_opts, path));
    let resp = (atoms::ok(), paths).encode(env);

    Ok(resp)
}

fn cf_handle<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let db_arc: ResourceArc<DBHandle> = args[0].decode()?;
    let db = db_arc.db.read().unwrap();

    let name: &str = args[1].decode()?;

    match db.cf_handle(name) {
        Some(cf) => Ok((
            atoms::ok(),
            ResourceArc::new(CFHandle {
                cf: cf,
                db: DatabaseRef::from(&db_arc.db),
            }),
        ).encode(env)),
        None => Ok((
            atoms::error(),
            format!("Could not find ColumnFamily named {}", name),
        ).encode(env)),
    }
}

fn put<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let db_arc: ResourceArc<DBHandle> = args[0].decode()?;
    let db = db_arc.deref().db.write().unwrap();

    let key: Binary = args[1].decode()?;
    let val: Binary = args[2].decode()?;

    let resp = if args[3].map_size()? > 0 {
        let write_opts = decode_write_options(env, args[2])?;
        db.put_opt(key.as_slice(), val.as_slice(), &write_opts)
    } else {
        db.put(key.as_slice(), val.as_slice())
    };

    handle_error!(env, resp);

    Ok(atoms::ok().encode(env))
}

fn put_cf<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let db_arc: ResourceArc<DBHandle> = args[0].decode()?;
    let db = db_arc.deref().db.write().unwrap();

    let cf_arc: ResourceArc<CFHandle> = args[1].decode()?;
    let cf = cf_arc.deref().cf;

    let key: Binary = args[2].decode()?;
    let val: Binary = args[3].decode()?;

    let resp = if args[4].map_size()? > 0 {
        let write_opts = decode_write_options(env, args[2])?;
        db.put_cf_opt(cf, key.as_slice(), val.as_slice(), &write_opts)
    } else {
        db.put_cf(cf, key.as_slice(), val.as_slice())
    };

    handle_error!(env, resp);

    Ok(atoms::ok().encode(env))
}

fn delete<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let db_arc: ResourceArc<DBHandle> = args[0].decode()?;
    let db = db_arc.deref().db.write().unwrap();

    let key: Binary = args[1].decode()?;

    let resp = if args[2].map_size()? > 0 {
        let write_opts = decode_write_options(env, args[2])?;
        db.delete_opt(key.as_slice(), &write_opts)
    } else {
        db.delete(key.as_slice())
    };

    handle_error!(env, resp);

    Ok(atoms::ok().encode(env))
}

fn delete_cf<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let db_arc: ResourceArc<DBHandle> = args[0].decode()?;
    let db = db_arc.deref().db.write().unwrap();

    let cf_arc: ResourceArc<CFHandle> = args[1].decode()?;
    let cf = cf_arc.deref().cf;

    let key: Binary = args[2].decode()?;

    let resp = if args[3].map_size()? > 0 {
        let write_opts = decode_write_options(env, args[3])?;
        db.delete_cf_opt(cf, key.as_slice(), &write_opts)
    } else {
        db.delete_cf(cf, key.as_slice())
    };

    handle_error!(env, resp);

    Ok(atoms::ok().encode(env))
}

fn get<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let key = args[1].decode::<Binary>()?.as_slice();

    let resp = args[0]
        .decode::<ResourceArc<DBHandle>>()
        .map(|db_arc| {
            let db = db_arc.db.read().unwrap();
            db.get(key)
        })
        .or_else(|_| {
            let snapshot_arc = args[0].decode::<ResourceArc<SnapshotHandle>>()?;
            Ok(snapshot_arc.get(key))
        })?;

    let val_option = handle_error!(env, resp);

    match val_option {
        Some(val) => {
            let mut bin = OwnedBinary::new(val.len()).unwrap();
            bin.as_mut_slice().write(&val).unwrap();

            Ok((atoms::ok(), bin.release(env).encode(env)).encode(env))
        }
        None => Ok(atoms::not_found().encode(env)),
    }
}

fn get_cf<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let cf = args[1].decode::<ResourceArc<CFHandle>>()?.cf;
    let key = args[2].decode::<Binary>()?.as_slice();

    let resp = args[0]
        .decode::<ResourceArc<DBHandle>>()
        .map(|db_arc| {
            let db = db_arc.db.read().unwrap();
            db.get_cf(cf, key)
        })
        .or_else(|_| {
            let snapshot_arc = args[0].decode::<ResourceArc<SnapshotHandle>>()?;
            Ok(snapshot_arc.get_cf(cf, key))
        })?;

    let val_option = handle_error!(env, resp);

    match val_option {
        Some(val) => {
            let mut bin = OwnedBinary::new(val.len()).unwrap();
            bin.as_mut_slice().write(&val).unwrap();

            Ok((atoms::ok(), bin.release(env).encode(env)).encode(env))
        }
        None => Ok(atoms::not_found().encode(env)),
    }
}

fn iterate<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let (iter, db_ref) = args[0]
        .decode::<ResourceArc<DBHandle>>()
        .and_then(|db_arc| {
            let db = db_arc.db.read().unwrap();
            Ok((
                db.iterator(decode_iterator_mode(args[1])?),
                DatabaseRef::from(&db_arc.db),
            ))
        })
        .or_else(|_| {
            let snapshot_arc = args[0].decode::<ResourceArc<SnapshotHandle>>()?;
            Ok((
                snapshot_arc
                    .snapshot
                    .snapshot
                    .as_ref()
                    .unwrap()
                    .iterator(decode_iterator_mode(args[1])?),
                DatabaseRef::from(&snapshot_arc.snapshot),
            ))
        })?;

    let resp = (
        atoms::ok(),
        ResourceArc::new(IteratorHandle {
            iter: RwLock::new(iter),
            db: db_ref,
        }),
    ).encode(env);

    Ok(resp)
}

fn iterate_cf<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let cf = args[1].decode::<ResourceArc<CFHandle>>()?.cf;

    let (iter_res, db_ref) = args[0]
        .decode::<ResourceArc<DBHandle>>()
        .and_then(|db_arc| {
            let db = db_arc.db.read().unwrap();
            Ok((
                db.iterator_cf(cf, decode_iterator_mode(args[2])?),
                DatabaseRef::from(&db_arc.db),
            ))
        })
        .or_else(|_| {
            let snapshot_arc = args[0].decode::<ResourceArc<SnapshotHandle>>()?;
            Ok((
                snapshot_arc
                    .snapshot
                    .snapshot
                    .as_ref()
                    .unwrap()
                    .iterator_cf(cf, decode_iterator_mode(args[2])?),
                DatabaseRef::from(&snapshot_arc.snapshot),
            ))
        })?;

    let iter = handle_error!(env, iter_res);

    let resp = (
        atoms::ok(),
        ResourceArc::new(IteratorHandle {
            iter: RwLock::new(iter),
            db: db_ref,
        }),
    ).encode(env);

    Ok(resp)
}

fn iterator_next<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let iter_arc: ResourceArc<IteratorHandle> = args[0].decode()?;
    let mut iter = iter_arc.iter.write().unwrap();

    match iter.next() {
        None => Ok(atoms::done().encode(env)),
        Some((key, val)) => {
            let mut enc_key = OwnedBinary::new(key.len()).unwrap();
            enc_key.as_mut_slice().write(&key).unwrap();

            let mut enc_val = OwnedBinary::new(val.len()).unwrap();
            enc_val.as_mut_slice().write(&val).unwrap();

            Ok((
                enc_key.release(env).encode(env),
                enc_val.release(env).encode(env),
            ).encode(env))
        }
    }
}

fn iterator_reset<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let iter_arc: ResourceArc<IteratorHandle> = args[0].decode()?;
    let mut iter = iter_arc.iter.write().unwrap();

    let iterator_mode = decode_iterator_mode(args[1])?;

    iter.set_mode(iterator_mode);

    Ok(atoms::ok().encode(env))
}

fn batch_write<'a>(env: Env<'a>, args: &[Term<'a>]) -> NifResult<Term<'a>> {
    let ops_iter: ListIterator = args[0].decode()?;
    let ops: Vec<BatchOperation> = try!(
        ops_iter
            .map(|x| x.decode())
            .collect::<NifResult<Vec<BatchOperation>>>()
    );

    let db_arc: ResourceArc<DBHandle> = args[1].decode()?;
    let db = db_arc.db.write().unwrap();

    let mut batch = WriteBatch::default();
    for op in ops {
        match op {
            BatchOperation::Put(key, val) => batch.put(key, val).unwrap(),
            BatchOperation::PutCf(cf, key, val) => batch.put_cf(cf, key, val).unwrap(),
            BatchOperation::Delete(key) => batch.delete(key).unwrap(),
            BatchOperation::DeleteCf(cf, key) => batch.delete_cf(cf, key).unwrap(),
        }
    }

    handle_error!(env, db.write(batch));
    Ok(atoms::ok().encode(env))
}

rustler_export_nifs!(
    "Elixir.Rox.Native",
    [
        ("open", 3, open),
        ("create_snapshot", 1, create_snapshot),
        ("create_cf", 3, create_cf),
        ("cf_handle", 2, cf_handle),
        ("put", 4, put),
        ("put_cf", 5, put_cf),
        ("delete", 3, delete),
        ("delete_cf", 4, delete_cf),
        ("count", 1, count),
        ("count_cf", 2, count_cf),
        ("list_cf", 2, list_cf),
        ("iterate", 2, iterate),
        ("iterate_cf", 3, iterate_cf),
        ("iterator_next", 1, iterator_next),
        ("iterator_reset", 2, iterator_reset),
        ("get", 3, get),
        ("get_cf", 4, get_cf),
        ("batch_write", 2, batch_write),
    ],
    Some(on_load)
);

fn on_load<'a>(env: Env<'a>, _load_info: Term<'a>) -> bool {
    resource_struct_init!(DBHandle, env);
    resource_struct_init!(CFHandle, env);
    resource_struct_init!(IteratorHandle, env);
    resource_struct_init!(SnapshotWrapper, env);
    resource_struct_init!(SnapshotHandle, env);
    true
}