<commit_before>use criterion::{criterion_group, criterion_main, Criterion};
use lazy_static::lazy_static;
use rosrust;
use std::process::Command;
use std::time;

mod util;

mod msg {
    rosrust::rosmsg_include!(roscpp_tutorials / TwoInts, rospy_tutorials / AddTwoInts);
}

fn global_init() -> util::ChildProcessTerminator {
    let roscore = util::run_roscore_for(util::Language::None, util::Feature::Benchmarks);
    rosrust::init("benchmarker");
    roscore
}

lazy_static! {
    static ref ROS_CORE: Option<util::ChildProcessTerminator> = Some(global_init());
}

fn setup() {
    assert!(ROS_CORE.is_some());
}

fn call_service(criterion: &mut Criterion) {
    setup();

    let namespace = format!("/namespaceat{}", line!());

    let _roscpp_service = util::ChildProcessTerminator::spawn(
        Command::new("rosrun")
            .arg("roscpp_tutorials")
            .arg("add_two_ints_server")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=roscpp_service")
            .arg("add_two_ints:=add_two_ints_cpp"),
    );

    let _roscpp_service = util::ChildProcessTerminator::spawn(
        Command::new("rosrun")
            .arg("rospy_tutorials")
            .arg("add_two_ints_server")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=rospy_service")
            .arg("add_two_ints:=add_two_ints_py"),
    );

    let _roscpp_service = util::ChildProcessTerminator::spawn_example_bench(
        "../examples/serviceclient",
        Command::new("cargo")
            .arg("run")
            .arg("--release")
            .arg("--bin")
            .arg("service")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=rosrust_service")
            .arg("add_two_ints:=add_two_ints_rust"),
    );

    let service_name_inline = format!("service_at_line_{}", line!());

    let _service_raii =
        rosrust::service::<msg::roscpp_tutorials::TwoInts, _>(&service_name_inline, move |req| {
            let sum = req.a + req.b;
            Ok(msg::roscpp_tutorials::TwoIntsRes { sum })
        })
        .unwrap();

    let service_name_cpp = format!("{}/add_two_ints_cpp", namespace);
    let service_name_py = format!("{}/add_two_ints_py", namespace);
    let service_name_rust = format!("{}/add_two_ints_rust", namespace);

    rosrust::wait_for_service(&service_name_cpp, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_py, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_rust, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_inline, Some(time::Duration::from_secs(10))).unwrap();

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_cpp).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call roscpp service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function("call roscpp service 50 times in parallel", move |b| {
        b.iter(|| {
            let requests = (0..50)
                .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                .collect::<Vec<_>>();
            for (idx, request) in requests.into_iter().enumerate() {
                assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
            }
        });
    });

    let client_original =
        rosrust::client::<msg::rospy_tutorials::AddTwoInts>(&service_name_py).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call rospy service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::rospy_tutorials::AddTwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_rust).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call rosrust service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function("call rosrust service 50 times in parallel", move |b| {
        b.iter(|| {
            let requests = (0..50)
                .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                .collect::<Vec<_>>();
            for (idx, request) in requests.into_iter().enumerate() {
                assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
            }
        });
    });

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_inline).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call inline rosrust service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function(
        "call inline rosrust service 50 times in parallel",
        move |b| {
            b.iter(|| {
                let requests = (0..50)
                    .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                    .collect::<Vec<_>>();
                for (idx, request) in requests.into_iter().enumerate() {
                    assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
                }
            });
        },
    );
}

criterion_group!(benches, call_service);
criterion_main!(benches);
<commit_msg>Add direct pubsub benchmark<commit_after>use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use crossbeam::unbounded;
use lazy_static::lazy_static;
use rosrust;
use std::process::Command;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time;

mod util;

mod msg {
    rosrust::rosmsg_include!(
        roscpp_tutorials / TwoInts,
        rospy_tutorials / AddTwoInts,
        std_msgs / String
    );
}

fn global_init() -> util::ChildProcessTerminator {
    let roscore = util::run_roscore_for(util::Language::None, util::Feature::Benchmarks);
    rosrust::init("benchmarker");
    roscore
}

lazy_static! {
    static ref ROS_CORE: Option<util::ChildProcessTerminator> = Some(global_init());
}

fn setup() {
    assert!(ROS_CORE.is_some());
}

fn subscribe_publish_directly(criterion: &mut Criterion) {
    setup();

    let topic = format!("/topicat{}", line!());

    let publisher = rosrust::publish::<msg::std_msgs::String>(&topic, 2000).unwrap();

    let (tx, rx) = unbounded();

    let _subscriber =
        rosrust::subscribe::<msg::std_msgs::String, _>(&topic, 2000, move |message| {
            tx.send(message.data).unwrap();
        })
        .unwrap();

    loop {
        publisher
            .send(msg::std_msgs::String { data: "".into() })
            .unwrap();
        if !rx.is_empty() {
            break;
        }
    }

    publisher
        .send(msg::std_msgs::String {
            data: "ready".into(),
        })
        .unwrap();

    while rx.recv().unwrap() != "ready" {}

    let inner_publisher = publisher.clone();
    let receiver = rx.clone();
    criterion.bench_function("send and receive single message directly", move |b| {
        let idx = AtomicUsize::new(0);

        b.iter_batched(
            || format!("{}:{}", line!(), idx.fetch_add(1, Ordering::SeqCst)),
            |data| {
                inner_publisher
                    .send(msg::std_msgs::String { data: data.clone() })
                    .unwrap();
                let response = receiver.recv().unwrap();
                assert_eq!(data, response);
            },
            BatchSize::SmallInput,
        );
    });

    let inner_publisher = publisher.clone();
    let receiver = rx.clone();
    criterion.bench_function("send and receive 100 messages directly", move |b| {
        let idx = AtomicUsize::new(0);

        b.iter_batched(
            || format!("{}:{}", line!(), idx.fetch_add(1, Ordering::SeqCst)),
            |data| {
                for _ in 0..100 {
                    inner_publisher
                        .send(msg::std_msgs::String { data: data.clone() })
                        .unwrap();
                }
                for _ in 0..100 {
                    let response = receiver.recv().unwrap();
                    assert_eq!(data, response);
                }
            },
            BatchSize::SmallInput,
        );
    });
}

fn call_service(criterion: &mut Criterion) {
    setup();

    let namespace = format!("/namespaceat{}", line!());

    let _roscpp_service = util::ChildProcessTerminator::spawn(
        Command::new("rosrun")
            .arg("roscpp_tutorials")
            .arg("add_two_ints_server")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=roscpp_service")
            .arg("add_two_ints:=add_two_ints_cpp"),
    );

    let _roscpp_service = util::ChildProcessTerminator::spawn(
        Command::new("rosrun")
            .arg("rospy_tutorials")
            .arg("add_two_ints_server")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=rospy_service")
            .arg("add_two_ints:=add_two_ints_py"),
    );

    let _roscpp_service = util::ChildProcessTerminator::spawn_example_bench(
        "../examples/serviceclient",
        Command::new("cargo")
            .arg("run")
            .arg("--release")
            .arg("--bin")
            .arg("service")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=rosrust_service")
            .arg("add_two_ints:=add_two_ints_rust"),
    );

    let service_name_inline = format!("service_at_line_{}", line!());

    let _service_raii =
        rosrust::service::<msg::roscpp_tutorials::TwoInts, _>(&service_name_inline, move |req| {
            let sum = req.a + req.b;
            Ok(msg::roscpp_tutorials::TwoIntsRes { sum })
        })
        .unwrap();

    let service_name_cpp = format!("{}/add_two_ints_cpp", namespace);
    let service_name_py = format!("{}/add_two_ints_py", namespace);
    let service_name_rust = format!("{}/add_two_ints_rust", namespace);

    rosrust::wait_for_service(&service_name_cpp, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_py, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_rust, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_inline, Some(time::Duration::from_secs(10))).unwrap();

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_cpp).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call roscpp service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function("call roscpp service 50 times in parallel", move |b| {
        b.iter(|| {
            let requests = (0..50)
                .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                .collect::<Vec<_>>();
            for (idx, request) in requests.into_iter().enumerate() {
                assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
            }
        });
    });

    let client_original =
        rosrust::client::<msg::rospy_tutorials::AddTwoInts>(&service_name_py).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call rospy service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::rospy_tutorials::AddTwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_rust).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call rosrust service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function("call rosrust service 50 times in parallel", move |b| {
        b.iter(|| {
            let requests = (0..50)
                .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                .collect::<Vec<_>>();
            for (idx, request) in requests.into_iter().enumerate() {
                assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
            }
        });
    });

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_inline).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call inline rosrust service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function(
        "call inline rosrust service 50 times in parallel",
        move |b| {
            b.iter(|| {
                let requests = (0..50)
                    .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                    .collect::<Vec<_>>();
                for (idx, request) in requests.into_iter().enumerate() {
                    assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
                }
            });
        },
    );
}

criterion_group!(benches, subscribe_publish_directly, call_service);
criterion_main!(benches);
<|endoftext|>


<commit_before><commit_msg>Remove unnecessary type declarations in iter example<commit_after><|endoftext|>


<commit_before><commit_msg>[cli] add support for txn params for script functions<commit_after><|endoftext|>


<commit_before>use kernel_micro::{include_file, MicroKernel};

/// A microkernel for Bash
///
/// The `kernel-bash.sh` script relies on `/dev/stdin`, `/dev/stdout`,
/// and `/dev/stderr` so this kernel is not available on Windows.
pub fn new() -> MicroKernel {
    MicroKernel::new(
        "bash-micro",
        &["Bash", "Shell"],
        true,
        false,
        false,
        ("bash", "*"),
        &["{{script}}"],
        include_file!("bash-kernel.sh"),
        &[],
        "{{name}}={{json}}",
        "echo ${{name}}",
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use kernel::{
        common::{
            eyre::{bail, Result},
            tokio,
        },
        stencila_schema::Node,
        KernelTrait,
    };
    use test_utils::{assert_json_eq, common::serde_json::json, skip_ci_os};

    /// Tests of basic functionality
    /// This test is replicated in all the microkernels.
    /// Other test should be written for language specific quirks and regressions.
    #[tokio::test]
    async fn basics() -> Result<()> {
        if skip_ci_os("windows", "test currently failing on windows") {
            return Ok(());
        }

        let mut kernel = new();
        match kernel.is_available().await {
            true => kernel.start_here().await?,
            false => return Ok(()),
        }

        // Assign a variable and output it
        let (outputs, messages) = kernel.exec("a=2\necho $a\n", None).await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [2]);

        // Syntax error
        let (outputs, messages) = kernel.exec("if", None).await?;
        assert!(messages[0]
            .error_message
            .ends_with("syntax error: unexpected end of file"));
        assert_json_eq!(outputs, json!([]));

        // Runtime error
        let (outputs, messages) = kernel.exec("foo", None).await?;
        assert!(messages[0]
            .error_message
            .ends_with("foo: command not found"));
        assert_json_eq!(outputs, json!([]));

        // Set and get another variable
        kernel.set("b", Node::Integer(3)).await?;
        let b = kernel.get("b").await?;
        assert_json_eq!(b, 3);

        // Use both variables
        let (outputs, messages) = kernel.exec("echo $a$b\n", None).await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [23]);

        Ok(())
    }

    #[tokio::test]
    async fn percent_escaping() -> Result<()> {
        if skip_ci_os("windows", "test currently failing on windows") {
            return Ok(());
        }

        let mut kernel = new();
        match kernel.is_available().await {
            true => kernel.start_here().await?,
            false => return Ok(()),
        }

        let (outputs, messages) = kernel.exec("date +%s", None).await?;
        assert_json_eq!(messages, json!([]));
        let timestamp = outputs.first().unwrap();
        match timestamp {
            Node::Integer(timestamp) => assert!(*timestamp > 1600000000),
            _ => bail!("Expected an integer"),
        }

        Ok(())
    }
}
<commit_msg>test(Bash kernel): Skip test on MacOS<commit_after>use kernel_micro::{include_file, MicroKernel};

/// A microkernel for Bash
///
/// The `kernel-bash.sh` script relies on `/dev/stdin`, `/dev/stdout`,
/// and `/dev/stderr` so this kernel is not available on Windows.
pub fn new() -> MicroKernel {
    MicroKernel::new(
        "bash-micro",
        &["Bash", "Shell"],
        true,
        false,
        false,
        ("bash", "*"),
        &["{{script}}"],
        include_file!("bash-kernel.sh"),
        &[],
        "{{name}}={{json}}",
        "echo ${{name}}",
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use kernel::{
        common::{
            eyre::{bail, Result},
            tokio,
        },
        stencila_schema::Node,
        KernelTrait,
    };
    use test_utils::{assert_json_eq, common::serde_json::json, skip_ci_os};

    /// Tests of basic functionality
    /// This test is replicated in all the microkernels.
    /// Other test should be written for language specific quirks and regressions.
    #[tokio::test]
    async fn basics() -> Result<()> {
        if skip_ci_os("windows", "test currently failing on Windows CI") {
            return Ok(());
        }
        if skip_ci_os("macos", "test currently failing on MacOS CI") {
            return Ok(());
        }

        let mut kernel = new();
        match kernel.is_available().await {
            true => kernel.start_here().await?,
            false => return Ok(()),
        }

        // Assign a variable and output it
        let (outputs, messages) = kernel.exec("a=2\necho $a\n", None).await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [2]);

        // Syntax error
        let (outputs, messages) = kernel.exec("if", None).await?;
        assert!(messages[0]
            .error_message
            .ends_with("syntax error: unexpected end of file"));
        assert_json_eq!(outputs, json!([]));

        // Runtime error
        let (outputs, messages) = kernel.exec("foo", None).await?;
        assert!(messages[0]
            .error_message
            .ends_with("foo: command not found"));
        assert_json_eq!(outputs, json!([]));

        // Set and get another variable
        kernel.set("b", Node::Integer(3)).await?;
        let b = kernel.get("b").await?;
        assert_json_eq!(b, 3);

        // Use both variables
        let (outputs, messages) = kernel.exec("echo $a$b\n", None).await?;
        assert_json_eq!(messages, json!([]));
        assert_json_eq!(outputs, [23]);

        Ok(())
    }

    #[tokio::test]
    async fn percent_escaping() -> Result<()> {
        if skip_ci_os("windows", "test currently failing on Windows CI") {
            return Ok(());
        }

        let mut kernel = new();
        match kernel.is_available().await {
            true => kernel.start_here().await?,
            false => return Ok(()),
        }

        let (outputs, messages) = kernel.exec("date +%s", None).await?;
        assert_json_eq!(messages, json!([]));
        let timestamp = outputs.first().unwrap();
        match timestamp {
            Node::Integer(timestamp) => assert!(*timestamp > 1600000000),
            _ => bail!("Expected an integer"),
        }

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


<commit_before>use std::{collections::BTreeMap, fmt, panic};

use quickcheck::{Arbitrary, Gen, RngCore};
use rand::distributions::{Distribution, Gamma};
use rand::{rngs::StdRng, Rng, SeedableRng};

use sled::*;

#[derive(Clone, Ord, PartialOrd, Eq, PartialEq)]
pub struct Key(pub Vec<u8>);

impl fmt::Debug for Key {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Key(vec!{:?})", self.0)
    }
}

struct SledGen {
    r: StdRng,
    size: usize,
}

impl Gen for SledGen {
    fn size(&self) -> usize {
        self.size
    }
}

impl RngCore for SledGen {
    fn next_u32(&mut self) -> u32 {
        self.r.gen::<u32>()
    }

    fn next_u64(&mut self) -> u64 {
        self.r.gen::<u64>()
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.r.fill_bytes(dest)
    }

    fn try_fill_bytes(
        &mut self,
        dest: &mut [u8],
    ) -> std::result::Result<(), rand::Error> {
        self.r.try_fill_bytes(dest)
    }
}

pub fn fuzz_then_shrink(buf: &[u8]) {
    let ops: Vec<Op> = buf
        .chunks(2)
        .map(|chunk| {
            let mut seed = [0u8; 32];
            seed[0..chunk.len()].copy_from_slice(chunk);
            let rng: StdRng = SeedableRng::from_seed(seed);
            let mut sled_rng = SledGen { r: rng, size: 2 };

            Op::arbitrary(&mut sled_rng)
        })
        .collect();

    println!("working on ops {:?}", ops);
    match panic::catch_unwind(move || {
        prop_tree_matches_btreemap(ops, 0, 100, false)
    }) {
        Ok(_) => {}
        Err(_e) => panic!("TODO"),
    }
}

impl Arbitrary for Key {
    fn arbitrary<G: Gen>(g: &mut G) -> Key {
        if g.gen::<bool>() {
            let gs = g.size();
            let gamma = Gamma::new(0.3, gs as f64);
            let v = gamma.sample(&mut rand::thread_rng());
            let len = if v > 30000.0 {
                30000
            } else if v < 1. && v > 0.0001 {
                1
            } else {
                v as usize
            };

            let space = g.gen_range(0, gs) + 1;

            let inner = (0..len)
                .map(|_| g.gen_range(0, space) as u8)
                .collect();

            Key(inner)
        } else {
            let len = g.gen_range(0, 2);
            let mut inner = vec![];

            for _ in 0..len {
                inner.push(g.gen::<u8>());
            }

            Key(inner)
        }
    }

    fn shrink(&self) -> Box<Iterator<Item = Key>> {
        // we only want to shrink on length, not byte values
        Box::new(
            self.0
                .len()
                .shrink()
                .zip(std::iter::repeat(self.0.clone()))
                .map(|(len, underlying)| {
                    Key(underlying[..len].to_vec())
                }),
        )
    }
}

#[derive(Debug, Clone)]
pub enum Op {
    Set(Key, u8),
    Merge(Key, u8),
    Get(Key),
    GetLt(Key),
    GetGt(Key),
    Del(Key),
    Cas(Key, u8, u8),
    Scan(Key, usize),
    Restart,
}

use self::Op::*;

impl Op {
    pub fn is_restart(&self) -> bool {
        if let Restart = self {
            true
        } else {
            false
        }
    }
}

impl Arbitrary for Op {
    fn arbitrary<G: Gen>(g: &mut G) -> Op {
        if g.gen_bool(1. / 10.) {
            return Restart;
        }

        let choice = g.gen_range(0, 8);

        match choice {
            0 => Set(Key::arbitrary(g), g.gen::<u8>()),
            1 => Merge(Key::arbitrary(g), g.gen::<u8>()),
            2 => Get(Key::arbitrary(g)),
            3 => GetLt(Key::arbitrary(g)),
            4 => GetGt(Key::arbitrary(g)),
            5 => Del(Key::arbitrary(g)),
            6 => Cas(Key::arbitrary(g), g.gen::<u8>(), g.gen::<u8>()),
            7 => Scan(Key::arbitrary(g), g.gen_range(0, 40)),
            _ => panic!("impossible choice"),
        }
    }

    fn shrink(&self) -> Box<Iterator<Item = Op>> {
        match *self {
            Set(ref k, v) => {
                Box::new(k.shrink().map(move |sk| Set(sk, v)))
            }
            Merge(ref k, v) => {
                Box::new(k.shrink().map(move |k| Merge(k, v)))
            }
            Get(ref k) => Box::new(k.shrink().map(Get)),
            GetLt(ref k) => Box::new(k.shrink().map(GetLt)),
            GetGt(ref k) => Box::new(k.shrink().map(GetGt)),
            Cas(ref k, old, new) => Box::new(Box::new(
                k.shrink().map(move |k| Cas(k, old, new)),
            )),
            Scan(ref k, len) => {
                Box::new(k.shrink().map(move |k| Scan(k, len)))
            }
            Del(ref k) => Box::new(k.shrink().map(Del)),
            Restart => Box::new(vec![].into_iter()),
        }
    }
}

fn bytes_to_u16(v: &[u8]) -> u16 {
    assert_eq!(v.len(), 2);
    (u16::from(v[0]) << 8) + u16::from(v[1])
}

fn u16_to_bytes(u: u16) -> Vec<u8> {
    let ret = vec![(u >> 8) as u8, u as u8];
    ret
}

// just adds up values as if they were u16's
fn test_merge_operator(
    _k: &[u8],
    old: Option<&[u8]>,
    to_merge: &[u8],
) -> Option<Vec<u8>> {
    let base = old.unwrap_or(&[0, 0]);
    let base_n = bytes_to_u16(base);
    let new_n = base_n + u16::from(to_merge[0]);
    let ret = u16_to_bytes(new_n);
    Some(ret)
}

pub fn prop_tree_matches_btreemap(
    ops: Vec<Op>,
    blink_node_exponent: u8,
    snapshot_after: u8,
    flusher: bool,
) -> bool {
    use self::*;
    let config = ConfigBuilder::new()
        .async_io(false)
        .temporary(true)
        .snapshot_after_ops(snapshot_after as usize + 1)
        .flush_every_ms(if flusher { Some(1) } else { None })
        .io_buf_size(10000)
        .blink_node_split_size(
            1 << std::cmp::min(blink_node_exponent, 20),
        )
        .cache_capacity(40)
        .cache_bits(0)
        .merge_operator(test_merge_operator)
        .build();

    let mut tree = sled::Tree::start(config.clone()).unwrap();
    let mut reference: BTreeMap<Key, u16> = BTreeMap::new();

    for op in ops.into_iter() {
        match op {
            Set(k, v) => {
                tree.set(&k.0, vec![0, v]).unwrap();
                reference.insert(k.clone(), u16::from(v));
            }
            Merge(k, v) => {
                tree.merge(&k.0, vec![v]).unwrap();
                let entry = reference.entry(k).or_insert(0u16);
                *entry += u16::from(v);
            }
            Get(k) => {
                let res1 = tree
                    .get(&*k.0)
                    .unwrap()
                    .map(|v| bytes_to_u16(&*v));
                let res2 = reference.get(&k).cloned();
                assert_eq!(res1, res2);
            }
            GetLt(k) => {
                let res1 = tree
                    .get_lt(&*k.0)
                    .unwrap()
                    .map(|v| bytes_to_u16(&*v.1));
                let res2 = reference
                    .iter()
                    .rev()
                    .filter(|(key, _)| **key < k)
                    .nth(0)
                    .map(|(_, v)| *v);
                assert_eq!(res1, res2);
            }
            GetGt(k) => {
                let res1 = tree
                    .get_gt(&*k.0)
                    .unwrap()
                    .map(|v| bytes_to_u16(&*v.1));
                let res2 = reference
                    .iter()
                    .filter(|(key, _)| **key > k)
                    .nth(0)
                    .map(|(_, v)| *v);
                assert_eq!(res1, res2);
            }
            Del(k) => {
                tree.del(&*k.0).unwrap();
                reference.remove(&k);
            }
            Cas(k, old, new) => {
                let tree_old = tree.get(&*k.0).unwrap();
                if let Some(old_tree) = tree_old {
                    if old_tree == *vec![0, old] {
                        tree.set(&k.0, vec![0, new]).unwrap();
                    }
                }

                let ref_old = reference.get(&k).cloned();
                if ref_old == Some(u16::from(old)) {
                    reference.insert(k, u16::from(new));
                }
            }
            Scan(k, len) => {
                let mut tree_iter = tree
                    .scan(&*k.0)
                    .take(len)
                    .map(|res| res.unwrap());
                let ref_iter = reference
                    .iter()
                    .filter(|&(ref rk, _rv)| **rk >= k)
                    .take(len)
                    .map(|(ref rk, ref rv)| (rk.0.clone(), **rv));

                for r in ref_iter {
                    let tree_next = tree_iter.next().unwrap();
                    let lhs = (tree_next.0, &*tree_next.1);
                    let rhs = (r.0.clone(), &*u16_to_bytes(r.1));
                    assert_eq!(
                        lhs, rhs,
                        "expected iteration over the Tree \
                         to match our BTreeMap model"
                    );
                }
            }
            Restart => {
                drop(tree);
                tree = sled::Tree::start(config.clone()).unwrap();
            }
        }
    }

    true
}

#[test]
fn test_fuzz_test() {
    let seed = [0; 32];
    fuzz_then_shrink(&seed);
    let seed = [0; 31];
    fuzz_then_shrink(&seed);
    let seed = [0; 33];
    fuzz_then_shrink(&seed);
    fuzz_then_shrink(&[]);
}
<commit_msg>remove nested box<commit_after>use std::{collections::BTreeMap, fmt, panic};

use quickcheck::{Arbitrary, Gen, RngCore};
use rand::distributions::{Distribution, Gamma};
use rand::{rngs::StdRng, Rng, SeedableRng};

use sled::*;

#[derive(Clone, Ord, PartialOrd, Eq, PartialEq)]
pub struct Key(pub Vec<u8>);

impl fmt::Debug for Key {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Key(vec!{:?})", self.0)
    }
}

struct SledGen {
    r: StdRng,
    size: usize,
}

impl Gen for SledGen {
    fn size(&self) -> usize {
        self.size
    }
}

impl RngCore for SledGen {
    fn next_u32(&mut self) -> u32 {
        self.r.gen::<u32>()
    }

    fn next_u64(&mut self) -> u64 {
        self.r.gen::<u64>()
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.r.fill_bytes(dest)
    }

    fn try_fill_bytes(
        &mut self,
        dest: &mut [u8],
    ) -> std::result::Result<(), rand::Error> {
        self.r.try_fill_bytes(dest)
    }
}

pub fn fuzz_then_shrink(buf: &[u8]) {
    let ops: Vec<Op> = buf
        .chunks(2)
        .map(|chunk| {
            let mut seed = [0u8; 32];
            seed[0..chunk.len()].copy_from_slice(chunk);
            let rng: StdRng = SeedableRng::from_seed(seed);
            let mut sled_rng = SledGen { r: rng, size: 2 };

            Op::arbitrary(&mut sled_rng)
        })
        .collect();

    println!("working on ops {:?}", ops);
    match panic::catch_unwind(move || {
        prop_tree_matches_btreemap(ops, 0, 100, false)
    }) {
        Ok(_) => {}
        Err(_e) => panic!("TODO"),
    }
}

impl Arbitrary for Key {
    fn arbitrary<G: Gen>(g: &mut G) -> Key {
        if g.gen::<bool>() {
            let gs = g.size();
            let gamma = Gamma::new(0.3, gs as f64);
            let v = gamma.sample(&mut rand::thread_rng());
            let len = if v > 30000.0 {
                30000
            } else if v < 1. && v > 0.0001 {
                1
            } else {
                v as usize
            };

            let space = g.gen_range(0, gs) + 1;

            let inner = (0..len)
                .map(|_| g.gen_range(0, space) as u8)
                .collect();

            Key(inner)
        } else {
            let len = g.gen_range(0, 2);
            let mut inner = vec![];

            for _ in 0..len {
                inner.push(g.gen::<u8>());
            }

            Key(inner)
        }
    }

    fn shrink(&self) -> Box<Iterator<Item = Key>> {
        // we only want to shrink on length, not byte values
        Box::new(
            self.0
                .len()
                .shrink()
                .zip(std::iter::repeat(self.0.clone()))
                .map(|(len, underlying)| {
                    Key(underlying[..len].to_vec())
                }),
        )
    }
}

#[derive(Debug, Clone)]
pub enum Op {
    Set(Key, u8),
    Merge(Key, u8),
    Get(Key),
    GetLt(Key),
    GetGt(Key),
    Del(Key),
    Cas(Key, u8, u8),
    Scan(Key, usize),
    Restart,
}

use self::Op::*;

impl Op {
    pub fn is_restart(&self) -> bool {
        if let Restart = self {
            true
        } else {
            false
        }
    }
}

impl Arbitrary for Op {
    fn arbitrary<G: Gen>(g: &mut G) -> Op {
        if g.gen_bool(1. / 10.) {
            return Restart;
        }

        let choice = g.gen_range(0, 8);

        match choice {
            0 => Set(Key::arbitrary(g), g.gen::<u8>()),
            1 => Merge(Key::arbitrary(g), g.gen::<u8>()),
            2 => Get(Key::arbitrary(g)),
            3 => GetLt(Key::arbitrary(g)),
            4 => GetGt(Key::arbitrary(g)),
            5 => Del(Key::arbitrary(g)),
            6 => Cas(Key::arbitrary(g), g.gen::<u8>(), g.gen::<u8>()),
            7 => Scan(Key::arbitrary(g), g.gen_range(0, 40)),
            _ => panic!("impossible choice"),
        }
    }

    fn shrink(&self) -> Box<Iterator<Item = Op>> {
        match *self {
            Set(ref k, v) => {
                Box::new(k.shrink().map(move |sk| Set(sk, v)))
            }
            Merge(ref k, v) => {
                Box::new(k.shrink().map(move |k| Merge(k, v)))
            }
            Get(ref k) => Box::new(k.shrink().map(Get)),
            GetLt(ref k) => Box::new(k.shrink().map(GetLt)),
            GetGt(ref k) => Box::new(k.shrink().map(GetGt)),
            Cas(ref k, old, new) => Box::new(
                k.shrink().map(move |k| Cas(k, old, new)),
            ),
            Scan(ref k, len) => {
                Box::new(k.shrink().map(move |k| Scan(k, len)))
            }
            Del(ref k) => Box::new(k.shrink().map(Del)),
            Restart => Box::new(vec![].into_iter()),
        }
    }
}

fn bytes_to_u16(v: &[u8]) -> u16 {
    assert_eq!(v.len(), 2);
    (u16::from(v[0]) << 8) + u16::from(v[1])
}

fn u16_to_bytes(u: u16) -> Vec<u8> {
    let ret = vec![(u >> 8) as u8, u as u8];
    ret
}

// just adds up values as if they were u16's
fn test_merge_operator(
    _k: &[u8],
    old: Option<&[u8]>,
    to_merge: &[u8],
) -> Option<Vec<u8>> {
    let base = old.unwrap_or(&[0, 0]);
    let base_n = bytes_to_u16(base);
    let new_n = base_n + u16::from(to_merge[0]);
    let ret = u16_to_bytes(new_n);
    Some(ret)
}

pub fn prop_tree_matches_btreemap(
    ops: Vec<Op>,
    blink_node_exponent: u8,
    snapshot_after: u8,
    flusher: bool,
) -> bool {
    use self::*;
    let config = ConfigBuilder::new()
        .async_io(false)
        .temporary(true)
        .snapshot_after_ops(snapshot_after as usize + 1)
        .flush_every_ms(if flusher { Some(1) } else { None })
        .io_buf_size(10000)
        .blink_node_split_size(
            1 << std::cmp::min(blink_node_exponent, 20),
        )
        .cache_capacity(40)
        .cache_bits(0)
        .merge_operator(test_merge_operator)
        .build();

    let mut tree = sled::Tree::start(config.clone()).unwrap();
    let mut reference: BTreeMap<Key, u16> = BTreeMap::new();

    for op in ops.into_iter() {
        match op {
            Set(k, v) => {
                tree.set(&k.0, vec![0, v]).unwrap();
                reference.insert(k.clone(), u16::from(v));
            }
            Merge(k, v) => {
                tree.merge(&k.0, vec![v]).unwrap();
                let entry = reference.entry(k).or_insert(0u16);
                *entry += u16::from(v);
            }
            Get(k) => {
                let res1 = tree
                    .get(&*k.0)
                    .unwrap()
                    .map(|v| bytes_to_u16(&*v));
                let res2 = reference.get(&k).cloned();
                assert_eq!(res1, res2);
            }
            GetLt(k) => {
                let res1 = tree
                    .get_lt(&*k.0)
                    .unwrap()
                    .map(|v| bytes_to_u16(&*v.1));
                let res2 = reference
                    .iter()
                    .rev()
                    .filter(|(key, _)| **key < k)
                    .nth(0)
                    .map(|(_, v)| *v);
                assert_eq!(res1, res2);
            }
            GetGt(k) => {
                let res1 = tree
                    .get_gt(&*k.0)
                    .unwrap()
                    .map(|v| bytes_to_u16(&*v.1));
                let res2 = reference
                    .iter()
                    .filter(|(key, _)| **key > k)
                    .nth(0)
                    .map(|(_, v)| *v);
                assert_eq!(res1, res2);
            }
            Del(k) => {
                tree.del(&*k.0).unwrap();
                reference.remove(&k);
            }
            Cas(k, old, new) => {
                let tree_old = tree.get(&*k.0).unwrap();
                if let Some(old_tree) = tree_old {
                    if old_tree == *vec![0, old] {
                        tree.set(&k.0, vec![0, new]).unwrap();
                    }
                }

                let ref_old = reference.get(&k).cloned();
                if ref_old == Some(u16::from(old)) {
                    reference.insert(k, u16::from(new));
                }
            }
            Scan(k, len) => {
                let mut tree_iter = tree
                    .scan(&*k.0)
                    .take(len)
                    .map(|res| res.unwrap());
                let ref_iter = reference
                    .iter()
                    .filter(|&(ref rk, _rv)| **rk >= k)
                    .take(len)
                    .map(|(ref rk, ref rv)| (rk.0.clone(), **rv));

                for r in ref_iter {
                    let tree_next = tree_iter.next().unwrap();
                    let lhs = (tree_next.0, &*tree_next.1);
                    let rhs = (r.0.clone(), &*u16_to_bytes(r.1));
                    assert_eq!(
                        lhs, rhs,
                        "expected iteration over the Tree \
                         to match our BTreeMap model"
                    );
                }
            }
            Restart => {
                drop(tree);
                tree = sled::Tree::start(config.clone()).unwrap();
            }
        }
    }

    true
}

#[test]
fn test_fuzz_test() {
    let seed = [0; 32];
    fuzz_then_shrink(&seed);
    let seed = [0; 31];
    fuzz_then_shrink(&seed);
    let seed = [0; 33];
    fuzz_then_shrink(&seed);
    fuzz_then_shrink(&[]);
}
<|endoftext|>


<commit_before>extern crate herder;
extern crate serde_json;

use herder::api;

#[test]
fn account_from_json() {
    let ojson = r#"{
                    "name": "MYNAME",
                    "website": "MYURI"
                 }"#;
    let application: api::v1::Application = serde_json::from_str(ojson).unwrap();
    assert_eq!(application, api::v1::Application {
                        name: String::from("MYNAME"),
                        website: Some(String::from("MYURI"))
    });
}
<commit_msg>tests json deserialization of Application, Attachment, Card entities<commit_after>extern crate herder;
extern crate serde_json;

use herder::api::v1::{
    Application,
    Attachment,
    Card
};

#[test]
fn account_deserialize_from_json() {
    let ojson = r#"{
                    "name": "MYNAME",
                    "website": "MYURI"
                 }"#;
    let application: Application = serde_json::from_str(ojson).unwrap();
    assert_eq!(application, Application {
                        name: String::from("MYNAME"),
                        website: Some(String::from("MYURI"))
    });
}

#[test]
fn attachment_deserialize_from_json() {
    let ojson = r#"{
                    "id": 1234,
                    "type": "video",
                    "url": "MYURL",
                    "remote_url": "MYURL",
                    "preview_url": "MYURL",
                    "text_url": "MYURL"
                 }"#;
    let attachment: Attachment = serde_json::from_str(ojson).unwrap();
    assert_eq!(attachment, Attachment {
                        attachment_id: 1234,
                        attachment_type: String::from("video"),
                        url: String::from("MYURL"),
                        remote_url: String::from("MYURL"),
                        preview_url: String::from("MYURL"),
                        text_url: String::from("MYURL")
    });
}

#[test]
fn card_deserialize_from_json() {
    let ojson = r#"{
                    "url": "MYURL",
                    "title": "My Title",
                    "description": "My description",
                    "image": "MYIMG.png"
                 }"#;
    let card: Card = serde_json::from_str(ojson).unwrap();
    assert_eq!(card, Card {
                        url: String::from("MYURL"),
                        title: String::from("My Title"),
                        description: String::from("My description"),
                        image: String::from("MYIMG.png")
    });
}
<|endoftext|>


<commit_before>//! Types for the [`m.room.power_levels`] event.
//!
//! [`m.room.power_levels`]: https://spec.matrix.org/v1.2/client-server-api/#mroompower_levels

use std::collections::BTreeMap;

use js_int::{int, Int};
use ruma_macros::EventContent;
use serde::{Deserialize, Serialize};

use crate::{
    events::RoomEventType,
    power_levels::{default_power_level, NotificationPowerLevels},
    UserId,
};

/// The content of an `m.room.power_levels` event.
///
/// Defines the power levels (privileges) of users in the room.
#[derive(Clone, Debug, Deserialize, Serialize, EventContent)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[ruma_event(type = "m.room.power_levels", kind = State)]
pub struct RoomPowerLevelsEventContent {
    /// The level required to ban a user.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default = "default_power_level", skip_serializing_if = "is_default_power_level")]
    #[ruma_event(skip_redaction)]
    pub ban: Int,

    /// The level required to send specific event types.
    ///
    /// This is a mapping from event type to power level required.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::btreemap_deserialize_v1_powerlevel_values")
    )]
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    #[ruma_event(skip_redaction)]
    pub events: BTreeMap<RoomEventType, Int>,

    /// The default level required to send message events.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default, skip_serializing_if = "crate::serde::is_default")]
    #[ruma_event(skip_redaction)]
    pub events_default: Int,

    /// The level required to invite a user.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default, skip_serializing_if = "crate::serde::is_default")]
    pub invite: Int,

    /// The level required to kick a user.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default = "default_power_level", skip_serializing_if = "is_default_power_level")]
    #[ruma_event(skip_redaction)]
    pub kick: Int,

    /// The level required to redact an event.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default = "default_power_level", skip_serializing_if = "is_default_power_level")]
    #[ruma_event(skip_redaction)]
    pub redact: Int,

    /// The default level required to send state events.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default = "default_power_level", skip_serializing_if = "is_default_power_level")]
    #[ruma_event(skip_redaction)]
    pub state_default: Int,

    /// The power levels for specific users.
    ///
    /// This is a mapping from `user_id` to power level for that user.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::btreemap_deserialize_v1_powerlevel_values")
    )]
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    #[ruma_event(skip_redaction)]
    pub users: BTreeMap<Box<UserId>, Int>,

    /// The default power level for every user in the room.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default, skip_serializing_if = "crate::serde::is_default")]
    #[ruma_event(skip_redaction)]
    pub users_default: Int,

    /// The power level requirements for specific notification types.
    ///
    /// This is a mapping from `key` to power level for that notifications key.
    #[serde(default, skip_serializing_if = "crate::serde::is_default")]
    pub notifications: NotificationPowerLevels,
}

impl RoomPowerLevelsEventContent {
    /// Creates a new `RoomPowerLevelsEventContent` with all-default values.
    pub fn new() -> Self {
        // events_default, users_default and invite having a default of 0 while the others have a
        // default of 50 is not an oversight, these defaults are from the Matrix specification.
        Self {
            ban: default_power_level(),
            events: BTreeMap::new(),
            events_default: int!(0),
            invite: int!(0),
            kick: default_power_level(),
            redact: default_power_level(),
            state_default: default_power_level(),
            users: BTreeMap::new(),
            users_default: int!(0),
            notifications: NotificationPowerLevels::default(),
        }
    }
}

impl Default for RoomPowerLevelsEventContent {
    fn default() -> Self {
        Self::new()
    }
}

/// Used with `#[serde(skip_serializing_if)]` to omit default power levels.
#[allow(clippy::trivially_copy_pass_by_ref)]
fn is_default_power_level(l: &Int) -> bool {
    *l == int!(50)
}

impl RoomPowerLevelsEvent {
    /// Obtain the effective power levels, regardless of whether this event is redacted.
    pub fn power_levels(&self) -> RoomPowerLevels {
        match self {
            Self::Original(ev) => ev.content.clone().into(),
            Self::Redacted(ev) => ev.content.clone().into(),
        }
    }
}

impl SyncRoomPowerLevelsEvent {
    /// Obtain the effective power levels, regardless of whether this event is redacted.
    pub fn power_levels(&self) -> RoomPowerLevels {
        match self {
            Self::Original(ev) => ev.content.clone().into(),
            Self::Redacted(ev) => ev.content.clone().into(),
        }
    }
}

/// The effective power levels of a room.
///
/// This struct contains the same fields as [`RoomPowerLevelsEventContent`] and be created from that
/// using a `From` trait implementation, but it is also implements
/// `From<`[`RedactedRoomPowerLevelsEventContent`]`>`, so can be used when wanting to inspect the
/// power levels of a room, regardless of whether the most recent power-levels event is redacted or
/// not.
#[derive(Clone, Debug)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
pub struct RoomPowerLevels {
    /// The level required to ban a user.
    pub ban: Int,

    /// The level required to send specific event types.
    ///
    /// This is a mapping from event type to power level required.
    pub events: BTreeMap<RoomEventType, Int>,

    /// The default level required to send message events.
    pub events_default: Int,

    /// The level required to invite a user.
    pub invite: Int,

    /// The level required to kick a user.
    pub kick: Int,

    /// The level required to redact an event.
    pub redact: Int,

    /// The default level required to send state events.
    pub state_default: Int,

    /// The power levels for specific users.
    ///
    /// This is a mapping from `user_id` to power level for that user.
    pub users: BTreeMap<Box<UserId>, Int>,

    /// The default power level for every user in the room.
    pub users_default: Int,

    /// The power level requirements for specific notification types.
    ///
    /// This is a mapping from `key` to power level for that notifications key.
    pub notifications: NotificationPowerLevels,
}

impl From<RoomPowerLevelsEventContent> for RoomPowerLevels {
    fn from(c: RoomPowerLevelsEventContent) -> Self {
        Self {
            ban: c.ban,
            events: c.events,
            events_default: c.events_default,
            invite: c.invite,
            kick: c.kick,
            redact: c.redact,
            state_default: c.state_default,
            users: c.users,
            users_default: c.users_default,
            notifications: c.notifications,
        }
    }
}

impl From<RedactedRoomPowerLevelsEventContent> for RoomPowerLevels {
    fn from(c: RedactedRoomPowerLevelsEventContent) -> Self {
        Self {
            ban: c.ban,
            events: c.events,
            events_default: c.events_default,
            invite: int!(0),
            kick: c.kick,
            redact: c.redact,
            state_default: c.state_default,
            users: c.users,
            users_default: c.users_default,
            notifications: NotificationPowerLevels::default(),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use crate::{event_id, room_id, user_id, MilliSecondsSinceUnixEpoch};
    use assign::assign;
    use js_int::{int, uint};
    use maplit::btreemap;
    use serde_json::{json, to_value as to_json_value};

    use super::{default_power_level, NotificationPowerLevels, RoomPowerLevelsEventContent};
    use crate::events::{OriginalStateEvent, StateUnsigned};

    #[test]
    fn serialization_with_optional_fields_as_none() {
        let default = default_power_level();

        let power_levels_event = OriginalStateEvent {
            content: RoomPowerLevelsEventContent {
                ban: default,
                events: BTreeMap::new(),
                events_default: int!(0),
                invite: int!(0),
                kick: default,
                redact: default,
                state_default: default,
                users: BTreeMap::new(),
                users_default: int!(0),
                notifications: NotificationPowerLevels::default(),
            },
            event_id: event_id!("$h29iv0s8:example.com").to_owned(),
            origin_server_ts: MilliSecondsSinceUnixEpoch(uint!(1)),
            room_id: room_id!("!n8f893n9:example.com").to_owned(),
            unsigned: StateUnsigned::default(),
            sender: user_id!("@carl:example.com").to_owned(),
            state_key: "".into(),
        };

        let actual = to_json_value(&power_levels_event).unwrap();
        let expected = json!({
            "content": {},
            "event_id": "$h29iv0s8:example.com",
            "origin_server_ts": 1,
            "room_id": "!n8f893n9:example.com",
            "sender": "@carl:example.com",
            "state_key": "",
            "type": "m.room.power_levels"
        });

        assert_eq!(actual, expected);
    }

    #[test]
    fn serialization_with_all_fields() {
        let user = user_id!("@carl:example.com");
        let power_levels_event = OriginalStateEvent {
            content: RoomPowerLevelsEventContent {
                ban: int!(23),
                events: btreemap! {
                    "m.dummy".into() => int!(23)
                },
                events_default: int!(23),
                invite: int!(23),
                kick: int!(23),
                redact: int!(23),
                state_default: int!(23),
                users: btreemap! {
                    user.to_owned() => int!(23)
                },
                users_default: int!(23),
                notifications: assign!(NotificationPowerLevels::new(), { room: int!(23) }),
            },
            event_id: event_id!("$h29iv0s8:example.com").to_owned(),
            origin_server_ts: MilliSecondsSinceUnixEpoch(uint!(1)),
            room_id: room_id!("!n8f893n9:example.com").to_owned(),
            unsigned: StateUnsigned {
                age: Some(int!(100)),
                prev_content: Some(RoomPowerLevelsEventContent {
                    // Make just one field different so we at least know they're two different
                    // objects.
                    ban: int!(42),
                    events: btreemap! {
                        "m.dummy".into() => int!(42)
                    },
                    events_default: int!(42),
                    invite: int!(42),
                    kick: int!(42),
                    redact: int!(42),
                    state_default: int!(42),
                    users: btreemap! {
                        user.to_owned() => int!(42)
                    },
                    users_default: int!(42),
                    notifications: assign!(NotificationPowerLevels::new(), { room: int!(42) }),
                }),
                ..StateUnsigned::default()
            },
            sender: user.to_owned(),
            state_key: "".into(),
        };

        let actual = to_json_value(&power_levels_event).unwrap();
        let expected = json!({
            "content": {
                "ban": 23,
                "events": {
                    "m.dummy": 23
                },
                "events_default": 23,
                "invite": 23,
                "kick": 23,
                "redact": 23,
                "state_default": 23,
                "users": {
                    "@carl:example.com": 23
                },
                "users_default": 23,
                "notifications": {
                    "room": 23
                }
            },
            "event_id": "$h29iv0s8:example.com",
            "origin_server_ts": 1,
            "room_id": "!n8f893n9:example.com",
            "sender": "@carl:example.com",
            "state_key": "",
            "type": "m.room.power_levels",
            "unsigned": {
                "age": 100,
                "prev_content": {
                    "ban": 42,
                    "events": {
                        "m.dummy": 42
                    },
                    "events_default": 42,
                    "invite": 42,
                    "kick": 42,
                    "redact": 42,
                    "state_default": 42,
                    "users": {
                        "@carl:example.com": 42
                    },
                    "users_default": 42,
                    "notifications": {
                        "room": 42
                    },
                },
            }
        });

        assert_eq!(actual, expected);
    }
}
<commit_msg>events: Add RoomPowerLevels::{for_user, max}<commit_after>//! Types for the [`m.room.power_levels`] event.
//!
//! [`m.room.power_levels`]: https://spec.matrix.org/v1.2/client-server-api/#mroompower_levels

use std::{cmp::max, collections::BTreeMap};

use js_int::{int, Int};
use ruma_macros::EventContent;
use serde::{Deserialize, Serialize};

use crate::{
    events::RoomEventType,
    power_levels::{default_power_level, NotificationPowerLevels},
    UserId,
};

/// The content of an `m.room.power_levels` event.
///
/// Defines the power levels (privileges) of users in the room.
#[derive(Clone, Debug, Deserialize, Serialize, EventContent)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
#[ruma_event(type = "m.room.power_levels", kind = State)]
pub struct RoomPowerLevelsEventContent {
    /// The level required to ban a user.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default = "default_power_level", skip_serializing_if = "is_default_power_level")]
    #[ruma_event(skip_redaction)]
    pub ban: Int,

    /// The level required to send specific event types.
    ///
    /// This is a mapping from event type to power level required.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::btreemap_deserialize_v1_powerlevel_values")
    )]
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    #[ruma_event(skip_redaction)]
    pub events: BTreeMap<RoomEventType, Int>,

    /// The default level required to send message events.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default, skip_serializing_if = "crate::serde::is_default")]
    #[ruma_event(skip_redaction)]
    pub events_default: Int,

    /// The level required to invite a user.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default, skip_serializing_if = "crate::serde::is_default")]
    pub invite: Int,

    /// The level required to kick a user.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default = "default_power_level", skip_serializing_if = "is_default_power_level")]
    #[ruma_event(skip_redaction)]
    pub kick: Int,

    /// The level required to redact an event.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default = "default_power_level", skip_serializing_if = "is_default_power_level")]
    #[ruma_event(skip_redaction)]
    pub redact: Int,

    /// The default level required to send state events.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default = "default_power_level", skip_serializing_if = "is_default_power_level")]
    #[ruma_event(skip_redaction)]
    pub state_default: Int,

    /// The power levels for specific users.
    ///
    /// This is a mapping from `user_id` to power level for that user.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::btreemap_deserialize_v1_powerlevel_values")
    )]
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    #[ruma_event(skip_redaction)]
    pub users: BTreeMap<Box<UserId>, Int>,

    /// The default power level for every user in the room.
    ///
    /// If you activate the `compat` feature, deserialization will work for stringified
    /// integers too.
    #[cfg_attr(
        feature = "compat",
        serde(deserialize_with = "crate::serde::deserialize_v1_powerlevel")
    )]
    #[serde(default, skip_serializing_if = "crate::serde::is_default")]
    #[ruma_event(skip_redaction)]
    pub users_default: Int,

    /// The power level requirements for specific notification types.
    ///
    /// This is a mapping from `key` to power level for that notifications key.
    #[serde(default, skip_serializing_if = "crate::serde::is_default")]
    pub notifications: NotificationPowerLevels,
}

impl RoomPowerLevelsEventContent {
    /// Creates a new `RoomPowerLevelsEventContent` with all-default values.
    pub fn new() -> Self {
        // events_default, users_default and invite having a default of 0 while the others have a
        // default of 50 is not an oversight, these defaults are from the Matrix specification.
        Self {
            ban: default_power_level(),
            events: BTreeMap::new(),
            events_default: int!(0),
            invite: int!(0),
            kick: default_power_level(),
            redact: default_power_level(),
            state_default: default_power_level(),
            users: BTreeMap::new(),
            users_default: int!(0),
            notifications: NotificationPowerLevels::default(),
        }
    }
}

impl Default for RoomPowerLevelsEventContent {
    fn default() -> Self {
        Self::new()
    }
}

/// Used with `#[serde(skip_serializing_if)]` to omit default power levels.
#[allow(clippy::trivially_copy_pass_by_ref)]
fn is_default_power_level(l: &Int) -> bool {
    *l == int!(50)
}

impl RoomPowerLevelsEvent {
    /// Obtain the effective power levels, regardless of whether this event is redacted.
    pub fn power_levels(&self) -> RoomPowerLevels {
        match self {
            Self::Original(ev) => ev.content.clone().into(),
            Self::Redacted(ev) => ev.content.clone().into(),
        }
    }
}

impl SyncRoomPowerLevelsEvent {
    /// Obtain the effective power levels, regardless of whether this event is redacted.
    pub fn power_levels(&self) -> RoomPowerLevels {
        match self {
            Self::Original(ev) => ev.content.clone().into(),
            Self::Redacted(ev) => ev.content.clone().into(),
        }
    }
}

/// The effective power levels of a room.
///
/// This struct contains the same fields as [`RoomPowerLevelsEventContent`] and be created from that
/// using a `From` trait implementation, but it is also implements
/// `From<`[`RedactedRoomPowerLevelsEventContent`]`>`, so can be used when wanting to inspect the
/// power levels of a room, regardless of whether the most recent power-levels event is redacted or
/// not.
#[derive(Clone, Debug)]
#[cfg_attr(not(feature = "unstable-exhaustive-types"), non_exhaustive)]
pub struct RoomPowerLevels {
    /// The level required to ban a user.
    pub ban: Int,

    /// The level required to send specific event types.
    ///
    /// This is a mapping from event type to power level required.
    pub events: BTreeMap<RoomEventType, Int>,

    /// The default level required to send message events.
    pub events_default: Int,

    /// The level required to invite a user.
    pub invite: Int,

    /// The level required to kick a user.
    pub kick: Int,

    /// The level required to redact an event.
    pub redact: Int,

    /// The default level required to send state events.
    pub state_default: Int,

    /// The power levels for specific users.
    ///
    /// This is a mapping from `user_id` to power level for that user.
    pub users: BTreeMap<Box<UserId>, Int>,

    /// The default power level for every user in the room.
    pub users_default: Int,

    /// The power level requirements for specific notification types.
    ///
    /// This is a mapping from `key` to power level for that notifications key.
    pub notifications: NotificationPowerLevels,
}

impl RoomPowerLevels {
    /// Get the power level of a specific user.
    pub fn for_user(&self, user_id: &UserId) -> Int {
        self.users.get(user_id).map_or(self.users_default, |pl| *pl)
    }

    /// Get the maximum power level of any user.
    pub fn max(&self) -> Int {
        self.users.values().fold(self.users_default, |max_pl, user_pl| max(max_pl, *user_pl))
    }
}

impl From<RoomPowerLevelsEventContent> for RoomPowerLevels {
    fn from(c: RoomPowerLevelsEventContent) -> Self {
        Self {
            ban: c.ban,
            events: c.events,
            events_default: c.events_default,
            invite: c.invite,
            kick: c.kick,
            redact: c.redact,
            state_default: c.state_default,
            users: c.users,
            users_default: c.users_default,
            notifications: c.notifications,
        }
    }
}

impl From<RedactedRoomPowerLevelsEventContent> for RoomPowerLevels {
    fn from(c: RedactedRoomPowerLevelsEventContent) -> Self {
        Self {
            ban: c.ban,
            events: c.events,
            events_default: c.events_default,
            invite: int!(0),
            kick: c.kick,
            redact: c.redact,
            state_default: c.state_default,
            users: c.users,
            users_default: c.users_default,
            notifications: NotificationPowerLevels::default(),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use crate::{event_id, room_id, user_id, MilliSecondsSinceUnixEpoch};
    use assign::assign;
    use js_int::{int, uint};
    use maplit::btreemap;
    use serde_json::{json, to_value as to_json_value};

    use super::{default_power_level, NotificationPowerLevels, RoomPowerLevelsEventContent};
    use crate::events::{OriginalStateEvent, StateUnsigned};

    #[test]
    fn serialization_with_optional_fields_as_none() {
        let default = default_power_level();

        let power_levels_event = OriginalStateEvent {
            content: RoomPowerLevelsEventContent {
                ban: default,
                events: BTreeMap::new(),
                events_default: int!(0),
                invite: int!(0),
                kick: default,
                redact: default,
                state_default: default,
                users: BTreeMap::new(),
                users_default: int!(0),
                notifications: NotificationPowerLevels::default(),
            },
            event_id: event_id!("$h29iv0s8:example.com").to_owned(),
            origin_server_ts: MilliSecondsSinceUnixEpoch(uint!(1)),
            room_id: room_id!("!n8f893n9:example.com").to_owned(),
            unsigned: StateUnsigned::default(),
            sender: user_id!("@carl:example.com").to_owned(),
            state_key: "".into(),
        };

        let actual = to_json_value(&power_levels_event).unwrap();
        let expected = json!({
            "content": {},
            "event_id": "$h29iv0s8:example.com",
            "origin_server_ts": 1,
            "room_id": "!n8f893n9:example.com",
            "sender": "@carl:example.com",
            "state_key": "",
            "type": "m.room.power_levels"
        });

        assert_eq!(actual, expected);
    }

    #[test]
    fn serialization_with_all_fields() {
        let user = user_id!("@carl:example.com");
        let power_levels_event = OriginalStateEvent {
            content: RoomPowerLevelsEventContent {
                ban: int!(23),
                events: btreemap! {
                    "m.dummy".into() => int!(23)
                },
                events_default: int!(23),
                invite: int!(23),
                kick: int!(23),
                redact: int!(23),
                state_default: int!(23),
                users: btreemap! {
                    user.to_owned() => int!(23)
                },
                users_default: int!(23),
                notifications: assign!(NotificationPowerLevels::new(), { room: int!(23) }),
            },
            event_id: event_id!("$h29iv0s8:example.com").to_owned(),
            origin_server_ts: MilliSecondsSinceUnixEpoch(uint!(1)),
            room_id: room_id!("!n8f893n9:example.com").to_owned(),
            unsigned: StateUnsigned {
                age: Some(int!(100)),
                prev_content: Some(RoomPowerLevelsEventContent {
                    // Make just one field different so we at least know they're two different
                    // objects.
                    ban: int!(42),
                    events: btreemap! {
                        "m.dummy".into() => int!(42)
                    },
                    events_default: int!(42),
                    invite: int!(42),
                    kick: int!(42),
                    redact: int!(42),
                    state_default: int!(42),
                    users: btreemap! {
                        user.to_owned() => int!(42)
                    },
                    users_default: int!(42),
                    notifications: assign!(NotificationPowerLevels::new(), { room: int!(42) }),
                }),
                ..StateUnsigned::default()
            },
            sender: user.to_owned(),
            state_key: "".into(),
        };

        let actual = to_json_value(&power_levels_event).unwrap();
        let expected = json!({
            "content": {
                "ban": 23,
                "events": {
                    "m.dummy": 23
                },
                "events_default": 23,
                "invite": 23,
                "kick": 23,
                "redact": 23,
                "state_default": 23,
                "users": {
                    "@carl:example.com": 23
                },
                "users_default": 23,
                "notifications": {
                    "room": 23
                }
            },
            "event_id": "$h29iv0s8:example.com",
            "origin_server_ts": 1,
            "room_id": "!n8f893n9:example.com",
            "sender": "@carl:example.com",
            "state_key": "",
            "type": "m.room.power_levels",
            "unsigned": {
                "age": 100,
                "prev_content": {
                    "ban": 42,
                    "events": {
                        "m.dummy": 42
                    },
                    "events_default": 42,
                    "invite": 42,
                    "kick": 42,
                    "redact": 42,
                    "state_default": 42,
                    "users": {
                        "@carl:example.com": 42
                    },
                    "users_default": 42,
                    "notifications": {
                        "room": 42
                    },
                },
            }
        });

        assert_eq!(actual, expected);
    }
}
<|endoftext|>


<commit_before>pub mod doc_id_set;
pub mod boolean_retrieval;
pub mod scoring;
pub mod plan;

use std::fmt;
use std::io::{Cursor, Read, Write};
use std::collections::HashMap;
use std::rc::Rc;

use abra::schema::{FieldRef, SchemaRead};
use abra::query::Query;
use abra::query::term_scorer::TermScorer;
use abra::collectors::{Collector, DocumentMatch};
use rocksdb::DBVector;
use byteorder::{ByteOrder, BigEndian};

use key_builder::KeyBuilder;
use term_dictionary::TermRef;
use super::{RocksDBIndexReader, DocRef};
use search::doc_id_set::DocIdSet;
use search::boolean_retrieval::BooleanQueryOp;
use search::scoring::{CombinatorScorer, ScoreFunctionOp};
use search::plan::SearchPlan;
use search::plan::boolean_query_builder::BooleanQueryBuilder;

impl<'a> RocksDBIndexReader<'a> {
    fn plan_query_combinator(&self, mut plan: &mut SearchPlan, queries: &Vec<Query>, join_op: BooleanQueryOp, score: bool, scorer: CombinatorScorer) {
        match queries.len() {
            0 => plan.boolean_query.push(BooleanQueryOp::PushEmpty),
            1 =>  self.plan_query(&mut plan, &queries[0], score),
            _ => {
                let mut query_iter = queries.iter();
                self.plan_query(&mut plan, query_iter.next().unwrap(), score);

                for query in query_iter {
                    self.plan_query(&mut plan, query, score);
                    plan.boolean_query.push(join_op.clone());
                }
            }
        }

        plan.score_function.push(ScoreFunctionOp::CombinatorScorer(queries.len() as u32, scorer));
    }

    fn plan_query(&self, mut plan: &mut SearchPlan, query: &Query, score: bool) {
        match *query {
            Query::MatchAll{ref score} => {
                plan.boolean_query.push(BooleanQueryOp::PushFull);
                plan.score_function.push(ScoreFunctionOp::Literal(*score));
            }
            Query::MatchNone => {
                plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                plan.score_function.push(ScoreFunctionOp::Literal(0.0f64));
            }
            Query::MatchTerm{ref field, ref term, ref matcher, ref scorer} => {
                // Get term
                let term_bytes = term.to_bytes();
                let term_ref = match self.store.term_dictionary.get(&term_bytes) {
                    Some(term_ref) => term_ref,
                    None => {
                        // Term doesn't exist, so will never match
                        plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                        return
                    }
                };

                // Get field
                let field_ref = match self.schema().get_field_by_name(field) {
                    Some(field_ref) => field_ref,
                    None => {
                        // Field doesn't exist, so will never match
                        plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                        return
                    }
                };

                let tag = plan.allocate_tag().unwrap_or(0);
                plan.boolean_query.push(BooleanQueryOp::PushTermDirectory(field_ref, term_ref, tag));
                plan.score_function.push(ScoreFunctionOp::TermScorer(field_ref, term_ref, scorer.clone(), tag));
            }
            Query::Conjunction{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::And, score, CombinatorScorer::Avg);
            }
            Query::Disjunction{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Avg);
            }
            Query::NDisjunction{ref queries, minimum_should_match} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Avg);  // FIXME
            }
            Query::DisjunctionMax{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Max);
            }
            Query::Filter{ref query, ref filter} => {
                self.plan_query(&mut plan, query, score);
                self.plan_query(&mut plan, filter, false);
                plan.boolean_query.push(BooleanQueryOp::And);
            }
            Query::Exclude{ref query, ref exclude} => {
                self.plan_query(&mut plan, query, score);
                self.plan_query(&mut plan, exclude, false);
                plan.boolean_query.push(BooleanQueryOp::AndNot);
            }
        }
    }

    fn search_chunk_boolean_phase(&self, plan: &SearchPlan, chunk: u32) -> (DocIdSet, HashMap<u8, DocIdSet>) {
        let mut tagged_docid_sets = HashMap::new();

        // Execute boolean query
        let mut stack = Vec::new();
        for op in plan.boolean_query.iter() {
            match *op {
                BooleanQueryOp::PushEmpty => {
                    stack.push(DocIdSet::new_filled(0));
                }
                BooleanQueryOp::PushFull => {
                    stack.push(DocIdSet::new_filled(65536));
                }
                BooleanQueryOp::PushTermDirectory(field_ref, term_ref, tag) => {
                    let kb = KeyBuilder::chunk_dir_list(chunk, field_ref.ord(), term_ref.ord());
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(docid_set)) => {
                            let data = DocIdSet::FromRDB(docid_set);
                            tagged_docid_sets.insert(tag, data.clone());
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::PushDeletionList => {
                    let kb = KeyBuilder::chunk_del_list(chunk);
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(deletion_list)) => {
                            let data = DocIdSet::FromRDB(deletion_list);
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::And => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.intersection(&b));
                }
                BooleanQueryOp::Or => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.union(&b));
                }
                BooleanQueryOp::AndNot => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.exclusion(&b));
                }
            }
        }

        if !stack.len() == 1 {
            // TODO: Error
        }
        let mut matches = stack.pop().unwrap();

        (matches, tagged_docid_sets)
    }

    fn score_doc(&self, doc_id: u16, tagged_docid_sets: &HashMap<u8, DocIdSet>, plan: &SearchPlan) -> f64 {
        // Execute score function
        let mut stack = Vec::new();
        for op in plan.score_function.iter() {
            match *op {
                ScoreFunctionOp::Literal(val) => stack.push(val),
                ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer, tag) => {
                    match tagged_docid_sets.get(&tag) {
                        Some(docid_set) => {
                            if docid_set.contains_doc(doc_id) {
                                stack.push(1.0f64);
                            } else {
                                stack.push(0.0f64);
                            }
                        }
                        None => stack.push(0.0f64)
                    }
                }
                ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                    let score = match *scorer {
                        CombinatorScorer::Avg => {
                            let mut total_score = 0.0f64;

                            for i in 0..num_vals {
                                total_score += stack.pop().expect("stack underflow");
                            }

                            total_score / num_vals as f64
                        }
                        CombinatorScorer::Max => {
                            let mut max_score = 0.0f64;

                            for i in 0..num_vals {
                                let score = stack.pop().expect("stack underflow");
                                if score > max_score {
                                    max_score = score
                                }
                            }

                            max_score
                        }
                    };

                    stack.push(score);
                }
            }
        }

        stack.pop().expect("stack underflow")
    }

    fn search_chunk<C: Collector>(&self, collector: &mut C, plan: &SearchPlan, chunk: u32) {
        let (matches, tagged_docid_sets) = self.search_chunk_boolean_phase(plan, chunk);

        // Score documents and pass to collector
        for doc in matches.iter() {
            let score = self.score_doc(doc, &tagged_docid_sets, plan);

            let doc_ref = DocRef(chunk, doc);
            let doc_match = DocumentMatch::new_scored(doc_ref.as_u64(), score);
            collector.collect(doc_match);
        }
    }

    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) {
        let mut plan = SearchPlan::new();
        self.plan_query(&mut plan, query, true);

        // Add operations to exclude deleted documents to boolean query
        plan.boolean_query.push(BooleanQueryOp::PushDeletionList);
        plan.boolean_query.push(BooleanQueryOp::AndNot);

        // Optimise boolean query
        let mut optimiser = BooleanQueryBuilder::new();
        for op in plan.boolean_query.iter() {
            optimiser.push_op(op);
        }
        let (boolean_query, boolean_query_is_negated) = optimiser.build();
        plan.boolean_query = boolean_query;
        plan.boolean_query_is_negated = boolean_query_is_negated;

        // Run query on each chunk
        for chunk in self.store.chunks.iter_active(&self.snapshot) {
            self.search_chunk(collector, &plan, chunk);
        }
    }
}
<commit_msg>Invert the results list if the boolean query is negated<commit_after>pub mod doc_id_set;
pub mod boolean_retrieval;
pub mod scoring;
pub mod plan;

use std::fmt;
use std::io::{Cursor, Read, Write};
use std::collections::HashMap;
use std::rc::Rc;

use abra::schema::{FieldRef, SchemaRead};
use abra::query::Query;
use abra::query::term_scorer::TermScorer;
use abra::collectors::{Collector, DocumentMatch};
use rocksdb::DBVector;
use byteorder::{ByteOrder, BigEndian};

use key_builder::KeyBuilder;
use term_dictionary::TermRef;
use super::{RocksDBIndexReader, DocRef};
use search::doc_id_set::DocIdSet;
use search::boolean_retrieval::BooleanQueryOp;
use search::scoring::{CombinatorScorer, ScoreFunctionOp};
use search::plan::SearchPlan;
use search::plan::boolean_query_builder::BooleanQueryBuilder;

impl<'a> RocksDBIndexReader<'a> {
    fn plan_query_combinator(&self, mut plan: &mut SearchPlan, queries: &Vec<Query>, join_op: BooleanQueryOp, score: bool, scorer: CombinatorScorer) {
        match queries.len() {
            0 => plan.boolean_query.push(BooleanQueryOp::PushEmpty),
            1 =>  self.plan_query(&mut plan, &queries[0], score),
            _ => {
                let mut query_iter = queries.iter();
                self.plan_query(&mut plan, query_iter.next().unwrap(), score);

                for query in query_iter {
                    self.plan_query(&mut plan, query, score);
                    plan.boolean_query.push(join_op.clone());
                }
            }
        }

        plan.score_function.push(ScoreFunctionOp::CombinatorScorer(queries.len() as u32, scorer));
    }

    fn plan_query(&self, mut plan: &mut SearchPlan, query: &Query, score: bool) {
        match *query {
            Query::MatchAll{ref score} => {
                plan.boolean_query.push(BooleanQueryOp::PushFull);
                plan.score_function.push(ScoreFunctionOp::Literal(*score));
            }
            Query::MatchNone => {
                plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                plan.score_function.push(ScoreFunctionOp::Literal(0.0f64));
            }
            Query::MatchTerm{ref field, ref term, ref matcher, ref scorer} => {
                // Get term
                let term_bytes = term.to_bytes();
                let term_ref = match self.store.term_dictionary.get(&term_bytes) {
                    Some(term_ref) => term_ref,
                    None => {
                        // Term doesn't exist, so will never match
                        plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                        return
                    }
                };

                // Get field
                let field_ref = match self.schema().get_field_by_name(field) {
                    Some(field_ref) => field_ref,
                    None => {
                        // Field doesn't exist, so will never match
                        plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                        return
                    }
                };

                let tag = plan.allocate_tag().unwrap_or(0);
                plan.boolean_query.push(BooleanQueryOp::PushTermDirectory(field_ref, term_ref, tag));
                plan.score_function.push(ScoreFunctionOp::TermScorer(field_ref, term_ref, scorer.clone(), tag));
            }
            Query::Conjunction{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::And, score, CombinatorScorer::Avg);
            }
            Query::Disjunction{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Avg);
            }
            Query::NDisjunction{ref queries, minimum_should_match} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Avg);  // FIXME
            }
            Query::DisjunctionMax{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Max);
            }
            Query::Filter{ref query, ref filter} => {
                self.plan_query(&mut plan, query, score);
                self.plan_query(&mut plan, filter, false);
                plan.boolean_query.push(BooleanQueryOp::And);
            }
            Query::Exclude{ref query, ref exclude} => {
                self.plan_query(&mut plan, query, score);
                self.plan_query(&mut plan, exclude, false);
                plan.boolean_query.push(BooleanQueryOp::AndNot);
            }
        }
    }

    fn search_chunk_boolean_phase(&self, plan: &SearchPlan, chunk: u32) -> (DocIdSet, HashMap<u8, DocIdSet>) {
        let mut tagged_docid_sets = HashMap::new();

        // Execute boolean query
        let mut stack = Vec::new();
        for op in plan.boolean_query.iter() {
            match *op {
                BooleanQueryOp::PushEmpty => {
                    stack.push(DocIdSet::new_filled(0));
                }
                BooleanQueryOp::PushFull => {
                    stack.push(DocIdSet::new_filled(65536));
                }
                BooleanQueryOp::PushTermDirectory(field_ref, term_ref, tag) => {
                    let kb = KeyBuilder::chunk_dir_list(chunk, field_ref.ord(), term_ref.ord());
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(docid_set)) => {
                            let data = DocIdSet::FromRDB(docid_set);
                            tagged_docid_sets.insert(tag, data.clone());
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::PushDeletionList => {
                    let kb = KeyBuilder::chunk_del_list(chunk);
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(deletion_list)) => {
                            let data = DocIdSet::FromRDB(deletion_list);
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::And => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.intersection(&b));
                }
                BooleanQueryOp::Or => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.union(&b));
                }
                BooleanQueryOp::AndNot => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.exclusion(&b));
                }
            }
        }

        if !stack.len() == 1 {
            // TODO: Error
        }
        let mut matches = stack.pop().unwrap();

        // Invert the list if the query is negated
        if plan.boolean_query_is_negated {
            let kb = KeyBuilder::chunk_stat(chunk, b"total_docs");
            let total_docs = match self.snapshot.get(&kb.key()) {
                Ok(Some(total_docs)) => BigEndian::read_i64(&total_docs) as u16,
                Ok(None) => 0,
                Err(e) => 0,  // FIXME
            };

            let all_docs = DocIdSet::new_filled(total_docs);
            matches = all_docs.exclusion(&matches);
        }

        (matches, tagged_docid_sets)
    }

    fn score_doc(&self, doc_id: u16, tagged_docid_sets: &HashMap<u8, DocIdSet>, plan: &SearchPlan) -> f64 {
        // Execute score function
        let mut stack = Vec::new();
        for op in plan.score_function.iter() {
            match *op {
                ScoreFunctionOp::Literal(val) => stack.push(val),
                ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer, tag) => {
                    match tagged_docid_sets.get(&tag) {
                        Some(docid_set) => {
                            if docid_set.contains_doc(doc_id) {
                                stack.push(1.0f64);
                            } else {
                                stack.push(0.0f64);
                            }
                        }
                        None => stack.push(0.0f64)
                    }
                }
                ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                    let score = match *scorer {
                        CombinatorScorer::Avg => {
                            let mut total_score = 0.0f64;

                            for i in 0..num_vals {
                                total_score += stack.pop().expect("stack underflow");
                            }

                            total_score / num_vals as f64
                        }
                        CombinatorScorer::Max => {
                            let mut max_score = 0.0f64;

                            for i in 0..num_vals {
                                let score = stack.pop().expect("stack underflow");
                                if score > max_score {
                                    max_score = score
                                }
                            }

                            max_score
                        }
                    };

                    stack.push(score);
                }
            }
        }

        stack.pop().expect("stack underflow")
    }

    fn search_chunk<C: Collector>(&self, collector: &mut C, plan: &SearchPlan, chunk: u32) {
        let (matches, tagged_docid_sets) = self.search_chunk_boolean_phase(plan, chunk);

        // Score documents and pass to collector
        for doc in matches.iter() {
            let score = self.score_doc(doc, &tagged_docid_sets, plan);

            let doc_ref = DocRef(chunk, doc);
            let doc_match = DocumentMatch::new_scored(doc_ref.as_u64(), score);
            collector.collect(doc_match);
        }
    }

    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) {
        let mut plan = SearchPlan::new();
        self.plan_query(&mut plan, query, true);

        // Add operations to exclude deleted documents to boolean query
        plan.boolean_query.push(BooleanQueryOp::PushDeletionList);
        plan.boolean_query.push(BooleanQueryOp::AndNot);

        // Optimise boolean query
        let mut optimiser = BooleanQueryBuilder::new();
        for op in plan.boolean_query.iter() {
            optimiser.push_op(op);
        }
        let (boolean_query, boolean_query_is_negated) = optimiser.build();
        plan.boolean_query = boolean_query;
        plan.boolean_query_is_negated = boolean_query_is_negated;

        // Run query on each chunk
        for chunk in self.store.chunks.iter_active(&self.snapshot) {
            self.search_chunk(collector, &plan, chunk);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Chao Lin & William Sergeant (Sorbonne University)

// 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.

pub use self::useless_chaining::*;

grammar! useless_chaining {

  // test1 = !(!"a") // &"a"
  // test2 = &(&"a") // &"a"
  // test3 = !(&"a") // !"a"
  // test4 = &(!"a") // !"a"
  // //test5 = ("a"*)* // infinite loop -> deja detectee
  // test6 = ("a"+)+ // "a"+
  // test7 = ("a"+)* // "a"+
  // //test8 = ("a"*)+ // infinite loop -> deja detectee
  //
  test9 = !"a"
  // test10 = !test9
  //
  test11 = &"a"
  // test12 = &test11
  //
  // test13 = !test11
  //
  // test14 = &test9
  //
  // test15 = "a"+
  // test16 = test15+
  //
  // test17 = test15*
  //
  // test18 = &test12
  // test19 = test16+
  //
  // test20 = ((("a")+)+)+
  // test21 = &(&(&(&("a"))))
  //
  // test22 = &"a" / !"b"
  // test23 = &test22
  //
  // test24 = &"a" "b"
  // test25 = &test24
  //
  // test26 = &"a" / &"b"
  // test27 = &test26
  //
  // test28 = &"a" / "b"
  // test29 = &test28
  //
  // test30 = &"a" &(!"b")
  // test31 = &test30
  //
  // test32 = &(&"a") / &(&"b")
  // test33 = &test32
}
<commit_msg>[test] Update useless chaining test<commit_after>// Copyright 2018 Chao Lin & William Sergeant (Sorbonne University)

// 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.

pub use self::useless_chaining::*;

grammar! useless_chaining {

  // test1 = !(!"a") // &"a"
  // test2 = &(&"a") // &"a"
  // test3 = !(&"a") // !"a"
  // test4 = &(!"a") // !"a"
  // //test5 = ("a"*)* // infinite loop -> deja detectee
  // test6 = ("a"+)+ // "a"+
  // test7 = ("a"+)* // "a"+
  // //test8 = ("a"*)+ // infinite loop -> deja detectee
  //
  test9 = !"a"
  // test10 = !test9
  //
  // test11 = &"a"
  // test12 = &test11
  //
  // test13 = !test11
  //
  // test14 = &test9
  //
  // test15 = "a"+
  // test16 = test15+
  //
  // test17 = test15*
  //
  // test18 = &test12
  // test19 = test16+
  //
  // test20 = ((("a")+)+)+
  // test21 = &(&(&(&("a"))))
  //
  // test22 = &"a" / !"b"
  // test23 = &test22
  //
  // test24 = &"a" "b"
  // test25 = &test24
  //
  // test26 = &"a" / &"b"
  // test27 = &test26
  //
  // test28 = &"a" / "b"
  // test29 = &test28
  //
  // test30 = &"a" &(!"b")
  // test31 = &test30
  //
  // test32 = &(&"a") / &(&"b")
  // test33 = &test32
}
<|endoftext|>