<commit_before>//! Licensed under the Apache License, Version 2.0
//! http://www.apache.org/licenses/LICENSE-2.0 or the MIT license
//! http://opensource.org/licenses/MIT, at your
//! option. This file may not be copied, modified, or distributed
//! except according to those terms.

#![feature(core, test)]

#[macro_use]
extern crate itertools;

extern crate test;

use std::fmt::Debug;
use std::iter::order;
use std::iter::IntoIterator;
use itertools::Itertools;
use itertools::Interleave;
use itertools::Zip;

use itertools as it;

fn assert_iters_equal<
    A: PartialEq + Debug,
    I: Iterator<Item=A>,
    J: Iterator<Item=A>>(mut it: I, mut jt: J)
{
    loop {
        let elti = it.next();
        let eltj = jt.next();
        assert_eq!(elti, eltj);
        if elti.is_none() { break; }
    }
}

#[test]
fn product2() {
    let s = "αβ";

    let mut prod = iproduct!(s.chars(), range(0, 2));
    assert!(prod.next() == Some(('α', 0)));
    assert!(prod.next() == Some(('α', 1)));
    assert!(prod.next() == Some(('β', 0)));
    assert!(prod.next() == Some(('β', 1)));
    assert!(prod.next() == None);
}

#[test]
fn product3() {
    let prod = iproduct!(range(0, 3), range(0, 2), range(0, 2));
    assert_eq!(prod.size_hint(), (12, Some(12)));
    let v = prod.collect_vec();
    for i in range(0,3) {
        for j in range(0, 2) {
            for k in range(0, 2) {
                assert!((i, j, k) == v[(i * 2 * 2 + j * 2 + k) as usize]);
            }
        }
    }
    for (_, _, _, _) in iproduct!(range(0, 3), range(0, 2), range(0, 2), range(0, 3)) {
        /* test compiles */
    }
}

#[test]
fn izip3() {
    let mut zip = Zip::new((range(0, 3), range(0, 2), range(0, 2i8)));
    for i in range(0, 2) {
        assert!((i as usize, i, i as i8) == zip.next().unwrap());
    }
    assert!(zip.next().is_none());

    let xs: [isize; 0] = [];
    let mut zip = Zip::new((range(0, 3), range(0, 2), range(0, 2i8), xs.iter()));
    assert!(zip.next().is_none());

    for (_, _, _, _) in Zip::new((range(0, 3), range(0, 2), range(0, 2), range(0, 3))) {
        /* test compiles */
    }
}

#[test]
fn write_to() {
    let xs = [7, 9, 8];
    let mut ys = [0; 5];
    let cnt = it::write(ys.iter_mut(), xs.iter().map(|x| *x));
    assert!(cnt == xs.len());
    assert!(ys == [7, 9, 8, 0, 0]);

    let cnt = it::write(ys.iter_mut(), range(0,10));
    assert!(cnt == ys.len());
    assert!(ys == [0, 1, 2, 3, 4]);
}

#[test]
fn interleave() {
    let xs: [u8; 0]  = [];
    let ys = [7u8, 9, 8, 10];
    let zs = [2u8, 77];
    let it = Interleave::new(xs.iter(), ys.iter());
    assert!(order::eq(it, ys.iter()));

    let rs = [7u8, 2, 9, 77, 8, 10];
    let it = Interleave::new(ys.iter(), zs.iter());
    assert!(order::eq(it, rs.iter()));
}

#[test]
fn times() {
    assert!(it::times(0).count() == 0);
    assert!(it::times(5).count() == 5);
}

#[test]
fn drain() {
    let xs = [1i32,2,3];
    let mut sum = 0;
    xs.iter().map(|elt| sum += *elt).drain();
    assert!(sum == 6);
}

#[test]
fn foreach() {
    let xs = [1i32, 2, 3];
    let mut sum = 0;
    xs.iter().foreach(|elt| sum += *elt);
    assert!(sum == 6);
}

#[test]
fn dropn() {
    let xs = [1, 2, 3];
    let mut it = xs.iter();
    assert!(it.dropn(2) == 2);
    assert!(it.next().is_some());
    assert!(it.next().is_none());
    let mut it = xs.iter();
    assert!(it.dropn(5) == 3);
    assert!(it.next().is_none());
}

#[test]
fn dropping() {
    let xs = [1, 2, 3];
    let mut it = xs.iter().dropping(2);
    assert!(it.next().is_some());
    assert!(it.next().is_none());
    let mut it = xs.iter().dropping(5);
    assert!(it.next().is_none());
}

#[test]
fn intersperse() {
    let xs = ["a", "", "b", "c"];
    let v: Vec<&str> = xs.iter().map(|x| x.clone()).intersperse(", ").collect();
    let text: String = v.concat();
    assert_eq!(text, "a, , b, c".to_string());

    let ys = [0, 1, 2, 3];
    let mut it = ys[..0].iter().map(|x| *x).intersperse(1);
    assert!(it.next() == None);
}

#[test]
fn linspace() {
    let mut iter = it::linspace::<f32>(0., 2., 3);
    assert_eq!(iter.next(), Some(0.0));
    assert_eq!(iter.next(), Some(1.0));
    assert_eq!(iter.next(), Some(2.0));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., -2., 4);
    assert_eq!(iter.next(), Some(0.));
    assert_eq!(iter.next(), Some(-0.666666666667));
    assert_eq!(iter.next(), Some(-1.333333333333));
    assert_eq!(iter.next(), Some(-2.));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., 1., 1);
    assert_eq!(iter.next(), Some(0.));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., 1., 0);
    assert_eq!(iter.next(), None);
}

#[test]
fn dedup() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let ys = [0, 1, 2, 1, 3];
    assert_iters_equal(ys.iter(), xs.iter().dedup());
    let xs = [0, 0, 0, 0, 0];
    let ys = [0];
    assert_iters_equal(ys.iter(), xs.iter().dedup());
}

#[test]
fn batching() {
    let xs = [0, 1, 2, 1, 3];
    let ys = [(0, 1), (2, 1)];

    // An iterator that gathers elements up in pairs
    let pit = xs.iter().cloned().batching(|mut it| {
               match it.next() {
                   None => None,
                   Some(x) => match it.next() {
                       None => None,
                       Some(y) => Some((x, y)),
                   }
               }
           });
    assert_iters_equal(pit, ys.iter().cloned());
}

#[test]
fn group_by() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let ans = vec![(0, vec![0]), (1, vec![1, 1, 1]),
                   (2, vec![2]), (1, vec![1]), (3, vec![3, 3])];

    let gb = xs.iter().cloned().group_by(|elt| *elt);
    assert_iters_equal(gb, ans.into_iter());
}

#[test]
fn put_back() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let mut pb = it::PutBack::new(xs.iter().cloned());
    pb.next();
    pb.put_back(1);
    pb.put_back(0);
    assert_iters_equal(pb, xs.iter().cloned());
}

#[test]
fn tee() {
    let xs  = [0, 1, 2, 3];
    let (mut t1, mut t2) = xs.iter().cloned().tee();
    assert_eq!(t1.next(), Some(0));
    assert_eq!(t2.next(), Some(0));
    assert_eq!(t1.next(), Some(1));
    assert_eq!(t1.next(), Some(2));
    assert_eq!(t1.next(), Some(3));
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), Some(1));
    assert_eq!(t2.next(), Some(2));
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), Some(3));
    assert_eq!(t2.next(), None);
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), None);

    let (t1, t2) = xs.iter().cloned().tee();
    assert_iters_equal(t1, xs.iter().cloned());
    assert_iters_equal(t2, xs.iter().cloned());

    let (t1, t2) = xs.iter().cloned().tee();
    assert_iters_equal(t1.zip(t2), xs.iter().cloned().zip(xs.iter().cloned()));
}

#[test]
fn rciter() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 5, 6];

    let mut r1 = xs.iter().cloned().into_rc();
    let mut r2 = r1.clone();
    assert_eq!(r1.next(), Some(0));
    assert_eq!(r2.next(), Some(1));
    let mut z = r1.zip(r2);
    assert_eq!(z.next(), Some((1, 1)));
    assert_eq!(z.next(), Some((2, 1)));
    assert_eq!(z.next(), Some((3, 5)));
    assert_eq!(z.next(), None);
}

#[test]
fn slice() {
    assert_iters_equal((0..10).slice(..3), 0..3);
    assert_iters_equal((0..10).slice(3..7), 3..7);
    assert_iters_equal((0..10).slice(3..27), 3..10);
    assert_iters_equal((0..10).slice(44..), 0..0);
}

#[test]
fn step() {
    assert_iters_equal((0..10).step(1), (0..10));
    assert_iters_equal((0..10).step(2), (0..10).filter(|x: &i32| *x % 2 == 0));
    assert_iters_equal((0..10).step(10), 0..1);
}

#[test]
fn trait_pointers() {
    /*
    struct ByRef<'r, I: ?Sized>(&'r mut I) where I: 'r;
    impl<'r, X, I: ?Sized> Iterator for ByRef<'r, I> where
        I: 'r + Iterator<Item=X>
    {
        type Item = X;
        fn next(&mut self) -> Option<X>
        {
            self.0.next()
        }
    }
    */
    let mut it = Box::new(0..10) as Box<Iterator<Item=i32>>;
    assert_eq!(it.next(), Some(0));

    let mut jt: &mut Iterator<Item=i32> = &mut *it;
    assert_eq!(jt.next(), Some(1));

    // they all crash/ICE at the moment.
    /*
    let mut r = ByRef(jt);
    //assert_eq!(r.next(), Some(2));

    //it.drain();
    //assert_eq!(it.dropn(2), 2);
    */
}

#[test]
fn merge() {
    assert_iters_equal((0..10).step(2).merge((1..10).step(2)), (0..10));
}

#[test]
fn to_string_join() {
    let many = [1, 2, 3];
    let one  = [1];
    let none: Vec<i32> = vec![];

    assert_eq!(many.iter().to_string_join(", "), "1, 2, 3");
    assert_eq!( one.iter().to_string_join(", "), "1");
    assert_eq!(none.iter().to_string_join(", "), "");
}

#[test]
fn multipeek() {
    let nums = vec![1u8,2,3,4,5];

    let multipeek = nums.iter().map(|&x| x).multipeek();
    assert_eq!(nums, multipeek.collect());

    let mut multipeek = nums.iter().map(|&x| x).multipeek();
    assert_eq!(multipeek.peek(), Some(&1));
    assert_eq!(multipeek.next(), Some(1));
    assert_eq!(multipeek.peek(), Some(&2));
    assert_eq!(multipeek.peek(), Some(&3));
    assert_eq!(multipeek.next(), Some(2));
    assert_eq!(multipeek.peek(), Some(&3));
    assert_eq!(multipeek.peek(), Some(&4));
    assert_eq!(multipeek.peek(), Some(&5));
    assert_eq!(multipeek.peek(), None);
    assert_eq!(multipeek.next(), Some(3));
    assert_eq!(multipeek.next(), Some(4));
    assert_eq!(multipeek.next(), Some(5));
    assert_eq!(multipeek.next(), None);
    assert_eq!(multipeek.peek(), None);

}
<commit_msg>forgot to remove this line..<commit_after>//! Licensed under the Apache License, Version 2.0
//! http://www.apache.org/licenses/LICENSE-2.0 or the MIT license
//! http://opensource.org/licenses/MIT, at your
//! option. This file may not be copied, modified, or distributed
//! except according to those terms.

#![feature(core, test)]

#[macro_use]
extern crate itertools;

extern crate test;

use std::fmt::Debug;
use std::iter::order;
use itertools::Itertools;
use itertools::Interleave;
use itertools::Zip;

use itertools as it;

fn assert_iters_equal<
    A: PartialEq + Debug,
    I: Iterator<Item=A>,
    J: Iterator<Item=A>>(mut it: I, mut jt: J)
{
    loop {
        let elti = it.next();
        let eltj = jt.next();
        assert_eq!(elti, eltj);
        if elti.is_none() { break; }
    }
}

#[test]
fn product2() {
    let s = "αβ";

    let mut prod = iproduct!(s.chars(), range(0, 2));
    assert!(prod.next() == Some(('α', 0)));
    assert!(prod.next() == Some(('α', 1)));
    assert!(prod.next() == Some(('β', 0)));
    assert!(prod.next() == Some(('β', 1)));
    assert!(prod.next() == None);
}

#[test]
fn product3() {
    let prod = iproduct!(range(0, 3), range(0, 2), range(0, 2));
    assert_eq!(prod.size_hint(), (12, Some(12)));
    let v = prod.collect_vec();
    for i in range(0,3) {
        for j in range(0, 2) {
            for k in range(0, 2) {
                assert!((i, j, k) == v[(i * 2 * 2 + j * 2 + k) as usize]);
            }
        }
    }
    for (_, _, _, _) in iproduct!(range(0, 3), range(0, 2), range(0, 2), range(0, 3)) {
        /* test compiles */
    }
}

#[test]
fn izip3() {
    let mut zip = Zip::new((range(0, 3), range(0, 2), range(0, 2i8)));
    for i in range(0, 2) {
        assert!((i as usize, i, i as i8) == zip.next().unwrap());
    }
    assert!(zip.next().is_none());

    let xs: [isize; 0] = [];
    let mut zip = Zip::new((range(0, 3), range(0, 2), range(0, 2i8), xs.iter()));
    assert!(zip.next().is_none());

    for (_, _, _, _) in Zip::new((range(0, 3), range(0, 2), range(0, 2), range(0, 3))) {
        /* test compiles */
    }
}

#[test]
fn write_to() {
    let xs = [7, 9, 8];
    let mut ys = [0; 5];
    let cnt = it::write(ys.iter_mut(), xs.iter().map(|x| *x));
    assert!(cnt == xs.len());
    assert!(ys == [7, 9, 8, 0, 0]);

    let cnt = it::write(ys.iter_mut(), range(0,10));
    assert!(cnt == ys.len());
    assert!(ys == [0, 1, 2, 3, 4]);
}

#[test]
fn interleave() {
    let xs: [u8; 0]  = [];
    let ys = [7u8, 9, 8, 10];
    let zs = [2u8, 77];
    let it = Interleave::new(xs.iter(), ys.iter());
    assert!(order::eq(it, ys.iter()));

    let rs = [7u8, 2, 9, 77, 8, 10];
    let it = Interleave::new(ys.iter(), zs.iter());
    assert!(order::eq(it, rs.iter()));
}

#[test]
fn times() {
    assert!(it::times(0).count() == 0);
    assert!(it::times(5).count() == 5);
}

#[test]
fn drain() {
    let xs = [1i32,2,3];
    let mut sum = 0;
    xs.iter().map(|elt| sum += *elt).drain();
    assert!(sum == 6);
}

#[test]
fn foreach() {
    let xs = [1i32, 2, 3];
    let mut sum = 0;
    xs.iter().foreach(|elt| sum += *elt);
    assert!(sum == 6);
}

#[test]
fn dropn() {
    let xs = [1, 2, 3];
    let mut it = xs.iter();
    assert!(it.dropn(2) == 2);
    assert!(it.next().is_some());
    assert!(it.next().is_none());
    let mut it = xs.iter();
    assert!(it.dropn(5) == 3);
    assert!(it.next().is_none());
}

#[test]
fn dropping() {
    let xs = [1, 2, 3];
    let mut it = xs.iter().dropping(2);
    assert!(it.next().is_some());
    assert!(it.next().is_none());
    let mut it = xs.iter().dropping(5);
    assert!(it.next().is_none());
}

#[test]
fn intersperse() {
    let xs = ["a", "", "b", "c"];
    let v: Vec<&str> = xs.iter().map(|x| x.clone()).intersperse(", ").collect();
    let text: String = v.concat();
    assert_eq!(text, "a, , b, c".to_string());

    let ys = [0, 1, 2, 3];
    let mut it = ys[..0].iter().map(|x| *x).intersperse(1);
    assert!(it.next() == None);
}

#[test]
fn linspace() {
    let mut iter = it::linspace::<f32>(0., 2., 3);
    assert_eq!(iter.next(), Some(0.0));
    assert_eq!(iter.next(), Some(1.0));
    assert_eq!(iter.next(), Some(2.0));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., -2., 4);
    assert_eq!(iter.next(), Some(0.));
    assert_eq!(iter.next(), Some(-0.666666666667));
    assert_eq!(iter.next(), Some(-1.333333333333));
    assert_eq!(iter.next(), Some(-2.));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., 1., 1);
    assert_eq!(iter.next(), Some(0.));
    assert_eq!(iter.next(), None);

    let mut iter = it::linspace::<f32>(0., 1., 0);
    assert_eq!(iter.next(), None);
}

#[test]
fn dedup() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let ys = [0, 1, 2, 1, 3];
    assert_iters_equal(ys.iter(), xs.iter().dedup());
    let xs = [0, 0, 0, 0, 0];
    let ys = [0];
    assert_iters_equal(ys.iter(), xs.iter().dedup());
}

#[test]
fn batching() {
    let xs = [0, 1, 2, 1, 3];
    let ys = [(0, 1), (2, 1)];

    // An iterator that gathers elements up in pairs
    let pit = xs.iter().cloned().batching(|mut it| {
               match it.next() {
                   None => None,
                   Some(x) => match it.next() {
                       None => None,
                       Some(y) => Some((x, y)),
                   }
               }
           });
    assert_iters_equal(pit, ys.iter().cloned());
}

#[test]
fn group_by() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let ans = vec![(0, vec![0]), (1, vec![1, 1, 1]),
                   (2, vec![2]), (1, vec![1]), (3, vec![3, 3])];

    let gb = xs.iter().cloned().group_by(|elt| *elt);
    assert_iters_equal(gb, ans.into_iter());
}

#[test]
fn put_back() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 3];
    let mut pb = it::PutBack::new(xs.iter().cloned());
    pb.next();
    pb.put_back(1);
    pb.put_back(0);
    assert_iters_equal(pb, xs.iter().cloned());
}

#[test]
fn tee() {
    let xs  = [0, 1, 2, 3];
    let (mut t1, mut t2) = xs.iter().cloned().tee();
    assert_eq!(t1.next(), Some(0));
    assert_eq!(t2.next(), Some(0));
    assert_eq!(t1.next(), Some(1));
    assert_eq!(t1.next(), Some(2));
    assert_eq!(t1.next(), Some(3));
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), Some(1));
    assert_eq!(t2.next(), Some(2));
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), Some(3));
    assert_eq!(t2.next(), None);
    assert_eq!(t1.next(), None);
    assert_eq!(t2.next(), None);

    let (t1, t2) = xs.iter().cloned().tee();
    assert_iters_equal(t1, xs.iter().cloned());
    assert_iters_equal(t2, xs.iter().cloned());

    let (t1, t2) = xs.iter().cloned().tee();
    assert_iters_equal(t1.zip(t2), xs.iter().cloned().zip(xs.iter().cloned()));
}

#[test]
fn rciter() {
    let xs = [0, 1, 1, 1, 2, 1, 3, 5, 6];

    let mut r1 = xs.iter().cloned().into_rc();
    let mut r2 = r1.clone();
    assert_eq!(r1.next(), Some(0));
    assert_eq!(r2.next(), Some(1));
    let mut z = r1.zip(r2);
    assert_eq!(z.next(), Some((1, 1)));
    assert_eq!(z.next(), Some((2, 1)));
    assert_eq!(z.next(), Some((3, 5)));
    assert_eq!(z.next(), None);
}

#[test]
fn slice() {
    assert_iters_equal((0..10).slice(..3), 0..3);
    assert_iters_equal((0..10).slice(3..7), 3..7);
    assert_iters_equal((0..10).slice(3..27), 3..10);
    assert_iters_equal((0..10).slice(44..), 0..0);
}

#[test]
fn step() {
    assert_iters_equal((0..10).step(1), (0..10));
    assert_iters_equal((0..10).step(2), (0..10).filter(|x: &i32| *x % 2 == 0));
    assert_iters_equal((0..10).step(10), 0..1);
}

#[test]
fn trait_pointers() {
    /*
    struct ByRef<'r, I: ?Sized>(&'r mut I) where I: 'r;
    impl<'r, X, I: ?Sized> Iterator for ByRef<'r, I> where
        I: 'r + Iterator<Item=X>
    {
        type Item = X;
        fn next(&mut self) -> Option<X>
        {
            self.0.next()
        }
    }
    */
    let mut it = Box::new(0..10) as Box<Iterator<Item=i32>>;
    assert_eq!(it.next(), Some(0));

    let mut jt: &mut Iterator<Item=i32> = &mut *it;
    assert_eq!(jt.next(), Some(1));

    // they all crash/ICE at the moment.
    /*
    let mut r = ByRef(jt);
    //assert_eq!(r.next(), Some(2));

    //it.drain();
    //assert_eq!(it.dropn(2), 2);
    */
}

#[test]
fn merge() {
    assert_iters_equal((0..10).step(2).merge((1..10).step(2)), (0..10));
}

#[test]
fn to_string_join() {
    let many = [1, 2, 3];
    let one  = [1];
    let none: Vec<i32> = vec![];

    assert_eq!(many.iter().to_string_join(", "), "1, 2, 3");
    assert_eq!( one.iter().to_string_join(", "), "1");
    assert_eq!(none.iter().to_string_join(", "), "");
}

#[test]
fn multipeek() {
    let nums = vec![1u8,2,3,4,5];

    let multipeek = nums.iter().map(|&x| x).multipeek();
    assert_eq!(nums, multipeek.collect());

    let mut multipeek = nums.iter().map(|&x| x).multipeek();
    assert_eq!(multipeek.peek(), Some(&1));
    assert_eq!(multipeek.next(), Some(1));
    assert_eq!(multipeek.peek(), Some(&2));
    assert_eq!(multipeek.peek(), Some(&3));
    assert_eq!(multipeek.next(), Some(2));
    assert_eq!(multipeek.peek(), Some(&3));
    assert_eq!(multipeek.peek(), Some(&4));
    assert_eq!(multipeek.peek(), Some(&5));
    assert_eq!(multipeek.peek(), None);
    assert_eq!(multipeek.next(), Some(3));
    assert_eq!(multipeek.next(), Some(4));
    assert_eq!(multipeek.next(), Some(5));
    assert_eq!(multipeek.next(), None);
    assert_eq!(multipeek.peek(), None);

}
<|endoftext|>


<commit_before>mod common;

use std::thread;

use rand::{thread_rng, Rng};

use sled::*;

use sled::{
    pin, BatchManifest, Log, LogKind, LogRead, Lsn, PageId, SEG_HEADER_LEN,
};

const PID: PageId = 4;
const REPLACE: LogKind = LogKind::Replace;

#[test]
fn log_writebatch() -> crate::Result<()> {
    common::setup_logger();
    let config = Config::new().temporary(true);
    let db = config.open()?;
    let log = &db.context.pagecache.log;

    let guard = pin();

    log.reserve(REPLACE, PID, &IVec::from(b"1"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"2"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"3"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"4"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"5"), &guard)?.complete()?;

    // simulate a torn batch by
    // writing an LSN higher than
    // is possible to recover into
    // a batch manifest before
    // some writes.
    let mut batch_res =
        log.reserve(REPLACE, PID, &BatchManifest::default(), &guard)?;
    log.reserve(REPLACE, PID, &IVec::from(b"6"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"7"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"8"), &guard)?.complete()?;
    let last_res = log.reserve(REPLACE, PID, &IVec::from(b"9"), &guard)?;
    let last_res_lsn = last_res.lsn();
    last_res.complete()?;
    batch_res.mark_writebatch(last_res_lsn + 50, &guard);
    batch_res.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"10"), &guard)?.complete()?;

    let mut iter = log.iter_from(0);

    assert!(iter.next().is_some());
    assert!(iter.next().is_some());
    assert!(iter.next().is_some());
    assert!(iter.next().is_some());
    assert!(iter.next().is_some());
    assert_eq!(iter.next(), None);

    Ok(())
}

#[test]
fn more_log_reservations_than_buffers() -> Result<()> {
    let config = Config::new().temporary(true).segment_size(128);

    let db = config.open()?;
    let log = &db.context.pagecache.log;
    let mut reservations = vec![];

    let total_seg_overhead = SEG_HEADER_LEN;
    let big_msg_overhead = MAX_MSG_HEADER_LEN + total_seg_overhead;
    let big_msg_sz = config.segment_size - big_msg_overhead;

    for _ in 0..=30 {
        let guard = pin();
        reservations.push(
            log.reserve(REPLACE, PID, &IVec::from(vec![0; big_msg_sz]), &guard)
                .unwrap(),
        )
    }
    for res in reservations.into_iter().rev() {
        // abort in reverse order
        res.abort()?;
    }

    log.flush()?;
    Ok(())
}

#[test]
fn non_contiguous_log_flush() -> Result<()> {
    let config = Config::new().temporary(true).segment_size(1024);

    let db = config.open()?;
    let log = &db.context.pagecache.log;

    let seg_overhead = SEG_HEADER_LEN;
    let buf_len = (config.segment_size / MINIMUM_ITEMS_PER_SEGMENT)
        - (MAX_MSG_HEADER_LEN + seg_overhead);

    let guard = pin();
    let res1 = log
        .reserve(REPLACE, PID, &IVec::from(vec![0; buf_len]), &guard)
        .unwrap();
    let res2 = log
        .reserve(REPLACE, PID, &IVec::from(vec![0; buf_len]), &guard)
        .unwrap();
    let lsn = res2.lsn();
    res2.abort()?;
    res1.abort()?;
    log.make_stable(lsn)?;
    Ok(())
}

#[test]
fn concurrent_logging() -> Result<()> {
    common::setup_logger();
    for _ in 0..10 {
        let config = Config::new()
            .temporary(true)
            .flush_every_ms(Some(50))
            .segment_size(256);

        let db = config.open()?;

        let db2 = db.clone();
        let db3 = db.clone();
        let db4 = db.clone();
        let db5 = db.clone();
        let db6 = db.clone();

        let seg_overhead = SEG_HEADER_LEN;
        let buf_len = (config.segment_size / MINIMUM_ITEMS_PER_SEGMENT)
            - (MAX_MSG_HEADER_LEN + seg_overhead);

        let t1 = thread::Builder::new()
            .name("c1".to_string())
            .spawn(move || {
                let log = &db.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![1; i % buf_len]);
                    let guard = pin();
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();

        let t2 = thread::Builder::new()
            .name("c2".to_string())
            .spawn(move || {
                let log = &db2.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![2; i % buf_len]);
                    let guard = pin();
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();

        let t3 = thread::Builder::new()
            .name("c3".to_string())
            .spawn(move || {
                let log = &db3.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![3; i % buf_len]);
                    let guard = pin();
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();

        let t4 = thread::Builder::new()
            .name("c4".to_string())
            .spawn(move || {
                let log = &db4.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![4; i % buf_len]);
                    let guard = pin();
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();
        let t5 = thread::Builder::new()
            .name("c5".to_string())
            .spawn(move || {
                let log = &db5.context.pagecache.log;
                for i in 0..1_000 {
                    let guard = pin();
                    let buf = IVec::from(vec![5; i % buf_len]);
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();

        let t6 = thread::Builder::new()
            .name("c6".to_string())
            .spawn(move || {
                let log = &db6.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![6; i % buf_len]);
                    let guard = pin();
                    let (lsn, _lid) = log
                        .reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                    log.make_stable(lsn).unwrap();
                }
            })
            .unwrap();

        t6.join().unwrap();
        t5.join().unwrap();
        t4.join().unwrap();
        t3.join().unwrap();
        t2.join().unwrap();
        t1.join().unwrap();
    }

    Ok(())
}

fn write(log: &Log) {
    let data_bytes = IVec::from(b"yoyoyoyo");
    let guard = pin();
    let (lsn, ptr) = log
        .reserve(REPLACE, PID, &data_bytes, &guard)
        .unwrap()
        .complete()
        .unwrap();
    let read_buf = log.read(PID, lsn, ptr).unwrap().into_data().unwrap();
    assert_eq!(
        *read_buf,
        *data_bytes.serialize(),
        "after writing data, it should be readable"
    );
}

fn abort(log: &Log) {
    let guard = pin();
    let res = log.reserve(REPLACE, PID, &IVec::from(&[0; 5]), &guard).unwrap();
    let (lsn, ptr) = res.abort().unwrap();
    match log.read(PID, lsn, ptr) {
        Ok(LogRead::Canceled(_)) => {}
        other => {
            panic!(
                "expected to successfully read \
                 aborted log message, instead read {:?}",
                other
            );
        }
    }
}

#[test]
fn log_aborts() {
    let config = Config::new().temporary(true);
    let db = config.open().unwrap();
    let log = &db.context.pagecache.log;
    write(&log);
    abort(&log);
    write(&log);
    abort(&log);
    write(&log);
    abort(&log);
}

#[test]
#[cfg(not(target_os = "fuchsia"))]
fn log_chunky_iterator() {
    common::setup_logger();
    let mut threads = vec![];
    for tn in 0..1 {
        let thread = thread::spawn(move || {
            let config = Config::new().temporary(true).segment_size(128);

            let db = config.open().unwrap();
            let log = &db.context.pagecache.log;

            let mut reference = vec![];

            let max_valid_size =
                config.segment_size - (MAX_MSG_HEADER_LEN + SEG_HEADER_LEN);

            for i in PID..1000 {
                let len = thread_rng().gen_range(0, max_valid_size * 2);
                let item = thread_rng().gen::<u8>();
                let buf = IVec::from(vec![item; len]);
                let abort = thread_rng().gen::<bool>();

                let pid = (tn * 10000) + i;

                let guard = pin();

                if abort {
                    let res = log
                        .reserve(REPLACE, pid, &buf, &guard)
                        .expect("should be able to reserve");
                    res.abort().unwrap();
                } else {
                    let res = log
                        .reserve(REPLACE, pid, &buf, &guard)
                        .expect("should be able to write reservation");
                    let ptr = res.pointer();
                    let (lsn, _) = res.complete().unwrap();
                    reference.push((REPLACE, pid, lsn, ptr));
                }
            }

            let mut ref_iter = reference.clone().into_iter();
            for _ in log.iter_from(SEG_HEADER_LEN as Lsn) {
                assert!(ref_iter.next().is_some());
            }

            for (_, pid, lsn, ptr) in reference.clone().into_iter() {
                assert!(log.read(pid, lsn, ptr).is_ok());
            }
        });
        threads.push(thread);
    }
    for thread in threads.into_iter() {
        thread.join().unwrap();
    }
}

#[test]
fn multi_segment_log_iteration() -> Result<()> {
    common::setup_logger();
    // ensure segments are being linked
    // ensure trailers are valid
    let config = Config::new().temporary(true).segment_size(512);

    let total_seg_overhead = SEG_HEADER_LEN;
    let big_msg_overhead = MAX_MSG_HEADER_LEN + total_seg_overhead;
    let big_msg_sz = (config.segment_size - big_msg_overhead) / 64;

    let db = config.open().unwrap();
    let log = &db.context.pagecache.log;

    for i in 0..48 {
        let buf = IVec::from(vec![i as u8; big_msg_sz * i]);
        let guard = pin();
        log.reserve(REPLACE, i as PageId, &buf, &guard)
            .unwrap()
            .complete()
            .unwrap();
    }
    log.flush()?;

    // start iterating just past the first segment header
    let mut iter = log.iter_from(SEG_HEADER_LEN as Lsn);

    for _ in 0..48 {
        iter.next().expect("expected to read another message");
    }

    Ok(())
}
<commit_msg>Avoid race condition in log test<commit_after>mod common;

use std::thread;

use rand::{thread_rng, Rng};

use sled::*;

use sled::{
    pin, BatchManifest, Log, LogKind, LogRead, Lsn, PageId, SEG_HEADER_LEN,
};

const PID: PageId = 4;
const REPLACE: LogKind = LogKind::Replace;

#[test]
fn log_writebatch() -> crate::Result<()> {
    common::setup_logger();
    let config = Config::new().temporary(true);
    let db = config.open()?;
    let log = &db.context.pagecache.log;

    let guard = pin();

    log.reserve(REPLACE, PID, &IVec::from(b"1"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"2"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"3"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"4"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"5"), &guard)?.complete()?;

    // simulate a torn batch by
    // writing an LSN higher than
    // is possible to recover into
    // a batch manifest before
    // some writes.
    let mut batch_res =
        log.reserve(REPLACE, PID, &BatchManifest::default(), &guard)?;
    log.reserve(REPLACE, PID, &IVec::from(b"6"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"7"), &guard)?.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"8"), &guard)?.complete()?;
    let last_res = log.reserve(REPLACE, PID, &IVec::from(b"9"), &guard)?;
    let last_res_lsn = last_res.lsn();
    last_res.complete()?;
    batch_res.mark_writebatch(last_res_lsn + 50, &guard);
    batch_res.complete()?;
    log.reserve(REPLACE, PID, &IVec::from(b"10"), &guard)?.complete()?;

    let mut iter = log.iter_from(0);

    assert!(iter.next().is_some());
    assert!(iter.next().is_some());
    assert!(iter.next().is_some());
    assert!(iter.next().is_some());
    assert!(iter.next().is_some());
    assert_eq!(iter.next(), None);

    Ok(())
}

#[test]
fn more_log_reservations_than_buffers() -> Result<()> {
    let config = Config::new().temporary(true).segment_size(128);

    let db = config.open()?;
    let log = &db.context.pagecache.log;
    let mut reservations = vec![];

    let total_seg_overhead = SEG_HEADER_LEN;
    let big_msg_overhead = MAX_MSG_HEADER_LEN + total_seg_overhead;
    let big_msg_sz = config.segment_size - big_msg_overhead;

    for _ in 0..=30 {
        let guard = pin();
        reservations.push(
            log.reserve(REPLACE, PID, &IVec::from(vec![0; big_msg_sz]), &guard)
                .unwrap(),
        )
    }
    for res in reservations.into_iter().rev() {
        // abort in reverse order
        res.abort()?;
    }

    log.flush()?;
    Ok(())
}

#[test]
fn non_contiguous_log_flush() -> Result<()> {
    let config = Config::new().temporary(true).segment_size(1024);

    let db = config.open()?;
    let log = &db.context.pagecache.log;

    let seg_overhead = SEG_HEADER_LEN;
    let buf_len = (config.segment_size / MINIMUM_ITEMS_PER_SEGMENT)
        - (MAX_MSG_HEADER_LEN + seg_overhead);

    let guard = pin();
    let res1 = log
        .reserve(REPLACE, PID, &IVec::from(vec![0; buf_len]), &guard)
        .unwrap();
    let res2 = log
        .reserve(REPLACE, PID, &IVec::from(vec![0; buf_len]), &guard)
        .unwrap();
    let lsn = res2.lsn();
    res2.abort()?;
    res1.abort()?;
    log.make_stable(lsn)?;
    Ok(())
}

#[test]
fn concurrent_logging() -> Result<()> {
    common::setup_logger();
    for _ in 0..10 {
        let config = Config::new()
            .temporary(true)
            .flush_every_ms(Some(50))
            .segment_size(256);

        let db = config.open()?;

        let db2 = db.clone();
        let db3 = db.clone();
        let db4 = db.clone();
        let db5 = db.clone();
        let db6 = db.clone();

        let seg_overhead = SEG_HEADER_LEN;
        let buf_len = (config.segment_size / MINIMUM_ITEMS_PER_SEGMENT)
            - (MAX_MSG_HEADER_LEN + seg_overhead);

        let t1 = thread::Builder::new()
            .name("c1".to_string())
            .spawn(move || {
                let log = &db.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![1; i % buf_len]);
                    let guard = pin();
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();

        let t2 = thread::Builder::new()
            .name("c2".to_string())
            .spawn(move || {
                let log = &db2.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![2; i % buf_len]);
                    let guard = pin();
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();

        let t3 = thread::Builder::new()
            .name("c3".to_string())
            .spawn(move || {
                let log = &db3.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![3; i % buf_len]);
                    let guard = pin();
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();

        let t4 = thread::Builder::new()
            .name("c4".to_string())
            .spawn(move || {
                let log = &db4.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![4; i % buf_len]);
                    let guard = pin();
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();
        let t5 = thread::Builder::new()
            .name("c5".to_string())
            .spawn(move || {
                let log = &db5.context.pagecache.log;
                for i in 0..1_000 {
                    let guard = pin();
                    let buf = IVec::from(vec![5; i % buf_len]);
                    log.reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                }
            })
            .unwrap();

        let t6 = thread::Builder::new()
            .name("c6".to_string())
            .spawn(move || {
                let log = &db6.context.pagecache.log;
                for i in 0..1_000 {
                    let buf = IVec::from(vec![6; i % buf_len]);
                    let guard = pin();
                    let (lsn, _lid) = log
                        .reserve(REPLACE, PID, &buf, &guard)
                        .unwrap()
                        .complete()
                        .unwrap();
                    log.make_stable(lsn).unwrap();
                }
            })
            .unwrap();

        t6.join().unwrap();
        t5.join().unwrap();
        t4.join().unwrap();
        t3.join().unwrap();
        t2.join().unwrap();
        t1.join().unwrap();
    }

    Ok(())
}

fn write(log: &Log) {
    let data_bytes = IVec::from(b"yoyoyoyo");
    let guard = pin();
    let (lsn, ptr) = log
        .reserve(REPLACE, PID, &data_bytes, &guard)
        .unwrap()
        .complete()
        .unwrap();
    let read_buf = log.read(PID, lsn, ptr).unwrap().into_data().unwrap();
    assert_eq!(
        *read_buf,
        *data_bytes.serialize(),
        "after writing data, it should be readable"
    );
}

fn abort(log: &Log) {
    let guard = pin();
    let res = log.reserve(REPLACE, PID, &IVec::from(&[0; 5]), &guard).unwrap();
    let (lsn, ptr) = res.abort().unwrap();
    match log.read(PID, lsn, ptr) {
        Ok(LogRead::Canceled(_)) => {}
        other => {
            panic!(
                "expected to successfully read \
                 aborted log message, instead read {:?}",
                other
            );
        }
    }
}

#[test]
fn log_aborts() {
    let config = Config::new().temporary(true);
    let db = config.open().unwrap();
    let log = &db.context.pagecache.log;
    write(&log);
    abort(&log);
    write(&log);
    abort(&log);
    write(&log);
    abort(&log);
}

#[test]
#[cfg(not(target_os = "fuchsia"))]
fn log_chunky_iterator() {
    common::setup_logger();
    let config =
        Config::new().flush_every_ms(None).temporary(true).segment_size(128);

    let db = config.open().unwrap();
    let log = &db.context.pagecache.log;

    let mut reference = vec![];

    let max_valid_size =
        config.segment_size - (MAX_MSG_HEADER_LEN + SEG_HEADER_LEN);

    for i in PID..1000 {
        let len = thread_rng().gen_range(0, max_valid_size * 2);
        let item = thread_rng().gen::<u8>();
        let buf = IVec::from(vec![item; len]);
        let abort = thread_rng().gen::<bool>();

        let pid = 10000 + i;

        let guard = pin();

        if abort {
            let res = log
                .reserve(REPLACE, pid, &buf, &guard)
                .expect("should be able to reserve");
            res.abort().unwrap();
        } else {
            let res = log
                .reserve(REPLACE, pid, &buf, &guard)
                .expect("should be able to write reservation");
            let ptr = res.pointer();
            let (lsn, _) = res.complete().unwrap();
            reference.push((REPLACE, pid, lsn, ptr));
        }
    }

    for (_, pid, lsn, ptr) in reference.clone().into_iter() {
        assert!(log.read(pid, lsn, ptr).is_ok());
    }
}

#[test]
fn multi_segment_log_iteration() -> Result<()> {
    common::setup_logger();
    // ensure segments are being linked
    // ensure trailers are valid
    let config = Config::new().temporary(true).segment_size(512);

    let total_seg_overhead = SEG_HEADER_LEN;
    let big_msg_overhead = MAX_MSG_HEADER_LEN + total_seg_overhead;
    let big_msg_sz = (config.segment_size - big_msg_overhead) / 64;

    let db = config.open().unwrap();
    let log = &db.context.pagecache.log;

    for i in 0..48 {
        let buf = IVec::from(vec![i as u8; big_msg_sz * i]);
        let guard = pin();
        log.reserve(REPLACE, i as PageId, &buf, &guard)
            .unwrap()
            .complete()
            .unwrap();
    }
    log.flush()?;

    // start iterating just past the first segment header
    let mut iter = log.iter_from(SEG_HEADER_LEN as Lsn);

    for _ in 0..48 {
        iter.next().expect("expected to read another message");
    }

    Ok(())
}
<|endoftext|>


<commit_before>// It's OK to destroy a heap when some GcRefs still exist. (!)
//
// Rust's safety rules ensure that the GcRefs won't still be accessible
// afterwards, so if this happens, it means the values were leaked.

extern crate cell_gc;

use cell_gc::Heap;
use std::mem;
use std::sync::Arc;

#[test]
fn test_leaked_ref() {
    let counted = Arc::new(false);
    assert_eq!(Arc::strong_count(&counted), 1);

    let mut heap = Heap::new();
    heap.enter(|hs| {
        let a = hs.alloc(counted.clone());
        assert_eq!(Arc::strong_count(&counted), 2);
        mem::forget(a);

        // The Arc in the heap remains pinned.
        hs.force_gc();
        assert_eq!(Arc::strong_count(&counted), 2);
    });

    // This destroys all remaining in-heap values, even though we leaked a
    // strong reference.
    drop(heap);

    assert_eq!(Arc::strong_count(&counted), 1);
}
<commit_msg>Make a comment on a test less sensational.<commit_after>// It's OK to destroy a heap when some GcRefs have not been dropped. (!)
//
// Rust's safety rules ensure that the GcRefs won't still be accessible
// afterwards, so if this happens, it means the values were leaked.

extern crate cell_gc;

use cell_gc::Heap;
use std::mem;
use std::sync::Arc;

#[test]
fn test_leaked_ref() {
    let counted = Arc::new(false);
    assert_eq!(Arc::strong_count(&counted), 1);

    let mut heap = Heap::new();
    heap.enter(|hs| {
        let a = hs.alloc(counted.clone());
        assert_eq!(Arc::strong_count(&counted), 2);
        mem::forget(a);

        // The Arc in the heap remains pinned.
        hs.force_gc();
        assert_eq!(Arc::strong_count(&counted), 2);
    });

    // This destroys all remaining in-heap values, even though we leaked a
    // strong reference.
    drop(heap);

    assert_eq!(Arc::strong_count(&counted), 1);
}
<|endoftext|>


<commit_before>use evalbotlib::{util, EvalService, Language};

use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;

use futures::StreamExt;
use log::{debug, error, info, warn};
use serde::{Deserialize, Serialize};
use telegram_bot::*;
use tokio::sync::RwLock;

static WHITELIST_FILENAME: &'static str = "tgwhitelist.toml";

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
struct TgCfg {
    owners: HashSet<i64>,
    msg_owner_id: Option<i64>,
    bot_id: String,
    lang_subst: HashMap<String, String>,
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
struct TgWhitelist {
    priv_enabled: bool,
    group_enabled: bool,
    allowed: HashSet<i64>,
    blocked: HashSet<i64>,
}

impl TgWhitelist {
    fn priv_ok(&self, id: ChatId) -> bool {
        (!self.priv_enabled || self.allowed.contains(&id.into()))
            && !self.blocked.contains(&id.into())
    }

    fn group_ok(&self, id: ChatId) -> bool {
        (!self.group_enabled || self.allowed.contains(&id.into()))
            && !self.blocked.contains(&id.into())
    }

    fn allow(&mut self, id: i64) {
        self.allowed.insert(id);
    }

    fn unallow(&mut self, id: i64) {
        self.allowed.remove(&id);
    }

    fn block(&mut self, id: i64) {
        self.blocked.insert(id);
    }

    fn unblock(&mut self, id: i64) {
        self.blocked.remove(&id);
    }

    async fn save<'a>(&'a self, path: &'static str) {
        drop(util::encode(self, path).await);
    }
}

struct TgSvc {
    config: TgCfg,
    whitelist: RwLock<TgWhitelist>,
    service: EvalService,
    api: Api,
    bot_user: User,
    username: String,
}

fn telegram_wrap_result(s: &str, group: bool) -> String {
    // FIXME configurable max-lines and max-bytes
    if s.is_empty() {
        "no output".to_owned()
    } else {
        let mut r = "
".to_owned();
        let input = s.as_bytes();
        let mut cut_input = String::from_utf8_lossy(&input[..512.min(input.len())]);
        if group {
            cut_input = Cow::Owned(cut_input.lines().take(10).collect::<Vec<_>>().join("\n"));
        }
        r.push_str(
            &cut_input
                .replace(
                    |c: char| c == '\u{FFFD}' || (c.is_control() && c != '\n' && c != '\t'),
                    "",
                )
                .replace('&', "&")
                .replace('<', "<")
                .replace('>', ">")
                .replace('"', """),
        );
        r.push_str("
");
        if cut_input.len() + 1 // we also cut off the trailing \n
            < input.len()
        {
            r.push_str("... (truncated)");
        }
        r
    }
}

fn is_from_owner(msg: &Message, tgsvc: &TgSvc) -> bool {
    tgsvc.is_owner(msg.from.to_user_id())
}

async fn verify_allowed(chat: &MessageChat, tgsvc: &Arc<TgSvc>) -> Result<(), ()> {
    let group = if let MessageChat::Private(_) = chat {
        false
    } else {
        true
    };
    let wl = tgsvc.whitelist.read().await;
    let chat_id = chat.id();
    if group && !wl.group_ok(chat_id) || !group && !wl.priv_ok(chat_id) {
        let send_message = tgsvc.api.send(SendMessage::new(
            &chat,
            format!(
                "You or this group is not on the whitelist. Seek help. ID: {}",
                chat_id
            ),
        ));
        let leave_group = if group {
            Some(tgsvc.api.send(LeaveChat::new(&chat)))
        } else {
            None
        };
        tokio::spawn(async move {
            drop(send_message.await.map_err(|e| {
                warn!("failed to send message: {}", e);
            }));
            if let Some(leave_group) = leave_group {
                drop(leave_group.await.map_err(|e| {
                    warn!("failed to leave group {}: {}", chat_id, e);
                }));
            }
        });
        Err(())
    } else {
        Ok(())
    }
}

async fn handle_update(update: Update, tgsvc: Arc<TgSvc>) -> Result<(), ()> {
    let message = match update {
        Update {
            kind: UpdateKind::Message(message),
            ..
        } => message,
        _ => return Ok(()),
    };

    match message.kind {
        MessageKind::GroupChatCreated => handle_group_created(message.chat, &tgsvc).await?,
        MessageKind::NewChatMembers { data } => {
            handle_new_chat_members(message.chat, data, &tgsvc).await?
        }
        MessageKind::Text { .. } => handle_message(message, &tgsvc).await?,
        _ => (),
    };

    Ok(())
}

async fn handle_new_chat_members(
    chat: MessageChat,
    new_members: Vec<User>,
    tgsvc: &Arc<TgSvc>,
) -> Result<(), ()> {
    if new_members.iter().any(|u| u.id == tgsvc.bot_user.id) {
        verify_allowed(&chat, tgsvc).await?;
    }

    Ok(())
}

async fn handle_group_created(chat: MessageChat, tgsvc: &Arc<TgSvc>) -> Result<(), ()> {
    verify_allowed(&chat, tgsvc).await?;
    Ok(())
}

async fn handle_message(message: Message, tgsvc: &Arc<TgSvc>) -> Result<(), ()> {
    let text = if let MessageKind::Text { ref data, .. } = message.kind {
        data
    } else {
        return Ok(());
    };

    if !text.starts_with("/") {
        return Ok(());
    }

    let (cmd, args) = {
        let first_tok = if let Some(first_tok) = text.split_whitespace().nth(0) {
            first_tok
        } else {
            return Ok(());
        };

        let mut split = first_tok.splitn(1, '@');
        let cmd = split.next();
        let user = split.next();

        if let Some(user) = user {
            if user != tgsvc.username {
                return Ok(());
            }
        }

        (
            cmd.expect("First split should exist"),
            text[first_tok.len()..].trim_start(),
        )
    };

    match cmd {
        "/privwl" => {
            handle_whitelist_toggle(tgsvc, &message, WhitelistToggleOp::TogglePrivate).await?
        }
        "/groupwl" => {
            handle_whitelist_toggle(tgsvc, &message, WhitelistToggleOp::ToggleGroup).await?
        }
        "/allow" => handle_whitelist_mod(tgsvc, &message, args, WhitelistModOp::Allow).await?,
        "/unallow" => handle_whitelist_mod(tgsvc, &message, args, WhitelistModOp::Unallow).await?,
        "/block" => handle_whitelist_mod(tgsvc, &message, args, WhitelistModOp::Block).await?,
        "/unblock" => handle_whitelist_mod(tgsvc, &message, args, WhitelistModOp::Unblock).await?,
        "/leave" => handle_leave(tgsvc, &message, args).await?,
        _ => {
            let (cmd, is_hash) = if cmd.ends_with('#') {
                (&cmd[1..cmd.len() - 1], true)
            } else {
                (&cmd[1..], false)
            };
            if let Some((_, lang)) = tgsvc.service.langs().find(|(name, _)| *name == cmd) {
                handle_eval(tgsvc, &message, args, lang, is_hash).await?;
            }
        }
    }

    Ok(())
}

async fn handle_eval(
    tgsvc: &Arc<TgSvc>,
    msg: &Message,
    args: &str,
    lang: &Arc<Language>,
    is_hash: bool,
) -> Result<(), ()> {
    verify_allowed(&msg.chat, tgsvc).await?;
    let no_limit = is_hash && is_from_owner(&msg, tgsvc);
    let is_group = if let MessageChat::Private(_) = msg.chat {
        false
    } else {
        true
    };
    let chat_id = msg.chat.id();
    let msg_id = msg.id;
    info!("({}) evaluating from {:?}: {:?}", msg_id, msg.from, args);
    let code = {
        let mut r = args.to_owned();
        r.push_str("\n");
        r
    };

    let eval_result = lang
        .eval(
            code,
            if no_limit { Some(0) } else { None },
            Some(format!("tg{}", chat_id)),
        )
        .await;
    let ok = eval_result.is_ok();
    info!(
        "({}) result: {:?}",
        msg_id,
        eval_result.as_ref().map(|x| x.as_str()).unwrap_or("")
    );
    let mut request = SendMessage::new(
        &msg.chat,
        eval_result
            .map(|r| telegram_wrap_result(&r, is_group))
            .unwrap_or_else(|e| e),
    );
    request.reply_to(msg);
    if ok {
        request.parse_mode(ParseMode::Html);
    }
    tokio::spawn(tgsvc.api.send(request));

    Ok(())
}

#[derive(Clone, Copy)]
enum WhitelistToggleOp {
    TogglePrivate,
    ToggleGroup,
}

#[derive(Clone, Copy)]
enum WhitelistModOp {
    Allow,
    Unallow,
    Block,
    Unblock,
}

async fn handle_whitelist_toggle(
    tgsvc: &Arc<TgSvc>,
    msg: &Message,
    op: WhitelistToggleOp,
) -> Result<(), ()> {
    if !is_from_owner(&msg, tgsvc) {
        return Ok(());
    }

    let mut wl = tgsvc.whitelist.write().await;
    let resp = match op {
        WhitelistToggleOp::TogglePrivate => {
            wl.priv_enabled = !wl.priv_enabled;
            wl.save(WHITELIST_FILENAME).await;
            format!("Private whitelist enabled: {}", wl.priv_enabled)
        }
        WhitelistToggleOp::ToggleGroup => {
            wl.group_enabled = !wl.group_enabled;
            wl.save(WHITELIST_FILENAME).await;
            format!("Group whitelist enabled: {}", wl.group_enabled)
        }
    };
    tokio::spawn(tgsvc.api.send(SendMessage::new(&msg.chat, resp)));
    Ok(())
}

async fn handle_whitelist_mod(
    tgsvc: &Arc<TgSvc>,
    msg: &Message,
    args: &str,
    op: WhitelistModOp,
) -> Result<(), ()> {
    if !is_from_owner(&msg, tgsvc) {
        return Ok(());
    }

    let arg = args
        .trim()
        .split_whitespace()
        .nth(0)
        .and_then(|arg| arg.parse().ok());
    let mut wl = tgsvc.whitelist.write().await;
    let resp = if let Some(id) = arg {
        match op {
            WhitelistModOp::Allow => {
                wl.allow(id);
                wl.save(WHITELIST_FILENAME).await;
                format!("Allowed {}", id)
            }
            WhitelistModOp::Unallow => {
                wl.unallow(id);
                wl.save(WHITELIST_FILENAME).await;
                format!("Unallowed {}", id)
            }
            WhitelistModOp::Block => {
                wl.block(id);
                wl.save(WHITELIST_FILENAME).await;
                format!("Blocked {}", id)
            }
            WhitelistModOp::Unblock => {
                wl.unblock(id);
                wl.save(WHITELIST_FILENAME).await;
                format!("Unblocked {}", id)
            }
        }
    } else {
        "Invalid ID".to_owned()
    };

    tokio::spawn(tgsvc.api.send(SendMessage::new(&msg.chat, resp)));
    Ok(())
}

async fn handle_leave(tgsvc: &Arc<TgSvc>, msg: &Message, args: &str) -> Result<(), ()> {
    if !is_from_owner(&msg, tgsvc) {
        return Ok(());
    }

    let arg: Option<i64> = args
        .trim()
        .split_whitespace()
        .nth(0)
        .and_then(|arg| arg.parse().ok());
    let (maybe_leave_chat, resp) = match arg {
        Some(id) => (Some(tgsvc.api.send(LeaveChat::new(ChatId::from(id)))), "OK"),
        None => (None, "Invalid ID"),
    };

    let send_response = tgsvc.api.send(SendMessage::new(&msg.chat, resp));

    tokio::spawn(async move {
        if let Some(fut) = maybe_leave_chat {
            drop(fut.await.map_err(|e| warn!("failed to leave group: {}", e)));
        }
        drop(
            send_response
                .await
                .map_err(|e| warn!("failed to send message: {}", e)),
        );
    });
    Ok(())
}

impl TgSvc {
    async fn run() -> Result<(), ()> {
        let cfg = util::decode::<TgCfg, _>("evalbot.tg.toml")
            .await
            .map(|cfg| {
                debug!("Loaded config: {:?}", cfg);
                cfg
            })
            .map_err(|e| error!("failed to read evalbot.tg.toml: {}", e))?;
        let wl = util::decode::<TgWhitelist, _>(WHITELIST_FILENAME)
            .await
            .or_else(|e| {
                warn!("failed to read whitelist: {}; using empty whitelist", e);
                Ok(TgWhitelist {
                    priv_enabled: false,
                    group_enabled: false,
                    allowed: HashSet::new(),
                    blocked: HashSet::new(),
                })
            })?;

        let api = Api::new(&cfg.bot_id);
        let bot_user = api
            .send(GetMe)
            .await
            .map_err(|e| error!("failed to get bot's user info: {}", e))?;
        let service = EvalService::from_toml_file("evalbot.toml")
            .await
            .map_err(|e| error!("failed to read evalbot.toml: {}", e))?;
        TgSvc {
            api: Api::new(&cfg.bot_id),
            config: cfg,
            whitelist: RwLock::new(wl),
            service: service,
            username: bot_user
                .username
                .as_ref()
                .expect("Bot must have username")
                .clone(),
            bot_user: bot_user,
        }
        .handle()
        .await;
        Ok(())
    }

    async fn handle(self) {
        let me = Arc::new(self);

        let mut stream = me.api.stream();
        stream
            .timeout(Duration::from_secs(35))
            .error_delay(Duration::from_secs(30));
        while let Some(update) = stream.next().await {
            match update {
                Ok(update) => {
                    tokio::spawn(handle_update(update, me.clone()));
                }
                Err(error) => error!("received error: {:?}", error),
            }
        }
    }

    fn is_owner(&self, id: UserId) -> bool {
        self.config.owners.contains(&id.into())
    }
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();
    drop(TgSvc::run().await);
}
<commit_msg>Split into 2 parts, not 1<commit_after>use evalbotlib::{util, EvalService, Language};

use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;

use futures::StreamExt;
use log::{debug, error, info, warn};
use serde::{Deserialize, Serialize};
use telegram_bot::*;
use tokio::sync::RwLock;

static WHITELIST_FILENAME: &'static str = "tgwhitelist.toml";

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
struct TgCfg {
    owners: HashSet<i64>,
    msg_owner_id: Option<i64>,
    bot_id: String,
    lang_subst: HashMap<String, String>,
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
struct TgWhitelist {
    priv_enabled: bool,
    group_enabled: bool,
    allowed: HashSet<i64>,
    blocked: HashSet<i64>,
}

impl TgWhitelist {
    fn priv_ok(&self, id: ChatId) -> bool {
        (!self.priv_enabled || self.allowed.contains(&id.into()))
            && !self.blocked.contains(&id.into())
    }

    fn group_ok(&self, id: ChatId) -> bool {
        (!self.group_enabled || self.allowed.contains(&id.into()))
            && !self.blocked.contains(&id.into())
    }

    fn allow(&mut self, id: i64) {
        self.allowed.insert(id);
    }

    fn unallow(&mut self, id: i64) {
        self.allowed.remove(&id);
    }

    fn block(&mut self, id: i64) {
        self.blocked.insert(id);
    }

    fn unblock(&mut self, id: i64) {
        self.blocked.remove(&id);
    }

    async fn save<'a>(&'a self, path: &'static str) {
        drop(util::encode(self, path).await);
    }
}

struct TgSvc {
    config: TgCfg,
    whitelist: RwLock<TgWhitelist>,
    service: EvalService,
    api: Api,
    bot_user: User,
    username: String,
}

fn telegram_wrap_result(s: &str, group: bool) -> String {
    // FIXME configurable max-lines and max-bytes
    if s.is_empty() {
        "no output".to_owned()
    } else {
        let mut r = "
".to_owned();
        let input = s.as_bytes();
        let mut cut_input = String::from_utf8_lossy(&input[..512.min(input.len())]);
        if group {
            cut_input = Cow::Owned(cut_input.lines().take(10).collect::<Vec<_>>().join("\n"));
        }
        r.push_str(
            &cut_input
                .replace(
                    |c: char| c == '\u{FFFD}' || (c.is_control() && c != '\n' && c != '\t'),
                    "",
                )
                .replace('&', "&")
                .replace('<', "<")
                .replace('>', ">")
                .replace('"', """),
        );
        r.push_str("
");
        if cut_input.len() + 1 // we also cut off the trailing \n
            < input.len()
        {
            r.push_str("... (truncated)");
        }
        r
    }
}

fn is_from_owner(msg: &Message, tgsvc: &TgSvc) -> bool {
    tgsvc.is_owner(msg.from.to_user_id())
}

async fn verify_allowed(chat: &MessageChat, tgsvc: &Arc<TgSvc>) -> Result<(), ()> {
    let group = if let MessageChat::Private(_) = chat {
        false
    } else {
        true
    };
    let wl = tgsvc.whitelist.read().await;
    let chat_id = chat.id();
    if group && !wl.group_ok(chat_id) || !group && !wl.priv_ok(chat_id) {
        let send_message = tgsvc.api.send(SendMessage::new(
            &chat,
            format!(
                "You or this group is not on the whitelist. Seek help. ID: {}",
                chat_id
            ),
        ));
        let leave_group = if group {
            Some(tgsvc.api.send(LeaveChat::new(&chat)))
        } else {
            None
        };
        tokio::spawn(async move {
            drop(send_message.await.map_err(|e| {
                warn!("failed to send message: {}", e);
            }));
            if let Some(leave_group) = leave_group {
                drop(leave_group.await.map_err(|e| {
                    warn!("failed to leave group {}: {}", chat_id, e);
                }));
            }
        });
        Err(())
    } else {
        Ok(())
    }
}

async fn handle_update(update: Update, tgsvc: Arc<TgSvc>) -> Result<(), ()> {
    let message = match update {
        Update {
            kind: UpdateKind::Message(message),
            ..
        } => message,
        _ => return Ok(()),
    };

    match message.kind {
        MessageKind::GroupChatCreated => handle_group_created(message.chat, &tgsvc).await?,
        MessageKind::NewChatMembers { data } => {
            handle_new_chat_members(message.chat, data, &tgsvc).await?
        }
        MessageKind::Text { .. } => handle_message(message, &tgsvc).await?,
        _ => (),
    };

    Ok(())
}

async fn handle_new_chat_members(
    chat: MessageChat,
    new_members: Vec<User>,
    tgsvc: &Arc<TgSvc>,
) -> Result<(), ()> {
    if new_members.iter().any(|u| u.id == tgsvc.bot_user.id) {
        verify_allowed(&chat, tgsvc).await?;
    }

    Ok(())
}

async fn handle_group_created(chat: MessageChat, tgsvc: &Arc<TgSvc>) -> Result<(), ()> {
    verify_allowed(&chat, tgsvc).await?;
    Ok(())
}

async fn handle_message(message: Message, tgsvc: &Arc<TgSvc>) -> Result<(), ()> {
    let text = if let MessageKind::Text { ref data, .. } = message.kind {
        data
    } else {
        return Ok(());
    };

    if !text.starts_with("/") {
        return Ok(());
    }

    let (cmd, args) = {
        let first_tok = if let Some(first_tok) = text.split_whitespace().nth(0) {
            first_tok
        } else {
            return Ok(());
        };

        let mut split = first_tok.splitn(2, '@');
        let cmd = split.next();
        let user = split.next();

        if let Some(user) = user {
            if user != tgsvc.username {
                return Ok(());
            }
        }

        (
            cmd.expect("First split should exist"),
            text[first_tok.len()..].trim_start(),
        )
    };

    match cmd {
        "/privwl" => {
            handle_whitelist_toggle(tgsvc, &message, WhitelistToggleOp::TogglePrivate).await?
        }
        "/groupwl" => {
            handle_whitelist_toggle(tgsvc, &message, WhitelistToggleOp::ToggleGroup).await?
        }
        "/allow" => handle_whitelist_mod(tgsvc, &message, args, WhitelistModOp::Allow).await?,
        "/unallow" => handle_whitelist_mod(tgsvc, &message, args, WhitelistModOp::Unallow).await?,
        "/block" => handle_whitelist_mod(tgsvc, &message, args, WhitelistModOp::Block).await?,
        "/unblock" => handle_whitelist_mod(tgsvc, &message, args, WhitelistModOp::Unblock).await?,
        "/leave" => handle_leave(tgsvc, &message, args).await?,
        _ => {
            let (cmd, is_hash) = if cmd.ends_with('#') {
                (&cmd[1..cmd.len() - 1], true)
            } else {
                (&cmd[1..], false)
            };
            if let Some((_, lang)) = tgsvc.service.langs().find(|(name, _)| *name == cmd) {
                handle_eval(tgsvc, &message, args, lang, is_hash).await?;
            }
        }
    }

    Ok(())
}

async fn handle_eval(
    tgsvc: &Arc<TgSvc>,
    msg: &Message,
    args: &str,
    lang: &Arc<Language>,
    is_hash: bool,
) -> Result<(), ()> {
    verify_allowed(&msg.chat, tgsvc).await?;
    let no_limit = is_hash && is_from_owner(&msg, tgsvc);
    let is_group = if let MessageChat::Private(_) = msg.chat {
        false
    } else {
        true
    };
    let chat_id = msg.chat.id();
    let msg_id = msg.id;
    info!("({}) evaluating from {:?}: {:?}", msg_id, msg.from, args);
    let code = {
        let mut r = args.to_owned();
        r.push_str("\n");
        r
    };

    let eval_result = lang
        .eval(
            code,
            if no_limit { Some(0) } else { None },
            Some(format!("tg{}", chat_id)),
        )
        .await;
    let ok = eval_result.is_ok();
    info!(
        "({}) result: {:?}",
        msg_id,
        eval_result.as_ref().map(|x| x.as_str()).unwrap_or("")
    );
    let mut request = SendMessage::new(
        &msg.chat,
        eval_result
            .map(|r| telegram_wrap_result(&r, is_group))
            .unwrap_or_else(|e| e),
    );
    request.reply_to(msg);
    if ok {
        request.parse_mode(ParseMode::Html);
    }
    tokio::spawn(tgsvc.api.send(request));

    Ok(())
}

#[derive(Clone, Copy)]
enum WhitelistToggleOp {
    TogglePrivate,
    ToggleGroup,
}

#[derive(Clone, Copy)]
enum WhitelistModOp {
    Allow,
    Unallow,
    Block,
    Unblock,
}

async fn handle_whitelist_toggle(
    tgsvc: &Arc<TgSvc>,
    msg: &Message,
    op: WhitelistToggleOp,
) -> Result<(), ()> {
    if !is_from_owner(&msg, tgsvc) {
        return Ok(());
    }

    let mut wl = tgsvc.whitelist.write().await;
    let resp = match op {
        WhitelistToggleOp::TogglePrivate => {
            wl.priv_enabled = !wl.priv_enabled;
            wl.save(WHITELIST_FILENAME).await;
            format!("Private whitelist enabled: {}", wl.priv_enabled)
        }
        WhitelistToggleOp::ToggleGroup => {
            wl.group_enabled = !wl.group_enabled;
            wl.save(WHITELIST_FILENAME).await;
            format!("Group whitelist enabled: {}", wl.group_enabled)
        }
    };
    tokio::spawn(tgsvc.api.send(SendMessage::new(&msg.chat, resp)));
    Ok(())
}

async fn handle_whitelist_mod(
    tgsvc: &Arc<TgSvc>,
    msg: &Message,
    args: &str,
    op: WhitelistModOp,
) -> Result<(), ()> {
    if !is_from_owner(&msg, tgsvc) {
        return Ok(());
    }

    let arg = args
        .trim()
        .split_whitespace()
        .nth(0)
        .and_then(|arg| arg.parse().ok());
    let mut wl = tgsvc.whitelist.write().await;
    let resp = if let Some(id) = arg {
        match op {
            WhitelistModOp::Allow => {
                wl.allow(id);
                wl.save(WHITELIST_FILENAME).await;
                format!("Allowed {}", id)
            }
            WhitelistModOp::Unallow => {
                wl.unallow(id);
                wl.save(WHITELIST_FILENAME).await;
                format!("Unallowed {}", id)
            }
            WhitelistModOp::Block => {
                wl.block(id);
                wl.save(WHITELIST_FILENAME).await;
                format!("Blocked {}", id)
            }
            WhitelistModOp::Unblock => {
                wl.unblock(id);
                wl.save(WHITELIST_FILENAME).await;
                format!("Unblocked {}", id)
            }
        }
    } else {
        "Invalid ID".to_owned()
    };

    tokio::spawn(tgsvc.api.send(SendMessage::new(&msg.chat, resp)));
    Ok(())
}

async fn handle_leave(tgsvc: &Arc<TgSvc>, msg: &Message, args: &str) -> Result<(), ()> {
    if !is_from_owner(&msg, tgsvc) {
        return Ok(());
    }

    let arg: Option<i64> = args
        .trim()
        .split_whitespace()
        .nth(0)
        .and_then(|arg| arg.parse().ok());
    let (maybe_leave_chat, resp) = match arg {
        Some(id) => (Some(tgsvc.api.send(LeaveChat::new(ChatId::from(id)))), "OK"),
        None => (None, "Invalid ID"),
    };

    let send_response = tgsvc.api.send(SendMessage::new(&msg.chat, resp));

    tokio::spawn(async move {
        if let Some(fut) = maybe_leave_chat {
            drop(fut.await.map_err(|e| warn!("failed to leave group: {}", e)));
        }
        drop(
            send_response
                .await
                .map_err(|e| warn!("failed to send message: {}", e)),
        );
    });
    Ok(())
}

impl TgSvc {
    async fn run() -> Result<(), ()> {
        let cfg = util::decode::<TgCfg, _>("evalbot.tg.toml")
            .await
            .map(|cfg| {
                debug!("Loaded config: {:?}", cfg);
                cfg
            })
            .map_err(|e| error!("failed to read evalbot.tg.toml: {}", e))?;
        let wl = util::decode::<TgWhitelist, _>(WHITELIST_FILENAME)
            .await
            .or_else(|e| {
                warn!("failed to read whitelist: {}; using empty whitelist", e);
                Ok(TgWhitelist {
                    priv_enabled: false,
                    group_enabled: false,
                    allowed: HashSet::new(),
                    blocked: HashSet::new(),
                })
            })?;

        let api = Api::new(&cfg.bot_id);
        let bot_user = api
            .send(GetMe)
            .await
            .map_err(|e| error!("failed to get bot's user info: {}", e))?;
        let service = EvalService::from_toml_file("evalbot.toml")
            .await
            .map_err(|e| error!("failed to read evalbot.toml: {}", e))?;
        TgSvc {
            api: Api::new(&cfg.bot_id),
            config: cfg,
            whitelist: RwLock::new(wl),
            service: service,
            username: bot_user
                .username
                .as_ref()
                .expect("Bot must have username")
                .clone(),
            bot_user: bot_user,
        }
        .handle()
        .await;
        Ok(())
    }

    async fn handle(self) {
        let me = Arc::new(self);

        let mut stream = me.api.stream();
        stream
            .timeout(Duration::from_secs(35))
            .error_delay(Duration::from_secs(30));
        while let Some(update) = stream.next().await {
            match update {
                Ok(update) => {
                    tokio::spawn(handle_update(update, me.clone()));
                }
                Err(error) => error!("received error: {:?}", error),
            }
        }
    }

    fn is_owner(&self, id: UserId) -> bool {
        self.config.owners.contains(&id.into())
    }
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();
    drop(TgSvc::run().await);
}
<|endoftext|>


<commit_before>use cairo;
use cairo_sys;
use glib::translate::*;
use libc;

use regex::{Captures, Regex};
use std::borrow::Cow;
use std::cell::RefCell;

use attributes::Attribute;
use drawing_ctx::DrawingCtx;
use handle::RsvgHandle;
use node::*;
use property_bag::PropertyBag;

pub struct NodeLink {
    link: RefCell<Option<String>>,
}

impl NodeLink {
    pub fn new() -> NodeLink {
        NodeLink {
            link: RefCell::new(None),
        }
    }
}

impl NodeTrait for NodeLink {
    fn set_atts(&self, _: &RsvgNode, _: *const RsvgHandle, pbag: &PropertyBag) -> NodeResult {
        for (_key, attr, value) in pbag.iter() {
            match attr {
                Attribute::XlinkHref => *self.link.borrow_mut() = Some(value.to_owned()),

                _ => (),
            }
        }

        Ok(())
    }

    fn draw(
        &self,
        node: &RsvgNode,
        cascaded: &CascadedValues,
        draw_ctx: &mut DrawingCtx,
        clipping: bool,
    ) {
        let link = self.link.borrow();

        let cascaded = CascadedValues::new(cascaded, node);
        let values = cascaded.get();

        draw_ctx.with_discrete_layer(node, values, clipping, &mut |dc| {
            if link.is_some() && link.as_ref().unwrap() != "" {
                const CAIRO_TAG_LINK: &str = "Link";

                let attributes = link.as_ref().map(|i| format!("uri='{}'", escape_value(i)));

                dc.get_cairo_context().with_tag(
                    CAIRO_TAG_LINK,
                    attributes.as_ref().map(|i| i.as_str()),
                    || {
                        node.draw_children(&cascaded, dc, clipping);
                    },
                )
            } else {
                node.draw_children(&cascaded, dc, clipping)
            }
        });
    }
}

/// escape quotes and backslashes with backslash
fn escape_value(value: &str) -> Cow<str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"['\\]").unwrap();
    }

    REGEX.replace_all(value, |caps: &Captures| {
        match caps.get(0).unwrap().as_str() {
            "'" => "\\'".to_owned(),
            "\\" => "\\\\".to_owned(),
            _ => unreachable!(),
        }
    })
}

extern "C" {
    fn cairo_tag_begin(
        cr: *mut cairo_sys::cairo_t,
        tag_name: *const libc::c_char,
        attibutes: *const libc::c_char,
    );
    fn cairo_tag_end(cr: *mut cairo_sys::cairo_t, tag_name: *const libc::c_char);
}

/// Bindings that aren't supported by `cairo-rs` for now
trait CairoTagging {
    fn tag_begin(&self, tag_name: &str, attributes: Option<&str>);
    fn tag_end(&self, tag_name: &str);
    fn with_tag<U, T>(&self, tag_name: &str, attributes: Option<&str>, f: U) -> T
    where
        U: Fn() -> T,
    {
        self.tag_begin(tag_name, attributes);
        let result = f();
        self.tag_end(tag_name);

        result
    }
}

impl CairoTagging for cairo::Context {
    fn tag_begin(&self, tag_name: &str, attributes: Option<&str>) {
        unsafe {
            cairo_tag_begin(
                self.to_glib_none().0,
                tag_name.to_glib_none().0,
                attributes.to_glib_none().0,
            );
        }
    }

    fn tag_end(&self, tag_name: &str) {
        unsafe {
            cairo_tag_end(self.to_glib_none().0, tag_name.to_glib_none().0);
        }
    }
}
<commit_msg>link.rs: Remove the with_tag() convenience function; use the primitives<commit_after>use cairo;
use cairo_sys;
use glib::translate::*;
use libc;

use regex::{Captures, Regex};
use std::borrow::Cow;
use std::cell::RefCell;

use attributes::Attribute;
use drawing_ctx::DrawingCtx;
use handle::RsvgHandle;
use node::*;
use property_bag::PropertyBag;

pub struct NodeLink {
    link: RefCell<Option<String>>,
}

impl NodeLink {
    pub fn new() -> NodeLink {
        NodeLink {
            link: RefCell::new(None),
        }
    }
}

impl NodeTrait for NodeLink {
    fn set_atts(&self, _: &RsvgNode, _: *const RsvgHandle, pbag: &PropertyBag) -> NodeResult {
        for (_key, attr, value) in pbag.iter() {
            match attr {
                Attribute::XlinkHref => *self.link.borrow_mut() = Some(value.to_owned()),

                _ => (),
            }
        }

        Ok(())
    }

    fn draw(
        &self,
        node: &RsvgNode,
        cascaded: &CascadedValues,
        draw_ctx: &mut DrawingCtx,
        clipping: bool,
    ) {
        let link = self.link.borrow();

        let cascaded = CascadedValues::new(cascaded, node);
        let values = cascaded.get();

        draw_ctx.with_discrete_layer(node, values, clipping, &mut |dc| {
            if link.is_some() && link.as_ref().unwrap() != "" {
                const CAIRO_TAG_LINK: &str = "Link";

                let attributes = link.as_ref().map(|i| format!("uri='{}'", escape_value(i)));

                let cr = dc.get_cairo_context();

                cr.tag_begin(CAIRO_TAG_LINK, attributes.as_ref().map(|i| i.as_str()));
                node.draw_children(&cascaded, dc, clipping);
                cr.tag_end(CAIRO_TAG_LINK);
            } else {
                node.draw_children(&cascaded, dc, clipping)
            }
        });
    }
}

/// escape quotes and backslashes with backslash
fn escape_value(value: &str) -> Cow<str> {
    lazy_static! {
        static ref REGEX: Regex = Regex::new(r"['\\]").unwrap();
    }

    REGEX.replace_all(value, |caps: &Captures| {
        match caps.get(0).unwrap().as_str() {
            "'" => "\\'".to_owned(),
            "\\" => "\\\\".to_owned(),
            _ => unreachable!(),
        }
    })
}

extern "C" {
    fn cairo_tag_begin(
        cr: *mut cairo_sys::cairo_t,
        tag_name: *const libc::c_char,
        attibutes: *const libc::c_char,
    );
    fn cairo_tag_end(cr: *mut cairo_sys::cairo_t, tag_name: *const libc::c_char);
}

/// Bindings that aren't supported by `cairo-rs` for now
trait CairoTagging {
    fn tag_begin(&self, tag_name: &str, attributes: Option<&str>);
    fn tag_end(&self, tag_name: &str);
}

impl CairoTagging for cairo::Context {
    fn tag_begin(&self, tag_name: &str, attributes: Option<&str>) {
        unsafe {
            cairo_tag_begin(
                self.to_glib_none().0,
                tag_name.to_glib_none().0,
                attributes.to_glib_none().0,
            );
        }
    }

    fn tag_end(&self, tag_name: &str) {
        unsafe {
            cairo_tag_end(self.to_glib_none().0, tag_name.to_glib_none().0);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2020-present, Gregory Szorc
// All rights reserved.
//
// This software may be modified and distributed under the terms
// of the BSD license. See the LICENSE file for details.

use crate::compression_dict::ZstdCompressionDict;
use crate::compression_parameters::{CCtxParams, ZstdCompressionParameters};
use crate::ZstdError;
use cpython::exc::ValueError;
use cpython::{py_class, PyErr, PyModule, PyObject, PyResult, Python, PythonObject};
use std::cell::RefCell;
use std::marker::PhantomData;

pub(crate) struct CCtx<'a>(*mut zstd_sys::ZSTD_CCtx, PhantomData<&'a ()>);

impl<'a> Drop for CCtx<'a> {
    fn drop(&mut self) {
        unsafe {
            zstd_sys::ZSTD_freeCCtx(self.0);
        }
    }
}

unsafe impl<'a> Send for CCtx<'a> {}

impl<'a> CCtx<'a> {
    fn new() -> Result<Self, &'static str> {
        let cctx = unsafe { zstd_sys::ZSTD_createCCtx() };
        if cctx.is_null() {
            return Err("could not allocate ZSTD_CCtx instance");
        }

        Ok(Self(cctx, PhantomData))
    }

    fn set_parameters(&self, params: &CCtxParams) -> Result<(), String> {
        let zresult = unsafe {
            zstd_sys::ZSTD_CCtx_setParametersUsingCCtxParams(self.0, params.get_raw_ptr())
        };
        if unsafe { zstd_sys::ZSTD_isError(zresult) } != 0 {
            return Err(zstd_safe::get_error_name(zresult).to_string());
        }

        Ok(())
    }

    pub fn memory_size(&self) -> usize {
        unsafe { zstd_sys::ZSTD_sizeof_CCtx(self.0 as *const _) }
    }
}

struct CompressorState<'params, 'cctx> {
    threads: i32,
    dict: Option<ZstdCompressionDict>,
    params: CCtxParams<'params>,
    cctx: CCtx<'cctx>,
}

impl<'params, 'cctx> CompressorState<'params, 'cctx> {
    pub(crate) fn setup_cctx(&self, py: Python) -> PyResult<()> {
        self.cctx
            .set_parameters(&self.params)
            .or_else(|msg| Err(PyErr::new::<ZstdError, _>(py, msg)))?;

        if let Some(dict) = &self.dict {
            dict.load_into_cctx(py, self.cctx.0)?;
        }

        Ok(())
    }
}

py_class!(class ZstdCompressor |py| {
    data state: RefCell<CompressorState<'static, 'static>>;

    def __new__(
        _cls,
        level: i32 = 3,
        dict_data: Option<ZstdCompressionDict> = None,
        compression_params: Option<ZstdCompressionParameters> = None,
        write_checksum: Option<bool> = None,
        write_content_size: Option<bool> = None,
        write_dict_id: Option<bool> = None,
        threads: i32 = 0
    ) -> PyResult<PyObject> {
        ZstdCompressor::new_impl(
            py,
            level,
            dict_data,
            compression_params,
            write_checksum,
            write_content_size,
            write_dict_id,
            threads,
        )
    }

    def memory_size(&self) -> PyResult<usize> {
        Ok(self.state(py).borrow().cctx.memory_size())
    }
});

impl ZstdCompressor {
    fn new_impl(
        py: Python,
        level: i32,
        dict_data: Option<ZstdCompressionDict>,
        compression_params: Option<ZstdCompressionParameters>,
        write_checksum: Option<bool>,
        write_content_size: Option<bool>,
        write_dict_id: Option<bool>,
        threads: i32,
    ) -> PyResult<PyObject> {
        if level > zstd_safe::max_c_level() {
            return Err(PyErr::new::<ValueError, _>(
                py,
                format!(
                    "level must be less than {}",
                    zstd_safe::max_c_level() as i32 + 1
                ),
            ));
        }

        let threads = if threads < 0 {
            num_cpus::get() as i32
        } else {
            threads
        };

        let cctx = CCtx::new().or_else(|msg| Err(PyErr::new::<ZstdError, _>(py, msg)))?;
        let params = CCtxParams::create(py)?;

        if let Some(ref compression_params) = compression_params {
            if write_checksum.is_some() {
                return Err(PyErr::new::<ValueError, _>(
                    py,
                    "cannot define compression_params and write_checksum",
                ));
            }
            if write_content_size.is_some() {
                return Err(PyErr::new::<ValueError, _>(
                    py,
                    "cannot define compression_params and write_content_size",
                ));
            }
            if write_dict_id.is_some() {
                return Err(PyErr::new::<ValueError, _>(
                    py,
                    "cannot define compression_params and write_dict_id",
                ));
            }
            if threads != 0 {
                return Err(PyErr::new::<ValueError, _>(
                    py,
                    "cannot define compression_params and threads",
                ));
            }

            params.apply_compression_parameters(py, compression_params)?;

        // TODO set parameters from CompressionParameters
        } else {
            params.set_parameter(
                py,
                zstd_sys::ZSTD_cParameter::ZSTD_c_compressionLevel,
                level,
            )?;
            params.set_parameter(
                py,
                zstd_sys::ZSTD_cParameter::ZSTD_c_contentSizeFlag,
                if write_content_size.unwrap_or(true) {
                    1
                } else {
                    0
                },
            )?;
            params.set_parameter(
                py,
                zstd_sys::ZSTD_cParameter::ZSTD_c_checksumFlag,
                if write_checksum.unwrap_or(false) {
                    1
                } else {
                    0
                },
            )?;
            params.set_parameter(
                py,
                zstd_sys::ZSTD_cParameter::ZSTD_c_dictIDFlag,
                if write_dict_id.unwrap_or(true) { 1 } else { 0 },
            )?;
            if threads != 0 {
                params.set_parameter(py, zstd_sys::ZSTD_cParameter::ZSTD_c_nbWorkers, threads)?;
            }
        }

        let state = CompressorState {
            threads,
            dict: dict_data,
            params,
            cctx,
        };

        state.setup_cctx(py)?;

        Ok(ZstdCompressor::create_instance(py, RefCell::new(state))?.into_object())
    }
}

pub(crate) fn init_module(py: Python, module: &PyModule) -> PyResult<()> {
    module.add(py, "ZstdCompressor", py.get_type::<ZstdCompressor>())?;

    Ok(())
}
<commit_msg>rust: implement ZstdCompressor.compress()<commit_after>// Copyright (c) 2020-present, Gregory Szorc
// All rights reserved.
//
// This software may be modified and distributed under the terms
// of the BSD license. See the LICENSE file for details.

use crate::compression_dict::ZstdCompressionDict;
use crate::compression_parameters::{CCtxParams, ZstdCompressionParameters};
use crate::ZstdError;
use cpython::buffer::PyBuffer;
use cpython::exc::ValueError;
use cpython::{py_class, PyBytes, PyErr, PyModule, PyObject, PyResult, Python, PythonObject};
use std::cell::RefCell;
use std::marker::PhantomData;

pub(crate) struct CCtx<'a>(*mut zstd_sys::ZSTD_CCtx, PhantomData<&'a ()>);

impl<'a> Drop for CCtx<'a> {
    fn drop(&mut self) {
        unsafe {
            zstd_sys::ZSTD_freeCCtx(self.0);
        }
    }
}

unsafe impl<'a> Send for CCtx<'a> {}
unsafe impl<'a> Sync for CCtx<'a> {}

impl<'a> CCtx<'a> {
    fn new() -> Result<Self, &'static str> {
        let cctx = unsafe { zstd_sys::ZSTD_createCCtx() };
        if cctx.is_null() {
            return Err("could not allocate ZSTD_CCtx instance");
        }

        Ok(Self(cctx, PhantomData))
    }

    fn set_parameters(&self, params: &CCtxParams) -> Result<(), String> {
        let zresult = unsafe {
            zstd_sys::ZSTD_CCtx_setParametersUsingCCtxParams(self.0, params.get_raw_ptr())
        };
        if unsafe { zstd_sys::ZSTD_isError(zresult) } != 0 {
            return Err(zstd_safe::get_error_name(zresult).to_string());
        }

        Ok(())
    }

    pub fn memory_size(&self) -> usize {
        unsafe { zstd_sys::ZSTD_sizeof_CCtx(self.0 as *const _) }
    }

    pub fn reset(&self) -> usize {
        unsafe {
            zstd_sys::ZSTD_CCtx_reset(
                self.0,
                zstd_sys::ZSTD_ResetDirective::ZSTD_reset_session_only,
            )
        }
    }

    pub fn set_pledged_source_size(&self, size: u64) -> Result<(), &'static str> {
        let zresult = unsafe { zstd_sys::ZSTD_CCtx_setPledgedSrcSize(self.0, size) };
        if unsafe { zstd_sys::ZSTD_isError(zresult) } != 0 {
            Err(zstd_safe::get_error_name(zresult))
        } else {
            Ok(())
        }
    }

    pub fn compress(&self, source: &[u8]) -> Result<Vec<u8>, &'static str> {
        self.reset();

        let dest_len = unsafe { zstd_sys::ZSTD_compressBound(source.len()) };

        let mut dest: Vec<u8> = Vec::with_capacity(dest_len);

        self.set_pledged_source_size(dest_len as _)?;

        let mut in_buffer = zstd_sys::ZSTD_inBuffer {
            src: source.as_ptr() as *const _,
            size: source.len(),
            pos: 0,
        };

        let mut out_buffer = zstd_sys::ZSTD_outBuffer {
            dst: dest.as_mut_ptr() as *mut _,
            size: dest.capacity(),
            pos: 0,
        };

        // By avoiding ZSTD_compress(), we don't necessarily write out content
        // size. This means the parameters to control frame parameters are honored.
        let zresult = unsafe {
            zstd_sys::ZSTD_compressStream2(
                self.0,
                &mut out_buffer as *mut _,
                &mut in_buffer as *mut _,
                zstd_sys::ZSTD_EndDirective::ZSTD_e_end,
            )
        };

        if unsafe { zstd_sys::ZSTD_isError(zresult) } != 0 {
            Err(zstd_safe::get_error_name(zresult))
        } else if zresult > 0 {
            Err("unexpected partial frame flush")
        } else {
            unsafe { dest.set_len(out_buffer.pos) }

            Ok(dest)
        }
    }
}

struct CompressorState<'params, 'cctx> {
    threads: i32,
    dict: Option<ZstdCompressionDict>,
    params: CCtxParams<'params>,
    cctx: CCtx<'cctx>,
}

impl<'params, 'cctx> CompressorState<'params, 'cctx> {
    pub(crate) fn setup_cctx(&self, py: Python) -> PyResult<()> {
        self.cctx
            .set_parameters(&self.params)
            .or_else(|msg| Err(PyErr::new::<ZstdError, _>(py, msg)))?;

        if let Some(dict) = &self.dict {
            dict.load_into_cctx(py, self.cctx.0)?;
        }

        Ok(())
    }
}

py_class!(class ZstdCompressor |py| {
    data state: RefCell<CompressorState<'static, 'static>>;

    def __new__(
        _cls,
        level: i32 = 3,
        dict_data: Option<ZstdCompressionDict> = None,
        compression_params: Option<ZstdCompressionParameters> = None,
        write_checksum: Option<bool> = None,
        write_content_size: Option<bool> = None,
        write_dict_id: Option<bool> = None,
        threads: i32 = 0
    ) -> PyResult<PyObject> {
        ZstdCompressor::new_impl(
            py,
            level,
            dict_data,
            compression_params,
            write_checksum,
            write_content_size,
            write_dict_id,
            threads,
        )
    }

    def memory_size(&self) -> PyResult<usize> {
        Ok(self.state(py).borrow().cctx.memory_size())
    }

    def compress(&self, data: PyObject) -> PyResult<PyBytes> {
        self.compress_impl(py, data)
    }
});

impl ZstdCompressor {
    fn new_impl(
        py: Python,
        level: i32,
        dict_data: Option<ZstdCompressionDict>,
        compression_params: Option<ZstdCompressionParameters>,
        write_checksum: Option<bool>,
        write_content_size: Option<bool>,
        write_dict_id: Option<bool>,
        threads: i32,
    ) -> PyResult<PyObject> {
        if level > zstd_safe::max_c_level() {
            return Err(PyErr::new::<ValueError, _>(
                py,
                format!(
                    "level must be less than {}",
                    zstd_safe::max_c_level() as i32 + 1
                ),
            ));
        }

        let threads = if threads < 0 {
            num_cpus::get() as i32
        } else {
            threads
        };

        let cctx = CCtx::new().or_else(|msg| Err(PyErr::new::<ZstdError, _>(py, msg)))?;
        let params = CCtxParams::create(py)?;

        if let Some(ref compression_params) = compression_params {
            if write_checksum.is_some() {
                return Err(PyErr::new::<ValueError, _>(
                    py,
                    "cannot define compression_params and write_checksum",
                ));
            }
            if write_content_size.is_some() {
                return Err(PyErr::new::<ValueError, _>(
                    py,
                    "cannot define compression_params and write_content_size",
                ));
            }
            if write_dict_id.is_some() {
                return Err(PyErr::new::<ValueError, _>(
                    py,
                    "cannot define compression_params and write_dict_id",
                ));
            }
            if threads != 0 {
                return Err(PyErr::new::<ValueError, _>(
                    py,
                    "cannot define compression_params and threads",
                ));
            }

            params.apply_compression_parameters(py, compression_params)?;

        // TODO set parameters from CompressionParameters
        } else {
            params.set_parameter(
                py,
                zstd_sys::ZSTD_cParameter::ZSTD_c_compressionLevel,
                level,
            )?;
            params.set_parameter(
                py,
                zstd_sys::ZSTD_cParameter::ZSTD_c_contentSizeFlag,
                if write_content_size.unwrap_or(true) {
                    1
                } else {
                    0
                },
            )?;
            params.set_parameter(
                py,
                zstd_sys::ZSTD_cParameter::ZSTD_c_checksumFlag,
                if write_checksum.unwrap_or(false) {
                    1
                } else {
                    0
                },
            )?;
            params.set_parameter(
                py,
                zstd_sys::ZSTD_cParameter::ZSTD_c_dictIDFlag,
                if write_dict_id.unwrap_or(true) { 1 } else { 0 },
            )?;
            if threads != 0 {
                params.set_parameter(py, zstd_sys::ZSTD_cParameter::ZSTD_c_nbWorkers, threads)?;
            }
        }

        let state = CompressorState {
            threads,
            dict: dict_data,
            params,
            cctx,
        };

        state.setup_cctx(py)?;

        Ok(ZstdCompressor::create_instance(py, RefCell::new(state))?.into_object())
    }

    fn compress_impl(&self, py: Python, data: PyObject) -> PyResult<PyBytes> {
        let state: std::cell::Ref<CompressorState> = self.state(py).borrow();

        let buffer = PyBuffer::get(py, &data)?;

        if !buffer.is_c_contiguous() || buffer.dimensions() > 1 {
            return Err(ZstdError::from_message(
                py,
                "data buffer should be contiguous and have at most one dimension",
            ));
        }

        let source: &[u8] =
            unsafe { std::slice::from_raw_parts(buffer.buf_ptr() as *const _, buffer.len_bytes()) };

        let cctx = &state.cctx;

        // TODO implement 0 copy via Py_SIZE().
        let data = py.allow_threads(|| cctx.compress(source)).or_else(|msg| {
            Err(ZstdError::from_message(
                py,
                format!("cannot compress: {}", msg).as_ref(),
            ))
        })?;

        Ok(PyBytes::new(py, &data))
    }
}

pub(crate) fn init_module(py: Python, module: &PyModule) -> PyResult<()> {
    module.add(py, "ZstdCompressor", py.get_type::<ZstdCompressor>())?;

    Ok(())
}
<|endoftext|>