<commit_before>use num::{Signed};

use super::*;

#[derive (Debug, Fail)]
pub enum Error <T: Integer> {
  #[fail (display = "out-of-bounds value in the x^{} term (coefficient is {}, but must have abs <= {})", power, coefficient, maximum_abs)]
  CoefficientOutOfBounds {
    power: usize,
    coefficient: T,
    maximum_abs: T,
  },
  #[fail (display = "out-of-bounds input to polynomial function (for the x^{} term, |{}*{}^{}| was {}, but must have abs <= {})", power, coefficient, input, power, value, maximum_abs)]
  TermOutOfBounds {
    power: usize,
    input: T,
    coefficient: T,
    value: T,
    maximum_abs: T,
  },
  #[fail (display = "out-of-bounds input to polynomial function (for the x^{} term, |{}*{}^{}| overflowed the type)", power, input, coefficient, power)]
  TermOverflowed {
    power: usize,
    input: T,
    coefficient: T,
  },
}

/// Approximately evaluate an integer polynomial at an input in the range [-0.5, 0.5].
///
/// The input is represented as an integer combined with a right-shift size.
/// Each coefficient, except the constant one, must obey coefficient.abs() <= T::max_value() >> shift.
/// The constant coefficient must obey coefficient.abs() <= T::max_value() - (T::max_value() >> shift),
/// i.e. be small enough that you can add one of the other coefficients to it without overflowing.
/// Given these conditions, the result is guaranteed to be strictly within 1 of the ideal result.
/// For instance, if the ideal result was 2.125, this function could return 2 or 3,
/// but if it was 2, this function can only return exactly 2.
pub fn evaluate_at_small_input <Coefficient: Copy, T: Integer + Signed + From <Coefficient>> (coefficients: & [Coefficient], input: T, shift: u32)->T {
  let half = (T::one() << shift) >> 1u32;
  assert!(-half <= input && input <= half, "inputs to evaluate_at_small_input must be in the range [-0.5, 0.5]");

  if coefficients.len () == 0 {return T::zero()}

  let mut result = T::zero();
  for coefficient in coefficients.iter().skip(1).rev() {
    result = shr_nicely_rounded ((result + (*coefficient).into())*input, shift);
  }
  result + coefficients [0].into()
}

/// Translate a polynomial horizontally, letting it take inputs relative to a given input instead of relative to 0.
///
/// This is equivalent to replacing its coefficients with its Taylor coefficients at the given input.
/// This is equivalent to translating the graph *left* by the given input amount.
///
/// For constant polynomials, this is a no-op.
///
/// To avoid overflow, each term, when evaluating at the given input,
/// must obey term.abs() <= T::max_value() >> (coefficients.len() - 1).
/// If this condition is broken, translate() returns Err and makes no changes.
pub fn translate <T: Integer + Signed> (coefficients: &mut [T], input: T)->Result <(), Error <T>> {
  translate_check (coefficients, input, |_| T::max_value())?;
  translate_unchecked (coefficients, input);
  Ok (())
}

/// Same as translate(), but assume no overflow will occur.
///
/// Useful in performance-critical code when you already know there won't be overflow,
/// such as in the range returned by conservative_safe_translation_range.
pub fn translate_unchecked <T: Integer + Signed> (coefficients: &mut [T], input: T) {
  coefficients.reverse();
  for index in 0..coefficients.len() {
    let coefficient = mem::replace (&mut coefficients [index], T::zero());
    for derivative in (1..(index + 1)).rev() {
      coefficients [derivative] = coefficients [derivative]*input + coefficients [derivative - 1]
    }
    coefficients [0] = coefficients [0]*input + coefficient
  }
}

pub fn within_bounds_check <T: Integer + Signed, MaximumFn: Fn (usize)->T> (coefficients: & [T], maximum: MaximumFn)->Result <(), Error <T>> {
  for (power, coefficient) in coefficients.iter().enumerate() {
    let max = maximum (power);
    if coefficient == &T::min_value() {return Err (Error::CoefficientOutOfBounds {power, coefficient: *coefficient, maximum_abs: max})}
    let magnitude = coefficient.abs();
    if magnitude > max {return Err (Error::CoefficientOutOfBounds {power, coefficient: *coefficient, maximum_abs: max})}
  }
  Ok (())
}

pub fn translate_check <T: Integer + Signed, MaximumFn: Fn (usize)->T> (mut coefficients: & [T], input: T, maximum: MaximumFn)->Result <(), Error <T>> {
  let original_len = coefficients.len();
  loop {
    if coefficients.len() == 0 {return Ok (());}
    if coefficients [coefficients.len() - 1] == T::zero() {coefficients = & coefficients [..(coefficients.len() - 1)];}
    else {break;}
  }
  let coefficients = coefficients;
  if input == T::zero() {within_bounds_check (coefficients, maximum)?; return Ok (());}
  let mut factor = T::one();
  let input = input.abs();
  // base the safety check on the original length, so that the bounds are consistent
  // for all polynomials with a given len(), rather than being dependent on the values in them
  let sum_safety_shift = original_len as u32 - 1;
  let mut running_maximum = T::max_value();
  for (power, coefficient) in coefficients.iter().enumerate() {
    if power > 0 { match factor.checked_mul (&input) {
      None => return Err (Error::TermOverflowed {power, input, coefficient:*coefficient}),
      Some (next_factor) => factor = next_factor,
    }}
    running_maximum = min (running_maximum, maximum (power).checked_shr (sum_safety_shift).unwrap_or (T::zero()));
    if coefficient == &T::min_value() {return Err (Error::TermOverflowed {power, input, coefficient:*coefficient});}
    let magnitude = coefficient.abs();
    match magnitude.checked_mul (&factor) {
      None => return Err (Error::TermOverflowed {power, input, coefficient:*coefficient}),
      Some (term) => if term > running_maximum {return Err (Error::TermOutOfBounds {power, input, coefficient:*coefficient, value: term, maximum_abs: running_maximum})},
    }
  }
  Ok (())
}

/// Find the range of inputs to which the polynomial can safely be translated.
///
/// Magnitudes <= result are safe to translate to; magnitudes > result aren't.
/// If even the initial value is out of bounds, this function returns -1.
///
/// The third argument can impose a stricter maximum on the resulting coefficients.
pub fn exact_safe_translation_range <T: Integer + Signed, MaximumFn: Fn (usize)->T> (coefficients: & [T], maximum: MaximumFn)->T {
  if within_bounds_check (coefficients, & maximum).is_err() {return -T::one();}
  // pick a min that's definitely legal and a max where max + 1 is definitely illegal
  let mut min = T::zero();
  let mut max = T::max_value();
  while min < max {
    let mid = mean_ceil (min, max);
    if translate_check (coefficients, mid, & maximum).is_ok() {
      min = mid;
    } else {
      max = mid - T::one();
    }
  }
  min
}

/// Find a range of inputs to which the polynomial can safely be translated.
///
/// Magnitudes <= result are safe to translate to.
/// Magnitudes >= result*2 aren't.
/// If even the initial value is out of bounds, this function returns -1.
///
/// The third argument can impose a stricter maximum on the resulting coefficients.
///
/// This function is much faster than exact_safe_translation_range, at the cost of being less precise.
pub fn conservative_safe_translation_range <T: Integer + Signed, MaximumFn: Fn (usize)->T> (coefficients: &[T], maximum: MaximumFn)->T {
  if coefficients.len() == 0 {return T::max_value()}
  if within_bounds_check (coefficients, & maximum).is_err() {return -T::one();}
  let sum_safety_shift = coefficients.len() as u32 - 1;
  let mut running_maximum = T::max_value();
  let mut result_shift = mem::size_of::<T>()*8 - 2;
  for (power, coefficient) in coefficients.iter().enumerate() {
    let magnitude = coefficient.abs();
    running_maximum = min (running_maximum, maximum (power).checked_shr (sum_safety_shift).unwrap_or (T::zero()));
    while magnitude > running_maximum.checked_shr((result_shift*power) as u32).unwrap_or (T::zero()) {
      if result_shift == 0 {return T::zero();}
      result_shift -= 1;
    }
  }
  T::one() << result_shift
}

#[cfg (test)]
mod tests {
  use super::*;
  use num::{Zero, One, FromPrimitive, Integer};
  use num::bigint::BigInt;
  use num::rational::{Ratio, BigRational};
  use rand::distributions::Distribution;
  use rand::distributions::range::{Range};
  use rand::{Rng, SeedableRng};

  fn evaluate_exactly <Coefficient: Copy, T: Integer + Signed + From <Coefficient>> (coefficients: & [Coefficient], input: T, shift: u32)->BigRational
    where BigInt: From <Coefficient> + From <T> {
    let mut result: BigRational = Ratio::zero();
    let input = Ratio::new(BigInt::from (input), BigInt::one() << shift as usize);
    for coefficient in coefficients.iter().rev() {
      result = result*&input + BigInt::from(*coefficient);
    }
    result
  }

  #[test]
  fn test_evaluate_at_small_input() {
    let mut generator =::rand::chacha::ChaChaRng::from_seed ([33; 32]) ;
    for shift in 0..5 {
      let shift = 1 << shift;
      let half = 1i64 << (shift - 1);
      let maximum = i64::max_value() >> shift;
      let constant_maximum = i64::max_value() - maximum;
      let range = Range::new_inclusive (- maximum, maximum) ;
      let constant_range = Range::new_inclusive (- constant_maximum, constant_maximum);
      let mut test_inputs = vec![- half, - half + 1, - 1, 0, 1, half - 1, half];
      if half > 2 {for _ in 0..6 {test_inputs.push (generator.gen_range (- half + 2, half - 1));}}
      for _ in 0..40 {
        let mut coefficients = vec![constant_range.sample (&mut generator)];
        while generator.gen() {coefficients.push (range.sample (&mut generator));}
        for input in test_inputs.iter() {
          let result = evaluate_at_small_input (& coefficients,*input, shift);

          let perfect_result = evaluate_exactly (& coefficients,*input, shift);
          let difference = Ratio::from_integer (FromPrimitive::from_i64 (result).unwrap()) - perfect_result;
          assert!(difference < Ratio::from_integer (FromPrimitive::from_i64 (1).unwrap()));
        }
      }
    }
  }

  quickcheck! {
    fn quickcheck_safe_translation_is_safe (coefficients_and_maxima: Vec<(i64, i64)>, input: i64)->bool {
      let mut coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      if translate_check (& coefficients, input, maxima).is_err() {return true}
      translate_unchecked (&mut coefficients, input);
      println!( "{:?}", coefficients);
      within_bounds_check (& coefficients, maxima).is_ok()
    }

    fn quickcheck_safe_translation_reverses_correctly (coefficients: Vec<i64>, input: i64)->bool {
      let mut coefficients = coefficients;
      let maxima = |_: usize | i64::max_value();
      if input == i64::min_value () {return true}
      let original_coefficients = coefficients.clone ();
      if translate_check (& coefficients, input, maxima).is_err() {return true}
      translate_unchecked (&mut coefficients, input);
      if translate_check (& coefficients, -input, maxima).is_err() {return true}
      println!( "{:?}", coefficients);
      translate_unchecked (&mut coefficients, -input);
      coefficients == original_coefficients
    }

    fn quickcheck_exact_safe_translation_range_is_safe (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      let range = exact_safe_translation_range (& coefficients, maxima);
      println!( "{:?}", range);
      range < 0 || translate_check (& coefficients, range, maxima).is_ok()
    }

    fn quickcheck_exact_safe_translation_range_is_maximal (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      let range = exact_safe_translation_range (& coefficients, maxima);
      println!( "{:?}", range);
      range == i64::max_value() || translate_check (& coefficients, range + 1, maxima).is_err()
    }

    fn quickcheck_conservative_safe_translation_range_is_safe (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      let exact_range = exact_safe_translation_range (& coefficients, maxima);
      let conservative_range = conservative_safe_translation_range (& coefficients, maxima);
      println!( "{:?}", (exact_range, conservative_range));
      conservative_range <= exact_range
    }

    fn quickcheck_conservative_safe_translation_range_is_within_half_maximal (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      let exact_range = exact_safe_translation_range (& coefficients, maxima);
      let conservative_range = conservative_safe_translation_range (& coefficients, maxima);
      println!( "{:?}", (exact_range, conservative_range));
      conservative_range == exact_range || conservative_range > exact_range >> 1
    }
  }
}
<commit_msg>Made a technical fix to the handling of overflowing inputs with 0 coefficients<commit_after>use num::{Signed};

use super::*;

#[derive (Debug, Fail)]
pub enum Error <T: Integer> {
  #[fail (display = "out-of-bounds value in the x^{} term (coefficient is {}, but must have abs <= {})", power, coefficient, maximum_abs)]
  CoefficientOutOfBounds {
    power: usize,
    coefficient: T,
    maximum_abs: T,
  },
  #[fail (display = "out-of-bounds input to polynomial function (for the x^{} term, |{}*{}^{}| was {}, but must have abs <= {})", power, coefficient, input, power, value, maximum_abs)]
  TermOutOfBounds {
    power: usize,
    input: T,
    coefficient: T,
    value: T,
    maximum_abs: T,
  },
  #[fail (display = "out-of-bounds input to polynomial function (for the x^{} term, |{}*{}^{}| overflowed the type)", power, input, coefficient, power)]
  TermOverflowed {
    power: usize,
    input: T,
    coefficient: T,
  },
}

/// Approximately evaluate an integer polynomial at an input in the range [-0.5, 0.5].
///
/// The input is represented as an integer combined with a right-shift size.
/// Each coefficient, except the constant one, must obey coefficient.abs() <= T::max_value() >> shift.
/// The constant coefficient must obey coefficient.abs() <= T::max_value() - (T::max_value() >> shift),
/// i.e. be small enough that you can add one of the other coefficients to it without overflowing.
/// Given these conditions, the result is guaranteed to be strictly within 1 of the ideal result.
/// For instance, if the ideal result was 2.125, this function could return 2 or 3,
/// but if it was 2, this function can only return exactly 2.
pub fn evaluate_at_small_input <Coefficient: Copy, T: Integer + Signed + From <Coefficient>> (coefficients: & [Coefficient], input: T, shift: u32)->T {
  let half = (T::one() << shift) >> 1u32;
  assert!(-half <= input && input <= half, "inputs to evaluate_at_small_input must be in the range [-0.5, 0.5]");

  if coefficients.len () == 0 {return T::zero()}

  let mut result = T::zero();
  for coefficient in coefficients.iter().skip(1).rev() {
    result = shr_nicely_rounded ((result + (*coefficient).into())*input, shift);
  }
  result + coefficients [0].into()
}

/// Translate a polynomial horizontally, letting it take inputs relative to a given input instead of relative to 0.
///
/// This is equivalent to replacing its coefficients with its Taylor coefficients at the given input.
/// This is equivalent to translating the graph *left* by the given input amount.
///
/// For constant polynomials, this is a no-op.
///
/// To avoid overflow, each term, when evaluating at the given input,
/// must obey term.abs() <= T::max_value() >> (coefficients.len() - 1).
/// If this condition is broken, translate() returns Err and makes no changes.
pub fn translate <T: Integer + Signed> (coefficients: &mut [T], input: T)->Result <(), Error <T>> {
  translate_check (coefficients, input, |_| T::max_value())?;
  translate_unchecked (coefficients, input);
  Ok (())
}

/// Same as translate(), but assume no overflow will occur.
///
/// Useful in performance-critical code when you already know there won't be overflow,
/// such as in the range returned by conservative_safe_translation_range.
pub fn translate_unchecked <T: Integer + Signed> (coefficients: &mut [T], input: T) {
  coefficients.reverse();
  for index in 0..coefficients.len() {
    let coefficient = mem::replace (&mut coefficients [index], T::zero());
    for derivative in (1..(index + 1)).rev() {
      coefficients [derivative] = coefficients [derivative]*input + coefficients [derivative - 1]
    }
    coefficients [0] = coefficients [0]*input + coefficient
  }
}

pub fn within_bounds_check <T: Integer + Signed, MaximumFn: Fn (usize)->T> (coefficients: & [T], maximum: MaximumFn)->Result <(), Error <T>> {
  for (power, coefficient) in coefficients.iter().enumerate() {
    let max = maximum (power);
    if coefficient == &T::min_value() {return Err (Error::CoefficientOutOfBounds {power, coefficient: *coefficient, maximum_abs: max})}
    let magnitude = coefficient.abs();
    if magnitude > max {return Err (Error::CoefficientOutOfBounds {power, coefficient: *coefficient, maximum_abs: max})}
  }
  Ok (())
}

pub fn translate_check <T: Integer + Signed, MaximumFn: Fn (usize)->T> (coefficients: & [T], input: T, maximum: MaximumFn)->Result <(), Error <T>> {
  if coefficients.len() == 0 {return Ok (());}
  if input == T::zero() {within_bounds_check (coefficients, maximum)?; return Ok (());}
  let mut factor = T::one();
  let input = input.abs();
  let sum_safety_shift = coefficients.len() as u32 - 1;
  let mut running_maximum = T::max_value();
  for (power, coefficient) in coefficients.iter().enumerate() {
    if power > 0 { match factor.checked_mul (&input) {
      // note: there's a slight weirdness here, where if the factor overflows
      // but the current coefficient is 0, we still advance to the next iteration,
      // and the factor is the "wrong" value on that iteration.
      // However, it just overflows again, leading to the correct result.
      None => if coefficient != &T::zero() {return Err (Error::TermOverflowed {power, input, coefficient:*coefficient})},
      Some (next_factor) => factor = next_factor,
    }}
    running_maximum = min (running_maximum, maximum (power).checked_shr (sum_safety_shift).unwrap_or (T::zero()));
    if coefficient == &T::min_value() {return Err (Error::TermOverflowed {power, input, coefficient:*coefficient});}
    let magnitude = coefficient.abs();
    match magnitude.checked_mul (&factor) {
      None => return Err (Error::TermOverflowed {power, input, coefficient:*coefficient}),
      Some (term) => if term > running_maximum {return Err (Error::TermOutOfBounds {power, input, coefficient:*coefficient, value: term, maximum_abs: running_maximum})},
    }
  }
  Ok (())
}

/// Find the range of inputs to which the polynomial can safely be translated.
///
/// Magnitudes <= result are safe to translate to; magnitudes > result aren't.
/// If even the initial value is out of bounds, this function returns -1.
///
/// The third argument can impose a stricter maximum on the resulting coefficients.
pub fn exact_safe_translation_range <T: Integer + Signed, MaximumFn: Fn (usize)->T> (coefficients: & [T], maximum: MaximumFn)->T {
  if within_bounds_check (coefficients, & maximum).is_err() {return -T::one();}
  // pick a min that's definitely legal and a max where max + 1 is definitely illegal
  let mut min = T::zero();
  let mut max = T::max_value();
  while min < max {
    let mid = mean_ceil (min, max);
    if translate_check (coefficients, mid, & maximum).is_ok() {
      min = mid;
    } else {
      max = mid - T::one();
    }
  }
  min
}

/// Find a range of inputs to which the polynomial can safely be translated.
///
/// Magnitudes <= result are safe to translate to.
/// Magnitudes >= result*2 aren't.
/// If even the initial value is out of bounds, this function returns -1.
///
/// The third argument can impose a stricter maximum on the resulting coefficients.
///
/// This function is much faster than exact_safe_translation_range, at the cost of being less precise.
pub fn conservative_safe_translation_range <T: Integer + Signed, MaximumFn: Fn (usize)->T> (coefficients: &[T], maximum: MaximumFn)->T {
  if coefficients.len() == 0 {return T::max_value()}
  if within_bounds_check (coefficients, & maximum).is_err() {return -T::one();}
  let sum_safety_shift = coefficients.len() as u32 - 1;
  let mut running_maximum = T::max_value();
  let mut result_shift = mem::size_of::<T>()*8 - 2;
  for (power, coefficient) in coefficients.iter().enumerate() {
    let magnitude = coefficient.abs();
    running_maximum = min (running_maximum, maximum (power).checked_shr (sum_safety_shift).unwrap_or (T::zero()));
    while magnitude > running_maximum.checked_shr((result_shift*power) as u32).unwrap_or (T::zero()) {
      if result_shift == 0 {return T::zero();}
      result_shift -= 1;
    }
  }
  T::one() << result_shift
}

#[cfg (test)]
mod tests {
  use super::*;
  use num::{Zero, One, FromPrimitive, Integer};
  use num::bigint::BigInt;
  use num::rational::{Ratio, BigRational};
  use rand::distributions::Distribution;
  use rand::distributions::range::{Range};
  use rand::{Rng, SeedableRng};

  fn evaluate_exactly <Coefficient: Copy, T: Integer + Signed + From <Coefficient>> (coefficients: & [Coefficient], input: T, shift: u32)->BigRational
    where BigInt: From <Coefficient> + From <T> {
    let mut result: BigRational = Ratio::zero();
    let input = Ratio::new(BigInt::from (input), BigInt::one() << shift as usize);
    for coefficient in coefficients.iter().rev() {
      result = result*&input + BigInt::from(*coefficient);
    }
    result
  }

  #[test]
  fn test_evaluate_at_small_input() {
    let mut generator =::rand::chacha::ChaChaRng::from_seed ([33; 32]) ;
    for shift in 0..5 {
      let shift = 1 << shift;
      let half = 1i64 << (shift - 1);
      let maximum = i64::max_value() >> shift;
      let constant_maximum = i64::max_value() - maximum;
      let range = Range::new_inclusive (- maximum, maximum) ;
      let constant_range = Range::new_inclusive (- constant_maximum, constant_maximum);
      let mut test_inputs = vec![- half, - half + 1, - 1, 0, 1, half - 1, half];
      if half > 2 {for _ in 0..6 {test_inputs.push (generator.gen_range (- half + 2, half - 1));}}
      for _ in 0..40 {
        let mut coefficients = vec![constant_range.sample (&mut generator)];
        while generator.gen() {coefficients.push (range.sample (&mut generator));}
        for input in test_inputs.iter() {
          let result = evaluate_at_small_input (& coefficients,*input, shift);

          let perfect_result = evaluate_exactly (& coefficients,*input, shift);
          let difference = Ratio::from_integer (FromPrimitive::from_i64 (result).unwrap()) - perfect_result;
          assert!(difference < Ratio::from_integer (FromPrimitive::from_i64 (1).unwrap()));
        }
      }
    }
  }

  quickcheck! {
    fn quickcheck_safe_translation_is_safe (coefficients_and_maxima: Vec<(i64, i64)>, input: i64)->bool {
      let mut coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      if translate_check (& coefficients, input, maxima).is_err() {return true}
      translate_unchecked (&mut coefficients, input);
      println!( "{:?}", coefficients);
      within_bounds_check (& coefficients, maxima).is_ok()
    }

    fn quickcheck_safe_translation_reverses_correctly (coefficients: Vec<i64>, input: i64)->bool {
      let mut coefficients = coefficients;
      let maxima = |_: usize | i64::max_value();
      if input == i64::min_value () {return true}
      let original_coefficients = coefficients.clone ();
      if translate_check (& coefficients, input, maxima).is_err() {return true}
      translate_unchecked (&mut coefficients, input);
      if translate_check (& coefficients, -input, maxima).is_err() {return true}
      println!( "{:?}", coefficients);
      translate_unchecked (&mut coefficients, -input);
      coefficients == original_coefficients
    }

    fn quickcheck_exact_safe_translation_range_is_safe (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      let range = exact_safe_translation_range (& coefficients, maxima);
      println!( "{:?}", range);
      range < 0 || translate_check (& coefficients, range, maxima).is_ok()
    }

    fn quickcheck_exact_safe_translation_range_is_maximal (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      let range = exact_safe_translation_range (& coefficients, maxima);
      println!( "{:?}", range);
      range == i64::max_value() || translate_check (& coefficients, range + 1, maxima).is_err()
    }

    fn quickcheck_conservative_safe_translation_range_is_safe (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      let exact_range = exact_safe_translation_range (& coefficients, maxima);
      let conservative_range = conservative_safe_translation_range (& coefficients, maxima);
      println!( "{:?}", (exact_range, conservative_range));
      conservative_range <= exact_range
    }

    fn quickcheck_conservative_safe_translation_range_is_within_half_maximal (coefficients_and_maxima: Vec<(i64, i64)>)->bool {
      let coefficients: Vec<_> = coefficients_and_maxima.iter().map (| & (coefficient,_) | coefficient).collect();
      let maxima = | index: usize | coefficients_and_maxima [index].1.checked_abs().unwrap_or (0);
      let exact_range = exact_safe_translation_range (& coefficients, maxima);
      let conservative_range = conservative_safe_translation_range (& coefficients, maxima);
      println!( "{:?}", (exact_range, conservative_range));
      conservative_range == exact_range || conservative_range > exact_range >> 1
    }
  }
}
<|endoftext|>


<commit_before>mod statistics;
mod planner;

use kite::doc_id_set::DocIdSet;
use kite::segment::Segment;
use kite::query::Query;
use kite::collectors::{Collector, DocumentMatch};
use byteorder::{ByteOrder, BigEndian};

use super::RocksDBIndexReader;
use search::statistics::{StatisticsReader, RocksDBStatisticsReader};
use search::planner::{SearchPlan, plan_query};
use search::planner::boolean_query::BooleanQueryOp;
use search::planner::score_function::{CombinatorScorer, ScoreFunctionOp};

fn search_segment_boolean_phase<S: Segment>(boolean_query: &Vec<BooleanQueryOp>, is_negated: bool, segment: &S) -> Result<DocIdSet, String> {
    // Execute boolean query
    let mut stack = Vec::new();
    for op in boolean_query.iter() {
        match *op {
            BooleanQueryOp::PushEmpty => {
                stack.push(DocIdSet::new_filled(0));
            }
            BooleanQueryOp::PushFull => {
                stack.push(DocIdSet::new_filled(65536));
            }
            BooleanQueryOp::PushTermDirectory(field_ref, term_ref) => {
                match try!(segment.load_term_directory(field_ref, term_ref)) {
                    Some(doc_id_set) => stack.push(doc_id_set),
                    None => stack.push(DocIdSet::new_filled(0)),
                }
            }
            BooleanQueryOp::PushDeletionList => {
                    match try!(segment.load_deletion_list()) {
                    Some(doc_id_set) => stack.push(doc_id_set),
                    None => stack.push(DocIdSet::new_filled(0)),
                }
            }
            BooleanQueryOp::And => {
                let b = stack.pop().expect("boolean query executor: stack underflow");
                let a = stack.pop().expect("boolean query executor: stack underflow");
                stack.push(a.intersection(&b));
            }
            BooleanQueryOp::Or => {
                let b = stack.pop().expect("boolean query executor: stack underflow");
                let a = stack.pop().expect("boolean query executor: stack underflow");
                stack.push(a.union(&b));
            }
            BooleanQueryOp::AndNot => {
                let b = stack.pop().expect("boolean query executor: stack underflow");
                let a = stack.pop().expect("boolean query executor: stack underflow");
                stack.push(a.exclusion(&b));
            }
        }
    }

    if !stack.len() == 1 {
        // This shouldn't be possible unless there's a bug in the planner
        panic!("boolean query executor: stack size too big ({})", stack.len());
    }
    let mut matches = stack.pop().unwrap();

    // Invert the list if the query is negated
    if is_negated {
        let total_docs = try!(segment.load_statistic(b"total_docs")).unwrap_or(0);
        let all_docs = DocIdSet::new_filled(total_docs as u32);
        matches = all_docs.exclusion(&matches);
    }

    Ok(matches)
}

fn score_doc<S: Segment, R: StatisticsReader>(doc_id: u16, score_function: &Vec<ScoreFunctionOp>, segment: &S, mut stats: &mut R) -> Result<f64, String> {
    // Execute score function
    let mut stack = Vec::new();
    for op in score_function.iter() {
        match *op {
            ScoreFunctionOp::Literal(val) => stack.push(val),
            ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer) => {
                // TODO: Check this isn't really slow
                match try!(segment.load_term_directory(field_ref, term_ref)) {
                    Some(doc_id_set) => {
                        if doc_id_set.contains_doc(doc_id) {
                            // Read field length
                            // TODO: we only need this for BM25
                            let field_length_raw = try!(segment.load_stored_field_value_raw(doc_id, field_ref, b"len"));
                            let field_length = match field_length_raw {
                                Some(value) => {
                                    let length_sqrt = (value[0] as f64) / 3.0 + 1.0;
                                    length_sqrt * length_sqrt
                                }
                                None => 1.0
                            };

                            // Read term frequency
                            let mut value_type = vec![b't', b'f'];
                            value_type.extend(term_ref.ord().to_string().as_bytes());
                            let term_frequency_raw = try!(segment.load_stored_field_value_raw(doc_id, field_ref, &value_type));
                            let term_frequency = match term_frequency_raw {
                                Some(value) => BigEndian::read_i64(&value),
                                None => 1,
                            };

                            let score = scorer.similarity_model.score(term_frequency as u32, field_length, try!(stats.total_tokens(field_ref)) as u64, try!(stats.total_docs(field_ref)) as u64, try!(stats.term_document_frequency(field_ref, term_ref)) as u64);
                            stack.push(score * scorer.boost);
                        } else {
                            stack.push(0.0f64);
                        }
                    }
                    None => stack.push(0.0f64),
                }
            }
            ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                let score = match *scorer {
                    CombinatorScorer::Avg => {
                        let mut total_score = 0.0f64;

                        for _ in 0..num_vals {
                            total_score += stack.pop().expect("document scorer: stack underflow");
                        }

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

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

                        max_score
                    }
                };

                stack.push(score);
            }
        }
    }

    if !stack.len() == 1 {
        // This shouldn't be possible unless there's a bug in the planner
        panic!("document scorer: stack size too big ({})", stack.len());
    }

    Ok(stack.pop().expect("document scorer: stack underflow"))
}

fn search_segment<C: Collector, S: Segment, R: StatisticsReader>(collector: &mut C, plan: &SearchPlan, segment: &S, mut stats: &mut R) -> Result<(), String> {
    let matches = try!(search_segment_boolean_phase(&plan.boolean_query, plan.boolean_query_is_negated, segment));

    // Score documents and pass to collector
    for doc in matches.iter() {
        let score = try!(score_doc(doc, &plan.score_function, segment, stats));

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

    Ok(())
}

impl<'a> RocksDBIndexReader<'a> {
    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) -> Result<(), String> {
        // Plan query
        let plan = plan_query(&self, query, collector.needs_score());

        // Initialise statistics reader
        let mut stats = RocksDBStatisticsReader::new(&self);

        // Run query on each segment
        for segment in self.store.segments.iter_active(&self) {
            try!(search_segment(collector, &plan, &segment, &mut stats));
        }

        Ok(())
    }
}
<commit_msg>Renamed search_segment_boolean_phase to run_boolean_query<commit_after>mod statistics;
mod planner;

use kite::doc_id_set::DocIdSet;
use kite::segment::Segment;
use kite::query::Query;
use kite::collectors::{Collector, DocumentMatch};
use byteorder::{ByteOrder, BigEndian};

use super::RocksDBIndexReader;
use search::statistics::{StatisticsReader, RocksDBStatisticsReader};
use search::planner::{SearchPlan, plan_query};
use search::planner::boolean_query::BooleanQueryOp;
use search::planner::score_function::{CombinatorScorer, ScoreFunctionOp};

fn run_boolean_query<S: Segment>(boolean_query: &Vec<BooleanQueryOp>, is_negated: bool, segment: &S) -> Result<DocIdSet, String> {
    // Execute boolean query
    let mut stack = Vec::new();
    for op in boolean_query.iter() {
        match *op {
            BooleanQueryOp::PushEmpty => {
                stack.push(DocIdSet::new_filled(0));
            }
            BooleanQueryOp::PushFull => {
                stack.push(DocIdSet::new_filled(65536));
            }
            BooleanQueryOp::PushTermDirectory(field_ref, term_ref) => {
                match try!(segment.load_term_directory(field_ref, term_ref)) {
                    Some(doc_id_set) => stack.push(doc_id_set),
                    None => stack.push(DocIdSet::new_filled(0)),
                }
            }
            BooleanQueryOp::PushDeletionList => {
                    match try!(segment.load_deletion_list()) {
                    Some(doc_id_set) => stack.push(doc_id_set),
                    None => stack.push(DocIdSet::new_filled(0)),
                }
            }
            BooleanQueryOp::And => {
                let b = stack.pop().expect("boolean query executor: stack underflow");
                let a = stack.pop().expect("boolean query executor: stack underflow");
                stack.push(a.intersection(&b));
            }
            BooleanQueryOp::Or => {
                let b = stack.pop().expect("boolean query executor: stack underflow");
                let a = stack.pop().expect("boolean query executor: stack underflow");
                stack.push(a.union(&b));
            }
            BooleanQueryOp::AndNot => {
                let b = stack.pop().expect("boolean query executor: stack underflow");
                let a = stack.pop().expect("boolean query executor: stack underflow");
                stack.push(a.exclusion(&b));
            }
        }
    }

    if !stack.len() == 1 {
        // This shouldn't be possible unless there's a bug in the planner
        panic!("boolean query executor: stack size too big ({})", stack.len());
    }
    let mut matches = stack.pop().unwrap();

    // Invert the list if the query is negated
    if is_negated {
        let total_docs = try!(segment.load_statistic(b"total_docs")).unwrap_or(0);
        let all_docs = DocIdSet::new_filled(total_docs as u32);
        matches = all_docs.exclusion(&matches);
    }

    Ok(matches)
}

fn score_doc<S: Segment, R: StatisticsReader>(doc_id: u16, score_function: &Vec<ScoreFunctionOp>, segment: &S, mut stats: &mut R) -> Result<f64, String> {
    // Execute score function
    let mut stack = Vec::new();
    for op in score_function.iter() {
        match *op {
            ScoreFunctionOp::Literal(val) => stack.push(val),
            ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer) => {
                // TODO: Check this isn't really slow
                match try!(segment.load_term_directory(field_ref, term_ref)) {
                    Some(doc_id_set) => {
                        if doc_id_set.contains_doc(doc_id) {
                            // Read field length
                            // TODO: we only need this for BM25
                            let field_length_raw = try!(segment.load_stored_field_value_raw(doc_id, field_ref, b"len"));
                            let field_length = match field_length_raw {
                                Some(value) => {
                                    let length_sqrt = (value[0] as f64) / 3.0 + 1.0;
                                    length_sqrt * length_sqrt
                                }
                                None => 1.0
                            };

                            // Read term frequency
                            let mut value_type = vec![b't', b'f'];
                            value_type.extend(term_ref.ord().to_string().as_bytes());
                            let term_frequency_raw = try!(segment.load_stored_field_value_raw(doc_id, field_ref, &value_type));
                            let term_frequency = match term_frequency_raw {
                                Some(value) => BigEndian::read_i64(&value),
                                None => 1,
                            };

                            let score = scorer.similarity_model.score(term_frequency as u32, field_length, try!(stats.total_tokens(field_ref)) as u64, try!(stats.total_docs(field_ref)) as u64, try!(stats.term_document_frequency(field_ref, term_ref)) as u64);
                            stack.push(score * scorer.boost);
                        } else {
                            stack.push(0.0f64);
                        }
                    }
                    None => stack.push(0.0f64),
                }
            }
            ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                let score = match *scorer {
                    CombinatorScorer::Avg => {
                        let mut total_score = 0.0f64;

                        for _ in 0..num_vals {
                            total_score += stack.pop().expect("document scorer: stack underflow");
                        }

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

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

                        max_score
                    }
                };

                stack.push(score);
            }
        }
    }

    if !stack.len() == 1 {
        // This shouldn't be possible unless there's a bug in the planner
        panic!("document scorer: stack size too big ({})", stack.len());
    }

    Ok(stack.pop().expect("document scorer: stack underflow"))
}

fn search_segment<C: Collector, S: Segment, R: StatisticsReader>(collector: &mut C, plan: &SearchPlan, segment: &S, mut stats: &mut R) -> Result<(), String> {
    let matches = try!(run_boolean_query(&plan.boolean_query, plan.boolean_query_is_negated, segment));

    // Score documents and pass to collector
    for doc in matches.iter() {
        let score = try!(score_doc(doc, &plan.score_function, segment, stats));

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

    Ok(())
}

impl<'a> RocksDBIndexReader<'a> {
    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) -> Result<(), String> {
        // Plan query
        let plan = plan_query(&self, query, collector.needs_score());

        // Initialise statistics reader
        let mut stats = RocksDBStatisticsReader::new(&self);

        // Run query on each segment
        for segment in self.store.segments.iter_active(&self) {
            try!(search_segment(collector, &plan, &segment, &mut stats));
        }

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


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

#![deny(
    non_camel_case_types,
    non_snake_case,
    path_statements,
    trivial_numeric_casts,
    unstable_features,
    unused_allocation,
    unused_import_braces,
    unused_imports,
    unused_mut,
    unused_qualifications,
    while_true,
)]

#[macro_use] extern crate log;
extern crate crypto;
extern crate itertools;
extern crate semver;
extern crate toml;
extern crate version;
extern crate walkdir;

#[macro_use] extern crate libimagstore;
#[macro_use] extern crate libimagerror;
#[macro_use] extern crate libimagutil;
extern crate libimagentrylist;

module_entry_path_mod!("ref");

pub mod error;
pub mod flags;
pub mod hasher;
pub mod hashers;
pub mod lister;
pub mod reference;
pub mod result;
<commit_msg>Remove unused macro_use import<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

#![deny(
    non_camel_case_types,
    non_snake_case,
    path_statements,
    trivial_numeric_casts,
    unstable_features,
    unused_allocation,
    unused_import_braces,
    unused_imports,
    unused_mut,
    unused_qualifications,
    while_true,
)]

#[macro_use] extern crate log;
extern crate crypto;
extern crate itertools;
extern crate semver;
extern crate toml;
extern crate version;
extern crate walkdir;

#[macro_use] extern crate libimagstore;
#[macro_use] extern crate libimagerror;
extern crate libimagutil;
extern crate libimagentrylist;

module_entry_path_mod!("ref");

pub mod error;
pub mod flags;
pub mod hasher;
pub mod hashers;
pub mod lister;
pub mod reference;
pub mod result;
<|endoftext|>


<commit_before>// #![feature(scoped)]
// #![feature(collections)]

extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::thread;
use std::mem;

use dataflow_join::graph::{GraphTrait, GraphMMap};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_static::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      list of host:port for workers
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    // let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
    //               else { panic!("invalid setting for workers: {}", args.get_str("<workers>")) };
    // println!("starting pagerank dataflow with {:?} worker{}", workers, if workers == 1 { "" } else { "s" });
    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_multi(communicators, source);
    }
    else if workers > 1 {
        println!("Initializing ProcessCommunicator");
        pagerank_multi(ProcessCommunicator::new_vector(workers), source);
    }
    else {
        println!("Initializing ThreadCommunicator");
        pagerank_multi(vec![ThreadCommunicator], source);
    };
}

fn pagerank_multi<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

fn pagerank<C>(communicator: C, filename: String, workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    {   // new scope avoids long borrow on root
        let mut builder = root.new_subgraph();

        // establish the beginnings of a loop,
        // 20 iterations, each time around += 1.
        let (helper, stream) = builder.loop_variable::<(u32, f32)>(RootTimestamp::new(20), Local(1));

        let graph = GraphMMap::<u32>::new(&filename);

        let nodes = graph.nodes();

        let mut src = vec![1.0; 1 + (nodes / peers as usize)];  // local rank accumulation
        let mut tmp = vec![1.0; 1 + (nodes / peers as usize)];  // local rank accumulation
        // let mut dst = vec![0.0; nodes];                         // local rank accumulation

        let mut buf = vec![];

        let mut start = time::precise_time_s();

        // from feedback, place an operator that
        // aggregates and broadcasts ranks along edges.
        let ranks = stream.enable(builder).unary_notify(

            Exchange::new(|x: &(u32, f32)| x.0 as u64),     // 1. how data should be exchanged
            format!("PageRank"),                            // 2. a tasteful, descriptive name
            vec![RootTimestamp::new(0)],                    // 3. indicate an initial capability
            move |input, output, iterator| {                // 4. provide the operator logic

                while let Some((iter, _)) = iterator.next() {

                    mem::swap(&mut src, &mut tmp);

                    // /---- should look familiar! ----\
                    for node in 0..src.len() {
                        src[node] = 0.15 + 0.85 * src[node];
                    }

                    for node in 0..src.len() {
                        let mut session = output.session(&iter);

                        let edges = graph.edges(index + peers * node);
                        let value = src[node] / edges.len() as f32;
                        for &b in edges {
                            session.give((b, value));
                        }

                        while let Some((iter, data)) = input.pull() {
                            iterator.notify_at(&iter);
                            buf.extend(data.drain_temp());
                        }

                        for (node, rank) in buf.drain_temp() {
                            tmp[node as usize / peers] += rank;
                        }
                    }
                    // \------ end familiar part ------/

                    println!("iteration {:?}: {}s", iter, time::precise_time_s() - start);
                    start = time::precise_time_s();
                }

                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        tmp[node as usize / peers] += rank;
                    }
                }
            }
        );

        // let local_index = index as usize % workers;
        // let mut acc = vec![0.0; 1 + (nodes / workers)];

        ranks
        // .unary_notify(
        //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
        //     format!("Aggregation"),
        //     vec![],
        //     move |input, output, iterator| {
        //         while let Some((iter, data)) = input.pull() {
        //             iterator.notify_at(&iter);
        //             for (node, rank) in data.drain_temp() {
        //                 acc[node as usize / workers] += rank;
        //             }
        //         }
        //
        //         while let Some((item, _)) = iterator.next() {
        //
        //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
        //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
        //
        //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
        //             assert!(acc.len() == (1 + (nodes/workers)));
        //         }
        //     }
        // )
        .connect_loop(helper);
    }

    while root.step() { }
}
<commit_msg>update exotic pagerank<commit_after>// #![feature(scoped)]
// #![feature(collections)]

extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::thread;
use std::mem;

use dataflow_join::graph::{GraphTrait, GraphMMap};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_static::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      list of host:port for workers
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    // let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
    //               else { panic!("invalid setting for workers: {}", args.get_str("<workers>")) };
    // println!("starting pagerank dataflow with {:?} worker{}", workers, if workers == 1 { "" } else { "s" });
    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_multi(communicators, source);
    }
    else if workers > 1 {
        println!("Initializing ProcessCommunicator");
        pagerank_multi(ProcessCommunicator::new_vector(workers), source);
    }
    else {
        println!("Initializing ThreadCommunicator");
        pagerank_multi(vec![ThreadCommunicator], source);
    };
}

fn pagerank_multi<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

fn pagerank<C>(communicator: C, filename: String, workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    {   // new scope avoids long borrow on root
        let mut builder = root.new_subgraph();

        // establish the beginnings of a loop,
        // 20 iterations, each time around += 1.
        let (helper, stream) = builder.loop_variable::<(u32, f32)>(RootTimestamp::new(20), Local(1));

        let graph = GraphMMap::<u32>::new(&filename);

        let nodes = graph.nodes();

        let mut src = vec![1.0; 1 + (nodes / peers as usize)];  // local rank accumulation
        let mut tmp = vec![1.0; 1 + (nodes / peers as usize)];  // local rank accumulation
        // let mut dst = vec![0.0; nodes];                         // local rank accumulation

        let mut buf = vec![];

        let mut start = time::precise_time_s();

        // from feedback, place an operator that
        // aggregates and broadcasts ranks along edges.
        let ranks = stream.enable(builder).unary_notify(

            Exchange::new(|x: &(u32, f32)| x.0 as u64),     // 1. how data should be exchanged
            format!("PageRank"),                            // 2. a tasteful, descriptive name
            vec![RootTimestamp::new(0)],                    // 3. indicate an initial capability
            move |input, output, iterator| {                // 4. provide the operator logic

                while let Some((iter, _)) = iterator.next() {

                    mem::swap(&mut src, &mut tmp);

                    // /---- should look familiar! ----\
                    for node in 0..src.len() {
                        src[node] = 0.15 + 0.85 * src[node];
                    }

                    let mut node = 0;
                    let mut read = 0;
                    let mut counter = 0;

                    while node < src.len() {

                        let mut session = output.session(&iter);
                        for _ in 0 .. std::cmp::min(1_000, src.len() - node) {

                            let edges = graph.edges(index + peers * node);
                            let value = src[node] / edges.len() as f32;
                            for &b in edges {
                                session.give((b, value));
                            }

                            counter += edges.len();
                            node += 1;
                        }

                        while let Some((iter, data)) = input.pull() {
                            iterator.notify_at(&iter);
                            read += data.len();
                            buf.extend(data.drain_temp());
                            if read > counter { break; }
                        }

                        for (node, rank) in buf.drain_temp() {
                            tmp[node as usize / peers] += rank;
                        }

                        // if (node % 100_000) == 0 {
                        //     println!("status: {} node, {} counter, {} read, \tdefecit: {}", node, counter, read, counter as i64 - read as i64);
                        // }
                    }
                    // \------ end familiar part ------/

                    println!("iteration {:?}: {}s", iter, time::precise_time_s() - start);
                    start = time::precise_time_s();
                }

                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        tmp[node as usize / peers] += rank;
                    }
                }
            }
        );

        // let local_index = index as usize % workers;
        // let mut acc = vec![0.0; 1 + (nodes / workers)];

        ranks
        // .unary_notify(
        //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
        //     format!("Aggregation"),
        //     vec![],
        //     move |input, output, iterator| {
        //         while let Some((iter, data)) = input.pull() {
        //             iterator.notify_at(&iter);
        //             for (node, rank) in data.drain_temp() {
        //                 acc[node as usize / workers] += rank;
        //             }
        //         }
        //
        //         while let Some((item, _)) = iterator.next() {
        //
        //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
        //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
        //
        //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
        //             assert!(acc.len() == (1 + (nodes/workers)));
        //         }
        //     }
        // )
        .connect_loop(helper);
    }

    while root.step() { }
}
<|endoftext|>


<commit_before>use crate::agents::*;
use crate::crypto::SigningAlgorithm;
use crate::utils::{
    agent::*,
    keys::{KeyPairExt, NamedKeyPair, SignError},
    pem::{self, ParsedKeyPair},
    BoxFuture,
};
use err_derive::Error;
use log::{info, warn};
use ring::{
    rand::SecureRandom,
    signature::{Ed25519KeyPair, RsaKeyPair},
};
use serde_json as json;
use std::ffi::OsString;
use std::io::{Error as IoError, ErrorKind as IoErrorKind};
use std::mem;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::{Arc, RwLock};
use std::time::{Duration, SystemTime};
use tokio::fs::{self, File};
use tokio::io::BufReader;
use tokio::process::Command;
use tokio::task;

#[derive(Debug, Error)]
pub enum RotatingKeysError {
    #[error(display = "could not open '{}': {}", path, error)]
    Open {
        path: String,
        #[error(source, nofrom)]
        error: IoError,
    },
    #[error(display = "could not create directory '{}': {}", path, error)]
    Mkdir {
        path: String,
        #[error(source, nofrom)]
        error: IoError,
    },
    #[error(
        display = "could not determine last modification time of '{}': {}",
        path,
        error
    )]
    StatMtime {
        path: String,
        #[error(source, nofrom)]
        error: IoError,
    },
    #[error(display = "could not parse file '{}': {}", path, error)]
    Parse {
        path: String,
        #[error(source, nofrom)]
        error: pem::ParseError,
    },
    #[error(
        display = "invalid type of key in '{}', want {}, found {}",
        path,
        want,
        found
    )]
    InvalidKeyType {
        path: String,
        want: &'static str,
        found: &'static str,
    },
    #[error(display = "expected exactly one key in '{}', found {}", path, found)]
    ExpectedOneKey { path: String, found: usize },
}

/// Struct that contains configuration we pass around.
struct RotateConfig {
    keys_ttl: Duration,
    generate_rsa_command: Vec<String>,
    rng: Box<dyn SecureRandom + Send + Sync>,
}

/// KeyManager where we rotating 3 keys of each type.
pub struct RotatingKeys {
    ed25519_keys: Option<KeySetHandle<Ed25519KeyPair>>,
    rsa_keys: Option<KeySetHandle<RsaKeyPair>>,
}

impl RotatingKeys {
    pub async fn new(
        keysdir: impl AsRef<Path>,
        keys_ttl: Duration,
        signing_algs: &[SigningAlgorithm],
        generate_rsa_command: Vec<String>,
        rng: impl SecureRandom + Send + Sync + 'static,
    ) -> Result<Self, RotatingKeysError> {
        info!(
            "Using rotating keys with a {}s interval and algorithms: {}",
            keys_ttl.as_secs(),
            SigningAlgorithm::format_list(signing_algs)
        );
        let rng = Box::new(rng);
        let config = Arc::new(RotateConfig {
            keys_ttl,
            generate_rsa_command,
            rng,
        });
        let ed25519_keys = if signing_algs.contains(&SigningAlgorithm::EdDsa) {
            Some(KeySet::from_subdir(keysdir.as_ref(), "ed25519", &config).await?)
        } else {
            None
        };
        let rsa_keys = if signing_algs.contains(&SigningAlgorithm::Rs256) {
            Some(KeySet::from_subdir(keysdir.as_ref(), "rsa", &config).await?)
        } else {
            None
        };
        Ok(RotatingKeys {
            ed25519_keys,
            rsa_keys,
        })
    }
}

impl Agent for RotatingKeys {}

impl Handler<SignJws> for RotatingKeys {
    fn handle(&mut self, message: SignJws, reply: ReplySender<SignJws>) {
        // TODO: Read-locking is blocking.
        let maybe_jws = match message.signing_alg {
            SigningAlgorithm::EdDsa => self.ed25519_keys.as_ref().map(|key_set| {
                let key_set = key_set.read().unwrap();
                key_set
                    .current
                    .sign_jws(&message.payload, &*key_set.config.rng)
            }),
            SigningAlgorithm::Rs256 => self.rsa_keys.as_ref().map(|key_set| {
                let key_set = key_set.read().unwrap();
                key_set
                    .current
                    .sign_jws(&message.payload, &*key_set.config.rng)
            }),
        };
        reply.send(
            maybe_jws.unwrap_or_else(|| Err(SignError::UnsupportedAlgorithm(message.signing_alg))),
        );
    }
}

impl Handler<GetPublicJwks> for RotatingKeys {
    fn handle(&mut self, _message: GetPublicJwks, reply: ReplySender<GetPublicJwks>) {
        // TODO: Read-locking is blocking.
        let mut jwks = vec![];
        if let Some(ref key_set) = self.ed25519_keys {
            jwks.append(&mut key_set.read().unwrap().public_jwks());
        }
        if let Some(ref key_set) = self.rsa_keys {
            jwks.append(&mut key_set.read().unwrap().public_jwks());
        }
        reply.send(jwks)
    }
}

impl KeyManagerSender for Addr<RotatingKeys> {}

/// Thread-safe handle to a KeySet.
type KeySetHandle<T> = Arc<RwLock<KeySet<T>>>;

/// A rotating set of 3 keys of a single type.
struct KeySet<T: KeyPairExt> {
    config: Arc<RotateConfig>,
    mtime: SystemTime,
    current: NamedKeyPair<T>,
    current_path: PathBuf,
    next: NamedKeyPair<T>,
    next_path: PathBuf,
    previous: Option<NamedKeyPair<T>>,
    previous_path: PathBuf,
}

impl<T> KeySet<T>
where
    T: KeyPairExt + GeneratedKeyPair + Send + Sync + 'static,
{
    /// Read key pairs of type `T` from a subdir of `keysdir`.
    async fn from_subdir(
        keysdir: impl AsRef<Path>,
        subdir: &str,
        config: &Arc<RotateConfig>,
    ) -> Result<KeySetHandle<T>, RotatingKeysError> {
        let mut dir = keysdir.as_ref().to_path_buf();
        dir.push(subdir);

        fs::create_dir_all(&dir)
            .await
            .map_err(|error| RotatingKeysError::Mkdir {
                path: dir.to_string_lossy().into_owned(),
                error,
            })?;

        let mut current_path = dir.clone();
        current_path.push("current.pem");
        let current = match read_key_file::<T>(&current_path).await? {
            Some(key_pair) => key_pair,
            None => T::generate(config.clone(), current_path.clone()).await,
        };

        // The last modification time of `next.pem`, along with `keys_ttl`,
        // is used to determine when the next rotation should happen.
        let mut next_path = dir.clone();
        next_path.push("next.pem");
        let (next, mtime) = if let Some(key_pair) = read_key_file::<T>(&next_path).await? {
            let mtime = fs::metadata(&next_path)
                .await
                .and_then(|meta| meta.modified())
                .map_err(|error| RotatingKeysError::StatMtime {
                    path: next_path.to_string_lossy().into_owned(),
                    error,
                })?;
            (key_pair, mtime)
        } else {
            (
                T::generate(config.clone(), next_path.clone()).await,
                SystemTime::now(),
            )
        };

        let mut previous_path = dir;
        previous_path.push("previous.pem");
        let previous = read_key_file::<T>(&previous_path).await?;

        let key_set = Arc::new(RwLock::new(KeySet {
            config: config.clone(),
            mtime,
            current,
            current_path,
            next,
            next_path,
            previous,
            previous_path,
        }));

        let initial_delay = Self::check_expiry(key_set.clone()).await;
        task::spawn(Self::rotate_loop(key_set.clone(), initial_delay));

        Ok(key_set)
    }

    /// The async loop that performs key rotation when necessary.
    async fn rotate_loop(handle: KeySetHandle<T>, mut delay: Duration) {
        loop {
            tokio::time::delay_for(delay).await;
            delay = Self::check_expiry(handle.clone()).await;
        }
    }

    /// Check if we should rotate keys, and return how long to delay the next check.
    async fn check_expiry(handle: KeySetHandle<T>) -> Duration {
        let read_handle = handle.clone();
        let (keys_ttl, mtime) = task::spawn_blocking(move || {
            let key_set = read_handle.read().expect("could not read-lock key set");
            (key_set.config.keys_ttl, key_set.mtime)
        })
        .await
        .unwrap();
        let delay = if let Ok(age) = SystemTime::now().duration_since(mtime) {
            keys_ttl.checked_sub(age)
        } else {
            warn!("Key set mtime is from the future, treating as bad input.");
            None
        };
        if let Some(delay) = delay {
            delay
        } else {
            Self::rotate(handle).await;
            keys_ttl
        }
    }

    /// Rotate keys in memory and on disk.
    ///
    /// If this fails, we panic, because it may happen at an arbitrary moment at run-time.
    async fn rotate(handle: KeySetHandle<T>) {
        let read_handle = handle.clone();
        let (config, next_path) = task::spawn_blocking(move || {
            let key_set = read_handle.read().expect("could not read-lock key set");
            std::fs::rename(&key_set.current_path, &key_set.previous_path)
                .expect("could not rename current key for rotation");
            std::fs::rename(&key_set.next_path, &key_set.current_path)
                .expect("could not rename next key for rotation");
            (key_set.config.clone(), key_set.next_path.clone())
        })
        .await
        .unwrap();
        let mut tmp = T::generate(config, next_path).await;
        task::spawn_blocking(move || {
            let mut key_set = handle.write().expect("could not write-lock key set");
            mem::swap(&mut key_set.next, &mut tmp);
            mem::swap(&mut key_set.current, &mut tmp);
            key_set.previous = Some(tmp);
        })
        .await
        .unwrap()
    }
}

impl<T: KeyPairExt> KeySet<T> {
    /// Get a list of JWKs containing public keys.
    fn public_jwks(&self) -> Vec<json::Value> {
        let mut list = vec![self.current.public_jwk(), self.next.public_jwk()];
        if let Some(ref previous) = self.previous {
            list.push(previous.public_jwk());
        }
        list
    }
}

/// Trait for key pair types we can generate.
trait GeneratedKeyPair: KeyPairExt + Sized {
    /// Generate a new key pair.
    ///
    /// If this fails, we panic, because it may happen at an arbitrary moment at run-time.
    ///
    /// This should log a message at info-level on success.
    ///
    /// Note that this function may block.
    fn generate(config: Arc<RotateConfig>, out_file: PathBuf) -> BoxFuture<NamedKeyPair<Self>>;

    /// Convert a ParsedKeyPair, if it is of the correct type.
    fn from_parsed(
        parsed: ParsedKeyPair,
        path: &Path,
    ) -> Result<NamedKeyPair<Self>, RotatingKeysError>;
}

impl GeneratedKeyPair for Ed25519KeyPair {
    fn generate(config: Arc<RotateConfig>, out_file: PathBuf) -> BoxFuture<NamedKeyPair<Self>> {
        Box::pin(async move {
            let doc =
                Self::generate_pkcs8(&*config.rng).expect("could not generate Ed25519 key pair");
            let key_pair =
                Self::from_pkcs8(doc.as_ref()).expect("could not parse generated Ed25519 key pair");
            let file = File::create(&out_file)
                .await
                .expect("could not open generated key pair output file");
            pem::write_pkcs8(&doc, file)
                .await
                .expect("could not write generated key pair to output file");
            info!("Generated new Ed25519 key: {:?}", out_file);
            key_pair.into()
        })
    }

    fn from_parsed(
        parsed: ParsedKeyPair,
        path: &Path,
    ) -> Result<NamedKeyPair<Self>, RotatingKeysError> {
        match parsed {
            ParsedKeyPair::Ed25519(inner) => Ok(inner.into()),
            other => Err(RotatingKeysError::InvalidKeyType {
                path: path.to_string_lossy().into_owned(),
                want: "Ed25519",
                found: other.kind(),
            }),
        }
    }
}

impl GeneratedKeyPair for RsaKeyPair {
    fn generate(config: Arc<RotateConfig>, out_file: PathBuf) -> BoxFuture<NamedKeyPair<Self>> {
        Box::pin(async move {
            let mut args: Vec<OsString> = config
                .generate_rsa_command
                .iter()
                .map(|arg| arg.into())
                .collect();
            let program = args.remove(0);
            let file = if let Some(part) = args.iter_mut().find(|part| part.as_os_str() == "{}") {
                *part = out_file.clone().into();
                None
            } else {
                let file = File::create(&out_file)
                    .await
                    .expect("could not open generated key pair output file");
                Some(file)
            };
            let mut command = Command::new(program);
            command.args(args).stdin(Stdio::null());
            if let Some(file) = file {
                command.stdout(file.into_std().await);
            }
            let status = command
                .status()
                .await
                .expect("Failed to run command to generate RSA key");
            if !status.success() {
                panic!("Command to generate RSA key failed with status {}", status);
            }
            let key_pair = read_key_file(&out_file)
                .await
                .expect("could not read generated RSA key file")
                .expect("generated RSA key file not found");
            info!("Generated new RSA key: {:?}", out_file);
            key_pair
        })
    }

    fn from_parsed(
        parsed: ParsedKeyPair,
        path: &Path,
    ) -> Result<NamedKeyPair<Self>, RotatingKeysError> {
        match parsed {
            ParsedKeyPair::Rsa(inner) => Ok(inner.into()),
            other => Err(RotatingKeysError::InvalidKeyType {
                path: path.to_string_lossy().into_owned(),
                want: "RSA",
                found: other.kind(),
            }),
        }
    }
}

/// Read a single key pair of type `T` from `path`.
///
/// Returns `None` when the file does not exist.
///
/// Note that this function may block.
async fn read_key_file<T: KeyPairExt + GeneratedKeyPair>(
    path: &Path,
) -> Result<Option<NamedKeyPair<T>>, RotatingKeysError> {
    let file = match File::open(path).await {
        Ok(file) => file,
        Err(error) => {
            if error.kind() == IoErrorKind::NotFound {
                return Ok(None);
            } else {
                return Err(RotatingKeysError::Open {
                    path: path.to_string_lossy().into_owned(),
                    error,
                });
            }
        }
    };
    let mut key_pairs = pem::parse_key_pairs(BufReader::new(file))
        .await
        .map_err(|error| RotatingKeysError::Parse {
            path: path.to_string_lossy().into_owned(),
            error,
        })?;
    if key_pairs.len() != 1 {
        return Err(RotatingKeysError::ExpectedOneKey {
            path: path.to_string_lossy().into_owned(),
            found: key_pairs.len(),
        });
    }
    let key_pair = key_pairs.pop().unwrap();
    Ok(Some(T::from_parsed(key_pair, path)?))
}
<commit_msg>Fix blocking RwLock calls for rotating keys<commit_after>use crate::agents::*;
use crate::crypto::SigningAlgorithm;
use crate::utils::{
    agent::*,
    keys::{KeyPairExt, NamedKeyPair, SignError},
    pem::{self, ParsedKeyPair},
    BoxFuture,
};
use err_derive::Error;
use log::{info, warn};
use ring::{
    rand::SecureRandom,
    signature::{Ed25519KeyPair, RsaKeyPair},
};
use serde_json as json;
use std::ffi::OsString;
use std::io::{Error as IoError, ErrorKind as IoErrorKind};
use std::mem;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use tokio::fs::{self, File};
use tokio::io::BufReader;
use tokio::process::Command;
use tokio::sync::RwLock;
use tokio::task;

#[derive(Debug, Error)]
pub enum RotatingKeysError {
    #[error(display = "could not open '{}': {}", path, error)]
    Open {
        path: String,
        #[error(source, nofrom)]
        error: IoError,
    },
    #[error(display = "could not create directory '{}': {}", path, error)]
    Mkdir {
        path: String,
        #[error(source, nofrom)]
        error: IoError,
    },
    #[error(
        display = "could not determine last modification time of '{}': {}",
        path,
        error
    )]
    StatMtime {
        path: String,
        #[error(source, nofrom)]
        error: IoError,
    },
    #[error(display = "could not parse file '{}': {}", path, error)]
    Parse {
        path: String,
        #[error(source, nofrom)]
        error: pem::ParseError,
    },
    #[error(
        display = "invalid type of key in '{}', want {}, found {}",
        path,
        want,
        found
    )]
    InvalidKeyType {
        path: String,
        want: &'static str,
        found: &'static str,
    },
    #[error(display = "expected exactly one key in '{}', found {}", path, found)]
    ExpectedOneKey { path: String, found: usize },
}

/// Struct that contains configuration we pass around.
struct RotateConfig {
    keys_ttl: Duration,
    generate_rsa_command: Vec<String>,
    rng: Box<dyn SecureRandom + Send + Sync>,
}

/// KeyManager where we rotating 3 keys of each type.
pub struct RotatingKeys {
    ed25519_keys: Option<KeySetHandle<Ed25519KeyPair>>,
    rsa_keys: Option<KeySetHandle<RsaKeyPair>>,
}

impl RotatingKeys {
    pub async fn new(
        keysdir: impl AsRef<Path>,
        keys_ttl: Duration,
        signing_algs: &[SigningAlgorithm],
        generate_rsa_command: Vec<String>,
        rng: impl SecureRandom + Send + Sync + 'static,
    ) -> Result<Self, RotatingKeysError> {
        info!(
            "Using rotating keys with a {}s interval and algorithms: {}",
            keys_ttl.as_secs(),
            SigningAlgorithm::format_list(signing_algs)
        );
        let rng = Box::new(rng);
        let config = Arc::new(RotateConfig {
            keys_ttl,
            generate_rsa_command,
            rng,
        });
        let ed25519_keys = if signing_algs.contains(&SigningAlgorithm::EdDsa) {
            Some(KeySet::from_subdir(keysdir.as_ref(), "ed25519", &config).await?)
        } else {
            None
        };
        let rsa_keys = if signing_algs.contains(&SigningAlgorithm::Rs256) {
            Some(KeySet::from_subdir(keysdir.as_ref(), "rsa", &config).await?)
        } else {
            None
        };
        Ok(RotatingKeys {
            ed25519_keys,
            rsa_keys,
        })
    }
}

impl Agent for RotatingKeys {}

impl Handler<SignJws> for RotatingKeys {
    fn handle(&mut self, message: SignJws, reply: ReplySender<SignJws>) {
        match message.signing_alg {
            SigningAlgorithm::EdDsa => {
                let key_set = self.ed25519_keys.clone();
                reply.later(move || async move {
                    let key_set = key_set
                        .ok_or_else(|| SignError::UnsupportedAlgorithm(message.signing_alg))?;
                    let key_set = key_set.read().await;
                    key_set
                        .current
                        .sign_jws(&message.payload, &*key_set.config.rng)
                });
            }
            SigningAlgorithm::Rs256 => {
                let key_set = self.rsa_keys.clone();
                reply.later(move || async move {
                    let key_set = key_set
                        .ok_or_else(|| SignError::UnsupportedAlgorithm(message.signing_alg))?;
                    let key_set = key_set.read().await;
                    key_set
                        .current
                        .sign_jws(&message.payload, &*key_set.config.rng)
                });
            }
        }
    }
}

impl Handler<GetPublicJwks> for RotatingKeys {
    fn handle(&mut self, _message: GetPublicJwks, reply: ReplySender<GetPublicJwks>) {
        let ed25519_keys = self.ed25519_keys.clone();
        let rsa_keys = self.rsa_keys.clone();
        reply.later(move || async move {
            let mut jwks = vec![];
            if let Some(key_set) = ed25519_keys {
                jwks.append(&mut key_set.read().await.public_jwks());
            }
            if let Some(key_set) = rsa_keys {
                jwks.append(&mut key_set.read().await.public_jwks());
            }
            jwks
        })
    }
}

impl KeyManagerSender for Addr<RotatingKeys> {}

/// Thread-safe handle to a KeySet.
type KeySetHandle<T> = Arc<RwLock<KeySet<T>>>;

/// A rotating set of 3 keys of a single type.
struct KeySet<T: KeyPairExt> {
    config: Arc<RotateConfig>,
    mtime: SystemTime,
    current: NamedKeyPair<T>,
    current_path: PathBuf,
    next: NamedKeyPair<T>,
    next_path: PathBuf,
    previous: Option<NamedKeyPair<T>>,
    previous_path: PathBuf,
}

impl<T> KeySet<T>
where
    T: KeyPairExt + GeneratedKeyPair + Send + Sync + 'static,
{
    /// Read key pairs of type `T` from a subdir of `keysdir`.
    async fn from_subdir(
        keysdir: impl AsRef<Path>,
        subdir: &str,
        config: &Arc<RotateConfig>,
    ) -> Result<KeySetHandle<T>, RotatingKeysError> {
        let mut dir = keysdir.as_ref().to_path_buf();
        dir.push(subdir);

        fs::create_dir_all(&dir)
            .await
            .map_err(|error| RotatingKeysError::Mkdir {
                path: dir.to_string_lossy().into_owned(),
                error,
            })?;

        let mut current_path = dir.clone();
        current_path.push("current.pem");
        let current = match read_key_file::<T>(&current_path).await? {
            Some(key_pair) => key_pair,
            None => T::generate(config.clone(), current_path.clone()).await,
        };

        // The last modification time of `next.pem`, along with `keys_ttl`,
        // is used to determine when the next rotation should happen.
        let mut next_path = dir.clone();
        next_path.push("next.pem");
        let (next, mtime) = if let Some(key_pair) = read_key_file::<T>(&next_path).await? {
            let mtime = fs::metadata(&next_path)
                .await
                .and_then(|meta| meta.modified())
                .map_err(|error| RotatingKeysError::StatMtime {
                    path: next_path.to_string_lossy().into_owned(),
                    error,
                })?;
            (key_pair, mtime)
        } else {
            (
                T::generate(config.clone(), next_path.clone()).await,
                SystemTime::now(),
            )
        };

        let mut previous_path = dir;
        previous_path.push("previous.pem");
        let previous = read_key_file::<T>(&previous_path).await?;

        let key_set = Arc::new(RwLock::new(KeySet {
            config: config.clone(),
            mtime,
            current,
            current_path,
            next,
            next_path,
            previous,
            previous_path,
        }));

        let initial_delay = Self::check_expiry(key_set.clone()).await;
        task::spawn(Self::rotate_loop(key_set.clone(), initial_delay));

        Ok(key_set)
    }

    /// The async loop that performs key rotation when necessary.
    async fn rotate_loop(handle: KeySetHandle<T>, mut delay: Duration) {
        loop {
            tokio::time::delay_for(delay).await;
            delay = Self::check_expiry(handle.clone()).await;
        }
    }

    /// Check if we should rotate keys, and return how long to delay the next check.
    async fn check_expiry(handle: KeySetHandle<T>) -> Duration {
        let read_handle = handle.clone();
        let (keys_ttl, mtime) = async {
            let key_set = read_handle.read().await;
            (key_set.config.keys_ttl, key_set.mtime)
        }
        .await;
        let delay = if let Ok(age) = SystemTime::now().duration_since(mtime) {
            keys_ttl.checked_sub(age)
        } else {
            warn!("Key set mtime is from the future, treating as bad input.");
            None
        };
        if let Some(delay) = delay {
            delay
        } else {
            Self::rotate(handle).await;
            keys_ttl
        }
    }

    /// Rotate keys in memory and on disk.
    ///
    /// If this fails, we panic, because it may happen at an arbitrary moment at run-time.
    async fn rotate(handle: KeySetHandle<T>) {
        let (config, next_path) = async {
            let key_set = handle.read().await;
            fs::rename(&key_set.current_path, &key_set.previous_path)
                .await
                .expect("could not rename current key for rotation");
            fs::rename(&key_set.next_path, &key_set.current_path)
                .await
                .expect("could not rename next key for rotation");
            (key_set.config.clone(), key_set.next_path.clone())
        }
        .await;
        let mut tmp = T::generate(config, next_path).await;
        let mut key_set = handle.write().await;
        mem::swap(&mut key_set.next, &mut tmp);
        mem::swap(&mut key_set.current, &mut tmp);
        key_set.previous = Some(tmp);
    }
}

impl<T: KeyPairExt> KeySet<T> {
    /// Get a list of JWKs containing public keys.
    fn public_jwks(&self) -> Vec<json::Value> {
        let mut list = vec![self.current.public_jwk(), self.next.public_jwk()];
        if let Some(ref previous) = self.previous {
            list.push(previous.public_jwk());
        }
        list
    }
}

/// Trait for key pair types we can generate.
trait GeneratedKeyPair: KeyPairExt + Sized {
    /// Generate a new key pair.
    ///
    /// If this fails, we panic, because it may happen at an arbitrary moment at run-time.
    ///
    /// This should log a message at info-level on success.
    ///
    /// Note that this function may block.
    fn generate(config: Arc<RotateConfig>, out_file: PathBuf) -> BoxFuture<NamedKeyPair<Self>>;

    /// Convert a ParsedKeyPair, if it is of the correct type.
    fn from_parsed(
        parsed: ParsedKeyPair,
        path: &Path,
    ) -> Result<NamedKeyPair<Self>, RotatingKeysError>;
}

impl GeneratedKeyPair for Ed25519KeyPair {
    fn generate(config: Arc<RotateConfig>, out_file: PathBuf) -> BoxFuture<NamedKeyPair<Self>> {
        Box::pin(async move {
            let doc =
                Self::generate_pkcs8(&*config.rng).expect("could not generate Ed25519 key pair");
            let key_pair =
                Self::from_pkcs8(doc.as_ref()).expect("could not parse generated Ed25519 key pair");
            let file = File::create(&out_file)
                .await
                .expect("could not open generated key pair output file");
            pem::write_pkcs8(&doc, file)
                .await
                .expect("could not write generated key pair to output file");
            info!("Generated new Ed25519 key: {:?}", out_file);
            key_pair.into()
        })
    }

    fn from_parsed(
        parsed: ParsedKeyPair,
        path: &Path,
    ) -> Result<NamedKeyPair<Self>, RotatingKeysError> {
        match parsed {
            ParsedKeyPair::Ed25519(inner) => Ok(inner.into()),
            other => Err(RotatingKeysError::InvalidKeyType {
                path: path.to_string_lossy().into_owned(),
                want: "Ed25519",
                found: other.kind(),
            }),
        }
    }
}

impl GeneratedKeyPair for RsaKeyPair {
    fn generate(config: Arc<RotateConfig>, out_file: PathBuf) -> BoxFuture<NamedKeyPair<Self>> {
        Box::pin(async move {
            let mut args: Vec<OsString> = config
                .generate_rsa_command
                .iter()
                .map(|arg| arg.into())
                .collect();
            let program = args.remove(0);
            let file = if let Some(part) = args.iter_mut().find(|part| part.as_os_str() == "{}") {
                *part = out_file.clone().into();
                None
            } else {
                let file = File::create(&out_file)
                    .await
                    .expect("could not open generated key pair output file");
                Some(file)
            };
            let mut command = Command::new(program);
            command.args(args).stdin(Stdio::null());
            if let Some(file) = file {
                command.stdout(file.into_std().await);
            }
            let status = command
                .status()
                .await
                .expect("Failed to run command to generate RSA key");
            if !status.success() {
                panic!("Command to generate RSA key failed with status {}", status);
            }
            let key_pair = read_key_file(&out_file)
                .await
                .expect("could not read generated RSA key file")
                .expect("generated RSA key file not found");
            info!("Generated new RSA key: {:?}", out_file);
            key_pair
        })
    }

    fn from_parsed(
        parsed: ParsedKeyPair,
        path: &Path,
    ) -> Result<NamedKeyPair<Self>, RotatingKeysError> {
        match parsed {
            ParsedKeyPair::Rsa(inner) => Ok(inner.into()),
            other => Err(RotatingKeysError::InvalidKeyType {
                path: path.to_string_lossy().into_owned(),
                want: "RSA",
                found: other.kind(),
            }),
        }
    }
}

/// Read a single key pair of type `T` from `path`.
///
/// Returns `None` when the file does not exist.
async fn read_key_file<T: KeyPairExt + GeneratedKeyPair>(
    path: &Path,
) -> Result<Option<NamedKeyPair<T>>, RotatingKeysError> {
    let file = match File::open(path).await {
        Ok(file) => file,
        Err(error) => {
            if error.kind() == IoErrorKind::NotFound {
                return Ok(None);
            } else {
                return Err(RotatingKeysError::Open {
                    path: path.to_string_lossy().into_owned(),
                    error,
                });
            }
        }
    };
    let mut key_pairs = pem::parse_key_pairs(BufReader::new(file))
        .await
        .map_err(|error| RotatingKeysError::Parse {
            path: path.to_string_lossy().into_owned(),
            error,
        })?;
    if key_pairs.len() != 1 {
        return Err(RotatingKeysError::ExpectedOneKey {
            path: path.to_string_lossy().into_owned(),
            found: key_pairs.len(),
        });
    }
    let key_pair = key_pairs.pop().unwrap();
    Ok(Some(T::from_parsed(key_pair, path)?))
}
<|endoftext|>


<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

//! [TVM](https://github.com/apache/incubator-tvm) is a compiler stack for deep learning systems.
//!
//! This crate provides an idiomatic Rust API for TVM runtime frontend.
//!
//! One particular use case is that given optimized deep learning model artifacts,
//! (compiled with TVM) which include a shared library
//! `lib.so`, `graph.json` and a byte-array `param.params`, one can load them
//! in Rust idiomatically to create a TVM Graph Runtime and
//! run the model for some inputs and get the
//! desired predictions *all in Rust*.
//!
//! Checkout the `examples` repository for more details.

pub use crate::{errors::*, function::Function, module::Module, ndarray::NDArray};

pub use tvm_rt::{Context, DataType, DeviceType};

pub use tvm_rt::context;
pub use tvm_rt::errors;
pub use tvm_rt::function;
pub use tvm_rt::module;
pub use tvm_rt::ndarray;
pub use tvm_rt::value;
pub mod ir;
#[cfg(feature = "python")]
pub mod python;
pub mod runtime;
pub mod transform;

pub use runtime::version;

#[macro_export]
macro_rules! export {
    ($($fn_name:expr),*) => {
        pub fn tvm_export(ns: &str) -> Result<(), tvm::Error> {
            $(
                let name = String::fromwe(ns) + ::std::stringify!($fn_name);
                tvm::runtime::function::register_override($fn_name, name, true)?;
            )*
            Ok(())
        }
    }
}

#[macro_export]
macro_rules! export_mod {
    ($ns:expr, $($mod_name:expr),*) => {
        pub fn tvm_mod_export() -> Result<(), tvm::Error> {
            $(
                $mod_name::tvm_export($ns)?;
            )*
            Ok(())
        }
    }
}
<commit_msg>Syntax error String::fromwe() should be String::from() (#6846)<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

//! [TVM](https://github.com/apache/incubator-tvm) is a compiler stack for deep learning systems.
//!
//! This crate provides an idiomatic Rust API for TVM runtime frontend.
//!
//! One particular use case is that given optimized deep learning model artifacts,
//! (compiled with TVM) which include a shared library
//! `lib.so`, `graph.json` and a byte-array `param.params`, one can load them
//! in Rust idiomatically to create a TVM Graph Runtime and
//! run the model for some inputs and get the
//! desired predictions *all in Rust*.
//!
//! Checkout the `examples` repository for more details.

pub use crate::{errors::*, function::Function, module::Module, ndarray::NDArray};

pub use tvm_rt::{Context, DataType, DeviceType};

pub use tvm_rt::context;
pub use tvm_rt::errors;
pub use tvm_rt::function;
pub use tvm_rt::module;
pub use tvm_rt::ndarray;
pub use tvm_rt::value;
pub mod ir;
#[cfg(feature = "python")]
pub mod python;
pub mod runtime;
pub mod transform;

pub use runtime::version;

#[macro_export]
macro_rules! export {
    ($($fn_name:expr),*) => {
        pub fn tvm_export(ns: &str) -> Result<(), tvm::Error> {
            $(
                let name = String::from(ns) + ::std::stringify!($fn_name);
                tvm::runtime::function::register_override($fn_name, name, true)?;
            )*
            Ok(())
        }
    }
}

#[macro_export]
macro_rules! export_mod {
    ($ns:expr, $($mod_name:expr),*) => {
        pub fn tvm_mod_export() -> Result<(), tvm::Error> {
            $(
                $mod_name::tvm_export($ns)?;
            )*
            Ok(())
        }
    }
}
<|endoftext|>