<commit_before>// Implement data structures for a Huffman encoding tree:
//   http://rosettacode.org/wiki/Huffman_coding

extern crate core;
use std::collections::HashMap;
use std::collections::priority_queue::PriorityQueue;

// Each HNode has a weight, representing the sum of the frequencies for all its
// children. It is either a leaf (containing a character), or a HTree
// (containing two children)
struct HNode {
    weight: int,
    item: HTreeOrHLeaf,
}

enum HTreeOrHLeaf {
    HTree(HTreeData),
    HLeaf(char),
}

struct HTreeData {
    left: Box<HNode>,
    right: Box<HNode>,
}

// Implementing comparison traits (Ord and all its dependencies) such that
// the HNode with the greatest weight is the smallest in a comparison. Basically
// reversing all the comparison operators.
impl Ord for HNode {
    fn cmp(&self, other: &HNode) -> Ordering {
        match self.weight.cmp(&other.weight) {
            Less    => Greater,
            Equal   => Equal,
            Greater => Less,
        }
    }
}

impl Eq for HNode {}
impl PartialEq for HNode {
    fn eq(&self, other: &HNode) -> bool {
        self.weight == other.weight
    }
}

impl PartialOrd for HNode {
    fn lt(&self, other: &HNode) -> bool {
        self.weight > other.weight
    }
}

// Takes a non-empty string (function will fail if string is empty) and computes
// the Huffman encoding tree for that string.
fn huffman_tree(input: &str) -> HNode {
    // 1. Loop through all the characters in that string, adding them to a HashMap
    //    of character to frequency.
    let mut freq = HashMap::new();
    for ch in input.chars() {
        freq.insert_or_update_with(ch, 1, |_k, v: &mut int| {*v += 1;});
    }

    // 2. For each (character, frequency) pair in the HashMap, add a Leaf to a
    //    PriorityQueue
    let mut queue = PriorityQueue::<HNode>::new();
    for (ch, freq) in freq.iter() {
        let newNode = HNode{
            weight: *freq,
            item: HLeaf(*ch),
        };
        queue.push(newNode);
    }

    // 3. Pop two items with the least weight from the queue, combine them into
    //    a tree as children. The parent node's weight is the sum of the
    //    children's weight. Continue until one item is left on the queue, and
    //    return that item.
    while queue.len() > 1 {
        let item1 = queue.pop().unwrap();
        let item2 = queue.pop().unwrap();
        let newNode = HNode {
            weight: item1.weight + item2.weight,
            item: HTree(HTreeData{
                left: box item1,
                right: box item2,
            }),
        };
        queue.push(newNode);
    }
    queue.pop().unwrap()
}

// Takes a Huffman Tree, traverse it and build a table with each character and
// its encoding string.
fn build_encoding_table(tree: &HNode,
                      table: &mut HashMap<char,String>,
                      startStr: &str) {
    match tree.item {
        HTree(ref data) => {
            build_encoding_table(data.left, table,
                               String::from_str(startStr).append("0").as_slice());
            build_encoding_table(data.right, table,
                               String::from_str(startStr).append("1").as_slice());
        },
        HLeaf(ch)   => {table.insert(ch, String::from_str(startStr));}
    };
}

// Attempts to construct a tree, and test that the construction is successful
//    7
//   ----
//  /    \
// 4:'4'  3
//      -----
//     /     \
//    2:'2'  1:'1'
#[test]
fn test_tree_construction() {
    let to_encode = "4444221";
    let tree = huffman_tree(to_encode);
    assert!(tree.weight == 7);
    let children = match tree.item {
        HTree(data) => data,
        HLeaf(_)    => fail!("Tree Missing Children!"),
    };
    let left  = &children.left;
    let right = &children.right;
    assert!(right.weight == 4);
    assert!(left.weight == 3);
    let rightChar = match right.item {
        HTree(_)  => fail!("Node is not Leaf Node!"),
        HLeaf(ch) => ch,
    };
    assert!(rightChar == '4');
    let children = match left.item {
        HTree(ref data) => data,
        HLeaf(_)    => fail!("Tree Missing Children!"),
    };
    let left  = &children.left;
    let right = &children.right;
    let leftChar = match left.item {
        HTree(_)  => fail!("Node is not Leaf Node!"),
        HLeaf(ch) => ch,
    };
    let rightChar = match right.item {
        HTree(_)  => fail!("Node is not Leaf Node!"),
        HLeaf(ch) => ch,
    };
    match (left.weight, right.weight) {
        (1, 2) => {
            assert!(leftChar == '1');
            assert!(rightChar == '2');
        },
        (2, 1) => {
            assert!(leftChar == '2');
            assert!(rightChar == '1');
        },
        (_, _) => {
            fail!("Incorrect Leaf Nodes");
        }
    };
}

#[test]
// Constructs a table:
//  '4': 1
//  '2': 01 OR 00
//  '1': 00    01
// And tests that the table was correctly constructed
fn test_table_construction() {
    let to_encode = "4444221";
    let tree = huffman_tree(to_encode);
    let mut table = HashMap::<char,String>::new();
    build_encoding_table(&tree, &mut table, "");
    let one  = table.get(&'1').as_slice();
    let two  = table.get(&'2').as_slice();
    let four = table.get(&'4').as_slice();
    assert!(four == "1");
    assert!((one == "01" && two == "00") ||
            (one == "00" && two == "01"));
}

#[cfg(not(test))]
fn main() {
    let to_encode = "this is an example for huffman encoding";
    let tree = huffman_tree(to_encode);
    let mut table = HashMap::<char,String>::new();
    build_encoding_table(&tree, &mut table, "");

    for (ch, encoding) in table.iter() {
        println!("{}: {}", *ch, encoding);
    }
}
<commit_msg>small fix to case of variables (Rust coding convention)<commit_after>// Implement data structures for a Huffman encoding tree:
//   http://rosettacode.org/wiki/Huffman_coding

extern crate core;
use std::collections::HashMap;
use std::collections::priority_queue::PriorityQueue;

// Each HNode has a weight, representing the sum of the frequencies for all its
// children. It is either a leaf (containing a character), or a HTree
// (containing two children)
struct HNode {
    weight: int,
    item: HTreeOrHLeaf,
}

enum HTreeOrHLeaf {
    HTree(HTreeData),
    HLeaf(char),
}

struct HTreeData {
    left: Box<HNode>,
    right: Box<HNode>,
}

// Implementing comparison traits (Ord and all its dependencies) such that
// the HNode with the greatest weight is the smallest in a comparison. Basically
// reversing all the comparison operators.
impl Ord for HNode {
    fn cmp(&self, other: &HNode) -> Ordering {
        match self.weight.cmp(&other.weight) {
            Less    => Greater,
            Equal   => Equal,
            Greater => Less,
        }
    }
}

impl Eq for HNode {}
impl PartialEq for HNode {
    fn eq(&self, other: &HNode) -> bool {
        self.weight == other.weight
    }
}

impl PartialOrd for HNode {
    fn lt(&self, other: &HNode) -> bool {
        self.weight > other.weight
    }
}

// Takes a non-empty string (function will fail if string is empty) and computes
// the Huffman encoding tree for that string.
fn huffman_tree(input: &str) -> HNode {
    // 1. Loop through all the characters in that string, adding them to a HashMap
    //    of character to frequency.
    let mut freq = HashMap::new();
    for ch in input.chars() {
        freq.insert_or_update_with(ch, 1, |_k, v: &mut int| {*v += 1;});
    }

    // 2. For each (character, frequency) pair in the HashMap, add a Leaf to a
    //    PriorityQueue
    let mut queue = PriorityQueue::<HNode>::new();
    for (ch, freq) in freq.iter() {
        let new_node = HNode{
            weight: *freq,
            item: HLeaf(*ch),
        };
        queue.push(new_node);
    }

    // 3. Pop two items with the least weight from the queue, combine them into
    //    a tree as children. The parent node's weight is the sum of the
    //    children's weight. Continue until one item is left on the queue, and
    //    return that item.
    while queue.len() > 1 {
        let item1 = queue.pop().unwrap();
        let item2 = queue.pop().unwrap();
        let new_node = HNode {
            weight: item1.weight + item2.weight,
            item: HTree(HTreeData{
                left: box item1,
                right: box item2,
            }),
        };
        queue.push(new_node);
    }
    queue.pop().unwrap()
}

// Takes a Huffman Tree, traverse it and build a table with each character and
// its encoding string.
fn build_encoding_table(tree: &HNode,
                      table: &mut HashMap<char,String>,
                      startStr: &str) {
    match tree.item {
        HTree(ref data) => {
            build_encoding_table(data.left, table,
                               String::from_str(startStr).append("0").as_slice());
            build_encoding_table(data.right, table,
                               String::from_str(startStr).append("1").as_slice());
        },
        HLeaf(ch)   => {table.insert(ch, String::from_str(startStr));}
    };
}

// Attempts to construct a tree, and test that the construction is successful
//    7
//   ----
//  /    \
// 4:'4'  3
//      -----
//     /     \
//    2:'2'  1:'1'
#[test]
fn test_tree_construction() {
    let to_encode = "4444221";
    let tree = huffman_tree(to_encode);
    assert!(tree.weight == 7);
    let children = match tree.item {
        HTree(data) => data,
        HLeaf(_)    => fail!("Tree Missing Children!"),
    };
    let left  = &children.left;
    let right = &children.right;
    assert!(right.weight == 4);
    assert!(left.weight == 3);
    let right_char = match right.item {
        HTree(_)  => fail!("Node is not Leaf Node!"),
        HLeaf(ch) => ch,
    };
    assert!(right_char == '4');
    let children = match left.item {
        HTree(ref data) => data,
        HLeaf(_)    => fail!("Tree Missing Children!"),
    };
    let left  = &children.left;
    let right = &children.right;
    let left_char = match left.item {
        HTree(_)  => fail!("Node is not Leaf Node!"),
        HLeaf(ch) => ch,
    };
    let right_char = match right.item {
        HTree(_)  => fail!("Node is not Leaf Node!"),
        HLeaf(ch) => ch,
    };
    match (left.weight, right.weight) {
        (1, 2) => {
            assert!(left_char == '1');
            assert!(right_char == '2');
        },
        (2, 1) => {
            assert!(left_char == '2');
            assert!(right_char == '1');
        },
        (_, _) => {
            fail!("Incorrect Leaf Nodes");
        }
    };
}

#[test]
// Constructs a table:
//  '4': 1
//  '2': 01 OR 00
//  '1': 00    01
// And tests that the table was correctly constructed
fn test_table_construction() {
    let to_encode = "4444221";
    let tree = huffman_tree(to_encode);
    let mut table = HashMap::<char,String>::new();
    build_encoding_table(&tree, &mut table, "");
    let one  = table.get(&'1').as_slice();
    let two  = table.get(&'2').as_slice();
    let four = table.get(&'4').as_slice();
    assert!(four == "1");
    assert!((one == "01" && two == "00") ||
            (one == "00" && two == "01"));
}

#[cfg(not(test))]
fn main() {
    let to_encode = "this is an example for huffman encoding";
    let tree = huffman_tree(to_encode);
    let mut table = HashMap::<char,String>::new();
    build_encoding_table(&tree, &mut table, "");

    for (ch, encoding) in table.iter() {
        println!("{}: {}", *ch, encoding);
    }
}
<|endoftext|>


<commit_before>//  Copyright (C) 2021 Pascal Engélibert
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use crate::*;

#[derive(Clone, Copy, Debug, Default, Eq, Ord, PartialEq, PartialOrd)]
pub struct BlockCursor {
    pub number: BlockNumber,
}
impl std::fmt::Display for BlockCursor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.number)
    }
}

impl FromStr for BlockCursor {
    type Err = WrongCursor;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(Self {
            number: s.parse().map_err(|_| WrongCursor)?,
        })
    }
}

impl DbsReader {
    pub fn block(
        &self,
        bc_db: &BcV2DbRo<FileBackend>,
        number: U32BE,
    ) -> KvResult<Option<duniter_dbs::BlockMetaV2>> {
        bc_db.blocks_meta().get(&number)
    }

    pub fn blocks(
        &self,
        bc_db: &BcV2DbRo<FileBackend>,
        page_info: PageInfo<BlockCursor>,
    ) -> KvResult<PagedData<Vec<(BlockCursor, duniter_dbs::BlockMetaV2)>>> {
        let last_block_number = bc_db
            .blocks_meta()
            .iter_rev(.., |it| it.values().next_res())?
            .ok_or_else(|| KvError::Custom("Empty blockchain".into()))?
            .number;

        let first_cursor_opt = if page_info.not_all() {
            Some(BlockCursor {
                number: BlockNumber(0),
            })
        } else {
            None
        };

        let last_cursor_opt = if page_info.not_all() {
            Some(BlockCursor {
                number: BlockNumber(last_block_number),
            })
        } else {
            None
        };

        let k_min = U32BE(if page_info.order {
            page_info.pos.map_or_else(|| 0, |pos| pos.number.0)
        } else {
            page_info.limit_opt.map_or_else(
                || 0,
                |limit| {
                    page_info
                        .pos
                        .map_or_else(last_block_number + 1, |pos| pos.number.0)
                        .saturating_sub(limit.get() as u32 - 1)
                },
            )
        });
        let k_max = U32BE(if page_info.order {
            page_info.limit_opt.map_or_else(
                || last_block_number + 1,
                |limit| {
                    page_info.pos.map_or_else(
                        || limit.get() as u32,
                        |pos| pos.number.0.saturating_add(limit.get() as u32),
                    )
                },
            )
        } else {
            page_info.pos.map_or_else(
                || last_block_number + 1,
                |pos| pos.number.0.saturating_add(1),
            )
        });

        let blocks: Vec<(BlockCursor, duniter_dbs::BlockMetaV2)> = if page_info.order {
            bc_db.blocks_meta().iter(k_min..k_max, blocks_inner)?
        } else {
            bc_db.blocks_meta().iter_rev(k_min..k_max, blocks_inner)?
        };

        Ok(PagedData {
            has_next_page: has_next_page(
                blocks
                    .iter()
                    .map(|(block_cursor, _block)| block_cursor.into()),
                last_cursor_opt,
                page_info,
                page_info.order,
            ),
            has_previous_page: has_previous_page(
                blocks
                    .iter()
                    .map(|(block_cursor, _block)| block_cursor.into()),
                first_cursor_opt,
                page_info,
                page_info.order,
            ),
            data: blocks,
        })
    }
}

fn blocks_inner<I>(blocks_iter: I) -> KvResult<Vec<(BlockCursor, duniter_dbs::BlockMetaV2)>>
where
    I: Iterator<Item = KvResult<(U32BE, BlockMetaV2)>>,
{
    blocks_iter
        .map(|block_res| {
            block_res.map(|block| {
                (
                    BlockCursor {
                        number: BlockNumber(block.0 .0),
                    },
                    block.1,
                )
            })
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use duniter_dbs::databases::bc_v2::BcV2DbWritable;
    use duniter_gva_db::GvaV1DbWritable;
    use std::num::NonZeroUsize;

    #[test]
    fn test_block() -> KvResult<()> {
        let bc_db = duniter_dbs::databases::bc_v2::BcV2Db::<Mem>::open(MemConf::default())?;
        let gva_db = duniter_gva_db::GvaV1Db::<Mem>::open(MemConf::default())?;
        let bc_db_ro = bc_db.get_ro_handler();
        let db_reader = create_dbs_reader(unsafe { std::mem::transmute(&gva_db.get_ro_handler()) });

        bc_db
            .blocks_meta_write()
            .upsert(U32BE(0), duniter_dbs::BlockMetaV2::default())?;

        assert_eq!(
            db_reader.block(&bc_db_ro, U32BE(0))?,
            Some(duniter_dbs::BlockMetaV2::default())
        );

        Ok(())
    }

    #[test]
    fn test_blocks() -> KvResult<()> {
        let bc_db = duniter_dbs::databases::bc_v2::BcV2Db::<Mem>::open(MemConf::default())?;
        let gva_db = duniter_gva_db::GvaV1Db::<Mem>::open(MemConf::default())?;
        let bc_db_ro = bc_db.get_ro_handler();
        let db_reader = create_dbs_reader(unsafe { std::mem::transmute(&gva_db.get_ro_handler()) });

        for i in 0..20 {
            bc_db.blocks_meta_write().upsert(
                U32BE(i),
                duniter_dbs::BlockMetaV2 {
                    number: i,
                    ..Default::default()
                },
            )?;
        }

        let blocks = db_reader.blocks(
            &bc_db_ro,
            PageInfo {
                pos: Some(BlockCursor {
                    number: BlockNumber(10),
                }),
                order: true,
                limit_opt: NonZeroUsize::new(3),
            },
        )?;

        assert_eq!(blocks.data.len(), 3);
        assert_eq!(blocks.data[0].1.number, 10);
        assert_eq!(blocks.data[1].1.number, 11);
        assert_eq!(blocks.data[2].1.number, 12);
        assert!(blocks.has_previous_page);
        assert!(blocks.has_next_page);

        let blocks = db_reader.blocks(
            &bc_db_ro,
            PageInfo {
                pos: Some(BlockCursor {
                    number: BlockNumber(10),
                }),
                order: false,
                limit_opt: NonZeroUsize::new(3),
            },
        )?;

        assert_eq!(blocks.data.len(), 3);
        assert_eq!(blocks.data[0].1.number, 10);
        assert_eq!(blocks.data[1].1.number, 9);
        assert_eq!(blocks.data[2].1.number, 8);
        assert!(blocks.has_previous_page);
        assert!(blocks.has_next_page);

        Ok(())
    }
}
<commit_msg>[fix] gva block compile error<commit_after>//  Copyright (C) 2021 Pascal Engélibert
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use crate::*;

#[derive(Clone, Copy, Debug, Default, Eq, Ord, PartialEq, PartialOrd)]
pub struct BlockCursor {
    pub number: BlockNumber,
}
impl std::fmt::Display for BlockCursor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.number)
    }
}

impl FromStr for BlockCursor {
    type Err = WrongCursor;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(Self {
            number: s.parse().map_err(|_| WrongCursor)?,
        })
    }
}

impl DbsReader {
    pub fn block(
        &self,
        bc_db: &BcV2DbRo<FileBackend>,
        number: U32BE,
    ) -> KvResult<Option<duniter_dbs::BlockMetaV2>> {
        bc_db.blocks_meta().get(&number)
    }

    pub fn blocks(
        &self,
        bc_db: &BcV2DbRo<FileBackend>,
        page_info: PageInfo<BlockCursor>,
    ) -> KvResult<PagedData<Vec<(BlockCursor, duniter_dbs::BlockMetaV2)>>> {
        let last_block_number = bc_db
            .blocks_meta()
            .iter_rev(.., |it| it.values().next_res())?
            .ok_or_else(|| KvError::Custom("Empty blockchain".into()))?
            .number;

        let first_cursor_opt = if page_info.not_all() {
            Some(BlockCursor {
                number: BlockNumber(0),
            })
        } else {
            None
        };

        let last_cursor_opt = if page_info.not_all() {
            Some(BlockCursor {
                number: BlockNumber(last_block_number),
            })
        } else {
            None
        };

        let k_min = U32BE(if page_info.order {
            page_info.pos.map_or_else(|| 0, |pos| pos.number.0)
        } else {
            page_info.limit_opt.map_or_else(
                || 0,
                |limit| {
                    page_info
                        .pos
                        .map_or_else(|| last_block_number + 1, |pos| pos.number.0)
                        .saturating_sub(limit.get() as u32 - 1)
                },
            )
        });
        let k_max = U32BE(if page_info.order {
            page_info.limit_opt.map_or_else(
                || last_block_number + 1,
                |limit| {
                    page_info.pos.map_or_else(
                        || limit.get() as u32,
                        |pos| pos.number.0.saturating_add(limit.get() as u32),
                    )
                },
            )
        } else {
            page_info.pos.map_or_else(
                || last_block_number + 1,
                |pos| pos.number.0.saturating_add(1),
            )
        });

        let blocks: Vec<(BlockCursor, duniter_dbs::BlockMetaV2)> = if page_info.order {
            bc_db.blocks_meta().iter(k_min..k_max, blocks_inner)?
        } else {
            bc_db.blocks_meta().iter_rev(k_min..k_max, blocks_inner)?
        };

        Ok(PagedData {
            has_next_page: has_next_page(
                blocks
                    .iter()
                    .map(|(block_cursor, _block)| block_cursor.into()),
                last_cursor_opt,
                page_info,
                page_info.order,
            ),
            has_previous_page: has_previous_page(
                blocks
                    .iter()
                    .map(|(block_cursor, _block)| block_cursor.into()),
                first_cursor_opt,
                page_info,
                page_info.order,
            ),
            data: blocks,
        })
    }
}

fn blocks_inner<I>(blocks_iter: I) -> KvResult<Vec<(BlockCursor, duniter_dbs::BlockMetaV2)>>
where
    I: Iterator<Item = KvResult<(U32BE, BlockMetaV2)>>,
{
    blocks_iter
        .map(|block_res| {
            block_res.map(|block| {
                (
                    BlockCursor {
                        number: BlockNumber(block.0 .0),
                    },
                    block.1,
                )
            })
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use duniter_dbs::databases::bc_v2::BcV2DbWritable;
    use duniter_gva_db::GvaV1DbWritable;
    use std::num::NonZeroUsize;

    #[test]
    fn test_block() -> KvResult<()> {
        let bc_db = duniter_dbs::databases::bc_v2::BcV2Db::<Mem>::open(MemConf::default())?;
        let gva_db = duniter_gva_db::GvaV1Db::<Mem>::open(MemConf::default())?;
        let bc_db_ro = bc_db.get_ro_handler();
        let db_reader = create_dbs_reader(unsafe { std::mem::transmute(&gva_db.get_ro_handler()) });

        bc_db
            .blocks_meta_write()
            .upsert(U32BE(0), duniter_dbs::BlockMetaV2::default())?;

        assert_eq!(
            db_reader.block(&bc_db_ro, U32BE(0))?,
            Some(duniter_dbs::BlockMetaV2::default())
        );

        Ok(())
    }

    #[test]
    fn test_blocks() -> KvResult<()> {
        let bc_db = duniter_dbs::databases::bc_v2::BcV2Db::<Mem>::open(MemConf::default())?;
        let gva_db = duniter_gva_db::GvaV1Db::<Mem>::open(MemConf::default())?;
        let bc_db_ro = bc_db.get_ro_handler();
        let db_reader = create_dbs_reader(unsafe { std::mem::transmute(&gva_db.get_ro_handler()) });

        for i in 0..20 {
            bc_db.blocks_meta_write().upsert(
                U32BE(i),
                duniter_dbs::BlockMetaV2 {
                    number: i,
                    ..Default::default()
                },
            )?;
        }

        let blocks = db_reader.blocks(
            &bc_db_ro,
            PageInfo {
                pos: Some(BlockCursor {
                    number: BlockNumber(10),
                }),
                order: true,
                limit_opt: NonZeroUsize::new(3),
            },
        )?;

        assert_eq!(blocks.data.len(), 3);
        assert_eq!(blocks.data[0].1.number, 10);
        assert_eq!(blocks.data[1].1.number, 11);
        assert_eq!(blocks.data[2].1.number, 12);
        assert!(blocks.has_previous_page);
        assert!(blocks.has_next_page);

        let blocks = db_reader.blocks(
            &bc_db_ro,
            PageInfo {
                pos: Some(BlockCursor {
                    number: BlockNumber(10),
                }),
                order: false,
                limit_opt: NonZeroUsize::new(3),
            },
        )?;

        assert_eq!(blocks.data.len(), 3);
        assert_eq!(blocks.data[0].1.number, 10);
        assert_eq!(blocks.data[1].1.number, 9);
        assert_eq!(blocks.data[2].1.number, 8);
        assert!(blocks.has_previous_page);
        assert!(blocks.has_next_page);

        Ok(())
    }
}
<|endoftext|>


<commit_before>use clap::Parser;
use dicom_core::{dicom_value, smallvec};
use dicom_core::{DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_dump::DumpOptions;
use dicom_encoding::transfer_syntax;
use dicom_object::{mem::InMemDicomObject, open_file, StandardDataDictionary};
use dicom_transfer_syntax_registry::{entries, TransferSyntaxRegistry};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use query::parse_queries;
use smallvec::smallvec;
use snafu::{prelude::*, ErrorCompat};
use std::io::{stderr, Read};
use std::path::PathBuf;
use tracing::{debug, error, info, warn, Level};
use transfer_syntax::TransferSyntaxIndex;

mod query;

/// DICOM C-FIND SCU
#[derive(Debug, Parser)]
struct App {
    /// socket address to FIND SCP (example: "127.0.0.1:1045")
    addr: String,
    /// a DICOM file representing the query object
    file: Option<PathBuf>,
    /// sequence of queries
    #[clap(short('q'))]
    query: Vec<String>,

    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling AE title
    #[clap(long = "calling-ae-title", default_value = "FIND-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[clap(long = "called-ae-title")]
    called_ae_title: Option<String>,
    /// the maximum PDU length
    #[clap(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,

    /// use patient root information model
    #[clap(long, conflicts_with = "study")]
    patient: bool,
    /// use study root information model (default)
    #[clap(long, conflicts_with = "patient")]
    study: bool,
}

fn report<E: 'static>(err: &E)
where
    E: std::error::Error,
{
    error!("{}", err);
    if let Some(source) = err.source() {
        eprintln!();
        eprintln!("Caused by:");
        for (i, e) in std::iter::successors(Some(source), |e| e.source()).enumerate() {
            eprintln!("   {}: {}", i, e);
        }
    }
}

fn report_backtrace<E: 'static>(err: &E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    let env_backtrace = std::env::var("RUST_BACKTRACE").unwrap_or_default();
    let env_lib_backtrace = std::env::var("RUST_LIB_BACKTRACE").unwrap_or_default();
    if env_lib_backtrace == "1" || (env_backtrace == "1" && env_lib_backtrace != "0") {
        if let Some(backtrace) = ErrorCompat::backtrace(&err) {
            eprintln!();
            eprintln!("Backtrace:");
            eprintln!("{}", backtrace);
        }
    }
}

fn report_with_backtrace<E: 'static>(err: E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    report(&err);
    report_backtrace(&err);
}

fn main() {
    run().unwrap_or_else(|err| {
        report_with_backtrace(err);
        std::process::exit(-2);
    });
}

#[derive(Debug, Snafu)]
enum Error {
    /// Could not initialize SCU
    InitScu {
        source: dicom_ul::association::client::Error,
    },

    /// Could not construct DICOM command
    CreateCommand { source: dicom_object::Error },

    /// Could not read DICOM command
    ReadCommand { source: dicom_object::Error },

    /// Could not dump DICOM output
    DumpOutput { source: std::io::Error },

    #[snafu(whatever, display("{}", message))]
    Other {
        message: String,
        #[snafu(source(from(Box<dyn std::error::Error + 'static>, Some)))]
        source: Option<Box<dyn std::error::Error + 'static>>,
    },
}

fn build_query(
    file: Option<PathBuf>,
    q: Vec<String>,
    patient: bool,
    study: bool,
    verbose: bool,
) -> Result<InMemDicomObject, Error> {
    match (file, q) {
        (Some(file), q) => {
            if !q.is_empty() {
                whatever!("Conflicted file with query terms");
            }

            if verbose {
                info!("Opening file '{}'...", file.display());
            }

            open_file(file)
                .context(CreateCommandSnafu)
                .map(|file| file.into_inner())
        }
        (None, q) => {
            if q.is_empty() {
                whatever!("Query not specified");
            }

            let mut obj =
                parse_queries(&q).whatever_context("Could not build query object from terms")?;

            // (0008,0052) CS QueryRetrieveLevel
            let level = match (patient, study) {
                (true, false) => "PATIENT",
                (false, true) |
                (false, false) => "STUDY",
                _ => unreachable!(),
            };
            obj.put(DataElement::new(
                tags::QUERY_RETRIEVE_LEVEL,
                VR::CS,
                PrimitiveValue::from(level),
            ));

            Ok(obj)
        }
    }
}

fn run() -> Result<(), Error> {
    let App {
        addr,
        file,
        verbose,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
        patient,
        study,
        query,
    } = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if verbose { Level::DEBUG } else { Level::INFO })
            .finish(),
    )
    .unwrap_or_else(|e| {
        report(&e);
    });

    let dcm_query = build_query(file, query, patient, study, verbose)?;

    let abstract_syntax = match (patient, study) {
        // Patient Root Query/Retrieve Information Model - FIND
        (true, false) => "1.2.840.10008.5.1.4.1.2.1.1",
        // Study Root Query/Retrieve Information Model – FIND (default)
        (false, false) | (false, true) => "1.2.840.10008.5.1.4.1.2.2.1",
        // Series
        _ => unreachable!("Unexpected flag combination"),
    };

    if verbose {
        info!("Establishing association with '{}'...", &addr);
    }

    let mut scu_opt = ClientAssociationOptions::new()
        .with_abstract_syntax(abstract_syntax)
        .calling_ae_title(calling_ae_title)
        .max_pdu_length(max_pdu_length);

    if let Some(called_ae_title) = called_ae_title {
        scu_opt = scu_opt.called_ae_title(called_ae_title);
    }

    let mut scu = scu_opt.establish_with(&addr).context(InitScuSnafu)?;

    if verbose {
        info!("Association established");
    }

    let pc_selected = if let Some(pc_selected) = scu.presentation_contexts().first() {
        pc_selected
    } else {
        error!("Could not choose a presentation context");
        let _ = scu.abort();
        std::process::exit(-2);
    };
    let pc_selected_id = pc_selected.id;

    let ts = if let Some(ts) = TransferSyntaxRegistry.get(&pc_selected.transfer_syntax) {
        ts
    } else {
        error!("Poorly negotiated transfer syntax");
        let _ = scu.abort();
        std::process::exit(-2);
    };

    if verbose {
        debug!("Transfer Syntax: {}", ts.name());
    }

    let cmd = find_req_command(abstract_syntax, 1);

    let mut cmd_data = Vec::with_capacity(128);
    cmd.write_dataset_with_ts(&mut cmd_data, &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased())
        .whatever_context("Failed to write command")?;

    let implicit_vr_le = entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

    let mut iod_data = Vec::with_capacity(128);
    dcm_query
        .write_dataset_with_ts(&mut iod_data, &implicit_vr_le)
        .whatever_context("failed to write identifier dataset")?;

    let nbytes = cmd_data.len() + iod_data.len();

    if verbose {
        debug!("Sending query ({} B)...", nbytes);
    }

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Command,
            is_last: true,
            data: cmd_data,
        }],
    };
    scu.send(&pdu).whatever_context("Could not send command")?;

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Data,
            is_last: true,
            data: iod_data,
        }],
    };
    scu.send(&pdu)
        .whatever_context("Could not send C-Find request")?;

    if verbose {
        debug!("Awaiting response...");
    }

    let mut i = 0;
    loop {
        let rsp_pdu = scu
            .receive()
            .whatever_context("Failed to receive response from remote node")?;

        match rsp_pdu {
            Pdu::PData { data } => {
                let data_value = &data[0];

                let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                    &data_value.data[..],
                    &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
                )
                .context(ReadCommandSnafu)?;
                if verbose {
                    eprintln!("Match #{} Response command:", i);
                    DumpOptions::new()
                        .dump_object_to(stderr(), &cmd_obj)
                        .context(DumpOutputSnafu)?;
                }
                let status = cmd_obj
                    .element(tags::STATUS)
                    .whatever_context("status code from response is missing")?
                    .to_int::<u16>()
                    .whatever_context("failed to read status code")?;
                if status == 0 {
                    if verbose {
                        debug!("Matching is complete");
                    }
                    if i == 0 {
                        info!("No results matching query");
                    }
                    break;
                } else if status == 0xFF00 || status == 0xFF01 {
                    if verbose {
                        debug!("Operation pending: {:x}", status);
                    }

                    // fetch DICOM data

                    let dcm = {
                        let mut rsp = scu.receive_pdata();
                        let mut response_data = Vec::new();
                        rsp.read_to_end(&mut response_data)
                            .whatever_context("Failed to read response data")?;

                        InMemDicomObject::read_dataset_with_ts(&response_data[..], &implicit_vr_le)
                            .whatever_context("Could not read response data set")?
                    };

                    println!("------------ Match #{} ------------", i);
                    DumpOptions::new()
                        .dump_object(&dcm)
                        .context(DumpOutputSnafu)?;
                    i += 1;
                } else {
                    warn!("Operation failed (status code {})", status);
                    break;
                }
            }

            pdu @ Pdu::Unknown { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::AssociationAC { .. }
            | pdu @ Pdu::AssociationRJ { .. }
            | pdu @ Pdu::ReleaseRQ
            | pdu @ Pdu::ReleaseRP
            | pdu @ Pdu::AbortRQ { .. } => {
                error!("Unexpected SCP response: {:?}", pdu);
                let _ = scu.abort();
                std::process::exit(-2);
            }
        }
    }
    let _ = scu.release();

    Ok(())
}

fn find_req_command(
    sop_class_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + even_len(sop_class_uid.len())   // SOP Class UID
            + 8 + 2 // command field
            + 8 + 2 // message ID
            + 8 + 2 // priority
            + 8 + 2, // data set type
        ),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        PrimitiveValue::from(sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        // 0020H: C-FIND-RQ message
        dicom_value!(U16, [0x0020]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        // medium
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    obj
}

fn even_len(l: usize) -> u32 {
    ((l + 1) & !1) as u32
}

#[cfg(test)]
mod tests {
    use super::even_len;
    #[test]
    fn test_even_len() {
        assert_eq!(even_len(0), 0);
        assert_eq!(even_len(1), 2);
        assert_eq!(even_len(2), 2);
        assert_eq!(even_len(3), 4);
        assert_eq!(even_len(4), 4);
        assert_eq!(even_len(5), 6);
    }
}
<commit_msg>[findscu] Add short versions of --patient and --study<commit_after>use clap::Parser;
use dicom_core::{dicom_value, smallvec};
use dicom_core::{DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_dump::DumpOptions;
use dicom_encoding::transfer_syntax;
use dicom_object::{mem::InMemDicomObject, open_file, StandardDataDictionary};
use dicom_transfer_syntax_registry::{entries, TransferSyntaxRegistry};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use query::parse_queries;
use smallvec::smallvec;
use snafu::{prelude::*, ErrorCompat};
use std::io::{stderr, Read};
use std::path::PathBuf;
use tracing::{debug, error, info, warn, Level};
use transfer_syntax::TransferSyntaxIndex;

mod query;

/// DICOM C-FIND SCU
#[derive(Debug, Parser)]
struct App {
    /// socket address to FIND SCP (example: "127.0.0.1:1045")
    addr: String,
    /// a DICOM file representing the query object
    file: Option<PathBuf>,
    /// sequence of queries
    #[clap(short('q'))]
    query: Vec<String>,

    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling AE title
    #[clap(long = "calling-ae-title", default_value = "FIND-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[clap(long = "called-ae-title")]
    called_ae_title: Option<String>,
    /// the maximum PDU length
    #[clap(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,

    /// use patient root information model
    #[clap(short = 'P', long, conflicts_with = "study")]
    patient: bool,
    /// use study root information model (default)
    #[clap(short = 'S', long, conflicts_with = "patient")]
    study: bool,
}

fn report<E: 'static>(err: &E)
where
    E: std::error::Error,
{
    error!("{}", err);
    if let Some(source) = err.source() {
        eprintln!();
        eprintln!("Caused by:");
        for (i, e) in std::iter::successors(Some(source), |e| e.source()).enumerate() {
            eprintln!("   {}: {}", i, e);
        }
    }
}

fn report_backtrace<E: 'static>(err: &E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    let env_backtrace = std::env::var("RUST_BACKTRACE").unwrap_or_default();
    let env_lib_backtrace = std::env::var("RUST_LIB_BACKTRACE").unwrap_or_default();
    if env_lib_backtrace == "1" || (env_backtrace == "1" && env_lib_backtrace != "0") {
        if let Some(backtrace) = ErrorCompat::backtrace(&err) {
            eprintln!();
            eprintln!("Backtrace:");
            eprintln!("{}", backtrace);
        }
    }
}

fn report_with_backtrace<E: 'static>(err: E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    report(&err);
    report_backtrace(&err);
}

fn main() {
    run().unwrap_or_else(|err| {
        report_with_backtrace(err);
        std::process::exit(-2);
    });
}

#[derive(Debug, Snafu)]
enum Error {
    /// Could not initialize SCU
    InitScu {
        source: dicom_ul::association::client::Error,
    },

    /// Could not construct DICOM command
    CreateCommand { source: dicom_object::Error },

    /// Could not read DICOM command
    ReadCommand { source: dicom_object::Error },

    /// Could not dump DICOM output
    DumpOutput { source: std::io::Error },

    #[snafu(whatever, display("{}", message))]
    Other {
        message: String,
        #[snafu(source(from(Box<dyn std::error::Error + 'static>, Some)))]
        source: Option<Box<dyn std::error::Error + 'static>>,
    },
}

fn build_query(
    file: Option<PathBuf>,
    q: Vec<String>,
    patient: bool,
    study: bool,
    verbose: bool,
) -> Result<InMemDicomObject, Error> {
    match (file, q) {
        (Some(file), q) => {
            if !q.is_empty() {
                whatever!("Conflicted file with query terms");
            }

            if verbose {
                info!("Opening file '{}'...", file.display());
            }

            open_file(file)
                .context(CreateCommandSnafu)
                .map(|file| file.into_inner())
        }
        (None, q) => {
            if q.is_empty() {
                whatever!("Query not specified");
            }

            let mut obj =
                parse_queries(&q).whatever_context("Could not build query object from terms")?;

            // (0008,0052) CS QueryRetrieveLevel
            let level = match (patient, study) {
                (true, false) => "PATIENT",
                (false, true) |
                (false, false) => "STUDY",
                _ => unreachable!(),
            };
            obj.put(DataElement::new(
                tags::QUERY_RETRIEVE_LEVEL,
                VR::CS,
                PrimitiveValue::from(level),
            ));

            Ok(obj)
        }
    }
}

fn run() -> Result<(), Error> {
    let App {
        addr,
        file,
        verbose,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
        patient,
        study,
        query,
    } = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if verbose { Level::DEBUG } else { Level::INFO })
            .finish(),
    )
    .unwrap_or_else(|e| {
        report(&e);
    });

    let dcm_query = build_query(file, query, patient, study, verbose)?;

    let abstract_syntax = match (patient, study) {
        // Patient Root Query/Retrieve Information Model - FIND
        (true, false) => "1.2.840.10008.5.1.4.1.2.1.1",
        // Study Root Query/Retrieve Information Model – FIND (default)
        (false, false) | (false, true) => "1.2.840.10008.5.1.4.1.2.2.1",
        // Series
        _ => unreachable!("Unexpected flag combination"),
    };

    if verbose {
        info!("Establishing association with '{}'...", &addr);
    }

    let mut scu_opt = ClientAssociationOptions::new()
        .with_abstract_syntax(abstract_syntax)
        .calling_ae_title(calling_ae_title)
        .max_pdu_length(max_pdu_length);

    if let Some(called_ae_title) = called_ae_title {
        scu_opt = scu_opt.called_ae_title(called_ae_title);
    }

    let mut scu = scu_opt.establish_with(&addr).context(InitScuSnafu)?;

    if verbose {
        info!("Association established");
    }

    let pc_selected = if let Some(pc_selected) = scu.presentation_contexts().first() {
        pc_selected
    } else {
        error!("Could not choose a presentation context");
        let _ = scu.abort();
        std::process::exit(-2);
    };
    let pc_selected_id = pc_selected.id;

    let ts = if let Some(ts) = TransferSyntaxRegistry.get(&pc_selected.transfer_syntax) {
        ts
    } else {
        error!("Poorly negotiated transfer syntax");
        let _ = scu.abort();
        std::process::exit(-2);
    };

    if verbose {
        debug!("Transfer Syntax: {}", ts.name());
    }

    let cmd = find_req_command(abstract_syntax, 1);

    let mut cmd_data = Vec::with_capacity(128);
    cmd.write_dataset_with_ts(&mut cmd_data, &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased())
        .whatever_context("Failed to write command")?;

    let implicit_vr_le = entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

    let mut iod_data = Vec::with_capacity(128);
    dcm_query
        .write_dataset_with_ts(&mut iod_data, &implicit_vr_le)
        .whatever_context("failed to write identifier dataset")?;

    let nbytes = cmd_data.len() + iod_data.len();

    if verbose {
        debug!("Sending query ({} B)...", nbytes);
    }

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Command,
            is_last: true,
            data: cmd_data,
        }],
    };
    scu.send(&pdu).whatever_context("Could not send command")?;

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Data,
            is_last: true,
            data: iod_data,
        }],
    };
    scu.send(&pdu)
        .whatever_context("Could not send C-Find request")?;

    if verbose {
        debug!("Awaiting response...");
    }

    let mut i = 0;
    loop {
        let rsp_pdu = scu
            .receive()
            .whatever_context("Failed to receive response from remote node")?;

        match rsp_pdu {
            Pdu::PData { data } => {
                let data_value = &data[0];

                let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                    &data_value.data[..],
                    &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
                )
                .context(ReadCommandSnafu)?;
                if verbose {
                    eprintln!("Match #{} Response command:", i);
                    DumpOptions::new()
                        .dump_object_to(stderr(), &cmd_obj)
                        .context(DumpOutputSnafu)?;
                }
                let status = cmd_obj
                    .element(tags::STATUS)
                    .whatever_context("status code from response is missing")?
                    .to_int::<u16>()
                    .whatever_context("failed to read status code")?;
                if status == 0 {
                    if verbose {
                        debug!("Matching is complete");
                    }
                    if i == 0 {
                        info!("No results matching query");
                    }
                    break;
                } else if status == 0xFF00 || status == 0xFF01 {
                    if verbose {
                        debug!("Operation pending: {:x}", status);
                    }

                    // fetch DICOM data

                    let dcm = {
                        let mut rsp = scu.receive_pdata();
                        let mut response_data = Vec::new();
                        rsp.read_to_end(&mut response_data)
                            .whatever_context("Failed to read response data")?;

                        InMemDicomObject::read_dataset_with_ts(&response_data[..], &implicit_vr_le)
                            .whatever_context("Could not read response data set")?
                    };

                    println!("------------ Match #{} ------------", i);
                    DumpOptions::new()
                        .dump_object(&dcm)
                        .context(DumpOutputSnafu)?;
                    i += 1;
                } else {
                    warn!("Operation failed (status code {})", status);
                    break;
                }
            }

            pdu @ Pdu::Unknown { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::AssociationAC { .. }
            | pdu @ Pdu::AssociationRJ { .. }
            | pdu @ Pdu::ReleaseRQ
            | pdu @ Pdu::ReleaseRP
            | pdu @ Pdu::AbortRQ { .. } => {
                error!("Unexpected SCP response: {:?}", pdu);
                let _ = scu.abort();
                std::process::exit(-2);
            }
        }
    }
    let _ = scu.release();

    Ok(())
}

fn find_req_command(
    sop_class_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + even_len(sop_class_uid.len())   // SOP Class UID
            + 8 + 2 // command field
            + 8 + 2 // message ID
            + 8 + 2 // priority
            + 8 + 2, // data set type
        ),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        PrimitiveValue::from(sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        // 0020H: C-FIND-RQ message
        dicom_value!(U16, [0x0020]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        // medium
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    obj
}

fn even_len(l: usize) -> u32 {
    ((l + 1) & !1) as u32
}

#[cfg(test)]
mod tests {
    use super::even_len;
    #[test]
    fn test_even_len() {
        assert_eq!(even_len(0), 0);
        assert_eq!(even_len(1), 2);
        assert_eq!(even_len(2), 2);
        assert_eq!(even_len(3), 4);
        assert_eq!(even_len(4), 4);
        assert_eq!(even_len(5), 6);
    }
}
<|endoftext|>


<commit_before>//! This module integrates the buddy heap allocator into the Rust runtime.
use spin::Mutex;
use core::ptr;

use ::Allocator;
use super::{BuddyHeapAllocator, FreeList};

static ALLOC: Mutex<Option<BuddyHeapAllocator<'static>>>
    = Mutex::new(None);

pub unsafe fn init_heap( start_addr: *mut u8
                       , free_lists: &'static mut [FreeList]
                       , heap_size: usize ) {
    trace!("init_heap() was called.");
    *(ALLOC.lock())
        = Some(BuddyHeapAllocator::new(start_addr, free_lists, heap_size));
}

#[no_mangle]
pub extern "C" fn __rust_allocate(size: usize, align: usize) -> *mut u8 {
    trace!("__rust_allocate() was called.");
    unsafe {
        ALLOC.lock().as_mut()
             .expect("Cannot allocate memory, no system allocator exists!")
             .allocate(size, align)
             .map(|blck| {
                 trace!("__rust_allocate: allocatedd {:?}", blck);
                 blck })
             .unwrap_or(ptr::null_mut())
            //  .expect("Memory could not be allocated; either the allocator is\
            //           out of memory, or the allocation request was invalid.")
    }
}

#[no_mangle]
pub extern "C" fn __rust_deallocate( ptr: *mut u8, old_size: usize
                                   , align: usize ) {
    unsafe {
        ALLOC.lock().as_mut()
             .expect("Cannot deallocate memory, no system allocator exists!")
             .deallocate(ptr, old_size, align)
    }
}

#[no_mangle]
pub extern "C" fn __rust_reallocate( ptr: *mut u8, old_size: usize
                                   , size: usize, align: usize )
                                   -> *mut u8 {
    unsafe {
        ALLOC.lock().as_mut()
             .expect("Cannot reallocate memory, no system allocator exists!")
             .reallocate(ptr, old_size, size, align)
             .unwrap_or(ptr::null_mut())
            //  .expect("Memory could not be reallocated; either the allocator\
            //          is out of memory, or the allocation request was invalid.")
     }
}

/// This is currently unsupported, so we just silently ignore it
/// and return the old size.
#[no_mangle]
pub extern "C" fn __rust_reallocate_inplace( ptr: *mut u8, old_size: usize
                                           , _size: usize, _align: usize )
                                           -> usize {
    old_size
}

#[no_mangle]
pub extern "C" fn __rust_usable_size(size: usize, _align: usize) -> usize {
    size
}
<commit_msg>[alloc] Alloc crate owns free lists now<commit_after>//! This module integrates the buddy heap allocator into the Rust runtime.
use spin::Mutex;
use core::ptr;

use ::Allocator;
use super::{BuddyHeapAllocator, FreeList};

static ALLOC: Mutex<Option<BuddyHeapAllocator<'static>>>
    = Mutex::new(None);

static mut KERNEL_FREE_LISTS: [FreeList; 17]
    // TODO: I really wish there was a less awful way to do this...
    = [ FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      , FreeList::new()
      ];

pub unsafe fn init_heap(start_addr: *mut u8, heap_size: usize ) {
    trace!("init_heap() was called.");
    *(ALLOC.lock())
        = Some(BuddyHeapAllocator::new( start_addr
                                      , &mut KERNEL_FREE_LISTS
                                      , heap_size));
}

#[no_mangle]
pub extern "C" fn __rust_allocate(size: usize, align: usize) -> *mut u8 {
    trace!("__rust_allocate() was called.");
    unsafe {
        ALLOC.lock().as_mut()
             .expect("Cannot allocate memory, no system allocator exists!")
             .allocate(size, align)
             .map(|blck| {
                 trace!("__rust_allocate: allocatedd {:?}", blck);
                 blck })
             .unwrap_or(ptr::null_mut())
            //  .expect("Memory could not be allocated; either the allocator is\
            //           out of memory, or the allocation request was invalid.")
    }
}

#[no_mangle]
pub extern "C" fn __rust_deallocate( ptr: *mut u8, old_size: usize
                                   , align: usize ) {
    unsafe {
        ALLOC.lock().as_mut()
             .expect("Cannot deallocate memory, no system allocator exists!")
             .deallocate(ptr, old_size, align)
    }
}

#[no_mangle]
pub extern "C" fn __rust_reallocate( ptr: *mut u8, old_size: usize
                                   , size: usize, align: usize )
                                   -> *mut u8 {
    unsafe {
        ALLOC.lock().as_mut()
             .expect("Cannot reallocate memory, no system allocator exists!")
             .reallocate(ptr, old_size, size, align)
             .unwrap_or(ptr::null_mut())
            //  .expect("Memory could not be reallocated; either the allocator\
            //          is out of memory, or the allocation request was invalid.")
     }
}

/// This is currently unsupported, so we just silently ignore it
/// and return the old size.
#[no_mangle]
pub extern "C" fn __rust_reallocate_inplace( ptr: *mut u8, old_size: usize
                                           , _size: usize, _align: usize )
                                           -> usize {
    old_size
}

#[no_mangle]
pub extern "C" fn __rust_usable_size(size: usize, _align: usize) -> usize {
    size
}
<|endoftext|>


<commit_before>use super::*;

extern crate test;
use test::Bencher;

#[test]
fn test_new_filled() {
    for i in 0..128 {
        let idx_buf = BitSet::new_filled(i);
        let elems: Vec<usize> = idx_buf.iter().collect();
        let expected: Vec<usize> = (0..i).collect();
        assert_eq!(elems, expected);
    }
}

#[test]
fn bitset_iter_works() {
    let mut bitset: BitSet<usize> = BitSet::new_empty(100);
    bitset.insert(1);
    bitset.insert(10);
    bitset.insert(19);
    bitset.insert(62);
    bitset.insert(63);
    bitset.insert(64);
    bitset.insert(65);
    bitset.insert(66);
    bitset.insert(99);
    assert_eq!(bitset.iter().collect::<Vec<_>>(), [1, 10, 19, 62, 63, 64, 65, 66, 99]);
}

#[test]
fn bitset_iter_works_2() {
    let mut bitset: BitSet<usize> = BitSet::new_empty(320);
    bitset.insert(0);
    bitset.insert(127);
    bitset.insert(191);
    bitset.insert(255);
    bitset.insert(319);
    assert_eq!(bitset.iter().collect::<Vec<_>>(), [0, 127, 191, 255, 319]);
}

#[test]
fn union_two_sets() {
    let mut set1: BitSet<usize> = BitSet::new_empty(65);
    let mut set2: BitSet<usize> = BitSet::new_empty(65);
    assert!(set1.insert(3));
    assert!(!set1.insert(3));
    assert!(set2.insert(5));
    assert!(set2.insert(64));
    assert!(set1.union(&set2));
    assert!(!set1.union(&set2));
    assert!(set1.contains(3));
    assert!(!set1.contains(4));
    assert!(set1.contains(5));
    assert!(!set1.contains(63));
    assert!(set1.contains(64));
}

#[test]
fn hybrid_bitset() {
    let mut sparse038: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    assert!(sparse038.is_empty());
    assert!(sparse038.insert(0));
    assert!(sparse038.insert(1));
    assert!(sparse038.insert(8));
    assert!(sparse038.insert(3));
    assert!(!sparse038.insert(3));
    assert!(sparse038.remove(1));
    assert!(!sparse038.is_empty());
    assert_eq!(sparse038.iter().collect::<Vec<_>>(), [0, 3, 8]);

    for i in 0..256 {
        if i == 0 || i == 3 || i == 8 {
            assert!(sparse038.contains(i));
        } else {
            assert!(!sparse038.contains(i));
        }
    }

    let mut sparse01358 = sparse038.clone();
    assert!(sparse01358.insert(1));
    assert!(sparse01358.insert(5));
    assert_eq!(sparse01358.iter().collect::<Vec<_>>(), [0, 1, 3, 5, 8]);

    let mut dense10 = HybridBitSet::new_empty(256);
    for i in 0..10 {
        assert!(dense10.insert(i));
    }
    assert!(!dense10.is_empty());
    assert_eq!(dense10.iter().collect::<Vec<_>>(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);

    let mut dense256 = HybridBitSet::new_empty(256);
    assert!(dense256.is_empty());
    dense256.insert_all();
    assert!(!dense256.is_empty());
    for i in 0..256 {
        assert!(dense256.contains(i));
    }

    assert!(sparse038.superset(&sparse038)); // sparse + sparse (self)
    assert!(sparse01358.superset(&sparse038)); // sparse + sparse
    assert!(dense10.superset(&sparse038)); // dense + sparse
    assert!(dense10.superset(&dense10)); // dense + dense (self)
    assert!(dense256.superset(&dense10)); // dense + dense

    let mut hybrid = sparse038;
    assert!(!sparse01358.union(&hybrid)); // no change
    assert!(hybrid.union(&sparse01358));
    assert!(hybrid.superset(&sparse01358) && sparse01358.superset(&hybrid));
    assert!(!dense10.union(&sparse01358));
    assert!(!dense256.union(&dense10));
    let mut dense = dense10;
    assert!(dense.union(&dense256));
    assert!(dense.superset(&dense256) && dense256.superset(&dense));
    assert!(hybrid.union(&dense256));
    assert!(hybrid.superset(&dense256) && dense256.superset(&hybrid));

    assert_eq!(dense256.iter().count(), 256);
    let mut dense0 = dense256;
    for i in 0..256 {
        assert!(dense0.remove(i));
    }
    assert!(!dense0.remove(0));
    assert!(dense0.is_empty());
}

#[test]
fn grow() {
    let mut set: GrowableBitSet<usize> = GrowableBitSet::with_capacity(65);
    for index in 0..65 {
        assert!(set.insert(index));
        assert!(!set.insert(index));
    }
    set.ensure(128);

    // Check if the bits set before growing are still set
    for index in 0..65 {
        assert!(set.contains(index));
    }

    // Check if the new bits are all un-set
    for index in 65..128 {
        assert!(!set.contains(index));
    }

    // Check that we can set all new bits without running out of bounds
    for index in 65..128 {
        assert!(set.insert(index));
        assert!(!set.insert(index));
    }
}

#[test]
fn matrix_intersection() {
    let mut matrix: BitMatrix<usize, usize> = BitMatrix::new(200, 200);

    // (*) Elements reachable from both 2 and 65.

    matrix.insert(2, 3);
    matrix.insert(2, 6);
    matrix.insert(2, 10); // (*)
    matrix.insert(2, 64); // (*)
    matrix.insert(2, 65);
    matrix.insert(2, 130);
    matrix.insert(2, 160); // (*)

    matrix.insert(64, 133);

    matrix.insert(65, 2);
    matrix.insert(65, 8);
    matrix.insert(65, 10); // (*)
    matrix.insert(65, 64); // (*)
    matrix.insert(65, 68);
    matrix.insert(65, 133);
    matrix.insert(65, 160); // (*)

    let intersection = matrix.intersect_rows(2, 64);
    assert!(intersection.is_empty());

    let intersection = matrix.intersect_rows(2, 65);
    assert_eq!(intersection, &[10, 64, 160]);
}

#[test]
fn matrix_iter() {
    let mut matrix: BitMatrix<usize, usize> = BitMatrix::new(64, 100);
    matrix.insert(3, 22);
    matrix.insert(3, 75);
    matrix.insert(2, 99);
    matrix.insert(4, 0);
    matrix.union_rows(3, 5);
    matrix.insert_all_into_row(6);

    let expected = [99];
    let mut iter = expected.iter();
    for i in matrix.iter(2) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    assert_eq!(matrix.count(3), expected.len());
    for i in matrix.iter(3) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [0];
    let mut iter = expected.iter();
    assert_eq!(matrix.count(4), expected.len());
    for i in matrix.iter(4) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    assert_eq!(matrix.count(5), expected.len());
    for i in matrix.iter(5) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    assert_eq!(matrix.count(6), 100);
    let mut count = 0;
    for (idx, i) in matrix.iter(6).enumerate() {
        assert_eq!(idx, i);
        count += 1;
    }
    assert_eq!(count, 100);

    if let Some(i) = matrix.iter(7).next() {
        panic!("expected no elements in row, but contains element {:?}", i);
    }
}

#[test]
fn sparse_matrix_iter() {
    let mut matrix: SparseBitMatrix<usize, usize> = SparseBitMatrix::new(100);
    matrix.insert(3, 22);
    matrix.insert(3, 75);
    matrix.insert(2, 99);
    matrix.insert(4, 0);
    matrix.union_rows(3, 5);

    let expected = [99];
    let mut iter = expected.iter();
    for i in matrix.iter(2) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    for i in matrix.iter(3) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [0];
    let mut iter = expected.iter();
    for i in matrix.iter(4) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    for i in matrix.iter(5) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());
}

/// Merge dense hybrid set into empty sparse hybrid set.
#[bench]
fn union_hybrid_sparse_empty_to_dense(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    for i in 0..10 {
        assert!(pre_dense.insert(i));
    }
    let pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}

/// Merge dense hybrid set into full hybrid set with same indices.
#[bench]
fn union_hybrid_sparse_full_to_dense(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    for i in 0..10 {
        assert!(pre_dense.insert(i));
    }
    let mut pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    for i in 0..SPARSE_MAX {
        assert!(pre_sparse.insert(i));
    }
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}

/// Merge dense hybrid set into full hybrid set with indices over the whole domain.
#[bench]
fn union_hybrid_sparse_domain_to_dense(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX * 64);
    for i in 0..10 {
        assert!(pre_dense.insert(i));
    }
    let mut pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX * 64);
    for i in 0..SPARSE_MAX {
        assert!(pre_sparse.insert(i * 64));
    }
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}

/// Merge dense hybrid set into empty hybrid set where the domain is very small.
#[bench]
fn union_hybrid_sparse_empty_small_domain(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX);
    for i in 0..SPARSE_MAX {
        assert!(pre_dense.insert(i));
    }
    let pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX);
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}

/// Merge dense hybrid set into full hybrid set where the domain is very small.
#[bench]
fn union_hybrid_sparse_full_small_domain(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX);
    for i in 0..SPARSE_MAX {
        assert!(pre_dense.insert(i));
    }
    let mut pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX);
    for i in 0..SPARSE_MAX {
        assert!(pre_sparse.insert(i));
    }
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}
<commit_msg>Rollup merge of #81682 - JulianKnodt:bit_set_iter_benchmarks, r=oli-obk<commit_after>use super::*;

extern crate test;
use std::hint::black_box;
use test::Bencher;

#[test]
fn test_new_filled() {
    for i in 0..128 {
        let idx_buf = BitSet::new_filled(i);
        let elems: Vec<usize> = idx_buf.iter().collect();
        let expected: Vec<usize> = (0..i).collect();
        assert_eq!(elems, expected);
    }
}

#[test]
fn bitset_iter_works() {
    let mut bitset: BitSet<usize> = BitSet::new_empty(100);
    bitset.insert(1);
    bitset.insert(10);
    bitset.insert(19);
    bitset.insert(62);
    bitset.insert(63);
    bitset.insert(64);
    bitset.insert(65);
    bitset.insert(66);
    bitset.insert(99);
    assert_eq!(bitset.iter().collect::<Vec<_>>(), [1, 10, 19, 62, 63, 64, 65, 66, 99]);
}

#[test]
fn bitset_iter_works_2() {
    let mut bitset: BitSet<usize> = BitSet::new_empty(320);
    bitset.insert(0);
    bitset.insert(127);
    bitset.insert(191);
    bitset.insert(255);
    bitset.insert(319);
    assert_eq!(bitset.iter().collect::<Vec<_>>(), [0, 127, 191, 255, 319]);
}

#[test]
fn union_two_sets() {
    let mut set1: BitSet<usize> = BitSet::new_empty(65);
    let mut set2: BitSet<usize> = BitSet::new_empty(65);
    assert!(set1.insert(3));
    assert!(!set1.insert(3));
    assert!(set2.insert(5));
    assert!(set2.insert(64));
    assert!(set1.union(&set2));
    assert!(!set1.union(&set2));
    assert!(set1.contains(3));
    assert!(!set1.contains(4));
    assert!(set1.contains(5));
    assert!(!set1.contains(63));
    assert!(set1.contains(64));
}

#[test]
fn hybrid_bitset() {
    let mut sparse038: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    assert!(sparse038.is_empty());
    assert!(sparse038.insert(0));
    assert!(sparse038.insert(1));
    assert!(sparse038.insert(8));
    assert!(sparse038.insert(3));
    assert!(!sparse038.insert(3));
    assert!(sparse038.remove(1));
    assert!(!sparse038.is_empty());
    assert_eq!(sparse038.iter().collect::<Vec<_>>(), [0, 3, 8]);

    for i in 0..256 {
        if i == 0 || i == 3 || i == 8 {
            assert!(sparse038.contains(i));
        } else {
            assert!(!sparse038.contains(i));
        }
    }

    let mut sparse01358 = sparse038.clone();
    assert!(sparse01358.insert(1));
    assert!(sparse01358.insert(5));
    assert_eq!(sparse01358.iter().collect::<Vec<_>>(), [0, 1, 3, 5, 8]);

    let mut dense10 = HybridBitSet::new_empty(256);
    for i in 0..10 {
        assert!(dense10.insert(i));
    }
    assert!(!dense10.is_empty());
    assert_eq!(dense10.iter().collect::<Vec<_>>(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);

    let mut dense256 = HybridBitSet::new_empty(256);
    assert!(dense256.is_empty());
    dense256.insert_all();
    assert!(!dense256.is_empty());
    for i in 0..256 {
        assert!(dense256.contains(i));
    }

    assert!(sparse038.superset(&sparse038)); // sparse + sparse (self)
    assert!(sparse01358.superset(&sparse038)); // sparse + sparse
    assert!(dense10.superset(&sparse038)); // dense + sparse
    assert!(dense10.superset(&dense10)); // dense + dense (self)
    assert!(dense256.superset(&dense10)); // dense + dense

    let mut hybrid = sparse038;
    assert!(!sparse01358.union(&hybrid)); // no change
    assert!(hybrid.union(&sparse01358));
    assert!(hybrid.superset(&sparse01358) && sparse01358.superset(&hybrid));
    assert!(!dense10.union(&sparse01358));
    assert!(!dense256.union(&dense10));
    let mut dense = dense10;
    assert!(dense.union(&dense256));
    assert!(dense.superset(&dense256) && dense256.superset(&dense));
    assert!(hybrid.union(&dense256));
    assert!(hybrid.superset(&dense256) && dense256.superset(&hybrid));

    assert_eq!(dense256.iter().count(), 256);
    let mut dense0 = dense256;
    for i in 0..256 {
        assert!(dense0.remove(i));
    }
    assert!(!dense0.remove(0));
    assert!(dense0.is_empty());
}

#[test]
fn grow() {
    let mut set: GrowableBitSet<usize> = GrowableBitSet::with_capacity(65);
    for index in 0..65 {
        assert!(set.insert(index));
        assert!(!set.insert(index));
    }
    set.ensure(128);

    // Check if the bits set before growing are still set
    for index in 0..65 {
        assert!(set.contains(index));
    }

    // Check if the new bits are all un-set
    for index in 65..128 {
        assert!(!set.contains(index));
    }

    // Check that we can set all new bits without running out of bounds
    for index in 65..128 {
        assert!(set.insert(index));
        assert!(!set.insert(index));
    }
}

#[test]
fn matrix_intersection() {
    let mut matrix: BitMatrix<usize, usize> = BitMatrix::new(200, 200);

    // (*) Elements reachable from both 2 and 65.

    matrix.insert(2, 3);
    matrix.insert(2, 6);
    matrix.insert(2, 10); // (*)
    matrix.insert(2, 64); // (*)
    matrix.insert(2, 65);
    matrix.insert(2, 130);
    matrix.insert(2, 160); // (*)

    matrix.insert(64, 133);

    matrix.insert(65, 2);
    matrix.insert(65, 8);
    matrix.insert(65, 10); // (*)
    matrix.insert(65, 64); // (*)
    matrix.insert(65, 68);
    matrix.insert(65, 133);
    matrix.insert(65, 160); // (*)

    let intersection = matrix.intersect_rows(2, 64);
    assert!(intersection.is_empty());

    let intersection = matrix.intersect_rows(2, 65);
    assert_eq!(intersection, &[10, 64, 160]);
}

#[test]
fn matrix_iter() {
    let mut matrix: BitMatrix<usize, usize> = BitMatrix::new(64, 100);
    matrix.insert(3, 22);
    matrix.insert(3, 75);
    matrix.insert(2, 99);
    matrix.insert(4, 0);
    matrix.union_rows(3, 5);
    matrix.insert_all_into_row(6);

    let expected = [99];
    let mut iter = expected.iter();
    for i in matrix.iter(2) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    assert_eq!(matrix.count(3), expected.len());
    for i in matrix.iter(3) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [0];
    let mut iter = expected.iter();
    assert_eq!(matrix.count(4), expected.len());
    for i in matrix.iter(4) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    assert_eq!(matrix.count(5), expected.len());
    for i in matrix.iter(5) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    assert_eq!(matrix.count(6), 100);
    let mut count = 0;
    for (idx, i) in matrix.iter(6).enumerate() {
        assert_eq!(idx, i);
        count += 1;
    }
    assert_eq!(count, 100);

    if let Some(i) = matrix.iter(7).next() {
        panic!("expected no elements in row, but contains element {:?}", i);
    }
}

#[test]
fn sparse_matrix_iter() {
    let mut matrix: SparseBitMatrix<usize, usize> = SparseBitMatrix::new(100);
    matrix.insert(3, 22);
    matrix.insert(3, 75);
    matrix.insert(2, 99);
    matrix.insert(4, 0);
    matrix.union_rows(3, 5);

    let expected = [99];
    let mut iter = expected.iter();
    for i in matrix.iter(2) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    for i in matrix.iter(3) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [0];
    let mut iter = expected.iter();
    for i in matrix.iter(4) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());

    let expected = [22, 75];
    let mut iter = expected.iter();
    for i in matrix.iter(5) {
        let j = *iter.next().unwrap();
        assert_eq!(i, j);
    }
    assert!(iter.next().is_none());
}

/// Merge dense hybrid set into empty sparse hybrid set.
#[bench]
fn union_hybrid_sparse_empty_to_dense(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    for i in 0..10 {
        assert!(pre_dense.insert(i));
    }
    let pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}

/// Merge dense hybrid set into full hybrid set with same indices.
#[bench]
fn union_hybrid_sparse_full_to_dense(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    for i in 0..10 {
        assert!(pre_dense.insert(i));
    }
    let mut pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(256);
    for i in 0..SPARSE_MAX {
        assert!(pre_sparse.insert(i));
    }
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}

/// Merge dense hybrid set into full hybrid set with indices over the whole domain.
#[bench]
fn union_hybrid_sparse_domain_to_dense(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX * 64);
    for i in 0..10 {
        assert!(pre_dense.insert(i));
    }
    let mut pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX * 64);
    for i in 0..SPARSE_MAX {
        assert!(pre_sparse.insert(i * 64));
    }
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}

/// Merge dense hybrid set into empty hybrid set where the domain is very small.
#[bench]
fn union_hybrid_sparse_empty_small_domain(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX);
    for i in 0..SPARSE_MAX {
        assert!(pre_dense.insert(i));
    }
    let pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX);
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}

/// Merge dense hybrid set into full hybrid set where the domain is very small.
#[bench]
fn union_hybrid_sparse_full_small_domain(b: &mut Bencher) {
    let mut pre_dense: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX);
    for i in 0..SPARSE_MAX {
        assert!(pre_dense.insert(i));
    }
    let mut pre_sparse: HybridBitSet<usize> = HybridBitSet::new_empty(SPARSE_MAX);
    for i in 0..SPARSE_MAX {
        assert!(pre_sparse.insert(i));
    }
    b.iter(|| {
        let dense = pre_dense.clone();
        let mut sparse = pre_sparse.clone();
        sparse.union(&dense);
    })
}

#[bench]
fn bench_insert(b: &mut Bencher) {
    let mut bs = BitSet::new_filled(99999usize);
    b.iter(|| {
        black_box(bs.insert(black_box(100u32)));
    });
}

#[bench]
fn bench_remove(b: &mut Bencher) {
    let mut bs = BitSet::new_filled(99999usize);
    b.iter(|| {
        black_box(bs.remove(black_box(100u32)));
    });
}

#[bench]
fn bench_iter(b: &mut Bencher) {
    let bs = BitSet::new_filled(99999usize);
    b.iter(|| {
        bs.iter().map(|b: usize| black_box(b)).for_each(drop);
    });
}

#[bench]
fn bench_intersect(b: &mut Bencher) {
    let mut ba: BitSet<u32> = BitSet::new_filled(99999usize);
    let bb = BitSet::new_filled(99999usize);
    b.iter(|| {
        ba.intersect(black_box(&bb));
    });
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 Chef Software Inc. and/or applicable contributors
//
// 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.

use serde_json::{self,
                 json,
                 Value as Json};
use std::{fs::{create_dir_all,
               read_to_string,
               File},
          io::Write,
          path::Path};
use toml::Value;

use crate::{common::{templating::TemplateRenderer,
                     ui::{Status,
                          UIWriter,
                          UI}},
            error::Result};

pub fn start(ui: &mut UI,
             template_path: &Path,
             default_toml_path: &Path,
             user_toml_path: Option<&Path>,
             mock_data_path: Option<&Path>,
             print: bool,
             render: bool,
             render_dir: &Path,
             quiet: bool)
             -> Result<()> {
    // Strip the file name out of our passed template
    let file_name = Path::new(template_path.file_name().expect("valid template file"));

    if !quiet {
        ui.begin(format!("Rendering: {} into: {} as: {}",
                         template_path.display(), render_dir.display(), file_name.display()))?;
        ui.br()?;
    }

    // read our template from file
    let template = read_to_string(&template_path)?;

    // create a "data" json struct
    let mut data = json!({});

    if !quiet {
        // import default.toml values, convert to JSON
        ui.begin(format!("Importing default.toml: {}", &default_toml_path.display()))?;
    }

    // we should always have a default.toml, would be nice to "autodiscover" based on package name,
    // for now assume we're working in the plan dir if --default-toml not passed
    let default_toml = read_to_string(&default_toml_path)?;

    // merge default into data struct
    merge(&mut data, toml_to_json(&default_toml)?);

    // import default.toml values, convert to JSON
    let user_toml = match user_toml_path {
        Some(path) => {
            if !quiet {
                // print helper message, maybe only print if '--verbose'? how?
                ui.begin(format!("Importing user.toml: {}", path.display()))?;
            }
            read_to_string(path)?
        }
        None => String::new(),
    };
    // merge default into data struct
    merge(&mut data, toml_to_json(&user_toml)?);

    // read mock data if provided
    let mock_data = match mock_data_path {
        Some(path) => {
            if !quiet {
                // print helper message, maybe only print if '--verbose'? how?
                ui.begin(format!("Importing override file: {}", path.display()))?;
            }
            read_to_string(path)?
        }
        // return an empty json block if '--mock-data' isn't defined.
        // this allows us to merge an empty JSON block
        None => "{}".to_string(),
    };
    // merge mock data into data
    merge(&mut data, serde_json::from_str(&mock_data)?);

    // create a template renderer
    let mut renderer = TemplateRenderer::new();
    // register our template
    renderer.register_template_string(&template, &template)
            .expect("Could not register template content");
    // render our JSON override in our template.
    let rendered_template = renderer.render(&template, &data)?;

    if print {
        if !quiet {
            ui.br()?;
            ui.warn(format!("###======== Rendered template: {}", &template_path.display()))?;
        }

        println!("{}", rendered_template);

        if !quiet {
            ui.warn(format!("========### End rendered template: {}", &template_path.display()))?;
        }
    }

    if render {
        // Render our template file
        create_with_template(ui, &render_dir, &file_name, &rendered_template, quiet)?;
    }

    if !quiet {
        ui.br()?;
    }
    Ok(())
}

fn toml_to_json(cfg: &str) -> Result<Json> {
    let toml_value = cfg.parse::<Value>()?;
    let toml_string = serde_json::to_string(&toml_value)?;
    let json = serde_json::from_str(&format!(r#"{{ "cfg": {} }}"#, &toml_string))?;
    Ok(json)
}

// merge two Json structs
fn merge(a: &mut Json, b: Json) {
    if let Json::Object(a_map) = a {
        if let Json::Object(b_map) = b {
            for (k, v) in b_map {
                merge(a_map.entry(k).or_insert(Json::Null), v);
            }
            return;
        }
    }
    *a = b;
}

fn create_with_template(ui: &mut UI,
                        render_dir: &std::path::Path,
                        file_name: &std::path::Path,
                        template: &str,
                        quiet: bool)
                        -> Result<()> {
    let path = Path::new(&render_dir).join(&file_name);
    if !quiet {
        ui.status(Status::Creating, format!("file: {}", path.display()))?;
    }

    create_dir_all(render_dir)?;

    // Write file to disk
    File::create(path).and_then(|mut file| file.write(template.as_bytes()))?;
    Ok(())
}
<commit_msg>rustfmt cleanup<commit_after>// Copyright (c) 2019 Chef Software Inc. and/or applicable contributors
//
// 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.

use serde_json::{self,
                 json,
                 Value as Json};
use std::{fs::{create_dir_all,
               read_to_string,
               File},
          io::Write,
          path::Path};
use toml::Value;

use crate::{common::{templating::TemplateRenderer,
                     ui::{Status,
                          UIWriter,
                          UI}},
            error::Result};

pub fn start(ui: &mut UI,
             template_path: &Path,
             default_toml_path: &Path,
             user_toml_path: Option<&Path>,
             mock_data_path: Option<&Path>,
             print: bool,
             render: bool,
             render_dir: &Path,
             quiet: bool)
             -> Result<()> {
    // Strip the file name out of our passed template
    let file_name = Path::new(template_path.file_name().expect("valid template file"));

    if !quiet {
        ui.begin(format!("Rendering: {} into: {} as: {}",
                         template_path.display(),
                         render_dir.display(),
                         file_name.display()))?;
        ui.br()?;
    }

    // read our template from file
    let template = read_to_string(&template_path)?;

    // create a "data" json struct
    let mut data = json!({});

    if !quiet {
        // import default.toml values, convert to JSON
        ui.begin(format!("Importing default.toml: {}", &default_toml_path.display()))?;
    }

    // we should always have a default.toml, would be nice to "autodiscover" based on package name,
    // for now assume we're working in the plan dir if --default-toml not passed
    let default_toml = read_to_string(&default_toml_path)?;

    // merge default into data struct
    merge(&mut data, toml_to_json(&default_toml)?);

    // import default.toml values, convert to JSON
    let user_toml = match user_toml_path {
        Some(path) => {
            if !quiet {
                // print helper message, maybe only print if '--verbose'? how?
                ui.begin(format!("Importing user.toml: {}", path.display()))?;
            }
            read_to_string(path)?
        }
        None => String::new(),
    };
    // merge default into data struct
    merge(&mut data, toml_to_json(&user_toml)?);

    // read mock data if provided
    let mock_data = match mock_data_path {
        Some(path) => {
            if !quiet {
                // print helper message, maybe only print if '--verbose'? how?
                ui.begin(format!("Importing override file: {}", path.display()))?;
            }
            read_to_string(path)?
        }
        // return an empty json block if '--mock-data' isn't defined.
        // this allows us to merge an empty JSON block
        None => "{}".to_string(),
    };
    // merge mock data into data
    merge(&mut data, serde_json::from_str(&mock_data)?);

    // create a template renderer
    let mut renderer = TemplateRenderer::new();
    // register our template
    renderer.register_template_string(&template, &template)
            .expect("Could not register template content");
    // render our JSON override in our template.
    let rendered_template = renderer.render(&template, &data)?;

    if print {
        if !quiet {
            ui.br()?;
            ui.warn(format!("###======== Rendered template: {}",
                            &template_path.display()))?;
        }

        println!("{}", rendered_template);

        if !quiet {
            ui.warn(format!("========### End rendered template: {}",
                            &template_path.display()))?;
        }
    }

    if render {
        // Render our template file
        create_with_template(ui, &render_dir, &file_name, &rendered_template, quiet)?;
    }

    if !quiet {
        ui.br()?;
    }
    Ok(())
}

fn toml_to_json(cfg: &str) -> Result<Json> {
    let toml_value = cfg.parse::<Value>()?;
    let toml_string = serde_json::to_string(&toml_value)?;
    let json = serde_json::from_str(&format!(r#"{{ "cfg": {} }}"#, &toml_string))?;
    Ok(json)
}

// merge two Json structs
fn merge(a: &mut Json, b: Json) {
    if let Json::Object(a_map) = a {
        if let Json::Object(b_map) = b {
            for (k, v) in b_map {
                merge(a_map.entry(k).or_insert(Json::Null), v);
            }
            return;
        }
    }
    *a = b;
}

fn create_with_template(ui: &mut UI,
                        render_dir: &std::path::Path,
                        file_name: &std::path::Path,
                        template: &str,
                        quiet: bool)
                        -> Result<()> {
    let path = Path::new(&render_dir).join(&file_name);
    if !quiet {
        ui.status(Status::Creating, format!("file: {}", path.display()))?;
    }

    create_dir_all(render_dir)?;

    // Write file to disk
    File::create(path).and_then(|mut file| file.write(template.as_bytes()))?;
    Ok(())
}
<|endoftext|>