<commit_before>//! Bounded package merge algorithm, based on the paper
//! "A Fast and Space-Economical Algorithm for Length-Limited Coding
//! Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

use std::mem;
use std::ptr::{null, null_mut};

use libc::{c_void, size_t};
use libc::funcs::c95::stdlib::{free, malloc};

/// Nodes forming chains. Also used to represent leaves.
struct Node {
    /// Total weight (symbol count) of this chain.
    weight: usize,
    /// Previous node(s) of this chain, or 0 if none.
    tail: *mut Node,
    /// Leaf symbol index, or number of leaves before this chain.
    count: i32,
    /// Tracking for garbage collection.
    in_use: bool,
}

/// Memory pool for nodes.
struct NodePool {
    /// The pool.
    nodes: *mut Node,
    /// Pointer to a possibly free node in the pool.
    next: *mut Node,
    /// Size of the memory pool.
    size: u32,
}

/// Initializes a chain node with the given values and marks it as in use.
unsafe fn init_node(weight: usize, count: i32, tail: *mut Node, node: *mut Node) {
    (*node).weight = weight;
    (*node).count = count;
    (*node).tail = tail;
    (*node).in_use = true;
}

/**
 * Finds a free location in the memory pool. Performs garbage collection if needed.
 * lists: If given, used to mark in-use nodes during garbage collection.
 * maxbits: Size of lists.
 * pool: Memory pool to get free node from.
 */
unsafe fn get_free_node(lists: *const [*mut Node; 2],
                        maxbits: i32,
                        pool: *mut NodePool)
                        -> *mut Node {
    loop {
        if (*pool).next >= (*pool).nodes.offset((*pool).size as isize) {
            // Garbage collection
            for i in 0..(*pool).size {
                (*(*pool).nodes.offset(i as isize)).in_use = false;
            }
            if !lists.is_null() {
                for i in 0..maxbits * 2 {
                    let mut node = (*lists.offset((i / 2) as isize))[(i % 2) as usize];
                    while !node.is_null() {
                        (*node).in_use = true;
                        node = (*node).tail;
                    }
                }
            }
            (*pool).next = (*pool).nodes;
        }
        if !(*(*pool).next).in_use {
            break;
        }
        (*pool).next = (*pool).next.offset(1);
    }
    let free_node = (*pool).next;
    (*pool).next = (*pool).next.offset(1);
    free_node
}

/**
 * Performs a Boundary Package-Merge step. Puts a new chain in the given list. The
 * new chain is, depending on the weights, a leaf or a combination of two chains
 * from the previous list.
 * lists: The lists of chains.
 * maxbits: Number of lists.
 * leaves: The leaves, one per symbol.
 * numsymbols: Number of leaves.
 * pool: the node memory pool.
 * index: The index of the list in which a new chain or leaf is required.
 * final: Whether this is the last time this function is called. If it is then it
 *   is no more needed to recursively call self.
 */
unsafe fn boundary_pm(lists: *mut [*mut Node; 2],
                      maxbits: i32,
                      leaves: &[Node],
                      numsymbols: i32,
                      pool: *mut NodePool,
                      index: i32,
                      is_final: bool) {
    let lastcount = (*(*lists.offset(index as isize))[1]).count; // Count of last chain of list.

    if index == 0 && lastcount >= numsymbols {
        return;
    }

    let newchain = get_free_node(lists, maxbits, pool);
    let oldchain = (*lists.offset(index as isize))[1];

    // These are set up before the recursive calls below, so that there is a list
    // pointing to the new node, to let the garbage collection know it's in use.
    (*lists.offset(index as isize))[0] = oldchain;
    (*lists.offset(index as isize))[1] = newchain;

    if index == 0 {
        // New leaf node in list 0.
        init_node(leaves[lastcount as usize].weight,
                  lastcount + 1,
                  null_mut(),
                  newchain);
    } else {
        let sum = (*(*lists.offset((index - 1) as isize))[0]).weight +
                  (*(*lists.offset((index - 1) as isize))[1]).weight;
        if lastcount < numsymbols && sum > leaves[lastcount as usize].weight {
            // New leaf inserted in list, so count is incremented.
            init_node(leaves[lastcount as usize].weight,
                      lastcount + 1,
                      (*oldchain).tail,
                      newchain);
        } else {
            init_node(sum,
                      lastcount,
                      (*lists.offset((index - 1) as isize))[1],
                      newchain);
            if !is_final {
                // Two lookahead chains of previous list used up, create new ones.
                boundary_pm(lists, maxbits, leaves, numsymbols, pool, index - 1, false);
                boundary_pm(lists, maxbits, leaves, numsymbols, pool, index - 1, false);
            }
        }
    }
}

/// Initializes each list with as lookahead chains the two leaves with lowest weights.
unsafe fn init_lists(pool: *mut NodePool,
                     leaves: &[Node],
                     maxbits: i32,
                     lists: *mut [*mut Node; 2]) {
    let node0 = get_free_node(null(), maxbits, pool);
    let node1 = get_free_node(null(), maxbits, pool);
    init_node(leaves[0].weight, 1, null_mut(), node0);
    init_node(leaves[1].weight, 2, null_mut(), node1);
    for i in 0..maxbits {
        (*lists.offset(i as isize))[0] = node0;
        (*lists.offset(i as isize))[1] = node1;
    }
}

/**
 * Converts result of boundary package-merge to the bitlengths. The result in the
 * last chain of the last list contains the amount of active leaves in each list.
 * chain: Chain to extract the bit length from (last chain from last list).
 */
unsafe fn extract_bit_lengths(chain: *const Node, leaves: &[Node], bitlengths: *mut u32) {
    let mut node = chain;
    while node != null() {
        for i in 0..(*node).count {
            *bitlengths.offset(leaves[i as usize].count as isize) =
                *bitlengths.offset(leaves[i as usize].count as isize) + 1;
        }
        node = (*node).tail;
    }
}

/// Comparator for sorting the leaves. Has the function signature for qsort.
extern "C" fn leaf_comparator(a: *const c_void, b: *const c_void) -> i32 {
    unsafe {
        let node_a: *const Node = a as *const Node;
        let node_b: *const Node = b as *const Node;
        (*node_a).weight as i32 - (*node_b).weight as i32
    }
}

#[link(name="c")]
extern {
    fn qsort(base: *mut c_void,
             nmemb: size_t,
             size: size_t,
             compar: extern "C" fn(*const c_void, *const c_void) -> i32);
}

pub unsafe fn length_limited_code_lengths(frequencies: *const usize,
                                          n: i32,
                                          maxbits: i32,
                                          bitlengths: *mut u32)
                                          -> bool {
    // One leaf per symbol. Only numsymbols leaves will be used.
    let mut leaves: Vec<Node> = Vec::with_capacity(n as usize);

    // Initialize all bitlengths at 0.
    for i in 0..n {
        *bitlengths.offset(i as isize) = 0;
    }

    // Count used symbols and place them in the leaves.
    for i in 0..n {
        if *frequencies.offset(i as isize) != 0 {
            leaves.push(Node {
                weight: *frequencies.offset(i as isize),
                count: i, // Index of symbol this leaf represents.
                tail: null_mut(),
                in_use: false,
            });
        }
    }
    // Amount of symbols with frequency > 0.
    let numsymbols = leaves.len();

    // Check special cases and error conditions.
    if (1 << maxbits) < numsymbols {
        return true; // Error, too few maxbits to represent symbols.
    }
    if numsymbols == 0 {
        return false; // No symbols at all. OK.
    }
    if numsymbols == 1 {
        *bitlengths.offset(leaves[0].count as isize) = 1;
        return false; // Only one symbol, give it bitlength 1, not 0. OK.
    }

    // Sort the leaves from lightest to heaviest.
    qsort(leaves.as_mut_ptr() as *mut c_void,
          numsymbols as size_t,
          mem::size_of::<Node>() as size_t,
          leaf_comparator);

    // Initialize node memory pool.
    let pool_size = 2 * maxbits * (maxbits + 1);
    let pool_nodes = malloc((pool_size as usize * mem::size_of::<Node>()) as size_t) as *mut Node;
    let mut pool = NodePool {
        size: pool_size as u32,
        nodes: pool_nodes,
        next: pool_nodes,
    };
    for i in 0..pool.size {
        (*pool.nodes.offset(i as isize)).in_use = false;
    }

    // Array of lists of chains. Each list requires only two lookahead chains at
    // a time, so each list is a array of two Node*'s.
    let mut lists: *mut [*mut Node; 2] = mem::uninitialized();
    lists = malloc((maxbits as usize * mem::size_of_val(&*lists)) as size_t) as *mut [*mut Node; 2];
    init_lists(&mut pool, &leaves, maxbits, lists);

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each BoundaryPM run creates one.
    let num_boundary_pm_runs = 2 * numsymbols - 4;
    for i in 0..num_boundary_pm_runs {
        let is_final = i == num_boundary_pm_runs - 1;
        boundary_pm(lists,
                    maxbits,
                    &leaves,
                    numsymbols as i32,
                    &mut pool,
                    maxbits - 1,
                    is_final);
    }

    extract_bit_lengths((*lists.offset(maxbits as isize - 1))[1],
                        &leaves,
                        bitlengths);

    free(lists as *mut c_void);
    free(pool.nodes as *mut c_void);
    false // OK.
}

#[cfg(test)]
mod tests {
    use std::mem::size_of;
    use libc::{c_void, size_t};

    use super::qsort;

    extern "C" fn comparator(a: *const c_void, b: *const c_void) -> i32 {
        unsafe { *(a as *const i32) - *(b as *const i32) }
    }

    #[test]
    fn test_qsort() {
        unsafe {
            let mut a: [i32; 5] = [3, 2, 4, 1, 5];
            qsort(a.as_mut_ptr() as *mut c_void,
                  a.len() as size_t,
                  size_of::<i32>() as size_t,
                  comparator);
            assert_eq!(a, [1, 2, 3, 4, 5]);
        }
    }
}
<commit_msg>Use a Vec for "lists"<commit_after>//! Bounded package merge algorithm, based on the paper
//! "A Fast and Space-Economical Algorithm for Length-Limited Coding
//! Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

use std::iter;
use std::mem;
use std::ptr::{null, null_mut};

use libc::{c_void, size_t};
use libc::funcs::c95::stdlib::{free, malloc};

/// Nodes forming chains. Also used to represent leaves.
struct Node {
    /// Total weight (symbol count) of this chain.
    weight: usize,
    /// Previous node(s) of this chain, or 0 if none.
    tail: *mut Node,
    /// Leaf symbol index, or number of leaves before this chain.
    count: i32,
    /// Tracking for garbage collection.
    in_use: bool,
}

/// Memory pool for nodes.
struct NodePool {
    /// The pool.
    nodes: *mut Node,
    /// Pointer to a possibly free node in the pool.
    next: *mut Node,
    /// Size of the memory pool.
    size: u32,
}

/// Initializes a chain node with the given values and marks it as in use.
unsafe fn init_node(weight: usize, count: i32, tail: *mut Node, node: *mut Node) {
    (*node).weight = weight;
    (*node).count = count;
    (*node).tail = tail;
    (*node).in_use = true;
}

/**
 * Finds a free location in the memory pool. Performs garbage collection if needed.
 * lists: If given, used to mark in-use nodes during garbage collection.
 * maxbits: Size of lists.
 * pool: Memory pool to get free node from.
 */
unsafe fn get_free_node(lists: Option<&Vec<[*mut Node; 2]>>,
                        maxbits: i32,
                        pool: *mut NodePool)
                        -> *mut Node {
    loop {
        if (*pool).next >= (*pool).nodes.offset((*pool).size as isize) {
            // Garbage collection
            for i in 0..(*pool).size {
                (*(*pool).nodes.offset(i as isize)).in_use = false;
            }
            if let Some(lists) = lists {
                for i in 0..maxbits * 2 {
                    let mut node = lists[(i / 2) as usize][(i % 2) as usize];
                    while !node.is_null() {
                        (*node).in_use = true;
                        node = (*node).tail;
                    }
                }
            }
            (*pool).next = (*pool).nodes;
        }
        if !(*(*pool).next).in_use {
            break;
        }
        (*pool).next = (*pool).next.offset(1);
    }
    let free_node = (*pool).next;
    (*pool).next = (*pool).next.offset(1);
    free_node
}

/**
 * Performs a Boundary Package-Merge step. Puts a new chain in the given list. The
 * new chain is, depending on the weights, a leaf or a combination of two chains
 * from the previous list.
 * lists: The lists of chains.
 * maxbits: Number of lists.
 * leaves: The leaves, one per symbol.
 * numsymbols: Number of leaves.
 * pool: the node memory pool.
 * index: The index of the list in which a new chain or leaf is required.
 * final: Whether this is the last time this function is called. If it is then it
 *   is no more needed to recursively call self.
 */
unsafe fn boundary_pm(lists: &mut Vec<[*mut Node; 2]>,
                      maxbits: i32,
                      leaves: &[Node],
                      numsymbols: i32,
                      pool: *mut NodePool,
                      index: i32,
                      is_final: bool) {
    let lastcount = (*lists[index as usize][1]).count; // Count of last chain of list.

    if index == 0 && lastcount >= numsymbols {
        return;
    }

    let newchain = get_free_node(Some(&lists), maxbits, pool);
    let oldchain = lists[index as usize][1];

    // These are set up before the recursive calls below, so that there is a list
    // pointing to the new node, to let the garbage collection know it's in use.
    lists[index as usize][0] = oldchain;
    lists[index as usize][1] = newchain;

    if index == 0 {
        // New leaf node in list 0.
        init_node(leaves[lastcount as usize].weight,
                  lastcount + 1,
                  null_mut(),
                  newchain);
    } else {
        let sum = (*lists[(index - 1) as usize][0]).weight +
                  (*lists[(index - 1) as usize][1]).weight;
        if lastcount < numsymbols && sum > leaves[lastcount as usize].weight {
            // New leaf inserted in list, so count is incremented.
            init_node(leaves[lastcount as usize].weight,
                      lastcount + 1,
                      (*oldchain).tail,
                      newchain);
        } else {
            init_node(sum,
                      lastcount,
                      lists[(index - 1) as usize][1],
                      newchain);
            if !is_final {
                // Two lookahead chains of previous list used up, create new ones.
                boundary_pm(lists, maxbits, leaves, numsymbols, pool, index - 1, false);
                boundary_pm(lists, maxbits, leaves, numsymbols, pool, index - 1, false);
            }
        }
    }
}

/// Initializes each list with as lookahead chains the two leaves with lowest weights.
unsafe fn init_lists(pool: *mut NodePool,
                     leaves: &[Node],
                     maxbits: i32) -> Vec<[*mut Node; 2]> {
    let node0 = get_free_node(None, maxbits, pool);
    let node1 = get_free_node(None, maxbits, pool);
    init_node(leaves[0].weight, 1, null_mut(), node0);
    init_node(leaves[1].weight, 2, null_mut(), node1);
    iter::repeat([node0, node1]).take(maxbits as usize).collect()
}

/**
 * Converts result of boundary package-merge to the bitlengths. The result in the
 * last chain of the last list contains the amount of active leaves in each list.
 * chain: Chain to extract the bit length from (last chain from last list).
 */
unsafe fn extract_bit_lengths(chain: *const Node, leaves: &[Node], bitlengths: *mut u32) {
    let mut node = chain;
    while node != null() {
        for i in 0..(*node).count {
            *bitlengths.offset(leaves[i as usize].count as isize) =
                *bitlengths.offset(leaves[i as usize].count as isize) + 1;
        }
        node = (*node).tail;
    }
}

/// Comparator for sorting the leaves. Has the function signature for qsort.
extern "C" fn leaf_comparator(a: *const c_void, b: *const c_void) -> i32 {
    unsafe {
        let node_a: *const Node = a as *const Node;
        let node_b: *const Node = b as *const Node;
        (*node_a).weight as i32 - (*node_b).weight as i32
    }
}

#[link(name="c")]
extern {
    fn qsort(base: *mut c_void,
             nmemb: size_t,
             size: size_t,
             compar: extern "C" fn(*const c_void, *const c_void) -> i32);
}

pub unsafe fn length_limited_code_lengths(frequencies: *const usize,
                                          n: i32,
                                          maxbits: i32,
                                          bitlengths: *mut u32)
                                          -> bool {
    // One leaf per symbol. Only numsymbols leaves will be used.
    let mut leaves: Vec<Node> = Vec::with_capacity(n as usize);

    // Initialize all bitlengths at 0.
    for i in 0..n {
        *bitlengths.offset(i as isize) = 0;
    }

    // Count used symbols and place them in the leaves.
    for i in 0..n {
        if *frequencies.offset(i as isize) != 0 {
            leaves.push(Node {
                weight: *frequencies.offset(i as isize),
                count: i, // Index of symbol this leaf represents.
                tail: null_mut(),
                in_use: false,
            });
        }
    }
    // Amount of symbols with frequency > 0.
    let numsymbols = leaves.len();

    // Check special cases and error conditions.
    if (1 << maxbits) < numsymbols {
        return true; // Error, too few maxbits to represent symbols.
    }
    if numsymbols == 0 {
        return false; // No symbols at all. OK.
    }
    if numsymbols == 1 {
        *bitlengths.offset(leaves[0].count as isize) = 1;
        return false; // Only one symbol, give it bitlength 1, not 0. OK.
    }

    // Sort the leaves from lightest to heaviest.
    qsort(leaves.as_mut_ptr() as *mut c_void,
          numsymbols as size_t,
          mem::size_of::<Node>() as size_t,
          leaf_comparator);

    // Initialize node memory pool.
    let pool_size = 2 * maxbits * (maxbits + 1);
    let pool_nodes = malloc((pool_size as usize * mem::size_of::<Node>()) as size_t) as *mut Node;
    let mut pool = NodePool {
        size: pool_size as u32,
        nodes: pool_nodes,
        next: pool_nodes,
    };
    for i in 0..pool.size {
        (*pool.nodes.offset(i as isize)).in_use = false;
    }

    // Array of lists of chains. Each list requires only two lookahead chains at
    // a time, so each list is a array of two Node*'s.
    let mut lists: Vec<[*mut Node; 2]> = init_lists(&mut pool, &leaves, maxbits);

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each BoundaryPM run creates one.
    let num_boundary_pm_runs = 2 * numsymbols - 4;
    for i in 0..num_boundary_pm_runs {
        let is_final = i == num_boundary_pm_runs - 1;
        boundary_pm(&mut lists,
                    maxbits,
                    &leaves,
                    numsymbols as i32,
                    &mut pool,
                    maxbits - 1,
                    is_final);
    }

    extract_bit_lengths(lists[maxbits as usize - 1][1], &leaves, bitlengths);
    free(pool.nodes as *mut c_void);
    false // OK.
}

#[cfg(test)]
mod tests {
    use std::mem::size_of;
    use libc::{c_void, size_t};

    use super::qsort;

    extern "C" fn comparator(a: *const c_void, b: *const c_void) -> i32 {
        unsafe { *(a as *const i32) - *(b as *const i32) }
    }

    #[test]
    fn test_qsort() {
        unsafe {
            let mut a: [i32; 5] = [3, 2, 4, 1, 5];
            qsort(a.as_mut_ptr() as *mut c_void,
                  a.len() as size_t,
                  size_of::<i32>() as size_t,
                  comparator);
            assert_eq!(a, [1, 2, 3, 4, 5]);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017-2019 int08h LLC
//
// 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.

#[cfg(feature = "awskms")]
pub mod inner {
    extern crate bytes;

    use std::collections::HashMap;
    use std::default::Default;
    use std::error::Error;
    use std::fmt;
    use std::fmt::Formatter;
    use std::str::FromStr;

    use rusoto_core::Region;
    use rusoto_kms::{DecryptRequest, EncryptRequest, Kms, KmsClient};
    use self::bytes::Bytes;

    use crate::kms::{EncryptedDEK, KmsError, KmsProvider, PlaintextDEK, AD, DEK_SIZE_BYTES};

    /// Amazon Web Services Key Management Service
    /// https://aws.amazon.com/kms/
    pub struct AwsKms {
        kms_client: KmsClient,
        key_id: String,
    }

    impl AwsKms {
        /// Create a new instance from the full ARN of a AWS KMS key. The ARN is expected
        /// to be of the form `arn:aws:kms:some-aws-region:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab`
        pub fn from_arn(arn: &str) -> Result<Self, KmsError> {
            let parts: Vec<&str> = arn.split(':').collect();

            if parts.len() != 6 {
                return Err(KmsError::InvalidConfiguration(format!(
                    "invalid KMS arn: too few parts {}",
                    parts.len()
                )));
            }

            let region_part = parts.get(3).expect("region is missing");
            let region = match Region::from_str(region_part) {
                Ok(r) => r,
                Err(e) => return Err(KmsError::InvalidConfiguration(e.description().to_string())),
            };

            Ok(AwsKms {
                kms_client: KmsClient::new(region),
                key_id: arn.to_string(),
            })
        }
    }

    impl KmsProvider for AwsKms {
        fn encrypt_dek(&self, plaintext_dek: &PlaintextDEK) -> Result<EncryptedDEK, KmsError> {
            if plaintext_dek.len() != DEK_SIZE_BYTES {
                return Err(KmsError::InvalidKey(format!(
                    "provided DEK wrong length: {}",
                    plaintext_dek.len()
                )));
            }

            let mut encrypt_req: EncryptRequest = Default::default();
            encrypt_req.key_id = self.key_id.clone();
            encrypt_req.plaintext = Bytes::from(plaintext_dek.as_slice());

            let mut enc_context = HashMap::new();
            enc_context.insert("AD".to_string(), AD.to_string());
            encrypt_req.encryption_context = Some(enc_context);

            match self.kms_client.encrypt(encrypt_req).sync() {
                Ok(result) => {
                    if let Some(ciphertext) = result.ciphertext_blob {
                        Ok(ciphertext.to_vec())
                    } else {
                        Err(KmsError::OperationFailed(
                            "no ciphertext despite successful response".to_string(),
                        ))
                    }
                }
                Err(e) => Err(KmsError::OperationFailed(e.description().to_string())),
            }
        }

        fn decrypt_dek(&self, encrypted_dek: &EncryptedDEK) -> Result<PlaintextDEK, KmsError> {
            let mut decrypt_req: DecryptRequest = Default::default();
            decrypt_req.ciphertext_blob = Bytes::from(encrypted_dek.as_slice());

            let mut dec_context = HashMap::new();
            dec_context.insert("AD".to_string(), AD.to_string());
            decrypt_req.encryption_context = Some(dec_context);

            match self.kms_client.decrypt(decrypt_req).sync() {
                Ok(result) => {
                    if let Some(plaintext_dek) = result.plaintext {
                        if plaintext_dek.len() == DEK_SIZE_BYTES {
                            Ok(plaintext_dek.to_vec())
                        } else {
                            Err(KmsError::InvalidKey(format!(
                                "decrypted DEK wrong length: {}",
                                plaintext_dek.len()
                            )))
                        }
                    } else {
                        Err(KmsError::OperationFailed(
                            "decrypted payload is empty".to_string(),
                        ))
                    }
                }
                Err(e) => Err(KmsError::OperationFailed(e.description().to_string())),
            }
        }
    }

    #[cfg(feature = "awskms")]
    impl fmt::Display for AwsKms {
        fn fmt(&self, f: &mut Formatter) -> fmt::Result {
            write!(f, "{}", self.key_id)
        }
    }
}
<commit_msg>Remove deprecated `std::error::Error::description` calls<commit_after>// Copyright 2017-2019 int08h LLC
//
// 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.

#[cfg(feature = "awskms")]
pub mod inner {
    extern crate bytes;

    use std::collections::HashMap;
    use std::default::Default;
    use std::error::Error;
    use std::fmt;
    use std::fmt::Formatter;
    use std::str::FromStr;

    use rusoto_core::Region;
    use rusoto_kms::{DecryptRequest, EncryptRequest, Kms, KmsClient};
    use self::bytes::Bytes;

    use crate::kms::{EncryptedDEK, KmsError, KmsProvider, PlaintextDEK, AD, DEK_SIZE_BYTES};

    /// Amazon Web Services Key Management Service
    /// https://aws.amazon.com/kms/
    pub struct AwsKms {
        kms_client: KmsClient,
        key_id: String,
    }

    impl AwsKms {
        /// Create a new instance from the full ARN of a AWS KMS key. The ARN is expected
        /// to be of the form `arn:aws:kms:some-aws-region:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab`
        pub fn from_arn(arn: &str) -> Result<Self, KmsError> {
            let parts: Vec<&str> = arn.split(':').collect();

            if parts.len() != 6 {
                return Err(KmsError::InvalidConfiguration(format!(
                    "invalid KMS arn: too few parts {}",
                    parts.len()
                )));
            }

            let region_part = parts.get(3).expect("region is missing");
            let region = match Region::from_str(region_part) {
                Ok(r) => r,
                Err(e) => return Err(KmsError::InvalidConfiguration(e.to_string())),
            };

            Ok(AwsKms {
                kms_client: KmsClient::new(region),
                key_id: arn.to_string(),
            })
        }
    }

    impl KmsProvider for AwsKms {
        fn encrypt_dek(&self, plaintext_dek: &PlaintextDEK) -> Result<EncryptedDEK, KmsError> {
            if plaintext_dek.len() != DEK_SIZE_BYTES {
                return Err(KmsError::InvalidKey(format!(
                    "provided DEK wrong length: {}",
                    plaintext_dek.len()
                )));
            }

            let mut encrypt_req: EncryptRequest = Default::default();
            encrypt_req.key_id = self.key_id.clone();
            encrypt_req.plaintext = Bytes::from(plaintext_dek.as_slice());

            let mut enc_context = HashMap::new();
            enc_context.insert("AD".to_string(), AD.to_string());
            encrypt_req.encryption_context = Some(enc_context);

            match self.kms_client.encrypt(encrypt_req).sync() {
                Ok(result) => {
                    if let Some(ciphertext) = result.ciphertext_blob {
                        Ok(ciphertext.to_vec())
                    } else {
                        Err(KmsError::OperationFailed(
                            "no ciphertext despite successful response".to_string(),
                        ))
                    }
                }
                Err(e) => Err(KmsError::OperationFailed(e.to_string())),
            }
        }

        fn decrypt_dek(&self, encrypted_dek: &EncryptedDEK) -> Result<PlaintextDEK, KmsError> {
            let mut decrypt_req: DecryptRequest = Default::default();
            decrypt_req.ciphertext_blob = Bytes::from(encrypted_dek.as_slice());

            let mut dec_context = HashMap::new();
            dec_context.insert("AD".to_string(), AD.to_string());
            decrypt_req.encryption_context = Some(dec_context);

            match self.kms_client.decrypt(decrypt_req).sync() {
                Ok(result) => {
                    if let Some(plaintext_dek) = result.plaintext {
                        if plaintext_dek.len() == DEK_SIZE_BYTES {
                            Ok(plaintext_dek.to_vec())
                        } else {
                            Err(KmsError::InvalidKey(format!(
                                "decrypted DEK wrong length: {}",
                                plaintext_dek.len()
                            )))
                        }
                    } else {
                        Err(KmsError::OperationFailed(
                            "decrypted payload is empty".to_string(),
                        ))
                    }
                }
                Err(e) => Err(KmsError::OperationFailed(e.to_string())),
            }
        }
    }

    #[cfg(feature = "awskms")]
    impl fmt::Display for AwsKms {
        fn fmt(&self, f: &mut Formatter) -> fmt::Result {
            write!(f, "{}", self.key_id)
        }
    }
}
<|endoftext|>


<commit_before>use metric::{Measured, Line, Grapheme, Metric};
use super::{NodeLink, LeafRepr };

use self::Value::*;

use std::cell::Cell;
use std::convert;
use std::default::Default;
use std::fmt;
use std::ops;

/// A lazily-evaluated field
#[derive(Clone)]
struct Lazy<T: Copy>(Cell<Option<T>>);

impl<T> Lazy<T>
where T: Copy {

    /// Get the value of the field if it has been computed
    ///
    /// # Returns
    /// - `Some(T)` if the value of this field has been computed
    /// - `None` if the field has yet to be computed
    #[inline]
    pub fn get(&self) -> Option<T> { self.0.get() }

    /// Get the value of the field, or compute it
    ///
    /// # Arguments
    /// - `f`: a function returning type `T` to use to calculate the value of
    ///        the field if it has not already been calculated
    ///
    /// # Returns
    /// - If the field has already been evaluated, the value of the field.
    /// - If the field has not been evaluated, the value of `f`
    #[inline]
    pub fn get_or_else<F>(&self, f: F) -> T
    where F: FnOnce() -> T {
        if let Some(value) = self.0.get() {
            value
        } else {
            let value = f();
            self.0.set(Some(value));
            value
        }
    }

    #[inline]
    pub fn new() -> Self {
        Lazy(Cell::new(None))
    }

}

impl<T> Default for Lazy<T>
where T: Copy {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> fmt::Debug for Lazy<T>
where T: fmt::Debug
    , T: Copy {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.0.get() { Some(value) => value.fmt(f)
                           , None => write!(f, "?")

        }
    }
}

macro_rules! lazy_field {
    ($method: ident, $field: ident, $ty:ty) => {
        #[inline] fn $method(&self) -> $ty {
            self.$field.get_or_else(|| { self.value.$method() })
        }

    }
}

/// A `Node`.
#[derive(Clone, Default)]
pub struct Node { len: Lazy<usize>
                , weight: Lazy<usize>
                , line_count: Lazy<Line>
                , line_weight: Lazy<Line>
                , grapheme_count: Lazy<Grapheme>
                , grapheme_weight: Lazy<Grapheme>
                , pub value: Value
                }

impl Node {
    pub fn new(value: Value) -> Self {
        Node { value: value, ..Default::default() }
    }

    pub fn spanning(&self, i: usize, span_len: usize) -> (&Node, usize)
    where Node: Measured<usize> {
        assert!(self.len() >= span_len);
        match **self {
            Branch { ref right, ref left } if < Node as Measured<usize>>::measure_weight(self) < i => {
                // if this node is a branch, and the weight is less than the
                // index, where the span begins, then the first index of the
                // span is on the right side
                let span_i = or_zero!(i, left.len());
                assert!(or_zero!(right.len(), span_i) >= span_len);
                right.spanning(span_i, span_len)
            }
          , Branch { ref left, .. }
            // if the left child is long enough to contain the entire span,
            // walk to the left child
            if or_zero!(left.len(), i) >= span_len => left.spanning(i, span_len)
         ,  Leaf(_) | Branch {..} =>
            // if this function has walked as far as a leaf node,
            // then that leaf must be the spanning node. return it;
            //
            // otherwise, if the node is a branch node and the span is longer
            // than the left child, then this node must be the minimum
            // spanning node
            (self, i)
        }
    }
}

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.strings()
            .fold(Ok(()), |r, string| r.and_then(|_| write!(f, "{}", string)))
    }
}

impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!( f, "Node {{{}{}{}{:#?} }}"
              , self.len.get().map(|l| format!("{} chars, ", l))
                    .unwrap_or_else(|| { String::new() })
            //   , self.weight.get().map(|w| format!("weight: {:?}, ", w))
            //         .unwrap_or_else(|| { String::new() })
              , self.grapheme_count.get()
                    .map(|w| format!("{} graphemes, ", w.0))
                    .unwrap_or_else(|| { String::new() })
            //   , self.grapheme_weight.get().map(|w| format!("weight: {:?}, ", w))
            //         .unwrap_or_else(|| { String::new() })
              , self.line_count.get()
                  .map(|w| format!("{} lines, ", w.0))
                  .unwrap_or_else(|| { String::new() })
            //   , self.line_weight.get().map(|w| format!(", weight: {:?}", w))
            //         .unwrap_or_else(|| { String::new() })
              , self.value
              )
    }
}

impl convert::Into<NodeLink> for Node {
    #[inline] fn into(self) -> NodeLink {
        NodeLink::new(self)
    }
}

impl ops::Deref for Node {
    type Target = Value;
    fn deref(&self) -> &Value { &self.value }
}

impl Measured<usize> for Node {

        #[inline] fn to_byte_index(&self, index: usize) -> Option<usize>  {
             Some(index)
        }

        lazy_field!(measure, len, usize);
        lazy_field!(measure_weight, weight, usize);

}

impl Measured<Grapheme> for Node {

        #[inline] fn to_byte_index(&self, index: Grapheme) -> Option<usize>  {
            self.value.to_byte_index(index)
        }

        lazy_field!(measure, grapheme_count, Grapheme);
        lazy_field!(measure_weight, grapheme_weight, Grapheme);

}

impl Measured<Line> for Node {

        #[inline] fn to_byte_index(&self, index: Line) -> Option<usize>  {
            self.value.to_byte_index(index)
        }

        lazy_field!(measure, line_count, Line);
        lazy_field!(measure_weight, line_weight, Line);

}

impl<M> ops::Index<M> for Node
where M: Metric
    , Node: Measured<M>
    , LeafRepr: Measured<M>
    {
    type Output = str;

    fn index(&self, i: M) -> &str {
        let len = self.measure();
        assert!( i < len
               , "Node::index: index {:?} out of bounds (length {:?})", i, len);
        match **self {
            Leaf(ref string) => {
                let idx = string.to_byte_index(i)
                                .expect("index out of bounds!");
                &string[idx..idx+1]
            }
          , Branch { ref right, .. } if len < i =>
                &right[i - len]
          , Branch { ref left, .. } => &left[i]
        }
    }
}

/// A `Node` in the `Rope`'s tree.
///
/// A `Node` is either a `Leaf` holding a `String`, or a
/// a `Branch` concatenating together two `Node`s.
#[derive(Clone)]
pub enum Value {
    /// A leaf node
    Leaf(LeafRepr)
  , /// A branch concatenating together `l`eft and `r`ight nodes.
    Branch { /// The left branch node
             left: NodeLink
           , /// The right branch node
             right: NodeLink }
}

impl Value {
    #[inline]
    pub fn new_branch(left: NodeLink, right: NodeLink) -> Self {
        Branch { left: left, right: right }
    }
}

impl<M> Measured<M> for Value
where M: Metric
    , LeafRepr: Measured<M>
    , Node: Measured<M>
{
    fn to_byte_index(&self, index: M) -> Option<usize> {
        unimplemented!()
    }

    fn measure(&self) -> M {
        match *self {
            Leaf(ref r) => r.measure()
          , Branch { ref left, ref right } =>
                left.measure() + right.measure()
        }
    }

    fn measure_weight(&self) -> M {
        match *self {
            Leaf(ref r) => r.measure_weight()
          , Branch { ref left, .. } => left.measure()
        }
    }

}

impl convert::Into<Node> for Value {
    #[inline] fn into(self) -> Node {
        Node::new(self)
    }
}

impl Default for Value {
    fn default() -> Self {
        Leaf(LeafRepr::default())
    }
}

impl fmt::Debug for Value {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Leaf(ref r) => write!(f, "Leaf({:?})", r)
          , Branch { ref left, ref right } => write!( f, "Branch({:?}, {:?})"
                                                    , left
                                                    , right)
        }
    }
}
<commit_msg>fix(rope internals): re-enable missing to_byte_index fn<commit_after>use metric::{Measured, Line, Grapheme, Metric};
use super::{NodeLink, LeafRepr };

use self::Value::*;

use std::cell::Cell;
use std::convert;
use std::default::Default;
use std::fmt;
use std::ops;

/// A lazily-evaluated field
#[derive(Clone)]
struct Lazy<T: Copy>(Cell<Option<T>>);

impl<T> Lazy<T>
where T: Copy {

    /// Get the value of the field if it has been computed
    ///
    /// # Returns
    /// - `Some(T)` if the value of this field has been computed
    /// - `None` if the field has yet to be computed
    #[inline]
    pub fn get(&self) -> Option<T> { self.0.get() }

    /// Get the value of the field, or compute it
    ///
    /// # Arguments
    /// - `f`: a function returning type `T` to use to calculate the value of
    ///        the field if it has not already been calculated
    ///
    /// # Returns
    /// - If the field has already been evaluated, the value of the field.
    /// - If the field has not been evaluated, the value of `f`
    #[inline]
    pub fn get_or_else<F>(&self, f: F) -> T
    where F: FnOnce() -> T {
        if let Some(value) = self.0.get() {
            value
        } else {
            let value = f();
            self.0.set(Some(value));
            value
        }
    }

    #[inline]
    pub fn new() -> Self {
        Lazy(Cell::new(None))
    }

}

impl<T> Default for Lazy<T>
where T: Copy {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> fmt::Debug for Lazy<T>
where T: fmt::Debug
    , T: Copy {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.0.get() { Some(value) => value.fmt(f)
                           , None => write!(f, "?")

        }
    }
}

macro_rules! lazy_field {
    ($method: ident, $field: ident, $ty:ty) => {
        #[inline] fn $method(&self) -> $ty {
            self.$field.get_or_else(|| { self.value.$method() })
        }

    }
}

/// A `Node`.
#[derive(Clone, Default)]
pub struct Node { len: Lazy<usize>
                , weight: Lazy<usize>
                , line_count: Lazy<Line>
                , line_weight: Lazy<Line>
                , grapheme_count: Lazy<Grapheme>
                , grapheme_weight: Lazy<Grapheme>
                , pub value: Value
                }

impl Node {
    pub fn new(value: Value) -> Self {
        Node { value: value, ..Default::default() }
    }

    pub fn spanning(&self, i: usize, span_len: usize) -> (&Node, usize)
    where Node: Measured<usize> {
        assert!(self.len() >= span_len);
        match **self {
            Branch { ref right, ref left } if < Node as Measured<usize>>::measure_weight(self) < i => {
                // if this node is a branch, and the weight is less than the
                // index, where the span begins, then the first index of the
                // span is on the right side
                let span_i = or_zero!(i, left.len());
                assert!(or_zero!(right.len(), span_i) >= span_len);
                right.spanning(span_i, span_len)
            }
          , Branch { ref left, .. }
            // if the left child is long enough to contain the entire span,
            // walk to the left child
            if or_zero!(left.len(), i) >= span_len => left.spanning(i, span_len)
         ,  Leaf(_) | Branch {..} =>
            // if this function has walked as far as a leaf node,
            // then that leaf must be the spanning node. return it;
            //
            // otherwise, if the node is a branch node and the span is longer
            // than the left child, then this node must be the minimum
            // spanning node
            (self, i)
        }
    }
}

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.strings()
            .fold(Ok(()), |r, string| r.and_then(|_| write!(f, "{}", string)))
    }
}

impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!( f, "Node {{{}{}{}{:#?} }}"
              , self.len.get().map(|l| format!("{} chars, ", l))
                    .unwrap_or_else(|| { String::new() })
            //   , self.weight.get().map(|w| format!("weight: {:?}, ", w))
            //         .unwrap_or_else(|| { String::new() })
              , self.grapheme_count.get()
                    .map(|w| format!("{} graphemes, ", w.0))
                    .unwrap_or_else(|| { String::new() })
            //   , self.grapheme_weight.get().map(|w| format!("weight: {:?}, ", w))
            //         .unwrap_or_else(|| { String::new() })
              , self.line_count.get()
                  .map(|w| format!("{} lines, ", w.0))
                  .unwrap_or_else(|| { String::new() })
            //   , self.line_weight.get().map(|w| format!(", weight: {:?}", w))
            //         .unwrap_or_else(|| { String::new() })
              , self.value
              )
    }
}

impl convert::Into<NodeLink> for Node {
    #[inline] fn into(self) -> NodeLink {
        NodeLink::new(self)
    }
}

impl ops::Deref for Node {
    type Target = Value;
    fn deref(&self) -> &Value { &self.value }
}

impl Measured<usize> for Node {

        #[inline] fn to_byte_index(&self, index: usize) -> Option<usize>  {
             Some(index)
        }

        lazy_field!(measure, len, usize);
        lazy_field!(measure_weight, weight, usize);

}

impl Measured<Grapheme> for Node {

        #[inline] fn to_byte_index(&self, index: Grapheme) -> Option<usize>  {
            self.value.to_byte_index(index)
        }

        lazy_field!(measure, grapheme_count, Grapheme);
        lazy_field!(measure_weight, grapheme_weight, Grapheme);

}

impl Measured<Line> for Node {

        #[inline] fn to_byte_index(&self, index: Line) -> Option<usize>  {
            self.value.to_byte_index(index)
        }

        lazy_field!(measure, line_count, Line);
        lazy_field!(measure_weight, line_weight, Line);

}

impl<M> ops::Index<M> for Node
where M: Metric
    , Node: Measured<M>
    , LeafRepr: Measured<M>
    {
    type Output = str;

    fn index(&self, i: M) -> &str {
        let len = self.measure();
        assert!( i < len
               , "Node::index: index {:?} out of bounds (length {:?})", i, len);
        match **self {
            Leaf(ref string) => {
                let idx = string.to_byte_index(i)
                                .expect("index out of bounds!");
                &string[idx..idx+1]
            }
          , Branch { ref right, .. } if len < i =>
                &right[i - len]
          , Branch { ref left, .. } => &left[i]
        }
    }
}

/// A `Node` in the `Rope`'s tree.
///
/// A `Node` is either a `Leaf` holding a `String`, or a
/// a `Branch` concatenating together two `Node`s.
#[derive(Clone)]
pub enum Value {
    /// A leaf node
    Leaf(LeafRepr)
  , /// A branch concatenating together `l`eft and `r`ight nodes.
    Branch { /// The left branch node
             left: NodeLink
           , /// The right branch node
             right: NodeLink }
}

impl Value {
    #[inline]
    pub fn new_branch(left: NodeLink, right: NodeLink) -> Self {
        Branch { left: left, right: right }
    }
}

impl<M> Measured<M> for Value
where M: Metric
    , LeafRepr: Measured<M>
    , Node: Measured<M>
{
    fn to_byte_index(&self, index: M) -> Option<usize> {
        match *self {
            Leaf(ref r) => r.to_byte_index(index)
          , Branch { ref left, ref right } =>
                left.to_byte_index(index)
                    .or_else(|| { right.to_byte_index(index)
                                       .map(|i| i + left.len() ) })
        }
    }

    fn measure(&self) -> M {
        match *self {
            Leaf(ref r) => r.measure()
          , Branch { ref left, ref right } =>
                left.measure() + right.measure()
        }
    }

    fn measure_weight(&self) -> M {
        match *self {
            Leaf(ref r) => r.measure_weight()
          , Branch { ref left, .. } => left.measure()
        }
    }

}

impl convert::Into<Node> for Value {
    #[inline] fn into(self) -> Node {
        Node::new(self)
    }
}

impl Default for Value {
    fn default() -> Self {
        Leaf(LeafRepr::default())
    }
}

impl fmt::Debug for Value {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Leaf(ref r) => write!(f, "Leaf({:?})", r)
          , Branch { ref left, ref right } => write!( f, "Branch({:?}, {:?})"
                                                    , left
                                                    , right)
        }
    }
}
<|endoftext|>


<commit_before>/// MPI-based method for summing many arrays

// Argument parsing
// logging
#[macro_use] extern crate log;
#[macro_use] extern crate clap;
extern crate ndarray;
extern crate env_logger;
// lastly, this library
extern crate cabarrus;
use cabarrus::errors::*;

pub fn main() {
    // Main can't return a Result, and the ? operator needs the enclosing function to return Result
    inner_main().expect("Could not recover. Exiting.");
}
pub fn inner_main() -> Result<()> {
    env_logger::init().unwrap();
    let args = app_from_crate!()
        .arg_from_usage("<size> 'how many workers there are'")
        .arg_from_usage("<rank> 'which worker am I, starting from 0'")
        .arg_from_usage("<addends>.. 'files containin matrices to add'")
        .arg_from_usage("<output> 'file in which to store the resulting matrix'")
        .get_matches();

    let mats = args.values_of("addends").unwrap();
    let size = value_t!(args, "size", usize).unwrap_or_else(|e| e.exit());
    let rank = value_t!(args, "rank", usize).unwrap_or_else(|e| e.exit());
    let chunksize = (mats.len() + size - 1) / size;
    info!("{} files, {} workers, chunks of {}", mats.len(), size, chunksize);

    let mut accum = None;
    for matname in mats.skip(rank*chunksize).take(chunksize) {
        info!("Reading {}..", matname);
        let ref matfile = cabarrus::numpy::open_matrix_mmap(matname)?;
        let ref mat = cabarrus::numpy::read_matrix_mmap(matfile)?;
        accum = Some(match accum {
            Some(mut acc) => {acc += mat; acc}
            None => mat.to_owned()
        });
    }
    if let Some(ref acc) = accum {
        cabarrus::numpy::write_matrix(args.value_of("output").unwrap(), acc)?;
    } else {
        info!("No matrices processed so nothing saved.");
    }
    Ok(())
}
<commit_msg>Sort the incoming matrices<commit_after>/// MPI-based method for summing many arrays

// Argument parsing
// logging
#[macro_use] extern crate log;
#[macro_use] extern crate clap;
extern crate ndarray;
extern crate env_logger;
// lastly, this library
extern crate cabarrus;
use cabarrus::errors::*;

pub fn main() {
    // Main can't return a Result, and the ? operator needs the enclosing function to return Result
    inner_main().expect("Could not recover. Exiting.");
}
pub fn inner_main() -> Result<()> {
    env_logger::init().unwrap();
    let args = app_from_crate!()
        .arg_from_usage("<size> 'how many workers there are'")
        .arg_from_usage("<rank> 'which worker am I, starting from 0'")
        .arg_from_usage("<addends>.. 'files containin matrices to add'")
        .arg_from_usage("<output> 'file in which to store the resulting matrix'")
        .get_matches();

    let mut mats : Vec<&str> = args.values_of("addends").unwrap().collect();
    mats.sort();

    let size = value_t!(args, "size", usize).unwrap_or_else(|e| e.exit());
    let rank = value_t!(args, "rank", usize).unwrap_or_else(|e| e.exit());
    let chunksize = (mats.len() + size - 1) / size;
    info!("{} files, {} workers, chunks of {}", mats.len(), size, chunksize);

    let mut accum = None;
    for matname in mats.into_iter().skip(rank*chunksize).take(chunksize) {
        info!("Reading {}..", matname);
        let ref matfile = cabarrus::numpy::open_matrix_mmap(matname)?;
        let ref mat = cabarrus::numpy::read_matrix_mmap(matfile)?;
        accum = Some(match accum {
            Some(mut acc) => {acc += mat; acc}
            None => mat.to_owned()
        });
    }
    if let Some(ref acc) = accum {
        cabarrus::numpy::write_matrix(args.value_of("output").unwrap(), acc)?;
    } else {
        info!("No matrices processed so nothing saved.");
    }
    Ok(())
}
<|endoftext|>


<commit_before>#![allow(unused_must_use)]

use formatter::RdbParseFormatter;
use std::io;
use std::str;
use serialize::json;

pub struct JSONFormatter {
    out: Box<Writer+'static>,
    is_first_db: bool,
    has_databases: bool,
    is_first_key_in_db: bool,
    elements_in_key: u32,
    element_index: u32
}

impl JSONFormatter {
    pub fn new() -> JSONFormatter {
        let out = Box::new(io::stdout());
        JSONFormatter {
            out: out,
            is_first_db: true,
            has_databases: false,
            is_first_key_in_db: true,
            elements_in_key: 0,
            element_index: 0
        }
    }
}

fn encode_to_ascii(value: &[u8]) -> String {
    let s = unsafe{str::from_utf8_unchecked(value)};
    json::encode(&s)
}

impl JSONFormatter {
    fn start_key(&mut self, length: u32) {
        if !self.is_first_key_in_db {
            self.out.write_str(",");
        }

        self.is_first_key_in_db = false;
        self.elements_in_key = length;
        self.element_index = 0;
    }

    fn end_key(&mut self) { }

    fn write_comma(&mut self) {
        if self.element_index > 0 {
            self.out.write_str(",");
        }
        self.element_index += 1;
    }

    fn write_key(&mut self, key: &[u8]) {
        self.out.write(encode_to_ascii(key).as_bytes());
    }
    fn write_value(&mut self, value: &[u8]) {
        self.out.write(encode_to_ascii(value).as_bytes());
    }
}

impl RdbParseFormatter for JSONFormatter {
    fn start_rdb(&mut self) {
        self.out.write_str("[");
    }

    fn end_rdb(&mut self) {
        if self.has_databases {
            self.out.write_str("}");
        }
        self.out.write_str("]\n");
    }

    fn start_database(&mut self, _db_number: u32) {
        if !self.is_first_db {
            self.out.write_str("},");
        }

        self.out.write_str("{");
        self.is_first_db = false;
        self.has_databases = true;
        self.is_first_key_in_db = true;
    }

    fn set(&mut self, key: &[u8], value: &[u8], _expiry: Option<u64>) {
        self.start_key(0);
        self.write_key(key);
        self.out.write_str(":");
        self.write_value(value);
    }

    fn start_hash(&mut self, key: &[u8], length: u32,
                  _expiry: Option<u64>, _info: Option<()>) {
        self.start_key(length);
        self.write_key(key.as_slice());
        self.out.write_str(":{");
        self.out.flush();
    }
    fn end_hash(&mut self, _key: &[u8]) {
        self.end_key();
        self.out.write_str("}");
        self.out.flush();
    }
    fn hash_element(&mut self, _key: &[u8], field: &[u8], value: &[u8]) {
        self.write_comma();
        self.write_key(field.as_slice());
        self.out.write_str(":");
        self.write_value(value.as_slice());
        self.out.flush();
    }

    fn start_set(&mut self, key: &[u8], cardinality: u32,
                 _expiry: Option<u64>, _info: Option<()>) {
        self.start_key(cardinality);
        self.write_key(key);
        self.out.write_str(":[");
        self.out.flush();
    }
    fn end_set(&mut self, _key: &[u8]) {
        self.end_key();
        self.out.write_str("]");
    }
    fn set_element(&mut self, _key: &[u8], member: &[u8]) {
        self.write_comma();
        self.write_value(member.as_slice());
    }

    fn start_list(&mut self, key: &[u8], length: u32,
                  _expiry: Option<u64>, _info: Option<()>) {
        self.start_key(length);
        self.write_key(key.as_slice());
        self.out.write_str(":[");
    }
    fn end_list(&mut self, _key: &[u8]) {
        self.end_key();
        self.out.write_str("]");
    }
    fn list_element(&mut self, _key: &[u8], value: &[u8]) {
        self.write_comma();
        self.write_value(value.as_slice());
    }

    fn start_sorted_set(&mut self, key: &[u8], length: u32,
                        _expiry: Option<u64>, _info: Option<()>) {
        self.start_key(length);
        self.write_key(key.as_slice());
        self.out.write_str(":{");
    }
    fn end_sorted_set(&mut self, _key: &[u8]) {
        self.end_key();
        self.out.write_str("}");
    }
    fn sorted_set_element(&mut self, _key: &[u8],
                          score: f64, member: &[u8]) {
        self.write_comma();
        self.write_key(member.as_slice());
        self.out.write_str(":");
        self.out.write_str(score.to_string().as_slice());
    }

}
<commit_msg>Pass slices as-is<commit_after>#![allow(unused_must_use)]

use formatter::RdbParseFormatter;
use std::io;
use std::str;
use serialize::json;

pub struct JSONFormatter {
    out: Box<Writer+'static>,
    is_first_db: bool,
    has_databases: bool,
    is_first_key_in_db: bool,
    elements_in_key: u32,
    element_index: u32
}

impl JSONFormatter {
    pub fn new() -> JSONFormatter {
        let out = Box::new(io::stdout());
        JSONFormatter {
            out: out,
            is_first_db: true,
            has_databases: false,
            is_first_key_in_db: true,
            elements_in_key: 0,
            element_index: 0
        }
    }
}

fn encode_to_ascii(value: &[u8]) -> String {
    let s = unsafe{str::from_utf8_unchecked(value)};
    json::encode(&s)
}

impl JSONFormatter {
    fn start_key(&mut self, length: u32) {
        if !self.is_first_key_in_db {
            self.out.write_str(",");
        }

        self.is_first_key_in_db = false;
        self.elements_in_key = length;
        self.element_index = 0;
    }

    fn end_key(&mut self) { }

    fn write_comma(&mut self) {
        if self.element_index > 0 {
            self.out.write_str(",");
        }
        self.element_index += 1;
    }

    fn write_key(&mut self, key: &[u8]) {
        self.out.write(encode_to_ascii(key).as_bytes());
    }
    fn write_value(&mut self, value: &[u8]) {
        self.out.write(encode_to_ascii(value).as_bytes());
    }
}

impl RdbParseFormatter for JSONFormatter {
    fn start_rdb(&mut self) {
        self.out.write_str("[");
    }

    fn end_rdb(&mut self) {
        if self.has_databases {
            self.out.write_str("}");
        }
        self.out.write_str("]\n");
    }

    fn start_database(&mut self, _db_number: u32) {
        if !self.is_first_db {
            self.out.write_str("},");
        }

        self.out.write_str("{");
        self.is_first_db = false;
        self.has_databases = true;
        self.is_first_key_in_db = true;
    }

    fn set(&mut self, key: &[u8], value: &[u8], _expiry: Option<u64>) {
        self.start_key(0);
        self.write_key(key);
        self.out.write_str(":");
        self.write_value(value);
    }

    fn start_hash(&mut self, key: &[u8], length: u32,
                  _expiry: Option<u64>, _info: Option<()>) {
        self.start_key(length);
        self.write_key(key);
        self.out.write_str(":{");
        self.out.flush();
    }
    fn end_hash(&mut self, _key: &[u8]) {
        self.end_key();
        self.out.write_str("}");
        self.out.flush();
    }
    fn hash_element(&mut self, _key: &[u8], field: &[u8], value: &[u8]) {
        self.write_comma();
        self.write_key(field);
        self.out.write_str(":");
        self.write_value(value);
        self.out.flush();
    }

    fn start_set(&mut self, key: &[u8], cardinality: u32,
                 _expiry: Option<u64>, _info: Option<()>) {
        self.start_key(cardinality);
        self.write_key(key);
        self.out.write_str(":[");
        self.out.flush();
    }
    fn end_set(&mut self, _key: &[u8]) {
        self.end_key();
        self.out.write_str("]");
    }
    fn set_element(&mut self, _key: &[u8], member: &[u8]) {
        self.write_comma();
        self.write_value(member);
    }

    fn start_list(&mut self, key: &[u8], length: u32,
                  _expiry: Option<u64>, _info: Option<()>) {
        self.start_key(length);
        self.write_key(key);
        self.out.write_str(":[");
    }
    fn end_list(&mut self, _key: &[u8]) {
        self.end_key();
        self.out.write_str("]");
    }
    fn list_element(&mut self, _key: &[u8], value: &[u8]) {
        self.write_comma();
        self.write_value(value);
    }

    fn start_sorted_set(&mut self, key: &[u8], length: u32,
                        _expiry: Option<u64>, _info: Option<()>) {
        self.start_key(length);
        self.write_key(key);
        self.out.write_str(":{");
    }
    fn end_sorted_set(&mut self, _key: &[u8]) {
        self.end_key();
        self.out.write_str("}");
    }
    fn sorted_set_element(&mut self, _key: &[u8],
                          score: f64, member: &[u8]) {
        self.write_comma();
        self.write_key(member);
        self.out.write_str(":");
        self.write_value(score.to_string().as_bytes());
    }

}
<|endoftext|>


<commit_before>use super::*;
use crate::norm::Norm;
use num_traits::One;
use std::iter::*;

pub struct Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    a: F,
    v: ArrayBase<S, Ix1>,
    ortho: Ortho,
}

impl<A, S, F, Ortho> Arnoldi<A, S, F, Ortho>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    pub fn new(a: F, mut v: ArrayBase<S, Ix1>, mut ortho: Ortho) -> Self {
        assert_eq!(ortho.len(), 0);
        assert!(ortho.tolerance() < One::one());
        // normalize before append because |v| may be smaller than ortho.tolerance()
        let norm = v.norm_l2();
        azip!(mut v(&mut v) in { *v = v.div_real(norm) });
        ortho.append(v.view());
        Arnoldi { a, v, ortho }
    }

    /// Iterate until convergent
    pub fn complete(self) -> (Q<A>, H<A>) {
        let q = self.ortho.get_q();
        let hs: Vec<Array1<A>> = self.collect();
        let n = hs.len();
        let mut h = Array2::zeros((n, n).f());
        for (i, hc) in hs.iter().enumerate() {
            let m = std::cmp::max(n, i + 1);
            for j in 0..m {
                h[(j, i)] = hc[j];
            }
        }
        (q, h)
    }

    /// Dimension of Krylov subspace
    pub fn krylov_dimension(&self) -> usize {
        self.ortho.len()
    }
}

impl<A, S, F, Ortho> Iterator for Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    type Item = Array1<A>;

    fn next(&mut self) -> Option<Self::Item> {
        (self.a)(&mut self.v);
        let result = self.ortho.div_append(&mut self.v);
        azip!(mut v(&mut self.v) in { *v = v.div_real(result.residual_norm()) });
        if result.is_dependent() {
            None
        } else {
            Some(result.into_coeff())
        }
    }
}

pub fn arnoldi_householder<A, S1, S2>(
    a: ArrayBase<S1, Ix2>,
    v: ArrayBase<S2, Ix1>,
    tol: A::Real,
) -> impl Iterator<Item = Array1<A>>
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let (n, m) = a.dim();
    assert_eq!(n, m, "Input matrix must be square");
    assert_eq!(m, v.len(), "Input matrix and vector sizes mismach");

    let householder = Householder::new(n, tol);
    Arnoldi::new(
        move |x| {
            let ax = a.dot(x);
            azip!(mut x(x), ax in { *x = ax });
        },
        v,
        householder,
    )
}

pub fn arnoldi_mgs<A, S1, S2>(
    a: ArrayBase<S1, Ix2>,
    v: ArrayBase<S2, Ix1>,
    tol: A::Real,
) -> impl Iterator<Item = Array1<A>>
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let (n, m) = a.dim();
    assert_eq!(n, m, "Input matrix must be square");
    assert_eq!(m, v.len(), "Input matrix and vector sizes mismach");

    let mgs = MGS::new(n, tol);
    Arnoldi::new(
        move |x| {
            let ax = a.dot(x);
            azip!(mut x(x), ax in { *x = ax });
        },
        v,
        mgs,
    )
}
<commit_msg>Save coefficients in Arnoldi struct<commit_after>use super::*;
use crate::norm::Norm;
use num_traits::One;
use std::iter::*;

pub struct Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    a: F,
    v: ArrayBase<S, Ix1>,
    ortho: Ortho,
    /// Coefficients
    h: Vec<Array1<A>>,
}

impl<A, S, F, Ortho> Arnoldi<A, S, F, Ortho>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    pub fn new(a: F, mut v: ArrayBase<S, Ix1>, mut ortho: Ortho) -> Self {
        assert_eq!(ortho.len(), 0);
        assert!(ortho.tolerance() < One::one());
        // normalize before append because |v| may be smaller than ortho.tolerance()
        let norm = v.norm_l2();
        azip!(mut v(&mut v) in { *v = v.div_real(norm) });
        ortho.append(v.view());
        Arnoldi {
            a,
            v,
            ortho,
            h: Vec::new(),
        }
    }

    /// Dimension of Krylov subspace
    pub fn dim(&self) -> usize {
        self.ortho.len()
    }

    /// Iterate until convergent
    pub fn complete(mut self) -> (Q<A>, H<A>) {
        for _ in &mut self {} // execute iteration until convergent
        let q = self.ortho.get_q();
        let n = self.dim();
        let mut h = Array2::zeros((n, n).f());
        for (i, hc) in self.h.iter().enumerate() {
            let m = std::cmp::max(n, i + 1);
            for j in 0..m {
                h[(j, i)] = hc[j];
            }
        }
        (q, h)
    }
}

impl<A, S, F, Ortho> Iterator for Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    type Item = Array1<A>;

    fn next(&mut self) -> Option<Self::Item> {
        (self.a)(&mut self.v);
        let result = self.ortho.div_append(&mut self.v);
        azip!(mut v(&mut self.v) in { *v = v.div_real(result.residual_norm()) });
        if result.is_dependent() {
            None
        } else {
            let coef = result.into_coeff();
            self.h.push(coef.clone());
            Some(coef)
        }
    }
}

/// Interpret a matrix as a linear operator
pub fn mul_mat<A, S1, S2>(a: ArrayBase<S1, Ix2>) -> impl Fn(&mut ArrayBase<S2, Ix1>)
where
    A: Scalar,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let (n, m) = a.dim();
    assert_eq!(n, m, "Input matrix must be square");
    move |x| {
        assert_eq!(m, x.len(), "Input matrix and vector sizes mismatch");
        let ax = a.dot(x);
        azip!(mut x(x), ax in { *x = ax });
    }
}

pub fn arnoldi_householder<A, S1, S2>(a: ArrayBase<S1, Ix2>, v: ArrayBase<S2, Ix1>, tol: A::Real) -> (Q<A>, H<A>)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let householder = Householder::new(v.len(), tol);
    Arnoldi::new(mul_mat(a), v, householder).complete()
}

pub fn arnoldi_mgs<A, S1, S2>(a: ArrayBase<S1, Ix2>, v: ArrayBase<S2, Ix1>, tol: A::Real) -> (Q<A>, H<A>)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let mgs = MGS::new(v.len(), tol);
    Arnoldi::new(mul_mat(a), v, mgs).complete()
}
<|endoftext|>


<commit_before>mod transaction_mapper;
mod address_mapper;
mod bundle_mapper;
mod record;
mod error;

pub use self::address_mapper::AddressMapper;
pub use self::bundle_mapper::BundleMapper;
pub use self::error::{Error, Result};
pub use self::record::{AddressRecord, BundleRecord, Record, TransactionRecord};
pub use self::transaction_mapper::TransactionMapper;

use mysql;
use std::cell::Cell;
use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockWriteGuard};

pub type Records<T> = BTreeMap<u64, Arc<Mutex<T>>>;
pub type Hashes = HashMap<String, u64>;
pub type Index = Option<Vec<u64>>;
pub type Garbage<'a, T> = HashMap<
  u64,
  Option<
    (
      MutexGuard<'a, T>,
      Vec<MutexGuard<'a, Index>>,
      Cell<Option<bool>>,
    ),
  >,
>;

pub trait Mapper: Sized {
  type Record: Record;

  fn new(conn: &mut mysql::Conn) -> Result<Self>;

  fn counter(&self) -> &Mutex<u64>;

  fn records(&self) -> &RwLock<Records<Self::Record>>;

  fn hashes(&self) -> &RwLock<Hashes>;

  fn indices(&self) -> &[RwLock<Records<Index>>];

  fn fill_indices(
    indices: &mut [RwLockWriteGuard<Records<Index>>],
    record: &Self::Record,
  );

  fn mark_garbage(garbage: &Garbage<Self::Record>);

  fn lock_indices(&self) -> Vec<RwLockWriteGuard<Records<Index>>> {
    self
      .indices()
      .iter()
      .map(|index| index.write().unwrap())
      .collect()
  }

  fn init_counter(
    conn: &mut mysql::Conn,
    query: &str,
  ) -> mysql::Result<Mutex<u64>> {
    conn
      .first(query)
      .and_then(|row| row.map_or_else(|| Ok(0), mysql::from_row_opt))
      .map(Mutex::new)
  }

  fn next_id(&self) -> u64 {
    debug!("Mutex check at line {}", line!());
    let mut counter = self.counter().lock().unwrap();
    debug!("Mutex check at line {}", line!());
    *counter += 1;
    *counter
  }

  fn current_id(&self) -> u64 {
    debug!("Mutex check at line {}", line!());
    let counter = *self.counter().lock().unwrap();
    debug!("Mutex check at line {}", line!());
    counter
  }

  fn fetch(
    &self,
    conn: &mut mysql::Conn,
    id: u64,
  ) -> Result<Arc<Mutex<Self::Record>>> {
    let cached = {
      debug!("Mutex check at line {}", line!());
      let records = self.records().read().unwrap();
      debug!("Mutex check at line {}", line!());
      records.get(&id).cloned()
    };
    cached.map(Ok).unwrap_or_else(|| {
      Self::Record::find_by_id(conn, id).map(|record| {
        debug!("Mutex check at line {}", line!());
        let mut records = self.records().write().unwrap();
        debug!("Mutex check at line {}", line!());
        records
          .entry(record.id())
          .or_insert_with(|| {
            debug!("Mutex check at line {}", line!());
            let mut hashes = self.hashes().write().unwrap();
            debug!("Mutex check at line {}", line!());
            let mut indices = self.lock_indices();
            debug!("Mutex check at line {}", line!());
            hashes.insert(record.hash().to_owned(), record.id());
            Self::fill_indices(&mut indices, &record);
            Arc::new(Mutex::new(record))
          })
          .clone()
      })
    })
  }

  fn fetch_by_hash<T>(
    &self,
    conn: &mut mysql::Conn,
    hash: &str,
    f: T,
  ) -> Result<Arc<Mutex<Self::Record>>>
  where
    T: FnOnce(u64) -> Result<Self::Record>,
  {
    let cached = {
      debug!("Mutex check at line {}", line!());
      let records = self.records().read().unwrap();
      debug!("Mutex check at line {}", line!());
      let hashes = self.hashes().read().unwrap();
      debug!("Mutex check at line {}", line!());
      hashes.get(hash).and_then(|id| records.get(id)).cloned()
    };
    cached.map(Ok).unwrap_or_else(|| {
      Self::Record::find_by_hash(conn, hash).and_then(|record| {
        debug!("Mutex check at line {}", line!());
        let mut records = self.records().write().unwrap();
        debug!("Mutex check at line {}", line!());
        let mut hashes = self.hashes().write().unwrap();
        debug!("Mutex check at line {}", line!());
        let mut indices = self.lock_indices();
        debug!("Mutex check at line {}", line!());
        record.map_or_else(|| f(self.next_id()), Ok).map(|record| {
          records
            .entry(record.id())
            .or_insert_with(|| {
              hashes.insert(record.hash().to_owned(), record.id());
              Self::fill_indices(&mut indices, &record);
              Arc::new(Mutex::new(record))
            })
            .clone()
        })
      })
    })
  }

  fn update(&self, conn: &mut mysql::Conn) -> Result<usize> {
    let mut counter = 0;
    let records = {
      debug!("Mutex check at line {}", line!());
      let records = self.records().read().unwrap();
      debug!("Mutex check at line {}", line!());
      records.values().cloned().collect::<Vec<_>>()
    };
    for record in records {
      debug!("Mutex check at line {}", line!());
      let mut record = record.lock().unwrap();
      debug!("Mutex check at line {}", line!());
      if !record.is_persisted() {
        continue;
      }
      if record.is_modified() {
        record.update(conn)?;
        counter += 1;
      }
      record.advance_generation();
    }
    Ok(counter)
  }

  fn prune(&self, generation_limit: usize) -> usize {
    debug!("Mutex check at line {}", line!());
    let mut records = self.records().write().unwrap();
    debug!("Mutex check at line {}", line!());
    let mut hashes = self.hashes().write().unwrap();
    debug!("Mutex check at line {}", line!());
    let mut indices = self.lock_indices();
    debug!("Mutex check at line {}", line!());
    let record_refs = records.values().cloned().collect::<Vec<_>>();
    let index_refs = indices
      .iter()
      .map(|i| i.iter().map(|(&k, v)| (k, v.clone())).collect())
      .collect::<Vec<BTreeMap<_, _>>>();
    let garbage: Garbage<Self::Record> = record_refs
      .iter()
      .map(|reference| {
        debug!("Mutex check at line {}", line!());
        let record = reference.lock().unwrap();
        debug!("Mutex check at line {}", line!());
        let id = record.id();
        let index_refs = index_refs
          .iter()
          .filter_map(|index| index.get(&id))
          .collect::<Vec<_>>();
        debug!("Mutex check at line {}", line!());
        let indices = index_refs
          .iter()
          .map(|index| index.lock().unwrap())
          .collect::<Vec<_>>();
        debug!("Mutex check at line {}", line!());
        if Arc::strong_count(reference) == 2 &&
          index_refs.iter().all(|index| Arc::strong_count(index) == 2) &&
          record.is_persisted() && !record.is_modified() &&
          record.generation() > generation_limit
        {
          (id, Some((record, indices, Cell::new(None))))
        } else {
          (id, None)
        }
      })
      .collect();
    Self::mark_garbage(&garbage);
    for value in garbage.values() {
      if let Some((ref record, _, ref mark)) = *value {
        if let Some(true) = mark.get() {
          records.remove(&record.id());
          hashes.remove(record.hash());
          for index in &mut indices {
            index.remove(&record.id());
          }
        }
      }
    }
    record_refs.len() - records.len()
  }
}
<commit_msg>Deadlock fix<commit_after>mod transaction_mapper;
mod address_mapper;
mod bundle_mapper;
mod record;
mod error;

pub use self::address_mapper::AddressMapper;
pub use self::bundle_mapper::BundleMapper;
pub use self::error::{Error, Result};
pub use self::record::{AddressRecord, BundleRecord, Record, TransactionRecord};
pub use self::transaction_mapper::TransactionMapper;

use mysql;
use std::cell::Cell;
use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockWriteGuard};

pub type Records<T> = BTreeMap<u64, Arc<Mutex<T>>>;
pub type Hashes = HashMap<String, u64>;
pub type Index = Option<Vec<u64>>;
pub type Garbage<'a, T> = HashMap<
  u64,
  Option<
    (
      MutexGuard<'a, T>,
      Vec<MutexGuard<'a, Index>>,
      Cell<Option<bool>>,
    ),
  >,
>;

pub trait Mapper: Sized {
  type Record: Record;

  fn new(conn: &mut mysql::Conn) -> Result<Self>;

  fn counter(&self) -> &Mutex<u64>;

  fn records(&self) -> &RwLock<Records<Self::Record>>;

  fn hashes(&self) -> &RwLock<Hashes>;

  fn indices(&self) -> &[RwLock<Records<Index>>];

  fn fill_indices(
    indices: &mut [RwLockWriteGuard<Records<Index>>],
    record: &Self::Record,
  );

  fn mark_garbage(garbage: &Garbage<Self::Record>);

  fn lock_indices(&self) -> Vec<RwLockWriteGuard<Records<Index>>> {
    self
      .indices()
      .iter()
      .map(|index| index.write().unwrap())
      .collect()
  }

  fn init_counter(
    conn: &mut mysql::Conn,
    query: &str,
  ) -> mysql::Result<Mutex<u64>> {
    conn
      .first(query)
      .and_then(|row| row.map_or_else(|| Ok(0), mysql::from_row_opt))
      .map(Mutex::new)
  }

  fn next_id(&self) -> u64 {
    debug!("Mutex check at line {}", line!());
    let mut counter = self.counter().lock().unwrap();
    debug!("Mutex check at line {}", line!());
    *counter += 1;
    *counter
  }

  fn current_id(&self) -> u64 {
    debug!("Mutex check at line {}", line!());
    let counter = *self.counter().lock().unwrap();
    debug!("Mutex check at line {}", line!());
    counter
  }

  fn fetch(
    &self,
    conn: &mut mysql::Conn,
    id: u64,
  ) -> Result<Arc<Mutex<Self::Record>>> {
    let cached = {
      debug!("Mutex check at line {}", line!());
      let records = self.records().read().unwrap();
      debug!("Mutex check at line {}", line!());
      records.get(&id).cloned()
    };
    cached.map(Ok).unwrap_or_else(|| {
      Self::Record::find_by_id(conn, id).map(|record| {
        debug!("Mutex check at line {}", line!());
        let mut records = self.records().write().unwrap();
        debug!("Mutex check at line {}", line!());
        records
          .entry(record.id())
          .or_insert_with(|| {
            debug!("Mutex check at line {}", line!());
            let mut hashes = self.hashes().write().unwrap();
            debug!("Mutex check at line {}", line!());
            let mut indices = self.lock_indices();
            debug!("Mutex check at line {}", line!());
            hashes.insert(record.hash().to_owned(), record.id());
            Self::fill_indices(&mut indices, &record);
            Arc::new(Mutex::new(record))
          })
          .clone()
      })
    })
  }

  fn fetch_by_hash<T>(
    &self,
    conn: &mut mysql::Conn,
    hash: &str,
    f: T,
  ) -> Result<Arc<Mutex<Self::Record>>>
  where
    T: FnOnce(u64) -> Result<Self::Record>,
  {
    let cached = {
      debug!("Mutex check at line {}", line!());
      let records = self.records().read().unwrap();
      debug!("Mutex check at line {}", line!());
      let hashes = self.hashes().read().unwrap();
      debug!("Mutex check at line {}", line!());
      hashes.get(hash).and_then(|id| records.get(id)).cloned()
    };
    cached.map(Ok).unwrap_or_else(|| {
      Self::Record::find_by_hash(conn, hash).and_then(|record| {
        debug!("Mutex check at line {}", line!());
        let mut records = self.records().write().unwrap();
        debug!("Mutex check at line {}", line!());
        let mut hashes = self.hashes().write().unwrap();
        debug!("Mutex check at line {}", line!());
        let mut indices = self.lock_indices();
        debug!("Mutex check at line {}", line!());
        record.map_or_else(|| f(self.next_id()), Ok).map(|record| {
          records
            .entry(record.id())
            .or_insert_with(|| {
              hashes.insert(record.hash().to_owned(), record.id());
              Self::fill_indices(&mut indices, &record);
              Arc::new(Mutex::new(record))
            })
            .clone()
        })
      })
    })
  }

  fn update(&self, conn: &mut mysql::Conn) -> Result<usize> {
    let mut counter = 0;
    let records = {
      debug!("Mutex check at line {}", line!());
      let records = self.records().read().unwrap();
      debug!("Mutex check at line {}", line!());
      records.values().cloned().collect::<Vec<_>>()
    };
    for record in records {
      debug!("Mutex check at line {}", line!());
      let mut record = record.lock().unwrap();
      debug!("Mutex check at line {}", line!());
      if !record.is_persisted() {
        continue;
      }
      if record.is_modified() {
        record.update(conn)?;
        counter += 1;
      }
      record.advance_generation();
    }
    Ok(counter)
  }

  fn prune(&self, generation_limit: usize) -> usize {
    debug!("Mutex check at line {}", line!());
    let mut records = self.records().write().unwrap();
    debug!("Mutex check at line {}", line!());
    let mut hashes = self.hashes().write().unwrap();
    debug!("Mutex check at line {}", line!());
    let mut indices = self.lock_indices();
    debug!("Mutex check at line {}", line!());
    let record_refs = records
      .iter()
      .map(|(&k, v)| (k, v.clone()))
      .collect::<BTreeMap<_, _>>();
    let index_refs = indices
      .iter()
      .map(|i| i.iter().map(|(&k, v)| (k, v.clone())).collect())
      .collect::<Vec<HashMap<_, _>>>();
    let garbage: Garbage<Self::Record> = record_refs
      .iter()
      .map(|(&id, reference)| {
        if let Ok(record) = reference.try_lock() {
          let index_refs = index_refs
            .iter()
            .filter_map(|index| index.get(&id))
            .collect::<Vec<_>>();
          let indices = index_refs
            .iter()
            .filter_map(|index| index.try_lock().ok())
            .collect::<Vec<_>>();
          if indices.len() == index_refs.len() &&
            Arc::strong_count(reference) == 2 &&
            index_refs.iter().all(|index| Arc::strong_count(index) == 2) &&
            record.is_persisted() && !record.is_modified() &&
            record.generation() > generation_limit
          {
            return (id, Some((record, indices, Cell::new(None))));
          }
        };
        (id, None)
      })
      .collect();
    Self::mark_garbage(&garbage);
    for value in garbage.values() {
      if let Some((ref record, _, ref mark)) = *value {
        if let Some(true) = mark.get() {
          records.remove(&record.id());
          hashes.remove(record.hash());
          for index in &mut indices {
            index.remove(&record.id());
          }
        }
      }
    }
    record_refs.len() - records.len()
  }
}
<|endoftext|>


<commit_before>use std::process::exit;
use std::{convert::Infallible, sync::mpsc};
use std::{io, thread, time};

use cpal::traits::{DeviceTrait, HostTrait};
use thiserror::Error;

use super::{Open, Sink};

#[derive(Debug, Error)]
pub enum RodioError {
    #[error("Rodio: no device available")]
    NoDeviceAvailable,
    #[error("Rodio: device \"{0}\" is not available")]
    DeviceNotAvailable(String),
    #[error("Rodio play error: {0}")]
    PlayError(#[from] rodio::PlayError),
    #[error("Rodio stream error: {0}")]
    StreamError(#[from] rodio::StreamError),
    #[error("Cannot get audio devices: {0}")]
    DevicesError(#[from] cpal::DevicesError),
}

pub struct RodioSink {
    rodio_sink: rodio::Sink,

    // will produce a TryRecvError on the receiver side when it is dropped.
    _close_tx: mpsc::SyncSender<Infallible>,
}

fn list_formats(device: &rodio::Device) {
    let default_fmt = match device.default_output_config() {
        Ok(fmt) => fmt,
        Err(e) => {
            warn!("Error getting default rodio::Sink config: {}", e);
            return;
        }
    };
    debug!("  Default config:");
    debug!("    {:?}", default_fmt);

    let mut output_configs = match device.supported_output_configs() {
        Ok(f) => f.peekable(),
        Err(e) => {
            warn!("Error getting supported rodio::Sink configs: {}", e);
            return;
        }
    };

    if output_configs.peek().is_some() {
        debug!("  Available configs:");
        for format in output_configs {
            debug!("    {:?}", format);
        }
    }
}

fn list_outputs_and_exit() -> ! {
    let default_device = get_default_device().unwrap();
    let default_device_name = default_device.name().expect("cannot get output name");
    println!("Default Audio Device:\n  {}", default_device_name);
    list_formats(&default_device);

    println!("Other Available Audio Devices:");
    for device in cpal::default_host()
        .output_devices()
        .expect("cannot get list of output devices")
    {
        let device_name = device.name().expect("cannot get output name");
        if device_name != default_device_name {
            println!("  {}", device_name);
            list_formats(&device);
        }
    }

    exit(0)
}

fn get_default_device() -> Result<rodio::Device, RodioError> {
    cpal::default_host()
        .default_output_device()
        .ok_or(RodioError::NoDeviceAvailable)
}

fn create_sink(device: Option<String>) -> Result<(rodio::Sink, rodio::OutputStream), RodioError> {
    let rodio_device = match device {
        Some(ask) if &ask == "?" => list_outputs_and_exit(),
        Some(device_name) => {
            cpal::default_host()
                .output_devices()?
                .find(|d| d.name().ok().map_or(false, |name| name == device_name)) // Ignore devices for which getting name fails
                .ok_or(RodioError::DeviceNotAvailable(device_name))?
        }
        None => get_default_device()?,
    };

    let name = rodio_device.name().ok();
    info!(
        "Using audio device: {}",
        name.as_deref().unwrap_or("(unknown name)")
    );

    let (stream, handle) = rodio::OutputStream::try_from_device(&rodio_device)?;
    let sink = rodio::Sink::try_new(&handle)?;
    Ok((sink, stream))
}

impl Open for RodioSink {
    fn open(device: Option<String>) -> RodioSink {
        debug!(
            "Using rodio sink with cpal host: {:?}",
            cpal::default_host().id().name()
        );

        let (sink_tx, sink_rx) = mpsc::sync_channel(1);
        let (close_tx, close_rx) = mpsc::sync_channel(1);

        std::thread::spawn(move || match create_sink(device) {
            Ok((sink, stream)) => {
                sink_tx.send(Ok(sink)).unwrap();

                close_rx.recv().unwrap_err(); // This will fail as soon as the sender is dropped
                debug!("drop rodio::OutputStream");
                drop(stream);
            }
            Err(e) => {
                sink_tx.send(Err(e)).unwrap();
            }
        });

        // Instead of the second `unwrap`, better error handling could be introduced
        let sink = sink_rx.recv().unwrap().unwrap();

        debug!("Rodio sink was created");
        RodioSink {
            rodio_sink: sink,
            _close_tx: close_tx,
        }
    }
}

impl Sink for RodioSink {
    fn start(&mut self) -> io::Result<()> {
        // More similar to an "unpause" than "play". Doesn't undo "stop".
        // self.rodio_sink.play();
        Ok(())
    }

    fn stop(&mut self) -> io::Result<()> {
        // This will immediately stop playback, but the sink is then unusable.
        // We just have to let the current buffer play till the end.
        // self.rodio_sink.stop();
        Ok(())
    }

    fn write(&mut self, data: &[i16]) -> io::Result<()> {
        let source = rodio::buffer::SamplesBuffer::new(2, 44100, data);
        self.rodio_sink.append(source);

        // Chunk sizes seem to be about 256 to 3000 ish items long.
        // Assuming they're on average 1628 then a half second buffer is:
        // 44100 elements --> about 27 chunks
        while self.rodio_sink.len() > 26 {
            // sleep and wait for rodio to drain a bit
            thread::sleep(time::Duration::from_millis(10));
        }
        Ok(())
    }
}
<commit_msg>Improved error handling in rodio backend<commit_after>use std::process::exit;
use std::{convert::Infallible, sync::mpsc};
use std::{io, thread, time};

use cpal::traits::{DeviceTrait, HostTrait};
use thiserror::Error;

use super::{Open, Sink};

#[derive(Debug, Error)]
pub enum RodioError {
    #[error("Rodio: no device available")]
    NoDeviceAvailable,
    #[error("Rodio: device \"{0}\" is not available")]
    DeviceNotAvailable(String),
    #[error("Rodio play error: {0}")]
    PlayError(#[from] rodio::PlayError),
    #[error("Rodio stream error: {0}")]
    StreamError(#[from] rodio::StreamError),
    #[error("Cannot get audio devices: {0}")]
    DevicesError(#[from] cpal::DevicesError),
}

pub struct RodioSink {
    rodio_sink: rodio::Sink,

    // will produce a TryRecvError on the receiver side when it is dropped.
    _close_tx: mpsc::SyncSender<Infallible>,
}

fn list_formats(device: &rodio::Device) {
    match device.default_output_config() {
        Ok(cfg) => {
            debug!("  Default config:");
            debug!("    {:?}", cfg);
        }
        Err(e) => {
            // Use loglevel debug, since even the output is only debug
            debug!("Error getting default rodio::Sink config: {}", e);
        }
    };

    match device.supported_output_configs() {
        Ok(mut cfgs) => {
            if let Some(first) = cfgs.next() {
                debug!("  Available configs:");
                debug!("    {:?}", first);
            } else {
                return;
            }

            for cfg in cfgs {
                debug!("    {:?}", cfg);
            }
        }
        Err(e) => {
            debug!("Error getting supported rodio::Sink configs: {}", e);
        }
    }
}

fn list_outputs() -> Result<(), cpal::DevicesError> {
    let mut default_device_name = None;

    if let Some(default_device) = get_default_device() {
        default_device_name = default_device.name().ok();
        println!(
            "Default Audio Device:\n  {}",
            default_device_name.as_deref().unwrap_or("[unknown name]")
        );

        list_formats(&default_device);

        println!("Other Available Audio Devices:");
    } else {
        warn!("No default device was found");
    }

    for device in cpal::default_host().output_devices()? {
        match device.name() {
            Ok(name) if Some(&name) == default_device_name.as_ref() => (),
            Ok(name) => {
                println!("  {}", name);
                list_formats(&device);
            }
            Err(e) => {
                warn!("Cannot get device name: {}", e);
                println!("   [unknown name]");
                list_formats(&device);
            }
        }
    }

    Ok(())
}

fn get_default_device() -> Option<rodio::Device> {
    cpal::default_host().default_output_device()
}

fn create_sink(device: Option<String>) -> Result<(rodio::Sink, rodio::OutputStream), RodioError> {
    let rodio_device = match device {
        Some(ask) if &ask == "?" => {
            let exit_code = match list_outputs() {
                Ok(()) => 0,
                Err(e) => {
                    error!("{}", e);
                    1
                }
            };
            exit(exit_code)
        }
        Some(device_name) => {
            cpal::default_host()
                .output_devices()?
                .find(|d| d.name().ok().map_or(false, |name| name == device_name)) // Ignore devices for which getting name fails
                .ok_or(RodioError::DeviceNotAvailable(device_name))?
        }
        None => get_default_device().ok_or(RodioError::NoDeviceAvailable)?,
    };

    let name = rodio_device.name().ok();
    info!(
        "Using audio device: {}",
        name.as_deref().unwrap_or("[unknown name]")
    );

    let (stream, handle) = rodio::OutputStream::try_from_device(&rodio_device)?;
    let sink = rodio::Sink::try_new(&handle)?;
    Ok((sink, stream))
}

impl Open for RodioSink {
    fn open(device: Option<String>) -> RodioSink {
        debug!(
            "Using rodio sink with cpal host: {:?}",
            cpal::default_host().id().name()
        );

        let (sink_tx, sink_rx) = mpsc::sync_channel(1);
        let (close_tx, close_rx) = mpsc::sync_channel(1);

        std::thread::spawn(move || match create_sink(device) {
            Ok((sink, stream)) => {
                sink_tx.send(Ok(sink)).unwrap();

                close_rx.recv().unwrap_err(); // This will fail as soon as the sender is dropped
                debug!("drop rodio::OutputStream");
                drop(stream);
            }
            Err(e) => {
                sink_tx.send(Err(e)).unwrap();
            }
        });

        // Instead of the second `unwrap`, better error handling could be introduced
        let sink = sink_rx.recv().unwrap().unwrap();

        debug!("Rodio sink was created");
        RodioSink {
            rodio_sink: sink,
            _close_tx: close_tx,
        }
    }
}

impl Sink for RodioSink {
    fn start(&mut self) -> io::Result<()> {
        // More similar to an "unpause" than "play". Doesn't undo "stop".
        // self.rodio_sink.play();
        Ok(())
    }

    fn stop(&mut self) -> io::Result<()> {
        // This will immediately stop playback, but the sink is then unusable.
        // We just have to let the current buffer play till the end.
        // self.rodio_sink.stop();
        Ok(())
    }

    fn write(&mut self, data: &[i16]) -> io::Result<()> {
        let source = rodio::buffer::SamplesBuffer::new(2, 44100, data);
        self.rodio_sink.append(source);

        // Chunk sizes seem to be about 256 to 3000 ish items long.
        // Assuming they're on average 1628 then a half second buffer is:
        // 44100 elements --> about 27 chunks
        while self.rodio_sink.len() > 26 {
            // sleep and wait for rodio to drain a bit
            thread::sleep(time::Duration::from_millis(10));
        }
        Ok(())
    }
}
<|endoftext|>