tor-browser

parallel.rs (5542B)

---

use super::ProblemSolver;
use std::ops::{Deref, DerefMut};

use futures::ready;
use std::future::Future;
use std::pin::Pin;

pub trait AsyncTester {
   type Result: Future<Output = Vec<bool>>;

   fn test_async(&self, query: Vec<(usize, usize)>) -> Self::Result;
}

pub struct ParallelProblemSolver<T>
where
   T: AsyncTester,
{
   solver: ProblemSolver,
   current_test: Option<(T::Result, Vec<usize>)>,
}

impl<T: AsyncTester> Deref for ParallelProblemSolver<T> {
   type Target = ProblemSolver;

   fn deref(&self) -> &Self::Target {
       &self.solver
   }
}

impl<T: AsyncTester> DerefMut for ParallelProblemSolver<T> {
   fn deref_mut(&mut self) -> &mut Self::Target {
       &mut self.solver
   }
}

impl<T: AsyncTester> ParallelProblemSolver<T> {
   pub fn new(width: usize, depth: usize) -> Self {
       Self {
           solver: ProblemSolver::new(width, depth),
           current_test: None,
       }
   }
}

type TestQuery = (Vec<(usize, usize)>, Vec<usize>);

impl<T: AsyncTester> ParallelProblemSolver<T> {
   pub fn try_generate_complete_candidate(&mut self) -> bool {
       while !self.is_complete() {
           while self.is_current_cell_missing() {
               if !self.try_advance_source() {
                   return false;
               }
           }
           if !self.try_advance_resource() {
               return false;
           }
       }
       true
   }

   fn try_generate_test_query(&mut self) -> Result<TestQuery, usize> {
       let mut test_cells = vec![];
       let query = self
           .solution
           .iter()
           .enumerate()
           .filter_map(|(res_idx, source_idx)| {
               let cell = self.cache[res_idx][*source_idx];
               match cell {
                   None => {
                       test_cells.push(res_idx);
                       Some(Ok((res_idx, *source_idx)))
                   }
                   Some(false) => Some(Err(res_idx)),
                   Some(true) => None,
               }
           })
           .collect::<Result<_, _>>()?;
       Ok((query, test_cells))
   }

   fn apply_test_result(
       &mut self,
       resources: Vec<bool>,
       testing_cells: Vec<usize>,
   ) -> Result<(), usize> {
       let mut first_missing = None;
       for (result, res_idx) in resources.into_iter().zip(testing_cells) {
           let source_idx = self.solution[res_idx];
           self.cache[res_idx][source_idx] = Some(result);
           if !result && first_missing.is_none() {
               first_missing = Some(res_idx);
           }
       }
       if let Some(idx) = first_missing {
           Err(idx)
       } else {
           Ok(())
       }
   }

   pub fn try_poll_next(
       mut self: std::pin::Pin<&mut Self>,
       cx: &mut std::task::Context<'_>,
       tester: &T,
       prefetch: bool,
   ) -> std::task::Poll<Result<Option<Vec<usize>>, usize>>
   where
       <T as AsyncTester>::Result: Unpin,
   {
       if self.width == 0 || self.depth == 0 {
           return Ok(None).into();
       }

       'outer: loop {
           if let Some((test, testing_cells)) = &mut self.current_test {
               let pinned = Pin::new(test);
               let set = ready!(pinned.poll(cx));
               let testing_cells = testing_cells.clone();

               if let Err(res_idx) = self.apply_test_result(set, testing_cells) {
                   self.idx = res_idx;
                   self.prune();
                   if !self.bail() {
                       if let Some(res_idx) = self.has_missing_cell() {
                           return Err(res_idx).into();
                       } else {
                           return Ok(None).into();
                       }
                   }
                   self.current_test = None;
                   continue 'outer;
               } else {
                   self.current_test = None;
                   if !prefetch {
                       self.dirty = true;
                   }
                   return Ok(Some(self.solution.clone())).into();
               }
           } else {
               if self.dirty {
                   if !self.bail() {
                       if let Some(res_idx) = self.has_missing_cell() {
                           return Err(res_idx).into();
                       } else {
                           return Ok(None).into();
                       }
                   }
                   self.dirty = false;
               }
               while self.try_generate_complete_candidate() {
                   match self.try_generate_test_query() {
                       Ok((query, testing_cells)) => {
                           self.current_test = Some((tester.test_async(query), testing_cells));
                           continue 'outer;
                       }
                       Err(res_idx) => {
                           self.idx = res_idx;
                           self.prune();
                           if !self.bail() {
                               if let Some(res_idx) = self.has_missing_cell() {
                                   return Err(res_idx).into();
                               } else {
                                   return Ok(None).into();
                               }
                           }
                       }
                   }
               }
               return Ok(None).into();
           }
       }
   }
}