<commit_before>use crate::compiler::{
    hir,
    lir::{Chunk, ChunkIndex, Instruction},
};
use itertools::Itertools;
use log::debug;
use std::collections::HashMap;

use super::value::Value;

/// A VM can execute some byte code.
pub struct Vm {
    pub(super) chunks: Vec<Chunk>,
    pub(super) status: Status,
    next_instruction: ByteCodePointer,
    pub(super) stack: Vec<StackEntry>,
    heap: HashMap<ObjectPointer, Object>,
    next_heap_address: ObjectPointer,
    stack_trace: Vec<hir::Id>,
}

#[derive(Clone)]
pub enum Status {
    Running,
    Done(Value),
    Panicked(Value),
}

/// Stack entries point either to instructions or to objects.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum StackEntry {
    ByteCode(ByteCodePointer),
    Object(ObjectPointer),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ByteCodePointer {
    chunk: ChunkIndex,
    instruction: usize,
}
pub type ObjectPointer = usize;

#[derive(Debug, Clone)] // TODO: remove Clone once it's no longer needed
pub struct Object {
    reference_count: usize,
    data: ObjectData,
}
#[derive(Clone, Debug)] // TODO: remove Clone once it's no longer needed
pub enum ObjectData {
    Int(u64),
    Text(String),
    Symbol(String),
    Struct(HashMap<ObjectPointer, ObjectPointer>),
    Closure {
        // TODO: This could later be just a vector of object pointers, but for
        // now we capture the whole stack.
        captured: Vec<StackEntry>,
        body: ChunkIndex,
    },
}

impl Vm {
    pub fn new(chunks: Vec<Chunk>) -> Self {
        Self {
            chunks: chunks.clone(),
            status: Status::Running,
            next_instruction: ByteCodePointer {
                chunk: chunks.len() - 1,
                instruction: 0,
            },
            stack: vec![],
            heap: HashMap::new(),
            next_heap_address: 0,
            stack_trace: vec![],
        }
    }

    pub fn status(&self) -> Status {
        self.status.clone()
    }

    fn get_from_stack(&self, offset: usize) -> StackEntry {
        self.stack[self.stack.len() - 1 - offset as usize].clone()
    }
    fn get_from_heap(&self, address: ObjectPointer) -> &Object {
        self.heap.get(&address).unwrap()
    }
    fn get_mut_from_heap(&mut self, address: ObjectPointer) -> &mut Object {
        self.heap.get_mut(&address).unwrap()
    }

    pub fn create_object(&mut self, object: ObjectData) -> ObjectPointer {
        let address = self.next_heap_address;
        self.heap.insert(
            address,
            Object {
                reference_count: 1,
                data: object,
            },
        );
        self.next_heap_address += 1;
        address
    }
    pub fn free_object(&mut self, address: ObjectPointer) {
        let object = self.heap.remove(&address).unwrap();
        assert_eq!(object.reference_count, 0);
        match object.data {
            ObjectData::Int(_) => {}
            ObjectData::Text(_) => {}
            ObjectData::Symbol(_) => {}
            ObjectData::Struct(entries) => {
                for (key, value) in entries {
                    self.drop(key);
                    self.drop(value);
                }
            }
            ObjectData::Closure { captured, .. } => {
                for entry in captured {
                    if let StackEntry::Object(address) = entry {
                        self.drop(address);
                    }
                }
            }
        }
    }

    fn dup(&mut self, address: ObjectPointer) {
        self.get_mut_from_heap(address).reference_count += 1;
    }
    fn drop(&mut self, address: ObjectPointer) {
        let object = self.get_mut_from_heap(address);
        object.reference_count -= 1;
        if object.reference_count == 0 {
            self.free_object(address);
        }
    }

    pub(super) fn import(&mut self, value: Value) -> ObjectPointer {
        let value = match value {
            Value::Int(int) => ObjectData::Int(int),
            Value::Text(text) => ObjectData::Text(text),
            Value::Symbol(symbol) => ObjectData::Symbol(symbol),
            Value::Struct(struct_) => {
                let mut entries = HashMap::new();
                for (key, value) in struct_ {
                    let key = self.import(key);
                    let value = self.import(value);
                    entries.insert(key, value);
                }
                ObjectData::Struct(entries)
            }
            Value::Closure { captured, body } => ObjectData::Closure { captured, body },
        };
        self.create_object(value)
    }
    pub(super) fn export(&mut self, address: ObjectPointer) -> Value {
        let value = self.export_helper(address);
        self.drop(address);
        value
    }
    fn export_helper(&self, address: ObjectPointer) -> Value {
        match &self.get_from_heap(address).data {
            ObjectData::Int(int) => Value::Int(*int),
            ObjectData::Text(text) => Value::Text(text.clone()),
            ObjectData::Symbol(symbol) => Value::Symbol(symbol.clone()),
            ObjectData::Struct(struct_) => {
                let mut entries = im::HashMap::new();
                for (key, value) in struct_ {
                    let key = self.export_helper(*key);
                    let value = self.export_helper(*value);
                    entries.insert(key, value);
                }
                Value::Struct(entries)
            }
            ObjectData::Closure { captured, body } => Value::Closure {
                captured: captured.clone(),
                body: *body,
            },
        }
    }

    pub fn run(&mut self, mut num_instructions: u16) {
        assert!(
            matches!(self.status, Status::Running),
            "Called Vm::run on a vm that is not ready to run."
        );
        while matches!(self.status, Status::Running) && num_instructions > 0 {
            num_instructions -= 1;
            let instruction = self.chunks[self.next_instruction.chunk].instructions
                [self.next_instruction.instruction]
                .clone();
            debug!("Executing instruction: {:?}", &instruction);
            self.next_instruction.instruction += 1;
            self.run_instruction(instruction);
            debug!(
                "Stack: {}",
                self.stack
                    .iter()
                    .map(|entry| match entry {
                        StackEntry::ByteCode(byte_code) => {
                            format!("->{}:{}", byte_code.chunk, byte_code.instruction)
                        }
                        StackEntry::Object(address) => format!("{}", address),
                    })
                    .join(", ")
            );

            if self.next_instruction.instruction
                >= self.chunks[self.next_instruction.chunk].instructions.len()
            {
                self.status = Status::Done(Value::nothing());
            }
        }
    }
    pub fn run_instruction(&mut self, instruction: Instruction) {
        match instruction {
            Instruction::CreateInt(int) => {
                let address = self.create_object(ObjectData::Int(int));
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::CreateText(text) => {
                let address = self.create_object(ObjectData::Text(text));
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::CreateSymbol(symbol) => {
                let address = self.create_object(ObjectData::Symbol(symbol));
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::CreateStruct { num_entries } => {
                let mut key_value_addresses = vec![];
                for _ in 0..(2 * num_entries) {
                    match self.stack.pop().unwrap() {
                        StackEntry::ByteCode(_) => panic!("Struct can only contain objects."),
                        StackEntry::Object(address) => key_value_addresses.push(address),
                    }
                }
                let mut entries = HashMap::new();
                for mut key_and_value in &key_value_addresses.into_iter().rev().chunks(2) {
                    let key = key_and_value.next().unwrap();
                    let value = key_and_value.next().unwrap();
                    assert_eq!(key_and_value.next(), None);
                    entries.insert(key, value);
                }
                let address = self.create_object(ObjectData::Struct(entries));
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::CreateClosure(chunk_index) => {
                let stack = self.stack.clone();
                for entry in &stack {
                    match entry {
                        StackEntry::ByteCode(_) => {}
                        StackEntry::Object(address) => {
                            self.dup(*address);
                        }
                    }
                }
                let address = self.create_object(ObjectData::Closure {
                    captured: stack,
                    body: chunk_index,
                });
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::PopMultipleBelowTop(n) => {
                let top = self.stack.pop().unwrap();
                for _ in 0..n {
                    if let StackEntry::Object(address) = self.stack.pop().unwrap() {
                        self.drop(address);
                    }
                }
                self.stack.push(top);
            }
            Instruction::PushFromStack(offset) => {
                let entry = self.get_from_stack(offset);
                if let StackEntry::Object(address) = &entry {
                    self.dup(*address);
                }
                self.stack.push(entry);
            }
            Instruction::Call => {
                let closure_address = match self.stack.pop().unwrap() {
                    StackEntry::ByteCode(_) => panic!(),
                    StackEntry::Object(address) => address,
                };
                let (captured, body) = match &self.get_from_heap(closure_address).data {
                    ObjectData::Closure { captured, body } => (captured.clone(), *body),
                    _ => panic!("Can't call non-closure."),
                };
                let num_args = self.chunks[body].num_args;
                let mut args = vec![];
                for _ in 0..num_args {
                    let address = match self.stack.pop().unwrap() {
                        StackEntry::ByteCode(_) => {
                            panic!("You can only pass objects as arguments.")
                        }
                        StackEntry::Object(address) => address,
                    };
                    args.push(address);
                }
                self.stack.push(StackEntry::ByteCode(self.next_instruction));
                self.stack.append(&mut captured.clone());
                for arg in args {
                    self.stack.push(StackEntry::Object(arg));
                }
                self.drop(closure_address);
                self.next_instruction = ByteCodePointer {
                    chunk: body,
                    instruction: 0,
                };
            }
            Instruction::Return => {
                let return_value = self.stack.pop().unwrap();
                let caller = match self.stack.pop().unwrap() {
                    StackEntry::ByteCode(address) => address,
                    StackEntry::Object(_) => panic!("Return with no caller on stack"),
                };
                self.stack.push(return_value);
                self.next_instruction = caller;
            }
            Instruction::Builtin(builtin_function) => {
                self.run_builtin_function(builtin_function);
            }
            Instruction::DebugValueEvaluated(_) => {}
            Instruction::DebugClosureEntered(hir_id) => self.stack_trace.push(hir_id),
            Instruction::DebugClosureExited => {
                self.stack_trace.pop().unwrap();
            }
            Instruction::Error(_) => {
                self.panic(
                    "The VM crashed because there was an error in previous compilation stages."
                        .to_string(),
                );
            }
        }
    }

    pub fn panic(&mut self, message: String) -> Value {
        self.status = Status::Panicked(Value::Text(message));
        Value::Symbol("Never".to_string())
    }

    pub fn current_stack_trace(&self) -> Vec<hir::Id> {
        self.stack_trace.clone()
    }
}
<commit_msg>In VM, log objects, not IDs<commit_after>use crate::compiler::{
    hir,
    lir::{Chunk, ChunkIndex, Instruction},
};
use itertools::Itertools;
use log::debug;
use std::collections::HashMap;

use super::value::Value;

/// A VM can execute some byte code.
pub struct Vm {
    pub(super) chunks: Vec<Chunk>,
    pub(super) status: Status,
    next_instruction: ByteCodePointer,
    pub(super) stack: Vec<StackEntry>,
    heap: HashMap<ObjectPointer, Object>,
    next_heap_address: ObjectPointer,
    stack_trace: Vec<hir::Id>,
}

#[derive(Clone)]
pub enum Status {
    Running,
    Done(Value),
    Panicked(Value),
}

/// Stack entries point either to instructions or to objects.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum StackEntry {
    ByteCode(ByteCodePointer),
    Object(ObjectPointer),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ByteCodePointer {
    chunk: ChunkIndex,
    instruction: usize,
}
pub type ObjectPointer = usize;

#[derive(Debug, Clone)] // TODO: remove Clone once it's no longer needed
pub struct Object {
    reference_count: usize,
    data: ObjectData,
}
#[derive(Clone, Debug)] // TODO: remove Clone once it's no longer needed
pub enum ObjectData {
    Int(u64),
    Text(String),
    Symbol(String),
    Struct(HashMap<ObjectPointer, ObjectPointer>),
    Closure {
        // TODO: This could later be just a vector of object pointers, but for
        // now we capture the whole stack.
        captured: Vec<StackEntry>,
        body: ChunkIndex,
    },
}

impl Vm {
    pub fn new(chunks: Vec<Chunk>) -> Self {
        Self {
            chunks: chunks.clone(),
            status: Status::Running,
            next_instruction: ByteCodePointer {
                chunk: chunks.len() - 1,
                instruction: 0,
            },
            stack: vec![],
            heap: HashMap::new(),
            next_heap_address: 0,
            stack_trace: vec![],
        }
    }

    pub fn status(&self) -> Status {
        self.status.clone()
    }

    fn get_from_stack(&self, offset: usize) -> StackEntry {
        self.stack[self.stack.len() - 1 - offset as usize].clone()
    }
    fn get_from_heap(&self, address: ObjectPointer) -> &Object {
        self.heap.get(&address).unwrap()
    }
    fn get_mut_from_heap(&mut self, address: ObjectPointer) -> &mut Object {
        self.heap.get_mut(&address).unwrap()
    }

    pub fn create_object(&mut self, object: ObjectData) -> ObjectPointer {
        let address = self.next_heap_address;
        self.heap.insert(
            address,
            Object {
                reference_count: 1,
                data: object,
            },
        );
        self.next_heap_address += 1;
        address
    }
    pub fn free_object(&mut self, address: ObjectPointer) {
        let object = self.heap.remove(&address).unwrap();
        assert_eq!(object.reference_count, 0);
        match object.data {
            ObjectData::Int(_) => {}
            ObjectData::Text(_) => {}
            ObjectData::Symbol(_) => {}
            ObjectData::Struct(entries) => {
                for (key, value) in entries {
                    self.drop(key);
                    self.drop(value);
                }
            }
            ObjectData::Closure { captured, .. } => {
                for entry in captured {
                    if let StackEntry::Object(address) = entry {
                        self.drop(address);
                    }
                }
            }
        }
    }

    fn dup(&mut self, address: ObjectPointer) {
        self.get_mut_from_heap(address).reference_count += 1;
    }
    fn drop(&mut self, address: ObjectPointer) {
        let object = self.get_mut_from_heap(address);
        object.reference_count -= 1;
        if object.reference_count == 0 {
            self.free_object(address);
        }
    }

    pub(super) fn import(&mut self, value: Value) -> ObjectPointer {
        let value = match value {
            Value::Int(int) => ObjectData::Int(int),
            Value::Text(text) => ObjectData::Text(text),
            Value::Symbol(symbol) => ObjectData::Symbol(symbol),
            Value::Struct(struct_) => {
                let mut entries = HashMap::new();
                for (key, value) in struct_ {
                    let key = self.import(key);
                    let value = self.import(value);
                    entries.insert(key, value);
                }
                ObjectData::Struct(entries)
            }
            Value::Closure { captured, body } => ObjectData::Closure { captured, body },
        };
        self.create_object(value)
    }
    pub(super) fn export(&mut self, address: ObjectPointer) -> Value {
        let value = self.export_helper(address);
        self.drop(address);
        value
    }
    fn export_helper(&self, address: ObjectPointer) -> Value {
        match &self.get_from_heap(address).data {
            ObjectData::Int(int) => Value::Int(*int),
            ObjectData::Text(text) => Value::Text(text.clone()),
            ObjectData::Symbol(symbol) => Value::Symbol(symbol.clone()),
            ObjectData::Struct(struct_) => {
                let mut entries = im::HashMap::new();
                for (key, value) in struct_ {
                    let key = self.export_helper(*key);
                    let value = self.export_helper(*value);
                    entries.insert(key, value);
                }
                Value::Struct(entries)
            }
            ObjectData::Closure { captured, body } => Value::Closure {
                captured: captured.clone(),
                body: *body,
            },
        }
    }

    pub fn run(&mut self, mut num_instructions: u16) {
        assert!(
            matches!(self.status, Status::Running),
            "Called Vm::run on a vm that is not ready to run."
        );
        while matches!(self.status, Status::Running) && num_instructions > 0 {
            num_instructions -= 1;
            let instruction = self.chunks[self.next_instruction.chunk].instructions
                [self.next_instruction.instruction]
                .clone();
            debug!("Executing instruction: {:?}", &instruction);
            self.next_instruction.instruction += 1;
            self.run_instruction(instruction);
            debug!(
                "Stack: {}",
                self.stack
                    .iter()
                    .map(|entry| match entry {
                        StackEntry::ByteCode(byte_code) => {
                            format!("here#{}:{}", byte_code.chunk, byte_code.instruction)
                        }
                        StackEntry::Object(address) => format!("{}", self.export_helper(*address)),
                    })
                    .join(", ")
            );

            if self.next_instruction.instruction
                >= self.chunks[self.next_instruction.chunk].instructions.len()
            {
                self.status = Status::Done(Value::nothing());
            }
        }
    }
    pub fn run_instruction(&mut self, instruction: Instruction) {
        match instruction {
            Instruction::CreateInt(int) => {
                let address = self.create_object(ObjectData::Int(int));
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::CreateText(text) => {
                let address = self.create_object(ObjectData::Text(text));
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::CreateSymbol(symbol) => {
                let address = self.create_object(ObjectData::Symbol(symbol));
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::CreateStruct { num_entries } => {
                let mut key_value_addresses = vec![];
                for _ in 0..(2 * num_entries) {
                    match self.stack.pop().unwrap() {
                        StackEntry::ByteCode(_) => panic!("Struct can only contain objects."),
                        StackEntry::Object(address) => key_value_addresses.push(address),
                    }
                }
                let mut entries = HashMap::new();
                for mut key_and_value in &key_value_addresses.into_iter().rev().chunks(2) {
                    let key = key_and_value.next().unwrap();
                    let value = key_and_value.next().unwrap();
                    assert_eq!(key_and_value.next(), None);
                    entries.insert(key, value);
                }
                let address = self.create_object(ObjectData::Struct(entries));
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::CreateClosure(chunk_index) => {
                let stack = self.stack.clone();
                for entry in &stack {
                    match entry {
                        StackEntry::ByteCode(_) => {}
                        StackEntry::Object(address) => {
                            self.dup(*address);
                        }
                    }
                }
                let address = self.create_object(ObjectData::Closure {
                    captured: stack,
                    body: chunk_index,
                });
                self.stack.push(StackEntry::Object(address));
            }
            Instruction::PopMultipleBelowTop(n) => {
                let top = self.stack.pop().unwrap();
                for _ in 0..n {
                    if let StackEntry::Object(address) = self.stack.pop().unwrap() {
                        self.drop(address);
                    }
                }
                self.stack.push(top);
            }
            Instruction::PushFromStack(offset) => {
                let entry = self.get_from_stack(offset);
                if let StackEntry::Object(address) = &entry {
                    self.dup(*address);
                }
                self.stack.push(entry);
            }
            Instruction::Call => {
                let closure_address = match self.stack.pop().unwrap() {
                    StackEntry::ByteCode(_) => panic!(),
                    StackEntry::Object(address) => address,
                };
                let (captured, body) = match &self.get_from_heap(closure_address).data {
                    ObjectData::Closure { captured, body } => (captured.clone(), *body),
                    _ => panic!("Can't call non-closure."),
                };
                let num_args = self.chunks[body].num_args;
                let mut args = vec![];
                for _ in 0..num_args {
                    let address = match self.stack.pop().unwrap() {
                        StackEntry::ByteCode(_) => {
                            panic!("You can only pass objects as arguments.")
                        }
                        StackEntry::Object(address) => address,
                    };
                    args.push(address);
                }
                self.stack.push(StackEntry::ByteCode(self.next_instruction));
                self.stack.append(&mut captured.clone());
                for arg in args {
                    self.stack.push(StackEntry::Object(arg));
                }
                self.drop(closure_address);
                self.next_instruction = ByteCodePointer {
                    chunk: body,
                    instruction: 0,
                };
            }
            Instruction::Return => {
                let return_value = self.stack.pop().unwrap();
                let caller = match self.stack.pop().unwrap() {
                    StackEntry::ByteCode(address) => address,
                    StackEntry::Object(_) => panic!("Return with no caller on stack"),
                };
                self.stack.push(return_value);
                self.next_instruction = caller;
            }
            Instruction::Builtin(builtin_function) => {
                self.run_builtin_function(builtin_function);
            }
            Instruction::DebugValueEvaluated(_) => {}
            Instruction::DebugClosureEntered(hir_id) => self.stack_trace.push(hir_id),
            Instruction::DebugClosureExited => {
                self.stack_trace.pop().unwrap();
            }
            Instruction::Error(_) => {
                self.panic(
                    "The VM crashed because there was an error in previous compilation stages."
                        .to_string(),
                );
            }
        }
    }

    pub fn panic(&mut self, message: String) -> Value {
        self.status = Status::Panicked(Value::Text(message));
        Value::Symbol("Never".to_string())
    }

    pub fn current_stack_trace(&self) -> Vec<hir::Id> {
        self.stack_trace.clone()
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Tyler Neely
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
extern crate libc;
use self::libc::{c_int,size_t};
use std::ffi::CString;
use std::mem;

use rocksdb_ffi::{self, DBCompressionType};
use merge_operator::{self, MergeFn, MergeOperatorCallback,
                     full_merge_callback, partial_merge_callback};
use comparator::{self, ComparatorCallback, compare_callback};

pub struct BlockBasedOptions {
    inner: rocksdb_ffi::DBBlockBasedTableOptions,
}

pub struct Options {
    pub inner: rocksdb_ffi::DBOptions,
}

pub struct WriteOptions {
    pub inner: rocksdb_ffi::DBWriteOptions,
}

impl Drop for Options {
    fn drop(&mut self) {
        unsafe {
            rocksdb_ffi::rocksdb_options_destroy(self.inner);
        }
    }
}

impl Drop for BlockBasedOptions {
    fn drop(&mut self) {
        unsafe {
            rocksdb_ffi::rocksdb_block_based_options_destroy(self.inner);
        }
    }
}

impl Drop for WriteOptions {
    fn drop(&mut self) {
        unsafe {
            rocksdb_ffi::rocksdb_writeoptions_destroy(self.inner);
        }
    }
}

impl BlockBasedOptions {
    pub fn set_block_size(&mut self, size: usize) {
        unsafe {
            rocksdb_ffi::rocksdb_block_based_options_set_block_size(self.inner,
                                                                    size);
        }
    }

    pub fn set_lru_cache(&mut self, size: size_t) {
        let cache = rocksdb_ffi::new_cache(size);
        unsafe {
            // because cache is wrapped in shared_ptr, so we don't need to call
            // rocksdb_cache_destroy explicitly.
            rocksdb_ffi::rocksdb_block_based_options_set_block_cache(self.inner, cache);
        }
    }

    pub fn set_bloom_filter(&mut self, bits_per_key: c_int, block_based: bool) {
        unsafe {
            let bloom = if block_based {
                rocksdb_ffi::rocksdb_filterpolicy_create_bloom(bits_per_key)
            } else {
                rocksdb_ffi::rocksdb_filterpolicy_create_bloom_full(bits_per_key)
            };

            rocksdb_ffi::rocksdb_block_based_options_set_filter_policy(self.inner,
                                                                       bloom);
        }
    }

    pub fn set_cache_index_and_filter_blocks(&mut self, v: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_block_based_options_set_cache_index_and_filter_blocks(self.inner,
                                                                                v as u8);
        }
    }
}

impl Default for BlockBasedOptions {
    fn default() -> BlockBasedOptions {
        let block_opts = unsafe {
            rocksdb_ffi::rocksdb_block_based_options_create()
        };
        if block_opts.is_null() {
            panic!("Could not create rocksdb block based options".to_string());
        }
        BlockBasedOptions { inner: block_opts }
    }
}

impl Options {
    pub fn increase_parallelism(&mut self, parallelism: i32) {
        unsafe {
            rocksdb_ffi::rocksdb_options_increase_parallelism(self.inner,
                                                              parallelism);
        }
    }

    pub fn optimize_level_style_compaction(&mut self,
                                           memtable_memory_budget: i32) {
        unsafe {
            rocksdb_ffi::rocksdb_options_optimize_level_style_compaction(
                self.inner, memtable_memory_budget);
        }
    }

    pub fn create_if_missing(&mut self, create_if_missing: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_create_if_missing(
                self.inner, create_if_missing);
        }
    }

    pub fn compression(&mut self, t: DBCompressionType) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_compression(self.inner, t);
        }
    }

    pub fn compression_per_level(&mut self, level_types: &[DBCompressionType]) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_compression_per_level(self.inner,
                                                                level_types.as_ptr(),
                                                                level_types.len() as size_t)
        }
    }

    pub fn add_merge_operator(&mut self,
                              name: &str,
                              merge_fn: MergeFn) {
        let cb = Box::new(MergeOperatorCallback {
            name: CString::new(name.as_bytes()).unwrap(),
            merge_fn: merge_fn,
        });

        unsafe {
            let mo = rocksdb_ffi::rocksdb_mergeoperator_create(
                mem::transmute(cb),
                merge_operator::destructor_callback,
                full_merge_callback,
                partial_merge_callback,
                None,
                merge_operator::name_callback);
            rocksdb_ffi::rocksdb_options_set_merge_operator(self.inner, mo);
        }
    }

    pub fn add_comparator(&mut self,
                          name: &str,
                          compare_fn: fn(&[u8], &[u8]) -> i32) {
        let cb = Box::new(ComparatorCallback {
            name: CString::new(name.as_bytes()).unwrap(),
            f: compare_fn,
        });

        unsafe {
            let cmp = rocksdb_ffi::rocksdb_comparator_create(
                mem::transmute(cb),
                comparator::destructor_callback,
                compare_callback,
                comparator::name_callback);
            rocksdb_ffi::rocksdb_options_set_comparator(self.inner, cmp);
        }
    }

    pub fn set_block_cache_size_mb(&mut self, cache_size: u64) {
        unsafe {
            rocksdb_ffi::rocksdb_options_optimize_for_point_lookup(self.inner,
                                                                   cache_size);
        }
    }

    pub fn set_max_open_files(&mut self, nfiles: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_open_files(self.inner, nfiles);
        }
    }

    pub fn set_use_fsync(&mut self, useit: bool) {
        unsafe {
            if useit {
                rocksdb_ffi::rocksdb_options_set_use_fsync(self.inner, 1)
            } else {
                rocksdb_ffi::rocksdb_options_set_use_fsync(self.inner, 0)
            }
        }
    }

    pub fn set_bytes_per_sync(&mut self, nbytes: u64) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_bytes_per_sync(self.inner, nbytes);
        }
    }

    pub fn set_disable_data_sync(&mut self, disable: bool) {
        unsafe {
            if disable {
                rocksdb_ffi::rocksdb_options_set_disable_data_sync(self.inner, 1)
            } else {
                rocksdb_ffi::rocksdb_options_set_disable_data_sync(self.inner, 0)
            }
        }
    }

    pub fn allow_os_buffer(&mut self, is_allow: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_allow_os_buffer(self.inner,
                                                             is_allow);
        }
    }

    pub fn set_table_cache_num_shard_bits(&mut self, nbits: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_table_cache_numshardbits(self.inner,
                                                                      nbits);
        }
    }

    pub fn set_min_write_buffer_number(&mut self, nbuf: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_min_write_buffer_number_to_merge(
                self.inner, nbuf);
        }
    }

    pub fn set_max_write_buffer_number(&mut self, nbuf: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_write_buffer_number(self.inner,
                                                                     nbuf);
        }
    }

    pub fn set_write_buffer_size(&mut self, size: usize) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_write_buffer_size(self.inner,
                                                               size);
        }
    }

    pub fn set_max_bytes_for_level_base(&mut self, size: u64) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_bytes_for_level_base(self.inner, size);
        }
    }

    pub fn set_max_bytes_for_level_multiplier(&mut self, mul: i32) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_bytes_for_level_multiplier(self.inner, mul);
        }
    }

    pub fn set_target_file_size_base(&mut self, size: u64) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_target_file_size_base(self.inner,
                                                                   size);
        }
    }

    pub fn set_min_write_buffer_number_to_merge(&mut self, to_merge: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_min_write_buffer_number_to_merge(
                self.inner, to_merge);
        }
    }

    pub fn set_level_zero_compaction_trigger(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_level0_file_num_compaction_trigger(
                self.inner, n);
        }
    }

    pub fn set_level_zero_slowdown_writes_trigger(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_level0_slowdown_writes_trigger(
                self.inner, n);
        }
    }

    pub fn set_level_zero_stop_writes_trigger(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_level0_stop_writes_trigger(
                self.inner, n);
        }
    }

    pub fn set_compaction_style(&mut self,
                                style: rocksdb_ffi::DBCompactionStyle) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_compaction_style(self.inner,
                                                              style);
        }
    }

    pub fn set_max_background_compactions(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_background_compactions(
                self.inner, n);
        }
    }

    pub fn set_max_background_flushes(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_background_flushes(self.inner,
                                                                    n);
        }
    }

    pub fn set_filter_deletes(&mut self, filter: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_filter_deletes(self.inner, filter);
        }
    }

    pub fn set_disable_auto_compactions(&mut self, disable: bool) {
        let c_bool = if disable {
            1
        } else {
            0
        };
        unsafe {
            rocksdb_ffi::rocksdb_options_set_disable_auto_compactions(self.inner, c_bool)
        }
    }

    pub fn set_block_based_table_factory(&mut self,
                                         factory: &BlockBasedOptions) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_block_based_table_factory(self.inner, factory.inner);
        }
    }

    pub fn set_report_bg_io_stats(&mut self, enable: bool) {
        unsafe {
            if enable {
                rocksdb_ffi::rocksdb_options_set_report_bg_io_stats(self.inner, 1);
            } else {
                rocksdb_ffi::rocksdb_options_set_report_bg_io_stats(self.inner, 0);
            }
        }
    }
}

impl Default for Options {
    fn default() -> Options {
        unsafe {
            let opts = rocksdb_ffi::rocksdb_options_create();
            if opts.is_null() {
                panic!("Could not create rocksdb options".to_string());
            }
            Options { inner: opts }
        }
    }
}

impl WriteOptions {
    pub fn new() -> WriteOptions {
        WriteOptions::default()
    }

    pub fn set_sync(&mut self, sync: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_writeoptions_set_sync(self.inner, sync);
        }
    }

    pub fn disable_wal(&mut self, disable: bool) {
        unsafe {
            if disable {
                rocksdb_ffi::rocksdb_writeoptions_disable_WAL(self.inner, 1);
            } else {
                rocksdb_ffi::rocksdb_writeoptions_disable_WAL(self.inner, 0);
            }
        }
    }
}

impl Default for WriteOptions {
    fn default() -> WriteOptions {
        let write_opts = unsafe { rocksdb_ffi::rocksdb_writeoptions_create() };
        if write_opts.is_null() {
            panic!("Could not create rocksdb write options".to_string());
        }
        WriteOptions { inner: write_opts }
    }
}
<commit_msg>add set_level_zero_file_num_compaction_trigger api<commit_after>// Copyright 2014 Tyler Neely
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
extern crate libc;
use self::libc::{c_int,size_t};
use std::ffi::CString;
use std::mem;

use rocksdb_ffi::{self, DBCompressionType};
use merge_operator::{self, MergeFn, MergeOperatorCallback,
                     full_merge_callback, partial_merge_callback};
use comparator::{self, ComparatorCallback, compare_callback};

pub struct BlockBasedOptions {
    inner: rocksdb_ffi::DBBlockBasedTableOptions,
}

pub struct Options {
    pub inner: rocksdb_ffi::DBOptions,
}

pub struct WriteOptions {
    pub inner: rocksdb_ffi::DBWriteOptions,
}

impl Drop for Options {
    fn drop(&mut self) {
        unsafe {
            rocksdb_ffi::rocksdb_options_destroy(self.inner);
        }
    }
}

impl Drop for BlockBasedOptions {
    fn drop(&mut self) {
        unsafe {
            rocksdb_ffi::rocksdb_block_based_options_destroy(self.inner);
        }
    }
}

impl Drop for WriteOptions {
    fn drop(&mut self) {
        unsafe {
            rocksdb_ffi::rocksdb_writeoptions_destroy(self.inner);
        }
    }
}

impl BlockBasedOptions {
    pub fn set_block_size(&mut self, size: usize) {
        unsafe {
            rocksdb_ffi::rocksdb_block_based_options_set_block_size(self.inner,
                                                                    size);
        }
    }

    pub fn set_lru_cache(&mut self, size: size_t) {
        let cache = rocksdb_ffi::new_cache(size);
        unsafe {
            // because cache is wrapped in shared_ptr, so we don't need to call
            // rocksdb_cache_destroy explicitly.
            rocksdb_ffi::rocksdb_block_based_options_set_block_cache(self.inner, cache);
        }
    }

    pub fn set_bloom_filter(&mut self, bits_per_key: c_int, block_based: bool) {
        unsafe {
            let bloom = if block_based {
                rocksdb_ffi::rocksdb_filterpolicy_create_bloom(bits_per_key)
            } else {
                rocksdb_ffi::rocksdb_filterpolicy_create_bloom_full(bits_per_key)
            };

            rocksdb_ffi::rocksdb_block_based_options_set_filter_policy(self.inner,
                                                                       bloom);
        }
    }

    pub fn set_cache_index_and_filter_blocks(&mut self, v: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_block_based_options_set_cache_index_and_filter_blocks(self.inner,
                                                                                v as u8);
        }
    }
}

impl Default for BlockBasedOptions {
    fn default() -> BlockBasedOptions {
        let block_opts = unsafe {
            rocksdb_ffi::rocksdb_block_based_options_create()
        };
        if block_opts.is_null() {
            panic!("Could not create rocksdb block based options".to_string());
        }
        BlockBasedOptions { inner: block_opts }
    }
}

impl Options {
    pub fn increase_parallelism(&mut self, parallelism: i32) {
        unsafe {
            rocksdb_ffi::rocksdb_options_increase_parallelism(self.inner,
                                                              parallelism);
        }
    }

    pub fn optimize_level_style_compaction(&mut self,
                                           memtable_memory_budget: i32) {
        unsafe {
            rocksdb_ffi::rocksdb_options_optimize_level_style_compaction(
                self.inner, memtable_memory_budget);
        }
    }

    pub fn create_if_missing(&mut self, create_if_missing: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_create_if_missing(
                self.inner, create_if_missing);
        }
    }

    pub fn compression(&mut self, t: DBCompressionType) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_compression(self.inner, t);
        }
    }

    pub fn compression_per_level(&mut self, level_types: &[DBCompressionType]) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_compression_per_level(self.inner,
                                                                level_types.as_ptr(),
                                                                level_types.len() as size_t)
        }
    }

    pub fn add_merge_operator(&mut self,
                              name: &str,
                              merge_fn: MergeFn) {
        let cb = Box::new(MergeOperatorCallback {
            name: CString::new(name.as_bytes()).unwrap(),
            merge_fn: merge_fn,
        });

        unsafe {
            let mo = rocksdb_ffi::rocksdb_mergeoperator_create(
                mem::transmute(cb),
                merge_operator::destructor_callback,
                full_merge_callback,
                partial_merge_callback,
                None,
                merge_operator::name_callback);
            rocksdb_ffi::rocksdb_options_set_merge_operator(self.inner, mo);
        }
    }

    pub fn add_comparator(&mut self,
                          name: &str,
                          compare_fn: fn(&[u8], &[u8]) -> i32) {
        let cb = Box::new(ComparatorCallback {
            name: CString::new(name.as_bytes()).unwrap(),
            f: compare_fn,
        });

        unsafe {
            let cmp = rocksdb_ffi::rocksdb_comparator_create(
                mem::transmute(cb),
                comparator::destructor_callback,
                compare_callback,
                comparator::name_callback);
            rocksdb_ffi::rocksdb_options_set_comparator(self.inner, cmp);
        }
    }

    pub fn set_block_cache_size_mb(&mut self, cache_size: u64) {
        unsafe {
            rocksdb_ffi::rocksdb_options_optimize_for_point_lookup(self.inner,
                                                                   cache_size);
        }
    }

    pub fn set_max_open_files(&mut self, nfiles: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_open_files(self.inner, nfiles);
        }
    }

    pub fn set_use_fsync(&mut self, useit: bool) {
        unsafe {
            if useit {
                rocksdb_ffi::rocksdb_options_set_use_fsync(self.inner, 1)
            } else {
                rocksdb_ffi::rocksdb_options_set_use_fsync(self.inner, 0)
            }
        }
    }

    pub fn set_bytes_per_sync(&mut self, nbytes: u64) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_bytes_per_sync(self.inner, nbytes);
        }
    }

    pub fn set_disable_data_sync(&mut self, disable: bool) {
        unsafe {
            if disable {
                rocksdb_ffi::rocksdb_options_set_disable_data_sync(self.inner, 1)
            } else {
                rocksdb_ffi::rocksdb_options_set_disable_data_sync(self.inner, 0)
            }
        }
    }

    pub fn allow_os_buffer(&mut self, is_allow: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_allow_os_buffer(self.inner,
                                                             is_allow);
        }
    }

    pub fn set_table_cache_num_shard_bits(&mut self, nbits: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_table_cache_numshardbits(self.inner,
                                                                      nbits);
        }
    }

    pub fn set_min_write_buffer_number(&mut self, nbuf: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_min_write_buffer_number_to_merge(
                self.inner, nbuf);
        }
    }

    pub fn set_max_write_buffer_number(&mut self, nbuf: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_write_buffer_number(self.inner,
                                                                     nbuf);
        }
    }

    pub fn set_write_buffer_size(&mut self, size: usize) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_write_buffer_size(self.inner,
                                                               size);
        }
    }

    pub fn set_max_bytes_for_level_base(&mut self, size: u64) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_bytes_for_level_base(self.inner, size);
        }
    }

    pub fn set_max_bytes_for_level_multiplier(&mut self, mul: i32) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_bytes_for_level_multiplier(self.inner, mul);
        }
    }

    pub fn set_target_file_size_base(&mut self, size: u64) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_target_file_size_base(self.inner,
                                                                   size);
        }
    }

    pub fn set_min_write_buffer_number_to_merge(&mut self, to_merge: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_min_write_buffer_number_to_merge(
                self.inner, to_merge);
        }
    }

    pub fn set_level_zero_file_num_compaction_trigger(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_level0_file_num_compaction_trigger(
                self.inner, n);
        }
    }

    pub fn set_level_zero_slowdown_writes_trigger(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_level0_slowdown_writes_trigger(
                self.inner, n);
        }
    }

    pub fn set_level_zero_stop_writes_trigger(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_level0_stop_writes_trigger(
                self.inner, n);
        }
    }

    pub fn set_compaction_style(&mut self,
                                style: rocksdb_ffi::DBCompactionStyle) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_compaction_style(self.inner,
                                                              style);
        }
    }

    pub fn set_max_background_compactions(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_background_compactions(
                self.inner, n);
        }
    }

    pub fn set_max_background_flushes(&mut self, n: c_int) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_max_background_flushes(self.inner,
                                                                    n);
        }
    }

    pub fn set_filter_deletes(&mut self, filter: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_filter_deletes(self.inner, filter);
        }
    }

    pub fn set_disable_auto_compactions(&mut self, disable: bool) {
        let c_bool = if disable {
            1
        } else {
            0
        };
        unsafe {
            rocksdb_ffi::rocksdb_options_set_disable_auto_compactions(self.inner, c_bool)
        }
    }

    pub fn set_block_based_table_factory(&mut self,
                                         factory: &BlockBasedOptions) {
        unsafe {
            rocksdb_ffi::rocksdb_options_set_block_based_table_factory(self.inner, factory.inner);
        }
    }

    pub fn set_report_bg_io_stats(&mut self, enable: bool) {
        unsafe {
            if enable {
                rocksdb_ffi::rocksdb_options_set_report_bg_io_stats(self.inner, 1);
            } else {
                rocksdb_ffi::rocksdb_options_set_report_bg_io_stats(self.inner, 0);
            }
        }
    }
}

impl Default for Options {
    fn default() -> Options {
        unsafe {
            let opts = rocksdb_ffi::rocksdb_options_create();
            if opts.is_null() {
                panic!("Could not create rocksdb options".to_string());
            }
            Options { inner: opts }
        }
    }
}

impl WriteOptions {
    pub fn new() -> WriteOptions {
        WriteOptions::default()
    }

    pub fn set_sync(&mut self, sync: bool) {
        unsafe {
            rocksdb_ffi::rocksdb_writeoptions_set_sync(self.inner, sync);
        }
    }

    pub fn disable_wal(&mut self, disable: bool) {
        unsafe {
            if disable {
                rocksdb_ffi::rocksdb_writeoptions_disable_WAL(self.inner, 1);
            } else {
                rocksdb_ffi::rocksdb_writeoptions_disable_WAL(self.inner, 0);
            }
        }
    }
}

impl Default for WriteOptions {
    fn default() -> WriteOptions {
        let write_opts = unsafe { rocksdb_ffi::rocksdb_writeoptions_create() };
        if write_opts.is_null() {
            panic!("Could not create rocksdb write options".to_string());
        }
        WriteOptions { inner: write_opts }
    }
}
<|endoftext|>


<commit_before>import rustc::syntax::ast;
import rustc::syntax::visit;

export extract;

#[doc = "Converts the Rust AST to the rustdoc document model"]
fn extract(
    crate: @ast::crate
) -> doc::cratedoc {
    ~{
        topmod: top_moddoc_from_crate(crate),
    }
}

fn top_moddoc_from_crate(
    crate: @ast::crate
) -> doc::moddoc {
    moddoc_from_mod(crate.node.module, "crate", crate.node.attrs)
}

fn moddoc_from_mod(
    module: ast::_mod,
    name: ast::ident,
    _attrs: [ast::attribute]

) -> doc::moddoc {
    ~{
        name: name,
        mods: doc::modlist(
            vec::filter_map(module.items) {|item|
                alt item.node {
                  ast::item_mod(m) {
                    some(moddoc_from_mod(m, item.ident, item.attrs))
                  }
                  _ {
                    none
                  }
                }
            }),
        fns: doc::fnlist(
            vec::filter_map(module.items) {|item|
                alt item.node {
                  ast::item_fn(decl, typarams, _) {
                    some(fndoc_from_fn(
                        decl, typarams, item.ident, item.attrs))
                  }
                  _ {
                    none
                  }
                }
            })
    }
}

fn fndoc_from_fn(
    _decl: ast::fn_decl,
    _typarams: [ast::ty_param],
    name: ast::ident,
    _attrs: [ast::attribute]
) -> doc::fndoc {
    ~{
        name: name,
        brief: "todo",
        desc: none,
        return: none,
        args: map::new_str_hash::<str>()
    }
}

#[cfg(test)]
mod tests {

    #[test]
    fn extract_empty_crate() {
        let source = ""; // empty crate
        let ast = parse::from_str(source);
        let doc = extract(ast);
        // FIXME #1535: These are boxed to prevent a crash
        assert ~doc.topmod.mods == ~doc::modlist([]);
        assert ~doc.topmod.fns == ~doc::fnlist([]);
    }

    #[test]
    fn extract_mods() {
        let source = "mod a { mod b { } mod c { } }";
        let ast = parse::from_str(source);
        let doc = extract(ast);
        assert doc.topmod.mods[0].name == "a";
        assert doc.topmod.mods[0].mods[0].name == "b";
        assert doc.topmod.mods[0].mods[1].name == "c";
    }

    #[test]
    fn extract_mods_deep() {
        let source = "mod a { mod b { mod c { } } }";
        let ast = parse::from_str(source);
        let doc = extract(ast);
        assert doc.topmod.mods[0].mods[0].mods[0].name == "c";
    }

    #[test]
    fn extract_fns() {
        let source =
            "fn a() { } \
             mod b { fn c() { } }";
        let ast = parse::from_str(source);
        let doc = extract(ast);
        assert doc.topmod.fns[0].name == "a";
        assert doc.topmod.mods[0].fns[0].name == "c";
    }
}<commit_msg>rustdoc: Wire up the doc extractor to attribute parser<commit_after>import rustc::syntax::ast;
import rustc::syntax::visit;

export extract;

#[doc = "Converts the Rust AST to the rustdoc document model"]
fn extract(
    crate: @ast::crate
) -> doc::cratedoc {
    ~{
        topmod: top_moddoc_from_crate(crate),
    }
}

fn top_moddoc_from_crate(
    crate: @ast::crate
) -> doc::moddoc {
    moddoc_from_mod(crate.node.module, "crate", crate.node.attrs)
}

fn moddoc_from_mod(
    module: ast::_mod,
    name: ast::ident,
    _attrs: [ast::attribute]

) -> doc::moddoc {
    ~{
        name: name,
        mods: doc::modlist(
            vec::filter_map(module.items) {|item|
                alt item.node {
                  ast::item_mod(m) {
                    some(moddoc_from_mod(m, item.ident, item.attrs))
                  }
                  _ {
                    none
                  }
                }
            }),
        fns: doc::fnlist(
            vec::filter_map(module.items) {|item|
                alt item.node {
                  ast::item_fn(decl, typarams, _) {
                    some(fndoc_from_fn(
                        decl, typarams, item.ident, item.attrs))
                  }
                  _ {
                    none
                  }
                }
            })
    }
}

fn fndoc_from_fn(
    _decl: ast::fn_decl,
    _typarams: [ast::ty_param],
    name: ast::ident,
    attrs: [ast::attribute]
) -> doc::fndoc {
    attr_parser::parse_fn(name, attrs)
}

#[cfg(test)]
mod tests {

    #[test]
    fn extract_empty_crate() {
        let source = ""; // empty crate
        let ast = parse::from_str(source);
        let doc = extract(ast);
        // FIXME #1535: These are boxed to prevent a crash
        assert ~doc.topmod.mods == ~doc::modlist([]);
        assert ~doc.topmod.fns == ~doc::fnlist([]);
    }

    #[test]
    fn extract_mods() {
        let source = "mod a { mod b { } mod c { } }";
        let ast = parse::from_str(source);
        let doc = extract(ast);
        assert doc.topmod.mods[0].name == "a";
        assert doc.topmod.mods[0].mods[0].name == "b";
        assert doc.topmod.mods[0].mods[1].name == "c";
    }

    #[test]
    fn extract_mods_deep() {
        let source = "mod a { mod b { mod c { } } }";
        let ast = parse::from_str(source);
        let doc = extract(ast);
        assert doc.topmod.mods[0].mods[0].mods[0].name == "c";
    }

    #[test]
    fn extract_fns() {
        let source =
            "fn a() { } \
             mod b { fn c() { } }";
        let ast = parse::from_str(source);
        let doc = extract(ast);
        assert doc.topmod.fns[0].name == "a";
        assert doc.topmod.mods[0].fns[0].name == "c";
    }
}<|endoftext|>


<commit_before>#[doc = "Converts the Rust AST to the rustdoc document model"];

import rustc::syntax::ast;

export from_srv, extract;

fn from_srv(
    srv: astsrv::srv,
    default_name: str
) -> doc::cratedoc {

    #[doc = "Use the AST service to create a document tree"];

    astsrv::exec(srv) {|ctxt|
        extract(ctxt.ast, default_name)
    }
}

fn extract(
    crate: @ast::crate,
    default_name: str
) -> doc::cratedoc {
    ~{
        topmod: top_moddoc_from_crate(crate, default_name),
    }
}

fn top_moddoc_from_crate(
    crate: @ast::crate,
    default_name: str
) -> doc::moddoc {
    moddoc_from_mod(crate.node.module, default_name)
}

fn moddoc_from_mod(
    module: ast::_mod,
    name: ast::ident
) -> doc::moddoc {
    ~{
        name: name,
        mods: doc::modlist(
            vec::filter_map(module.items) {|item|
                alt item.node {
                  ast::item_mod(m) {
                    some(moddoc_from_mod(m, item.ident))
                  }
                  _ {
                    none
                  }
                }
            }),
        fns: doc::fnlist(
            vec::filter_map(module.items) {|item|
                alt item.node {
                  ast::item_fn(decl, _, _) {
                    some(fndoc_from_fn(
                        decl, item.ident, item.id))
                  }
                  _ {
                    none
                  }
                }
            })
    }
}

fn fndoc_from_fn(
    _decl: ast::fn_decl,
    name: ast::ident,
    id: ast::node_id
) -> doc::fndoc {
    ~{
        id: id,
        name: name,
        brief: none,
        desc: none,
        return: none,
        args: []
    }
}

#[cfg(test)]
mod tests {

    #[test]
    fn extract_empty_crate() {
        let source = ""; // empty crate
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        // FIXME #1535: These are boxed to prevent a crash
        assert ~doc.topmod.mods == ~doc::modlist([]);
        assert ~doc.topmod.fns == ~doc::fnlist([]);
    }

    #[test]
    fn extract_mods() {
        let source = "mod a { mod b { } mod c { } }";
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        assert doc.topmod.mods[0].name == "a";
        assert doc.topmod.mods[0].mods[0].name == "b";
        assert doc.topmod.mods[0].mods[1].name == "c";
    }

    #[test]
    fn extract_mods_deep() {
        let source = "mod a { mod b { mod c { } } }";
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        assert doc.topmod.mods[0].mods[0].mods[0].name == "c";
    }

    #[test]
    fn extract_fns() {
        let source =
            "fn a() { } \
             mod b { fn c() { } }";
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        assert doc.topmod.fns[0].name == "a";
        assert doc.topmod.mods[0].fns[0].name == "c";
    }

    #[test]
    fn extract_should_set_fn_ast_id() {
        let source = "fn a() { }";
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        assert doc.topmod.fns[0].id != 0;
    }

    #[test]
    fn extract_should_use_default_crate_name() {
        let source = "";
        let ast = parse::from_str(source);
        let doc = extract(ast, "burp");
        assert doc.topmod.name == "burp";
    }

    #[test]
    fn extract_from_seq_srv() {
        let source = "";
        let srv = astsrv::mk_srv_from_str(source);
        let doc = from_srv(srv, "name");
        assert doc.topmod.name == "name";
    }
}<commit_msg>rustdoc: Extract function argument names from the AST<commit_after>#[doc = "Converts the Rust AST to the rustdoc document model"];

import rustc::syntax::ast;

export from_srv, extract;

fn from_srv(
    srv: astsrv::srv,
    default_name: str
) -> doc::cratedoc {

    #[doc = "Use the AST service to create a document tree"];

    astsrv::exec(srv) {|ctxt|
        extract(ctxt.ast, default_name)
    }
}

fn extract(
    crate: @ast::crate,
    default_name: str
) -> doc::cratedoc {
    ~{
        topmod: top_moddoc_from_crate(crate, default_name),
    }
}

fn top_moddoc_from_crate(
    crate: @ast::crate,
    default_name: str
) -> doc::moddoc {
    moddoc_from_mod(crate.node.module, default_name)
}

fn moddoc_from_mod(
    module: ast::_mod,
    name: ast::ident
) -> doc::moddoc {
    ~{
        name: name,
        mods: doc::modlist(
            vec::filter_map(module.items) {|item|
                alt item.node {
                  ast::item_mod(m) {
                    some(moddoc_from_mod(m, item.ident))
                  }
                  _ {
                    none
                  }
                }
            }),
        fns: doc::fnlist(
            vec::filter_map(module.items) {|item|
                alt item.node {
                  ast::item_fn(decl, _, _) {
                    some(fndoc_from_fn(
                        decl, item.ident, item.id))
                  }
                  _ {
                    none
                  }
                }
            })
    }
}

fn fndoc_from_fn(
    decl: ast::fn_decl,
    name: ast::ident,
    id: ast::node_id
) -> doc::fndoc {
    ~{
        id: id,
        name: name,
        brief: none,
        desc: none,
        return: none,
        args: argdocs_from_args(decl.inputs)
    }
}

#[test]
fn should_extract_fn_args() {
    let source = "fn a(b: int, c: int) { }";
    let ast = parse::from_str(source);
    let doc = extract(ast, "");
    let fn_ = doc.topmod.fns[0];
    assert tuple::first(fn_.args[0]) == "b";
    assert tuple::first(fn_.args[1]) == "c";
}

fn argdocs_from_args(args: [ast::arg]) -> [(str, str)] {
    vec::map(args, argdoc_from_arg)
}

fn argdoc_from_arg(arg: ast::arg) -> (str, str) {
    (arg.ident, "")
}

#[cfg(test)]
mod tests {

    #[test]
    fn extract_empty_crate() {
        let source = ""; // empty crate
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        // FIXME #1535: These are boxed to prevent a crash
        assert ~doc.topmod.mods == ~doc::modlist([]);
        assert ~doc.topmod.fns == ~doc::fnlist([]);
    }

    #[test]
    fn extract_mods() {
        let source = "mod a { mod b { } mod c { } }";
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        assert doc.topmod.mods[0].name == "a";
        assert doc.topmod.mods[0].mods[0].name == "b";
        assert doc.topmod.mods[0].mods[1].name == "c";
    }

    #[test]
    fn extract_mods_deep() {
        let source = "mod a { mod b { mod c { } } }";
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        assert doc.topmod.mods[0].mods[0].mods[0].name == "c";
    }

    #[test]
    fn extract_fns() {
        let source =
            "fn a() { } \
             mod b { fn c() { } }";
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        assert doc.topmod.fns[0].name == "a";
        assert doc.topmod.mods[0].fns[0].name == "c";
    }

    #[test]
    fn extract_should_set_fn_ast_id() {
        let source = "fn a() { }";
        let ast = parse::from_str(source);
        let doc = extract(ast, "");
        assert doc.topmod.fns[0].id != 0;
    }

    #[test]
    fn extract_should_use_default_crate_name() {
        let source = "";
        let ast = parse::from_str(source);
        let doc = extract(ast, "burp");
        assert doc.topmod.name == "burp";
    }

    #[test]
    fn extract_from_seq_srv() {
        let source = "";
        let srv = astsrv::mk_srv_from_str(source);
        let doc = from_srv(srv, "name");
        assert doc.topmod.name == "name";
    }
}<|endoftext|>


<commit_before>use std::rand::RngUtil;
use std::rand;

use sdl;

pub fn main() {
	do sdl::start {
		sdl::init([sdl::InitVideo]);
		sdl::wm::set_caption("rust-sdl demo - video", "rust-sdl");

		let mut rng = rand::rng();
		let screen = match sdl::video::set_video_mode(800, 600, 32, [sdl::video::HWSurface],
			                                                        [sdl::video::DoubleBuf]) {
			Ok(screen) => screen,
			Err(err) => fail!(fmt!("failed to set video mode: %s", err))
		};

		// Note: You'll want to put this and the flip call inside the main loop
		// but we don't as to not startle epileptics
		for i in range(0, 10) {
			for j in range(0, 10) {
				screen.fill_rect(Some(sdl::Rect {
					x: (i as i16) * 800 / 10,
					y: (j as i16) * 600 / 10,
					w: 800 / 10,
					h: 600 / 10
				}), rng.gen::<sdl::video::Color>());
			}
		}

		screen.flip();

		'main : loop {
			'event : loop {
				match sdl::event::poll_event() {
					sdl::event::QuitEvent => break 'main,
					sdl::event::NoEvent => break 'event,
					sdl::event::KeyEvent(k, _, _, _)
					    if k == sdl::event::EscapeKey
					        => break 'main,
					_ => {}
				}
			}
		}

		sdl::quit();
	}
}
<commit_msg>Removed unneccesary call to SDL::Start<commit_after>use std::rand::RngUtil;
use std::rand;
use std::rt;

use sdl;

pub fn main() {
    sdl::init([sdl::InitVideo]);
    sdl::wm::set_caption("rust-sdl demo - video", "rust-sdl");

    let mut rng = rand::rng();
    let screen = match sdl::video::set_video_mode(800, 600, 32, [sdl::video::HWSurface],
                                                                [sdl::video::DoubleBuf]) {
        Ok(screen) => screen,
        Err(err) => fail!(fmt!("failed to set video mode: %s", err))
    };

    // Note: You'll want to put this and the flip call inside the main loop
    // but we don't as to not startle epileptics
    for i in range(0, 10) {
        for j in range(0, 10) {
            screen.fill_rect(Some(sdl::Rect {
                x: (i as i16) * 800 / 10,
                y: (j as i16) * 600 / 10,
                w: 800 / 10,
                h: 600 / 10
            }), rng.gen::<sdl::video::Color>());
        }
    }

    screen.flip();

    'main : loop {
        'event : loop {
            match sdl::event::poll_event() {
                sdl::event::QuitEvent => break 'main,
                sdl::event::NoEvent => break 'event,
                sdl::event::KeyEvent(k, _, _, _)
                    if k == sdl::event::EscapeKey
                        => break 'main,
                _ => {}
            }
        }
    }

    sdl::quit();
}
<|endoftext|>


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

use rustc_serialize::json::Json;

use search::term::Term;
use search::token::Token;
use search::mapping::Mapping;

#[derive(Debug)]
pub struct Document {
    pub key: String,
    pub fields: BTreeMap<String, Vec<Token>>,
}

#[derive(Debug)]
pub struct DocumentSource {
    pub key: String,
    pub data: Json,
}

impl DocumentSource {
    pub fn prepare(&self, mapping: &Mapping) -> Document {
        let mut fields = BTreeMap::new();
        let mut all_field_tokens: Vec<Token> = Vec::new();

        for (field_name, field_value) in self.data.as_object().unwrap() {
            if *field_value == Json::Null {
                // Treat null like a missing field
                continue;
            }

            match mapping.fields.get(field_name) {
                Some(field_mapping) => {
                    let value = field_mapping.process_value_for_index(field_value.clone());

                    match value {
                        Some(value) => {
                            // Copy the field's value into the _all field
                            if field_mapping.is_in_all {
                                // TODO: Should the positions be updated?
                                all_field_tokens.extend(value.iter().cloned());
                            }

                            // Insert the field
                            fields.insert(field_name.clone(), value);
                        }
                        None => {
                            // TODO: Should probably be an error
                            warn!("Unprocessable value: {}", field_value);
                        }
                    }
                }
                None => {
                    // No mapping found, just insert the value as-is
                    // TODO: This should probably be an error
                    if let Some(term) = Term::from_json(field_value) {
                        fields.insert(field_name.clone(), vec![Token{term: term, position: 1}]);
                    }
                }
            }
        }

        // Insert _all field
        fields.insert("_all".to_owned(), all_field_tokens);

        Document {
            key: self.key.clone(),
            fields: fields,
        }
    }
}
<commit_msg>Index _all field properly<commit_after>use std::collections::BTreeMap;

use rustc_serialize::json::Json;

use search::term::Term;
use search::token::Token;
use search::mapping::{Mapping, FieldMapping, FieldType};

#[derive(Debug)]
pub struct Document {
    pub key: String,
    pub fields: BTreeMap<String, Vec<Token>>,
}

#[derive(Debug)]
pub struct DocumentSource {
    pub key: String,
    pub data: Json,
}

impl DocumentSource {
    pub fn prepare(&self, mapping: &Mapping) -> Document {
        let mut fields = BTreeMap::new();
        let mut all_field_strings: Vec<String> = Vec::new();

        for (field_name, field_value) in self.data.as_object().unwrap() {
            if *field_value == Json::Null {
                // Treat null like a missing field
                continue;
            }

            match mapping.fields.get(field_name) {
                Some(field_mapping) => {
                    let value = field_mapping.process_value_for_index(field_value.clone());

                    match value {
                        Some(value) => {
                            // Copy the field's value into the _all field
                            if field_mapping.is_in_all {
                                if let Json::String(ref string) = *field_value {
                                    all_field_strings.push(string.clone());
                                }
                            }

                            // Insert the field
                            fields.insert(field_name.clone(), value);
                        }
                        None => {
                            // TODO: Should probably be an error
                            warn!("Unprocessable value: {}", field_value);
                        }
                    }
                }
                None => {
                    // No mapping found, just insert the value as-is
                    // TODO: This should probably be an error
                    if let Some(term) = Term::from_json(field_value) {
                        fields.insert(field_name.clone(), vec![Token{term: term, position: 1}]);
                    }
                }
            }
        }

        // Insert _all field
        let all_field_mapping = FieldMapping::default();
        let all_field_strings_json = Json::String(all_field_strings.join(" "));
        let all_field_value = all_field_mapping.process_value_for_index(all_field_strings_json.clone());

        match all_field_value {
            Some(all_field_value) => {
                fields.insert("_all".to_owned(), all_field_value);
            }
            None => {
                // TODO: Should probably be an error
                warn!("Unprocessable value: {}", all_field_strings_json);
            }
        }

        Document {
            key: self.key.clone(),
            fields: fields,
        }
    }
}
<|endoftext|>


<commit_before>//! This module provides a token source (`GetToken`) that obtains tokens for service accounts.
//! Service accounts are usually used by software (i.e., non-human actors) to get access to
//! resources. Currently, this module only works with RS256 JWTs, which makes it at least suitable for
//! authentication with Google services.
//!
//! Resources:
//! - [Using OAuth 2.0 for Server to Server
//! Applications](https://developers.google.com/identity/protocols/OAuth2ServiceAccount)
//! - [JSON Web Tokens](https://jwt.io/)
//!
//! Copyright (c) 2016 Google Inc (lewinb@google.com).
//!

use std::borrow::BorrowMut;
use std::default::Default;
use std::error;
use std::io::Read;
use std::result;
use std::str;

use crate::authenticator::GetToken;
use crate::storage::{hash_scopes, MemoryStorage, TokenStorage};
use crate::types::{StringError, Token};

use hyper::header;
use url::form_urlencoded;

#[cfg(not(feature = "no-openssl"))]
use openssl::{
    hash::MessageDigest,
    pkey::{PKey, Private},
    rsa::Padding,
    sign::Signer,
};

#[cfg(feature = "no-openssl")]
use rustls::{
    self,
    internal::pemfile,
    sign::{self, SigningKey},
    PrivateKey,
};
#[cfg(feature = "no-openssl")]
use std::io;

use base64;
use chrono;
use hyper;
use serde_json;

const GRANT_TYPE: &'static str = "urn:ietf:params:oauth:grant-type:jwt-bearer";
const GOOGLE_RS256_HEAD: &'static str = "{\"alg\":\"RS256\",\"typ\":\"JWT\"}";

// Encodes s as Base64
fn encode_base64<T: AsRef<[u8]>>(s: T) -> String {
    base64::encode_config(s.as_ref(), base64::URL_SAFE)
}

#[cfg(not(feature = "no-openssl"))]
fn decode_rsa_key(pem_pkcs8: &str) -> Result<PKey<Private>, Box<error::Error>> {
    let private = pem_pkcs8.to_string().replace("\\n", "\n").into_bytes();
    Ok(PKey::private_key_from_pem(&private)?)
}

#[cfg(feature = "no-openssl")]
fn decode_rsa_key(pem_pkcs8: &str) -> Result<PrivateKey, Box<error::Error>> {
    let private = pem_pkcs8.to_string().replace("\\n", "\n").into_bytes();
    let mut private_reader: &[u8] = private.as_ref();
    let private_keys = pemfile::pkcs8_private_keys(&mut private_reader);

    if let Ok(pk) = private_keys {
        if pk.len() > 0 {
            Ok(pk[0].clone())
        } else {
            Err(Box::new(io::Error::new(
                io::ErrorKind::InvalidInput,
                "Not enough private keys in PEM",
            )))
        }
    } else {
        Err(Box::new(io::Error::new(
            io::ErrorKind::InvalidInput,
            "Error reading key from PEM",
        )))
    }
}

/// JSON schema of secret service account key. You can obtain the key from
/// the Cloud Console at https://console.cloud.google.com/.
///
/// You can use `helpers::service_account_key_from_file()` as a quick way to read a JSON client
/// secret into a ServiceAccountKey.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ServiceAccountKey {
    #[serde(rename = "type")]
    pub key_type: Option<String>,
    pub project_id: Option<String>,
    pub private_key_id: Option<String>,
    pub private_key: Option<String>,
    pub client_email: Option<String>,
    pub client_id: Option<String>,
    pub auth_uri: Option<String>,
    pub token_uri: Option<String>,
    pub auth_provier_x509_cert_url: Option<String>,
    pub client_x509_cert_url: Option<String>,
}

#[derive(Serialize, Debug)]
struct Claims {
    iss: String,
    aud: String,
    exp: i64,
    iat: i64,
    sub: Option<String>,
    scope: String,
}

struct JWT {
    header: String,
    claims: Claims,
}

impl JWT {
    fn new(claims: Claims) -> JWT {
        JWT {
            header: GOOGLE_RS256_HEAD.to_string(),
            claims: claims,
        }
    }
    // Encodes the first two parts (header and claims) to base64 and assembles them into a form
    // ready to be signed.
    fn encode_claims(&self) -> String {
        let mut head = encode_base64(&self.header);
        let claims = encode_base64(serde_json::to_string(&self.claims).unwrap());

        head.push_str(".");
        head.push_str(&claims);
        head
    }

    #[cfg(not(feature = "no-openssl"))]
    fn sign(&self, private_key: &str) -> Result<String, Box<error::Error>> {
        let mut jwt_head = self.encode_claims();
        let key = decode_rsa_key(private_key)?;
        let mut signer = Signer::new(MessageDigest::sha256(), &key)?;
        signer.set_rsa_padding(Padding::PKCS1)?;
        signer.update(jwt_head.as_bytes())?;
        let signature = signer.sign_to_vec()?;
        let signature_b64 = encode_base64(signature);

        jwt_head.push_str(".");
        jwt_head.push_str(&signature_b64);

        Ok(jwt_head)
    }

    #[cfg(feature = "no-openssl")]
    fn sign(&self, private_key: &str) -> Result<String, Box<error::Error>> {
        let mut jwt_head = self.encode_claims();
        let key = decode_rsa_key(private_key)?;
        let signing_key = sign::RSASigningKey::new(&key)
            .map_err(|_| io::Error::new(io::ErrorKind::Other, "Couldn't initialize signer"))?;
        let signer = signing_key
            .choose_scheme(&[rustls::SignatureScheme::RSA_PKCS1_SHA256])
            .ok_or(io::Error::new(
                io::ErrorKind::Other,
                "Couldn't choose signing scheme",
            ))?;
        let signature = signer.sign(jwt_head.as_bytes())?;
        let signature_b64 = encode_base64(signature);

        jwt_head.push_str(".");
        jwt_head.push_str(&signature_b64);

        Ok(jwt_head)
    }
}

fn init_claims_from_key<'a, I, T>(key: &ServiceAccountKey, scopes: I) -> Claims
where
    T: AsRef<str> + 'a,
    I: IntoIterator<Item = &'a T>,
{
    let iat = chrono::Utc::now().timestamp();
    let expiry = iat + 3600 - 5; // Max validity is 1h.

    let mut scopes_string = scopes.into_iter().fold(String::new(), |mut acc, sc| {
        acc.push_str(sc.as_ref());
        acc.push_str(" ");
        acc
    });
    scopes_string.pop();

    Claims {
        iss: key.client_email.clone().unwrap(),
        aud: key.token_uri.clone().unwrap(),
        exp: expiry,
        iat: iat,
        sub: None,
        scope: scopes_string,
    }
}

/// See "Additional claims" at https://developers.google.com/identity/protocols/OAuth2ServiceAccount
#[allow(dead_code)]
fn set_sub_claim(mut claims: Claims, sub: String) -> Claims {
    claims.sub = Some(sub);
    claims
}

/// A token source (`GetToken`) yielding OAuth tokens for services that use ServiceAccount authorization.
/// This token source caches token and automatically renews expired ones.
pub struct ServiceAccountAccess<C> {
    client: C,
    key: ServiceAccountKey,
    cache: MemoryStorage,
    sub: Option<String>,
}

/// This is the schema of the server's response.
#[derive(Deserialize, Debug)]
struct TokenResponse {
    access_token: Option<String>,
    token_type: Option<String>,
    expires_in: Option<i64>,
}

impl TokenResponse {
    fn to_oauth_token(self) -> Token {
        let expires_ts = chrono::Utc::now().timestamp() + self.expires_in.unwrap_or(0);

        Token {
            access_token: self.access_token.unwrap(),
            token_type: self.token_type.unwrap(),
            refresh_token: String::new(),
            expires_in: self.expires_in,
            expires_in_timestamp: Some(expires_ts),
        }
    }
}

impl<'a, C> ServiceAccountAccess<C>
where
    C: BorrowMut<hyper::Client>,
{
    /// Returns a new `ServiceAccountAccess` token source.
    #[allow(dead_code)]
    pub fn new(key: ServiceAccountKey, client: C) -> ServiceAccountAccess<C> {
        ServiceAccountAccess {
            client: client,
            key: key,
            cache: MemoryStorage::default(),
            sub: None,
        }
    }

    pub fn with_sub(key: ServiceAccountKey, client: C, sub: String) -> ServiceAccountAccess<C> {
        ServiceAccountAccess {
            client: client,
            key: key,
            cache: MemoryStorage::default(),
            sub: Some(sub),
        }
    }

    fn request_token(&mut self, scopes: &Vec<&str>) -> result::Result<Token, Box<error::Error>> {
        let mut claims = init_claims_from_key(&self.key, scopes);
        claims.sub = self.sub.clone();
        let signed = JWT::new(claims).sign(self.key.private_key.as_ref().unwrap())?;

        let body = form_urlencoded::Serializer::new(String::new())
            .extend_pairs(vec![
                ("grant_type".to_string(), GRANT_TYPE.to_string()),
                ("assertion".to_string(), signed),
            ])
            .finish();

        let mut response = String::new();
        let mut result = self
            .client
            .borrow_mut()
            .post(self.key.token_uri.as_ref().unwrap())
            .body(&body)
            .header(header::ContentType(
                "application/x-www-form-urlencoded".parse().unwrap(),
            ))
            .send()?;

        result.read_to_string(&mut response)?;

        let token: Result<TokenResponse, serde_json::error::Error> =
            serde_json::from_str(&response);

        match token {
            Err(e) => return Err(Box::new(e)),
            Ok(token) => {
                if token.access_token.is_none()
                    || token.token_type.is_none()
                    || token.expires_in.is_none()
                {
                    Err(Box::new(StringError::new(
                        "Token response lacks fields".to_string(),
                        Some(&format!("{:?}", token)),
                    )))
                } else {
                    Ok(token.to_oauth_token())
                }
            }
        }
    }
}

impl<C: BorrowMut<hyper::Client>> GetToken for ServiceAccountAccess<C> {
    fn token<'b, I, T>(&mut self, scopes: I) -> result::Result<Token, Box<error::Error>>
    where
        T: AsRef<str> + Ord + 'b,
        I: IntoIterator<Item = &'b T>,
    {
        let (hash, scps) = hash_scopes(scopes);

        if let Some(token) = self.cache.get(hash, &scps)? {
            if !token.expired() {
                return Ok(token);
            }
        }

        let token = self.request_token(&scps)?;
        let _ = self.cache.set(hash, &scps, Some(token.clone()));

        Ok(token)
    }

    fn api_key(&mut self) -> Option<String> {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::authenticator::GetToken;
    use crate::helper::service_account_key_from_file;
    use hyper;
    use hyper::net::HttpsConnector;
    use hyper_native_tls::NativeTlsClient;

    // This is a valid but deactivated key.
    const TEST_PRIVATE_KEY_PATH: &'static str = "examples/Sanguine-69411a0c0eea.json";

    // Uncomment this test to verify that we can successfully obtain tokens.
    //#[test]
    #[allow(dead_code)]
    fn test_service_account_e2e() {
        let key = service_account_key_from_file(TEST_PRIVATE_KEY_PATH).unwrap();
        let client =
            hyper::Client::with_connector(HttpsConnector::new(NativeTlsClient::new().unwrap()));
        let mut acc = ServiceAccountAccess::new(key, client);
        println!(
            "{:?}",
            acc.token(vec![&"https://www.googleapis.com/auth/pubsub"])
                .unwrap()
        );
    }

    #[test]
    fn test_jwt_initialize_claims() {
        let key = service_account_key_from_file(TEST_PRIVATE_KEY_PATH).unwrap();
        let scopes = vec!["scope1", "scope2", "scope3"];
        let claims = super::init_claims_from_key(&key, &scopes);

        assert_eq!(
            claims.iss,
            "oauth2-public-test@sanguine-rhythm-105020.iam.gserviceaccount.com".to_string()
        );
        assert_eq!(claims.scope, "scope1 scope2 scope3".to_string());
        assert_eq!(
            claims.aud,
            "https://accounts.google.com/o/oauth2/token".to_string()
        );
        assert!(claims.exp > 1000000000);
        assert!(claims.iat < claims.exp);
        assert_eq!(claims.exp - claims.iat, 3595);
    }

    #[test]
    fn test_jwt_sign() {
        let key = service_account_key_from_file(TEST_PRIVATE_KEY_PATH).unwrap();
        let scopes = vec!["scope1", "scope2", "scope3"];
        let claims = super::init_claims_from_key(&key, &scopes);
        let jwt = super::JWT::new(claims);
        let signature = jwt.sign(key.private_key.as_ref().unwrap());

        assert!(signature.is_ok());

        let signature = signature.unwrap();
        assert_eq!(
            signature.split(".").nth(0).unwrap(),
            "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9"
        );
    }
}
<commit_msg>fix typo<commit_after>//! This module provides a token source (`GetToken`) that obtains tokens for service accounts.
//! Service accounts are usually used by software (i.e., non-human actors) to get access to
//! resources. Currently, this module only works with RS256 JWTs, which makes it at least suitable for
//! authentication with Google services.
//!
//! Resources:
//! - [Using OAuth 2.0 for Server to Server
//! Applications](https://developers.google.com/identity/protocols/OAuth2ServiceAccount)
//! - [JSON Web Tokens](https://jwt.io/)
//!
//! Copyright (c) 2016 Google Inc (lewinb@google.com).
//!

use std::borrow::BorrowMut;
use std::default::Default;
use std::error;
use std::io::Read;
use std::result;
use std::str;

use crate::authenticator::GetToken;
use crate::storage::{hash_scopes, MemoryStorage, TokenStorage};
use crate::types::{StringError, Token};

use hyper::header;
use url::form_urlencoded;

#[cfg(not(feature = "no-openssl"))]
use openssl::{
    hash::MessageDigest,
    pkey::{PKey, Private},
    rsa::Padding,
    sign::Signer,
};

#[cfg(feature = "no-openssl")]
use rustls::{
    self,
    internal::pemfile,
    sign::{self, SigningKey},
    PrivateKey,
};
#[cfg(feature = "no-openssl")]
use std::io;

use base64;
use chrono;
use hyper;
use serde_json;

const GRANT_TYPE: &'static str = "urn:ietf:params:oauth:grant-type:jwt-bearer";
const GOOGLE_RS256_HEAD: &'static str = "{\"alg\":\"RS256\",\"typ\":\"JWT\"}";

// Encodes s as Base64
fn encode_base64<T: AsRef<[u8]>>(s: T) -> String {
    base64::encode_config(s.as_ref(), base64::URL_SAFE)
}

#[cfg(not(feature = "no-openssl"))]
fn decode_rsa_key(pem_pkcs8: &str) -> Result<PKey<Private>, Box<error::Error>> {
    let private = pem_pkcs8.to_string().replace("\\n", "\n").into_bytes();
    Ok(PKey::private_key_from_pem(&private)?)
}

#[cfg(feature = "no-openssl")]
fn decode_rsa_key(pem_pkcs8: &str) -> Result<PrivateKey, Box<error::Error>> {
    let private = pem_pkcs8.to_string().replace("\\n", "\n").into_bytes();
    let mut private_reader: &[u8] = private.as_ref();
    let private_keys = pemfile::pkcs8_private_keys(&mut private_reader);

    if let Ok(pk) = private_keys {
        if pk.len() > 0 {
            Ok(pk[0].clone())
        } else {
            Err(Box::new(io::Error::new(
                io::ErrorKind::InvalidInput,
                "Not enough private keys in PEM",
            )))
        }
    } else {
        Err(Box::new(io::Error::new(
            io::ErrorKind::InvalidInput,
            "Error reading key from PEM",
        )))
    }
}

/// JSON schema of secret service account key. You can obtain the key from
/// the Cloud Console at https://console.cloud.google.com/.
///
/// You can use `helpers::service_account_key_from_file()` as a quick way to read a JSON client
/// secret into a ServiceAccountKey.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ServiceAccountKey {
    #[serde(rename = "type")]
    pub key_type: Option<String>,
    pub project_id: Option<String>,
    pub private_key_id: Option<String>,
    pub private_key: Option<String>,
    pub client_email: Option<String>,
    pub client_id: Option<String>,
    pub auth_uri: Option<String>,
    pub token_uri: Option<String>,
    pub auth_provider_x509_cert_url: Option<String>,
    pub client_x509_cert_url: Option<String>,
}

#[derive(Serialize, Debug)]
struct Claims {
    iss: String,
    aud: String,
    exp: i64,
    iat: i64,
    sub: Option<String>,
    scope: String,
}

struct JWT {
    header: String,
    claims: Claims,
}

impl JWT {
    fn new(claims: Claims) -> JWT {
        JWT {
            header: GOOGLE_RS256_HEAD.to_string(),
            claims: claims,
        }
    }
    // Encodes the first two parts (header and claims) to base64 and assembles them into a form
    // ready to be signed.
    fn encode_claims(&self) -> String {
        let mut head = encode_base64(&self.header);
        let claims = encode_base64(serde_json::to_string(&self.claims).unwrap());

        head.push_str(".");
        head.push_str(&claims);
        head
    }

    #[cfg(not(feature = "no-openssl"))]
    fn sign(&self, private_key: &str) -> Result<String, Box<error::Error>> {
        let mut jwt_head = self.encode_claims();
        let key = decode_rsa_key(private_key)?;
        let mut signer = Signer::new(MessageDigest::sha256(), &key)?;
        signer.set_rsa_padding(Padding::PKCS1)?;
        signer.update(jwt_head.as_bytes())?;
        let signature = signer.sign_to_vec()?;
        let signature_b64 = encode_base64(signature);

        jwt_head.push_str(".");
        jwt_head.push_str(&signature_b64);

        Ok(jwt_head)
    }

    #[cfg(feature = "no-openssl")]
    fn sign(&self, private_key: &str) -> Result<String, Box<error::Error>> {
        let mut jwt_head = self.encode_claims();
        let key = decode_rsa_key(private_key)?;
        let signing_key = sign::RSASigningKey::new(&key)
            .map_err(|_| io::Error::new(io::ErrorKind::Other, "Couldn't initialize signer"))?;
        let signer = signing_key
            .choose_scheme(&[rustls::SignatureScheme::RSA_PKCS1_SHA256])
            .ok_or(io::Error::new(
                io::ErrorKind::Other,
                "Couldn't choose signing scheme",
            ))?;
        let signature = signer.sign(jwt_head.as_bytes())?;
        let signature_b64 = encode_base64(signature);

        jwt_head.push_str(".");
        jwt_head.push_str(&signature_b64);

        Ok(jwt_head)
    }
}

fn init_claims_from_key<'a, I, T>(key: &ServiceAccountKey, scopes: I) -> Claims
where
    T: AsRef<str> + 'a,
    I: IntoIterator<Item = &'a T>,
{
    let iat = chrono::Utc::now().timestamp();
    let expiry = iat + 3600 - 5; // Max validity is 1h.

    let mut scopes_string = scopes.into_iter().fold(String::new(), |mut acc, sc| {
        acc.push_str(sc.as_ref());
        acc.push_str(" ");
        acc
    });
    scopes_string.pop();

    Claims {
        iss: key.client_email.clone().unwrap(),
        aud: key.token_uri.clone().unwrap(),
        exp: expiry,
        iat: iat,
        sub: None,
        scope: scopes_string,
    }
}

/// See "Additional claims" at https://developers.google.com/identity/protocols/OAuth2ServiceAccount
#[allow(dead_code)]
fn set_sub_claim(mut claims: Claims, sub: String) -> Claims {
    claims.sub = Some(sub);
    claims
}

/// A token source (`GetToken`) yielding OAuth tokens for services that use ServiceAccount authorization.
/// This token source caches token and automatically renews expired ones.
pub struct ServiceAccountAccess<C> {
    client: C,
    key: ServiceAccountKey,
    cache: MemoryStorage,
    sub: Option<String>,
}

/// This is the schema of the server's response.
#[derive(Deserialize, Debug)]
struct TokenResponse {
    access_token: Option<String>,
    token_type: Option<String>,
    expires_in: Option<i64>,
}

impl TokenResponse {
    fn to_oauth_token(self) -> Token {
        let expires_ts = chrono::Utc::now().timestamp() + self.expires_in.unwrap_or(0);

        Token {
            access_token: self.access_token.unwrap(),
            token_type: self.token_type.unwrap(),
            refresh_token: String::new(),
            expires_in: self.expires_in,
            expires_in_timestamp: Some(expires_ts),
        }
    }
}

impl<'a, C> ServiceAccountAccess<C>
where
    C: BorrowMut<hyper::Client>,
{
    /// Returns a new `ServiceAccountAccess` token source.
    #[allow(dead_code)]
    pub fn new(key: ServiceAccountKey, client: C) -> ServiceAccountAccess<C> {
        ServiceAccountAccess {
            client: client,
            key: key,
            cache: MemoryStorage::default(),
            sub: None,
        }
    }

    pub fn with_sub(key: ServiceAccountKey, client: C, sub: String) -> ServiceAccountAccess<C> {
        ServiceAccountAccess {
            client: client,
            key: key,
            cache: MemoryStorage::default(),
            sub: Some(sub),
        }
    }

    fn request_token(&mut self, scopes: &Vec<&str>) -> result::Result<Token, Box<error::Error>> {
        let mut claims = init_claims_from_key(&self.key, scopes);
        claims.sub = self.sub.clone();
        let signed = JWT::new(claims).sign(self.key.private_key.as_ref().unwrap())?;

        let body = form_urlencoded::Serializer::new(String::new())
            .extend_pairs(vec![
                ("grant_type".to_string(), GRANT_TYPE.to_string()),
                ("assertion".to_string(), signed),
            ])
            .finish();

        let mut response = String::new();
        let mut result = self
            .client
            .borrow_mut()
            .post(self.key.token_uri.as_ref().unwrap())
            .body(&body)
            .header(header::ContentType(
                "application/x-www-form-urlencoded".parse().unwrap(),
            ))
            .send()?;

        result.read_to_string(&mut response)?;

        let token: Result<TokenResponse, serde_json::error::Error> =
            serde_json::from_str(&response);

        match token {
            Err(e) => return Err(Box::new(e)),
            Ok(token) => {
                if token.access_token.is_none()
                    || token.token_type.is_none()
                    || token.expires_in.is_none()
                {
                    Err(Box::new(StringError::new(
                        "Token response lacks fields".to_string(),
                        Some(&format!("{:?}", token)),
                    )))
                } else {
                    Ok(token.to_oauth_token())
                }
            }
        }
    }
}

impl<C: BorrowMut<hyper::Client>> GetToken for ServiceAccountAccess<C> {
    fn token<'b, I, T>(&mut self, scopes: I) -> result::Result<Token, Box<error::Error>>
    where
        T: AsRef<str> + Ord + 'b,
        I: IntoIterator<Item = &'b T>,
    {
        let (hash, scps) = hash_scopes(scopes);

        if let Some(token) = self.cache.get(hash, &scps)? {
            if !token.expired() {
                return Ok(token);
            }
        }

        let token = self.request_token(&scps)?;
        let _ = self.cache.set(hash, &scps, Some(token.clone()));

        Ok(token)
    }

    fn api_key(&mut self) -> Option<String> {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::authenticator::GetToken;
    use crate::helper::service_account_key_from_file;
    use hyper;
    use hyper::net::HttpsConnector;
    use hyper_native_tls::NativeTlsClient;

    // This is a valid but deactivated key.
    const TEST_PRIVATE_KEY_PATH: &'static str = "examples/Sanguine-69411a0c0eea.json";

    // Uncomment this test to verify that we can successfully obtain tokens.
    //#[test]
    #[allow(dead_code)]
    fn test_service_account_e2e() {
        let key = service_account_key_from_file(TEST_PRIVATE_KEY_PATH).unwrap();
        let client =
            hyper::Client::with_connector(HttpsConnector::new(NativeTlsClient::new().unwrap()));
        let mut acc = ServiceAccountAccess::new(key, client);
        println!(
            "{:?}",
            acc.token(vec![&"https://www.googleapis.com/auth/pubsub"])
                .unwrap()
        );
    }

    #[test]
    fn test_jwt_initialize_claims() {
        let key = service_account_key_from_file(TEST_PRIVATE_KEY_PATH).unwrap();
        let scopes = vec!["scope1", "scope2", "scope3"];
        let claims = super::init_claims_from_key(&key, &scopes);

        assert_eq!(
            claims.iss,
            "oauth2-public-test@sanguine-rhythm-105020.iam.gserviceaccount.com".to_string()
        );
        assert_eq!(claims.scope, "scope1 scope2 scope3".to_string());
        assert_eq!(
            claims.aud,
            "https://accounts.google.com/o/oauth2/token".to_string()
        );
        assert!(claims.exp > 1000000000);
        assert!(claims.iat < claims.exp);
        assert_eq!(claims.exp - claims.iat, 3595);
    }

    #[test]
    fn test_jwt_sign() {
        let key = service_account_key_from_file(TEST_PRIVATE_KEY_PATH).unwrap();
        let scopes = vec!["scope1", "scope2", "scope3"];
        let claims = super::init_claims_from_key(&key, &scopes);
        let jwt = super::JWT::new(claims);
        let signature = jwt.sign(key.private_key.as_ref().unwrap());

        assert!(signature.is_ok());

        let signature = signature.unwrap();
        assert_eq!(
            signature.split(".").nth(0).unwrap(),
            "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9"
        );
    }
}
<|endoftext|>