<commit_before>use oppai_field::field::{euclidean, wave_diag, Field, Pos};
use oppai_field::player::Player;
use std::{
  collections::HashSet,
  ops::Index,
  sync::atomic::{AtomicBool, Ordering},
};

#[derive(Debug, Clone)]
struct Trajectory {
  points: Vec<Pos>,
  hash: u64,
}

impl Trajectory {
  pub fn new(points: Vec<Pos>, hash: u64) -> Trajectory {
    Trajectory { points, hash }
  }

  pub fn points(&self) -> &Vec<Pos> {
    &self.points
  }

  pub fn hash(&self) -> u64 {
    self.hash
  }

  pub fn len(&self) -> usize {
    self.points.len()
  }
}

pub struct TrajectoriesPruning {
  rebuild_trajectories: bool,
  cur_trajectories: Vec<Trajectory>,
  enemy_trajectories: Vec<Trajectory>,
  moves: Vec<Pos>,
}

impl TrajectoriesPruning {
  fn add_trajectory(field: &Field, trajectories: &mut Vec<Trajectory>, points: &[Pos], player: Player) {
    for &pos in points {
      if !field.cell(pos).is_bound() || field.number_near_groups(pos, player) < 2 {
        return;
      }
    }
    let zobrist = field.zobrist();
    let mut hash = 0u64;
    for &pos in points {
      hash ^= zobrist.get_hash(pos);
    }
    for trajectory in trajectories.iter() {
      if trajectory.hash() == hash {
        return;
      }
    }
    let trajectory = Trajectory::new(points.to_vec(), hash);
    trajectories.push(trajectory);
  }

  fn next_moves(
    field: &Field,
    start_pos: Pos,
    player: Player,
    empty_board: &mut Vec<u32>,
    marks: &mut Vec<Pos>,
  ) -> Vec<Pos> {
    let mut moves = Vec::new();
    wave_diag(field.width(), start_pos, |pos| {
      if empty_board[pos] != 0 {
        return false;
      }
      empty_board[pos] = 1;
      let cell = field.cell(pos);
      if cell.is_players_point(player) {
        marks.push(pos);
        true
      } else {
        if cell.is_putting_allowed() && !cell.is_players_empty_base(player) {
          moves.push(pos);
        } else {
          marks.push(pos);
        }
        false
      }
    });
    for &pos in &moves {
      empty_board[pos] = 0;
    }
    moves
  }

  fn build_trajectories_rec(
    field: &mut Field,
    trajectories: &mut Vec<Trajectory>,
    player: Player,
    cur_depth: u32,
    depth: u32,
    empty_board: &mut Vec<u32>,
    last_pos: Pos,
    moves: Vec<Pos>,
    should_stop: &AtomicBool,
  ) {
    for pos in moves {
      if should_stop.load(Ordering::Relaxed) {
        break;
      }
      if field.number_near_points(pos, player) >= 3 {
        continue;
      }
      let cell = field.cell(pos);
      if cell.is_players_empty_base(player.next()) {
        field.put_point(pos, player);
        if field.get_delta_score(player) > 0 {
          TrajectoriesPruning::add_trajectory(
            field,
            trajectories,
            field
              .points_seq()
              .index(field.moves_count() - cur_depth as usize..field.moves_count()),
            player,
          );
        }
        field.undo();
      } else {
        field.put_point(pos, player);
        if field.get_delta_score(player) > 0 {
          TrajectoriesPruning::add_trajectory(
            field,
            trajectories,
            field
              .points_seq()
              .index(field.moves_count() - cur_depth as usize..field.moves_count()),
            player,
          );
        } else if depth > 0 {
          let mut marks = Vec::new();
          let mut next_moves = TrajectoriesPruning::next_moves(field, pos, player, empty_board, &mut marks);
          if last_pos != 0 {
            next_moves.retain(|&next_pos| euclidean(field.width(), last_pos, next_pos) > 2);
          }
          TrajectoriesPruning::build_trajectories_rec(
            field,
            trajectories,
            player,
            cur_depth + 1,
            depth - 1,
            empty_board,
            pos,
            next_moves,
            should_stop,
          );
          for mark_pos in marks {
            empty_board[mark_pos] = 0;
          }
        }
        field.undo();
      }
    }
  }

  fn build_trajectories(
    field: &mut Field,
    trajectories: &mut Vec<Trajectory>,
    player: Player,
    depth: u32,
    empty_board: &mut Vec<u32>,
    should_stop: &AtomicBool,
  ) {
    if depth == 0 {
      return;
    }

    let mut marks = Vec::new();
    for pos in field.points_seq().clone() {
      if field.cell(pos).get_player() != player {
        continue;
      }

      if should_stop.load(Ordering::Relaxed) {
        break;
      }

      let moves = TrajectoriesPruning::next_moves(field, pos, player, empty_board, &mut marks);

      TrajectoriesPruning::build_trajectories_rec(
        field,
        trajectories,
        player,
        1,
        depth - 1,
        empty_board,
        0,
        moves,
        should_stop,
      );
    }

    for pos in marks {
      empty_board[pos] = 0;
    }
  }

  fn project(trajectories: &[Trajectory], empty_board: &mut Vec<u32>) {
    for &pos in trajectories.iter().flat_map(|trajectory| trajectory.points().iter()) {
      empty_board[pos] += 1;
    }
  }

  fn project_length(trajectories: &[Trajectory], empty_board: &mut Vec<u32>) {
    for trajectory in trajectories {
      let len = trajectory.len() as u32;
      for &pos in trajectory.points() {
        if empty_board[pos] == 0 || empty_board[pos] > len {
          empty_board[pos] = len;
        }
      }
    }
  }

  fn deproject(trajectories: &[Trajectory], empty_board: &mut Vec<u32>) {
    for &pos in trajectories.iter().flat_map(|trajectory| trajectory.points().iter()) {
      empty_board[pos] = 0;
    }
  }

  fn exclude_unnecessary_trajectories(trajectories: &mut Vec<Trajectory>, empty_board: &mut Vec<u32>) -> bool {
    let mut need_exclude = false;
    trajectories.retain(|trajectory| {
      let single_count = trajectory.points().iter().filter(|&&pos| empty_board[pos] == 1).count();
      if single_count > 1 {
        for &pos in trajectory.points() {
          empty_board[pos] -= 1;
        }
        need_exclude = true;
        false
      } else {
        true
      }
    });
    need_exclude
  }

  fn calculate_moves(
    trajectories1: &mut Vec<Trajectory>,
    trajectories2: &mut Vec<Trajectory>,
    empty_board: &mut Vec<u32>,
  ) -> Vec<Pos> {
    TrajectoriesPruning::project(trajectories1, empty_board);
    TrajectoriesPruning::project(trajectories2, empty_board);
    while TrajectoriesPruning::exclude_unnecessary_trajectories(trajectories1, empty_board)
      || TrajectoriesPruning::exclude_unnecessary_trajectories(trajectories2, empty_board)
    {}
    let mut result_set = HashSet::new();
    for &pos in trajectories1
      .iter()
      .chain(trajectories2.iter())
      .flat_map(|trajectory| trajectory.points().iter())
    {
      result_set.insert(pos);
    }
    let mut result = result_set.into_iter().collect::<Vec<Pos>>();
    result.sort_unstable_by(|&pos1, &pos2| empty_board[pos2].cmp(&empty_board[pos1]));
    TrajectoriesPruning::deproject(trajectories1, empty_board);
    TrajectoriesPruning::deproject(trajectories2, empty_board);
    TrajectoriesPruning::project_length(trajectories1, empty_board);
    TrajectoriesPruning::project_length(trajectories2, empty_board);
    result.sort_by(|&pos1, &pos2| empty_board[pos1].cmp(&empty_board[pos2]));
    TrajectoriesPruning::deproject(trajectories1, empty_board);
    TrajectoriesPruning::deproject(trajectories2, empty_board);
    result
  }

  #[inline]
  pub fn empty(rebuild_trajectories: bool) -> TrajectoriesPruning {
    TrajectoriesPruning {
      rebuild_trajectories,
      cur_trajectories: Vec::with_capacity(0),
      enemy_trajectories: Vec::with_capacity(0),
      moves: Vec::with_capacity(0),
    }
  }

  pub fn new(
    rebuild_trajectories: bool,
    field: &mut Field,
    player: Player,
    depth: u32,
    empty_board: &mut Vec<u32>,
    should_stop: &AtomicBool,
  ) -> TrajectoriesPruning {
    if depth == 0 {
      return TrajectoriesPruning::empty(rebuild_trajectories);
    }
    let mut cur_trajectories = Vec::new();
    let mut enemy_trajectories = Vec::new();
    TrajectoriesPruning::build_trajectories(
      field,
      &mut cur_trajectories,
      player,
      (depth + 1) / 2,
      empty_board,
      should_stop,
    );
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(rebuild_trajectories);
    }
    TrajectoriesPruning::build_trajectories(
      field,
      &mut enemy_trajectories,
      player.next(),
      depth / 2,
      empty_board,
      should_stop,
    );
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(rebuild_trajectories);
    }
    let moves = TrajectoriesPruning::calculate_moves(&mut cur_trajectories, &mut enemy_trajectories, empty_board);
    TrajectoriesPruning {
      rebuild_trajectories,
      cur_trajectories,
      enemy_trajectories,
      moves,
    }
  }

  fn last_pos_trajectory(field: &Field, player: Player, depth: u32, last_pos: Pos) -> Option<Trajectory> {
    let mut points = Vec::with_capacity(4);
    let mut hash = 0;
    for &pos in &field.directions(last_pos) {
      if field.cell(pos).is_putting_allowed() {
        let mut neighbors_count = 0;
        for &neighbor in &field.directions(pos) {
          if field.cell(neighbor).is_players_point(player) {
            neighbors_count += 1;
          }
        }
        if neighbors_count < 3 {
          points.push(pos);
          hash ^= field.zobrist().get_hash(pos);
        }
      } else if !field.cell(pos).is_players_point(player) {
        return None;
      }
    }
    if points.len() as u32 <= (depth + 1) / 2 {
      Some(Trajectory::new(points, hash))
    } else {
      None
    }
  }

  pub fn next(
    &self,
    field: &mut Field,
    player: Player,
    depth: u32,
    empty_board: &mut Vec<u32>,
    last_pos: Pos,
    should_stop: &AtomicBool,
  ) -> TrajectoriesPruning {
    if depth == 0 {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let mut cur_trajectories = Vec::new();
    let mut enemy_trajectories = Vec::new();
    if self.rebuild_trajectories {
      TrajectoriesPruning::build_trajectories(
        field,
        &mut cur_trajectories,
        player,
        (depth + 1) / 2,
        empty_board,
        should_stop,
      );
    } else {
      for trajectory in &self.enemy_trajectories {
        if trajectory
          .points()
          .iter()
          .all(|&pos| field.cell(pos).is_putting_allowed())
        {
          cur_trajectories.push(trajectory.clone());
        }
      }
      if let Some(new_cur_trajectory) = TrajectoriesPruning::last_pos_trajectory(field, player, depth, last_pos) {
        cur_trajectories.push(new_cur_trajectory);
      }
    }
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let enemy_depth = depth / 2;
    if enemy_depth > 0 {
      for trajectory in &self.cur_trajectories {
        let len = trajectory.len() as u32;
        let contains_pos = trajectory.points().contains(&last_pos);
        if (len <= enemy_depth || len == enemy_depth + 1 && contains_pos)
          && trajectory
            .points()
            .iter()
            .all(|&pos| field.cell(pos).is_putting_allowed() || pos == last_pos)
        {
          let new_trajectory = if contains_pos {
            if len == 1 {
              continue;
            }
            Trajectory::new(
              trajectory
                .points
                .iter()
                .cloned()
                .filter(|&pos| pos != last_pos)
                .collect(),
              trajectory.hash() ^ field.zobrist().get_hash(last_pos),
            )
          } else {
            Trajectory::new(trajectory.points.clone(), trajectory.hash())
          };
          enemy_trajectories.push(new_trajectory);
        }
      }
    }
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let moves = TrajectoriesPruning::calculate_moves(&mut cur_trajectories, &mut enemy_trajectories, empty_board);
    TrajectoriesPruning {
      rebuild_trajectories: self.rebuild_trajectories,
      cur_trajectories,
      enemy_trajectories,
      moves,
    }
  }

  pub fn dec_and_swap(&self, depth: u32, empty_board: &mut Vec<u32>, should_stop: &AtomicBool) -> TrajectoriesPruning {
    if depth == 0 {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let mut cur_trajectories = Vec::new();
    let mut enemy_trajectories = Vec::new();
    for trajectory in &self.enemy_trajectories {
      cur_trajectories.push(trajectory.clone());
    }
    let enemy_depth = depth / 2;
    if enemy_depth > 0 {
      for trajectory in self
        .cur_trajectories
        .iter()
        .filter(|trajectory| trajectory.len() as u32 <= enemy_depth)
      {
        enemy_trajectories.push(trajectory.clone());
      }
    }
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let moves = TrajectoriesPruning::calculate_moves(&mut cur_trajectories, &mut enemy_trajectories, empty_board);
    TrajectoriesPruning {
      rebuild_trajectories: self.rebuild_trajectories,
      cur_trajectories,
      enemy_trajectories,
      moves,
    }
  }

  pub fn inc(
    &self,
    field: &mut Field,
    player: Player,
    depth: u32,
    empty_board: &mut Vec<u32>,
    should_stop: &AtomicBool,
  ) -> TrajectoriesPruning {
    let mut cur_trajectories = Vec::new();
    let mut enemy_trajectories = Vec::new();
    if depth % 2 == 0 {
      TrajectoriesPruning::build_trajectories(
        field,
        &mut enemy_trajectories,
        player.next(),
        depth / 2,
        empty_board,
        should_stop,
      );
      if should_stop.load(Ordering::Relaxed) {
        return TrajectoriesPruning::empty(self.rebuild_trajectories);
      }
      for trajectory in &self.cur_trajectories {
        cur_trajectories.push(trajectory.clone());
      }
    } else {
      TrajectoriesPruning::build_trajectories(
        field,
        &mut cur_trajectories,
        player,
        (depth + 1) / 2,
        empty_board,
        should_stop,
      );
      if should_stop.load(Ordering::Relaxed) {
        return TrajectoriesPruning::empty(self.rebuild_trajectories);
      }
      for trajectory in &self.enemy_trajectories {
        enemy_trajectories.push(trajectory.clone());
      }
    }
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let moves = TrajectoriesPruning::calculate_moves(&mut cur_trajectories, &mut enemy_trajectories, empty_board);
    TrajectoriesPruning {
      rebuild_trajectories: self.rebuild_trajectories,
      cur_trajectories,
      enemy_trajectories,
      moves,
    }
  }

  pub fn moves(&self) -> &Vec<Pos> {
    &self.moves
  }

  pub fn moves_mut(&mut self) -> &mut Vec<Pos> {
    &mut self.moves
  }
}
<commit_msg>Use iterators more.<commit_after>use oppai_field::field::{euclidean, wave_diag, Field, Pos};
use oppai_field::player::Player;
use std::{
  collections::HashSet,
  ops::Index,
  sync::atomic::{AtomicBool, Ordering},
};

#[derive(Debug, Clone)]
struct Trajectory {
  points: Vec<Pos>,
  hash: u64,
}

impl Trajectory {
  pub fn new(points: Vec<Pos>, hash: u64) -> Trajectory {
    Trajectory { points, hash }
  }

  pub fn points(&self) -> &Vec<Pos> {
    &self.points
  }

  pub fn hash(&self) -> u64 {
    self.hash
  }

  pub fn len(&self) -> usize {
    self.points.len()
  }
}

pub struct TrajectoriesPruning {
  rebuild_trajectories: bool,
  cur_trajectories: Vec<Trajectory>,
  enemy_trajectories: Vec<Trajectory>,
  moves: Vec<Pos>,
}

impl TrajectoriesPruning {
  fn add_trajectory(field: &Field, trajectories: &mut Vec<Trajectory>, points: &[Pos], player: Player) {
    for &pos in points {
      if !field.cell(pos).is_bound() || field.number_near_groups(pos, player) < 2 {
        return;
      }
    }
    let zobrist = field.zobrist();
    let mut hash = 0u64;
    for &pos in points {
      hash ^= zobrist.get_hash(pos);
    }
    for trajectory in trajectories.iter() {
      if trajectory.hash() == hash {
        return;
      }
    }
    let trajectory = Trajectory::new(points.to_vec(), hash);
    trajectories.push(trajectory);
  }

  fn next_moves(
    field: &Field,
    start_pos: Pos,
    player: Player,
    empty_board: &mut Vec<u32>,
    marks: &mut Vec<Pos>,
  ) -> Vec<Pos> {
    let mut moves = Vec::new();
    wave_diag(field.width(), start_pos, |pos| {
      if empty_board[pos] != 0 {
        return false;
      }
      empty_board[pos] = 1;
      let cell = field.cell(pos);
      if cell.is_players_point(player) {
        marks.push(pos);
        true
      } else {
        if cell.is_putting_allowed() && !cell.is_players_empty_base(player) {
          moves.push(pos);
        } else {
          marks.push(pos);
        }
        false
      }
    });
    for &pos in &moves {
      empty_board[pos] = 0;
    }
    moves
  }

  fn build_trajectories_rec(
    field: &mut Field,
    trajectories: &mut Vec<Trajectory>,
    player: Player,
    cur_depth: u32,
    depth: u32,
    empty_board: &mut Vec<u32>,
    last_pos: Pos,
    moves: Vec<Pos>,
    should_stop: &AtomicBool,
  ) {
    for pos in moves {
      if should_stop.load(Ordering::Relaxed) {
        break;
      }
      if field.number_near_points(pos, player) >= 3 {
        continue;
      }
      let cell = field.cell(pos);
      if cell.is_players_empty_base(player.next()) {
        field.put_point(pos, player);
        if field.get_delta_score(player) > 0 {
          TrajectoriesPruning::add_trajectory(
            field,
            trajectories,
            field
              .points_seq()
              .index(field.moves_count() - cur_depth as usize..field.moves_count()),
            player,
          );
        }
        field.undo();
      } else {
        field.put_point(pos, player);
        if field.get_delta_score(player) > 0 {
          TrajectoriesPruning::add_trajectory(
            field,
            trajectories,
            field
              .points_seq()
              .index(field.moves_count() - cur_depth as usize..field.moves_count()),
            player,
          );
        } else if depth > 0 {
          let mut marks = Vec::new();
          let mut next_moves = TrajectoriesPruning::next_moves(field, pos, player, empty_board, &mut marks);
          if last_pos != 0 {
            next_moves.retain(|&next_pos| euclidean(field.width(), last_pos, next_pos) > 2);
          }
          TrajectoriesPruning::build_trajectories_rec(
            field,
            trajectories,
            player,
            cur_depth + 1,
            depth - 1,
            empty_board,
            pos,
            next_moves,
            should_stop,
          );
          for mark_pos in marks {
            empty_board[mark_pos] = 0;
          }
        }
        field.undo();
      }
    }
  }

  fn build_trajectories(
    field: &mut Field,
    player: Player,
    depth: u32,
    empty_board: &mut Vec<u32>,
    should_stop: &AtomicBool,
  ) -> Vec<Trajectory> {
    let mut trajectories = Vec::new();

    if depth == 0 {
      return trajectories;
    }

    let mut marks = Vec::new();
    for pos in field.points_seq().clone() {
      if field.cell(pos).get_player() != player {
        continue;
      }

      if should_stop.load(Ordering::Relaxed) {
        break;
      }

      let moves = TrajectoriesPruning::next_moves(field, pos, player, empty_board, &mut marks);

      TrajectoriesPruning::build_trajectories_rec(
        field,
        &mut trajectories,
        player,
        1,
        depth - 1,
        empty_board,
        0,
        moves,
        should_stop,
      );
    }

    for pos in marks {
      empty_board[pos] = 0;
    }

    trajectories
  }

  fn project(trajectories: &[Trajectory], empty_board: &mut Vec<u32>) {
    for &pos in trajectories.iter().flat_map(|trajectory| trajectory.points().iter()) {
      empty_board[pos] += 1;
    }
  }

  fn project_length(trajectories: &[Trajectory], empty_board: &mut Vec<u32>) {
    for trajectory in trajectories {
      let len = trajectory.len() as u32;
      for &pos in trajectory.points() {
        if empty_board[pos] == 0 || empty_board[pos] > len {
          empty_board[pos] = len;
        }
      }
    }
  }

  fn deproject(trajectories: &[Trajectory], empty_board: &mut Vec<u32>) {
    for &pos in trajectories.iter().flat_map(|trajectory| trajectory.points().iter()) {
      empty_board[pos] = 0;
    }
  }

  fn exclude_unnecessary_trajectories(trajectories: &mut Vec<Trajectory>, empty_board: &mut Vec<u32>) -> bool {
    let mut need_exclude = false;
    trajectories.retain(|trajectory| {
      let single_count = trajectory.points().iter().filter(|&&pos| empty_board[pos] == 1).count();
      if single_count > 1 {
        for &pos in trajectory.points() {
          empty_board[pos] -= 1;
        }
        need_exclude = true;
        false
      } else {
        true
      }
    });
    need_exclude
  }

  fn calculate_moves(
    trajectories1: &mut Vec<Trajectory>,
    trajectories2: &mut Vec<Trajectory>,
    empty_board: &mut Vec<u32>,
  ) -> Vec<Pos> {
    TrajectoriesPruning::project(trajectories1, empty_board);
    TrajectoriesPruning::project(trajectories2, empty_board);
    while TrajectoriesPruning::exclude_unnecessary_trajectories(trajectories1, empty_board)
      || TrajectoriesPruning::exclude_unnecessary_trajectories(trajectories2, empty_board)
    {}
    let mut result_set = HashSet::new();
    for &pos in trajectories1
      .iter()
      .chain(trajectories2.iter())
      .flat_map(|trajectory| trajectory.points().iter())
    {
      result_set.insert(pos);
    }
    let mut result = result_set.into_iter().collect::<Vec<Pos>>();
    result.sort_unstable_by(|&pos1, &pos2| empty_board[pos2].cmp(&empty_board[pos1]));
    TrajectoriesPruning::deproject(trajectories1, empty_board);
    TrajectoriesPruning::deproject(trajectories2, empty_board);
    TrajectoriesPruning::project_length(trajectories1, empty_board);
    TrajectoriesPruning::project_length(trajectories2, empty_board);
    result.sort_by(|&pos1, &pos2| empty_board[pos1].cmp(&empty_board[pos2]));
    TrajectoriesPruning::deproject(trajectories1, empty_board);
    TrajectoriesPruning::deproject(trajectories2, empty_board);
    result
  }

  #[inline]
  pub fn empty(rebuild_trajectories: bool) -> TrajectoriesPruning {
    TrajectoriesPruning {
      rebuild_trajectories,
      cur_trajectories: Vec::with_capacity(0),
      enemy_trajectories: Vec::with_capacity(0),
      moves: Vec::with_capacity(0),
    }
  }

  pub fn new(
    rebuild_trajectories: bool,
    field: &mut Field,
    player: Player,
    depth: u32,
    empty_board: &mut Vec<u32>,
    should_stop: &AtomicBool,
  ) -> TrajectoriesPruning {
    if depth == 0 {
      return TrajectoriesPruning::empty(rebuild_trajectories);
    }
    let mut cur_trajectories =
      TrajectoriesPruning::build_trajectories(field, player, (depth + 1) / 2, empty_board, should_stop);
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(rebuild_trajectories);
    }
    let mut enemy_trajectories =
      TrajectoriesPruning::build_trajectories(field, player.next(), depth / 2, empty_board, should_stop);
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(rebuild_trajectories);
    }
    let moves = TrajectoriesPruning::calculate_moves(&mut cur_trajectories, &mut enemy_trajectories, empty_board);
    TrajectoriesPruning {
      rebuild_trajectories,
      cur_trajectories,
      enemy_trajectories,
      moves,
    }
  }

  fn last_pos_trajectory(field: &Field, player: Player, depth: u32, last_pos: Pos) -> Option<Trajectory> {
    let mut points = Vec::with_capacity(4);
    let mut hash = 0;
    for &pos in &field.directions(last_pos) {
      if field.cell(pos).is_putting_allowed() {
        let mut neighbors_count = 0;
        for &neighbor in &field.directions(pos) {
          if field.cell(neighbor).is_players_point(player) {
            neighbors_count += 1;
          }
        }
        if neighbors_count < 3 {
          points.push(pos);
          hash ^= field.zobrist().get_hash(pos);
        }
      } else if !field.cell(pos).is_players_point(player) {
        return None;
      }
    }
    if points.len() as u32 <= (depth + 1) / 2 {
      Some(Trajectory::new(points, hash))
    } else {
      None
    }
  }

  pub fn next(
    &self,
    field: &mut Field,
    player: Player,
    depth: u32,
    empty_board: &mut Vec<u32>,
    last_pos: Pos,
    should_stop: &AtomicBool,
  ) -> TrajectoriesPruning {
    if depth == 0 {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let mut cur_trajectories = if self.rebuild_trajectories {
      TrajectoriesPruning::build_trajectories(field, player, (depth + 1) / 2, empty_board, should_stop)
    } else {
      self
        .enemy_trajectories
        .iter()
        .filter(|trajectory| {
          trajectory
            .points()
            .iter()
            .all(|&pos| field.cell(pos).is_putting_allowed())
        })
        .cloned()
        .chain(TrajectoriesPruning::last_pos_trajectory(field, player, depth, last_pos).into_iter())
        .collect()
    };
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let enemy_depth = depth / 2;
    let mut enemy_trajectories = if enemy_depth > 0 {
      self
        .cur_trajectories
        .iter()
        .filter_map(|trajectory| {
          let len = trajectory.len() as u32;
          let contains_pos = trajectory.points().contains(&last_pos);
          if (len <= enemy_depth || len == enemy_depth + 1 && contains_pos)
            && trajectory
              .points()
              .iter()
              .all(|&pos| field.cell(pos).is_putting_allowed() || pos == last_pos)
          {
            let new_trajectory = if contains_pos {
              if len == 1 {
                return None;
              }
              Trajectory::new(
                trajectory
                  .points
                  .iter()
                  .cloned()
                  .filter(|&pos| pos != last_pos)
                  .collect(),
                trajectory.hash() ^ field.zobrist().get_hash(last_pos),
              )
            } else {
              trajectory.clone()
            };
            Some(new_trajectory)
          } else {
            None
          }
        })
        .collect()
    } else {
      Vec::new()
    };
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let moves = TrajectoriesPruning::calculate_moves(&mut cur_trajectories, &mut enemy_trajectories, empty_board);
    TrajectoriesPruning {
      rebuild_trajectories: self.rebuild_trajectories,
      cur_trajectories,
      enemy_trajectories,
      moves,
    }
  }

  pub fn dec_and_swap(&self, depth: u32, empty_board: &mut Vec<u32>, should_stop: &AtomicBool) -> TrajectoriesPruning {
    if depth == 0 {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let mut cur_trajectories = self.enemy_trajectories.clone();
    let enemy_depth = depth / 2;
    let mut enemy_trajectories = if enemy_depth > 0 {
      self
        .cur_trajectories
        .iter()
        .filter(|trajectory| trajectory.len() as u32 <= enemy_depth)
        .cloned()
        .collect()
    } else {
      Vec::new()
    };
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let moves = TrajectoriesPruning::calculate_moves(&mut cur_trajectories, &mut enemy_trajectories, empty_board);
    TrajectoriesPruning {
      rebuild_trajectories: self.rebuild_trajectories,
      cur_trajectories,
      enemy_trajectories,
      moves,
    }
  }

  pub fn inc(
    &self,
    field: &mut Field,
    player: Player,
    depth: u32,
    empty_board: &mut Vec<u32>,
    should_stop: &AtomicBool,
  ) -> TrajectoriesPruning {
    let (mut cur_trajectories, mut enemy_trajectories) = if depth % 2 == 0 {
      let enemy_trajectories =
        TrajectoriesPruning::build_trajectories(field, player.next(), depth / 2, empty_board, should_stop);
      if should_stop.load(Ordering::Relaxed) {
        return TrajectoriesPruning::empty(self.rebuild_trajectories);
      }
      (self.cur_trajectories.clone(), enemy_trajectories)
    } else {
      let cur_trajectories =
        TrajectoriesPruning::build_trajectories(field, player, (depth + 1) / 2, empty_board, should_stop);
      if should_stop.load(Ordering::Relaxed) {
        return TrajectoriesPruning::empty(self.rebuild_trajectories);
      }
      (cur_trajectories, self.enemy_trajectories.clone())
    };
    if should_stop.load(Ordering::Relaxed) {
      return TrajectoriesPruning::empty(self.rebuild_trajectories);
    }
    let moves = TrajectoriesPruning::calculate_moves(&mut cur_trajectories, &mut enemy_trajectories, empty_board);
    TrajectoriesPruning {
      rebuild_trajectories: self.rebuild_trajectories,
      cur_trajectories,
      enemy_trajectories,
      moves,
    }
  }

  pub fn moves(&self) -> &Vec<Pos> {
    &self.moves
  }

  pub fn moves_mut(&mut self) -> &mut Vec<Pos> {
    &mut self.moves
  }
}
<|endoftext|>


<commit_before>use std::sync::Arc;

use salsa::Database;
use ra_db::{FilesDatabase, CrateGraph};
use relative_path::RelativePath;
use test_utils::assert_eq_text;

use crate::{
    self as hir,
    db::HirDatabase,
    mock::MockDatabase,
};

fn item_map(fixture: &str) -> (Arc<hir::ItemMap>, hir::ModuleId) {
    let (db, pos) = MockDatabase::with_position(fixture);
    let source_root = db.file_source_root(pos.file_id);
    let module = hir::source_binder::module_from_position(&db, pos)
        .unwrap()
        .unwrap();
    let module_id = module.module_id;
    (db.item_map(source_root).unwrap(), module_id)
}

fn check_module_item_map(map: &hir::ItemMap, module_id: hir::ModuleId, expected: &str) {
    let mut lines = map.per_module[&module_id]
        .items
        .iter()
        .map(|(name, res)| format!("{}: {}", name, dump_resolution(res)))
        .collect::<Vec<_>>();
    lines.sort();
    let actual = lines.join("\n");
    let expected = expected
        .trim()
        .lines()
        .map(|it| it.trim())
        .collect::<Vec<_>>()
        .join("\n");
    assert_eq_text!(&actual, &expected);

    fn dump_resolution(resolution: &hir::Resolution) -> &'static str {
        match (
            resolution.def_id.types.is_some(),
            resolution.def_id.values.is_some(),
        ) {
            (true, true) => "t v",
            (true, false) => "t",
            (false, true) => "v",
            (false, false) => "_",
        }
    }
}

#[test]
fn item_map_smoke_test() {
    let (item_map, module_id) = item_map(
        "
        //- /lib.rs
        mod foo;

        use crate::foo::bar::Baz;
        <|>

        //- /foo/mod.rs
        pub mod bar;

        //- /foo/bar.rs
        pub struct Baz;
    ",
    );
    check_module_item_map(
        &item_map,
        module_id,
        "
            Baz: t v
            foo: t
        ",
    );
}

#[test]
fn item_map_contains_items_from_expansions() {
    let (item_map, module_id) = item_map(
        "
        //- /lib.rs
        mod foo;

        use crate::foo::bar::Baz;
        <|>

        //- /foo/mod.rs
        pub mod bar;

        //- /foo/bar.rs
        salsa::query_group! {
            trait Baz {}
        }
    ",
    );
    check_module_item_map(
        &item_map,
        module_id,
        "
            Baz: t
            foo: t
        ",
    );
}

#[test]
fn item_map_using_self() {
    let (item_map, module_id) = item_map(
        "
            //- /lib.rs
            mod foo;
            use crate::foo::bar::Baz::{self};
            <|>
            //- /foo/mod.rs
            pub mod bar;
            //- /foo/bar.rs
            pub struct Baz;
        ",
    );
    check_module_item_map(
        &item_map,
        module_id,
        "
            Baz: t v
            foo: t
        ",
    );
}

#[test]
fn item_map_across_crates() {
    let (mut db, sr) = MockDatabase::with_files(
        "
        //- /main.rs
        use test_crate::Baz;

        //- /lib.rs
        pub struct Baz;
    ",
    );
    let main_id = sr.files[RelativePath::new("/main.rs")];
    let lib_id = sr.files[RelativePath::new("/lib.rs")];

    let mut crate_graph = CrateGraph::default();
    let main_crate = crate_graph.add_crate_root(main_id);
    let lib_crate = crate_graph.add_crate_root(lib_id);
    crate_graph.add_dep(main_crate, "test_crate".into(), lib_crate);

    db.set_crate_graph(crate_graph);

    let source_root = db.file_source_root(main_id);
    let module = hir::source_binder::module_from_file_id(&db, main_id)
        .unwrap()
        .unwrap();
    let module_id = module.module_id;
    let item_map = db.item_map(source_root).unwrap();

    check_module_item_map(
        &item_map,
        module_id,
        "
            Baz: t v
            test_crate: t
        ",
    );
}

#[test]
fn typing_inside_a_function_should_not_invalidate_item_map() {
    let (mut db, pos) = MockDatabase::with_position(
        "
        //- /lib.rs
        mod foo;

        use crate::foo::bar::Baz;

        //- /foo/mod.rs
        pub mod bar;

        //- /foo/bar.rs
        <|>
        salsa::query_group! {
            trait Baz {
                fn foo() -> i32 { 1 + 1 }
            }
        }
        ",
    );
    let source_root = db.file_source_root(pos.file_id);
    {
        let events = db.log_executed(|| {
            db.item_map(source_root).unwrap();
        });
        assert!(format!("{:?}", events).contains("item_map"))
    }

    let new_text = "
        salsa::query_group! {
            trait Baz {
                fn foo() -> i32 { 92 }
            }
        }
    "
    .to_string();

    db.query_mut(ra_db::FileTextQuery)
        .set(pos.file_id, Arc::new(new_text));

    {
        let events = db.log_executed(|| {
            db.item_map(source_root).unwrap();
        });
        assert!(
            !format!("{:?}", events).contains("item_map"),
            "{:#?}",
            events
        )
    }
}

#[test]
fn typing_inside_a_function_inside_a_macro_should_not_invalidate_item_map() {
    let (mut db, pos) = MockDatabase::with_position(
        "
        //- /lib.rs
        mod foo;<|>

        use crate::foo::bar::Baz;

        fn foo() -> i32 {
            1 + 1
        }
        //- /foo/mod.rs
        pub mod bar;

        //- /foo/bar.rs
        pub struct Baz;
    ",
    );
    let source_root = db.file_source_root(pos.file_id);
    {
        let events = db.log_executed(|| {
            db.item_map(source_root).unwrap();
        });
        assert!(format!("{:?}", events).contains("item_map"))
    }

    let new_text = "
        mod foo;

        use crate::foo::bar::Baz;

        fn foo() -> i32 { 92 }
    "
    .to_string();

    db.query_mut(ra_db::FileTextQuery)
        .set(pos.file_id, Arc::new(new_text));

    {
        let events = db.log_executed(|| {
            db.item_map(source_root).unwrap();
        });
        assert!(
            !format!("{:?}", events).contains("item_map"),
            "{:#?}",
            events
        )
    }
}
<commit_msg>fix the test<commit_after>use std::sync::Arc;

use salsa::Database;
use ra_db::{FilesDatabase, CrateGraph};
use relative_path::RelativePath;
use test_utils::assert_eq_text;

use crate::{
    self as hir,
    db::HirDatabase,
    mock::MockDatabase,
};

fn item_map(fixture: &str) -> (Arc<hir::ItemMap>, hir::ModuleId) {
    let (db, pos) = MockDatabase::with_position(fixture);
    let source_root = db.file_source_root(pos.file_id);
    let module = hir::source_binder::module_from_position(&db, pos)
        .unwrap()
        .unwrap();
    let module_id = module.def_id.loc(&db).module_id;
    (db.item_map(source_root).unwrap(), module_id)
}

fn check_module_item_map(map: &hir::ItemMap, module_id: hir::ModuleId, expected: &str) {
    let mut lines = map.per_module[&module_id]
        .items
        .iter()
        .map(|(name, res)| format!("{}: {}", name, dump_resolution(res)))
        .collect::<Vec<_>>();
    lines.sort();
    let actual = lines.join("\n");
    let expected = expected
        .trim()
        .lines()
        .map(|it| it.trim())
        .collect::<Vec<_>>()
        .join("\n");
    assert_eq_text!(&actual, &expected);

    fn dump_resolution(resolution: &hir::Resolution) -> &'static str {
        match (
            resolution.def_id.types.is_some(),
            resolution.def_id.values.is_some(),
        ) {
            (true, true) => "t v",
            (true, false) => "t",
            (false, true) => "v",
            (false, false) => "_",
        }
    }
}

#[test]
fn item_map_smoke_test() {
    let (item_map, module_id) = item_map(
        "
        //- /lib.rs
        mod foo;

        use crate::foo::bar::Baz;
        <|>

        //- /foo/mod.rs
        pub mod bar;

        //- /foo/bar.rs
        pub struct Baz;
    ",
    );
    check_module_item_map(
        &item_map,
        module_id,
        "
            Baz: t v
            foo: t
        ",
    );
}

#[test]
fn item_map_contains_items_from_expansions() {
    let (item_map, module_id) = item_map(
        "
        //- /lib.rs
        mod foo;

        use crate::foo::bar::Baz;
        <|>

        //- /foo/mod.rs
        pub mod bar;

        //- /foo/bar.rs
        salsa::query_group! {
            trait Baz {}
        }
    ",
    );
    check_module_item_map(
        &item_map,
        module_id,
        "
            Baz: t
            foo: t
        ",
    );
}

#[test]
fn item_map_using_self() {
    let (item_map, module_id) = item_map(
        "
            //- /lib.rs
            mod foo;
            use crate::foo::bar::Baz::{self};
            <|>
            //- /foo/mod.rs
            pub mod bar;
            //- /foo/bar.rs
            pub struct Baz;
        ",
    );
    check_module_item_map(
        &item_map,
        module_id,
        "
            Baz: t v
            foo: t
        ",
    );
}

#[test]
fn item_map_across_crates() {
    let (mut db, sr) = MockDatabase::with_files(
        "
        //- /main.rs
        use test_crate::Baz;

        //- /lib.rs
        pub struct Baz;
    ",
    );
    let main_id = sr.files[RelativePath::new("/main.rs")];
    let lib_id = sr.files[RelativePath::new("/lib.rs")];

    let mut crate_graph = CrateGraph::default();
    let main_crate = crate_graph.add_crate_root(main_id);
    let lib_crate = crate_graph.add_crate_root(lib_id);
    crate_graph.add_dep(main_crate, "test_crate".into(), lib_crate);

    db.set_crate_graph(crate_graph);

    let source_root = db.file_source_root(main_id);
    let module = hir::source_binder::module_from_file_id(&db, main_id)
        .unwrap()
        .unwrap();
    let module_id = module.def_id.loc(&db).module_id;
    let item_map = db.item_map(source_root).unwrap();

    check_module_item_map(
        &item_map,
        module_id,
        "
            Baz: t v
            test_crate: t
        ",
    );
}

#[test]
fn typing_inside_a_function_should_not_invalidate_item_map() {
    let (mut db, pos) = MockDatabase::with_position(
        "
        //- /lib.rs
        mod foo;

        use crate::foo::bar::Baz;

        //- /foo/mod.rs
        pub mod bar;

        //- /foo/bar.rs
        <|>
        salsa::query_group! {
            trait Baz {
                fn foo() -> i32 { 1 + 1 }
            }
        }
        ",
    );
    let source_root = db.file_source_root(pos.file_id);
    {
        let events = db.log_executed(|| {
            db.item_map(source_root).unwrap();
        });
        assert!(format!("{:?}", events).contains("item_map"))
    }

    let new_text = "
        salsa::query_group! {
            trait Baz {
                fn foo() -> i32 { 92 }
            }
        }
    "
    .to_string();

    db.query_mut(ra_db::FileTextQuery)
        .set(pos.file_id, Arc::new(new_text));

    {
        let events = db.log_executed(|| {
            db.item_map(source_root).unwrap();
        });
        assert!(
            !format!("{:?}", events).contains("item_map"),
            "{:#?}",
            events
        )
    }
}

#[test]
fn typing_inside_a_function_inside_a_macro_should_not_invalidate_item_map() {
    let (mut db, pos) = MockDatabase::with_position(
        "
        //- /lib.rs
        mod foo;<|>

        use crate::foo::bar::Baz;

        fn foo() -> i32 {
            1 + 1
        }
        //- /foo/mod.rs
        pub mod bar;

        //- /foo/bar.rs
        pub struct Baz;
    ",
    );
    let source_root = db.file_source_root(pos.file_id);
    {
        let events = db.log_executed(|| {
            db.item_map(source_root).unwrap();
        });
        assert!(format!("{:?}", events).contains("item_map"))
    }

    let new_text = "
        mod foo;

        use crate::foo::bar::Baz;

        fn foo() -> i32 { 92 }
    "
    .to_string();

    db.query_mut(ra_db::FileTextQuery)
        .set(pos.file_id, Arc::new(new_text));

    {
        let events = db.log_executed(|| {
            db.item_map(source_root).unwrap();
        });
        assert!(
            !format!("{:?}", events).contains("item_map"),
            "{:#?}",
            events
        )
    }
}
<|endoftext|>


<commit_before>//! `LineIndex` maps flat `TextSize` offsets into `(Line, Column)`
//! representation.
use std::iter;

use ra_syntax::{TextRange, TextSize};
use rustc_hash::FxHashMap;
use superslice::Ext;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LineIndex {
    /// Offset the the beginning of each line, zero-based
    pub(crate) newlines: Vec<TextSize>,
    /// List of non-ASCII characters on each line
    pub(crate) utf16_lines: FxHashMap<u32, Vec<Utf16Char>>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct LineCol {
    /// Zero-based
    pub line: u32,
    /// Zero-based
    pub col_utf16: u32,
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub(crate) struct Utf16Char {
    /// Start offset of a character inside a line, zero-based
    pub(crate) start: TextSize,
    /// End offset of a character inside a line, zero-based
    pub(crate) end: TextSize,
}

impl Utf16Char {
    fn len(&self) -> TextSize {
        self.end - self.start
    }
}

impl LineIndex {
    pub fn new(text: &str) -> LineIndex {
        let mut utf16_lines = FxHashMap::default();
        let mut utf16_chars = Vec::new();

        let mut newlines = vec![0.into()];
        let mut curr_row = 0.into();
        let mut curr_col = 0.into();
        let mut line = 0;
        for c in text.chars() {
            let c_len = TextSize::of(c);
            curr_row += c_len;
            if c == '\n' {
                newlines.push(curr_row);

                // Save any utf-16 characters seen in the previous line
                if !utf16_chars.is_empty() {
                    utf16_lines.insert(line, utf16_chars);
                    utf16_chars = Vec::new();
                }

                // Prepare for processing the next line
                curr_col = 0.into();
                line += 1;
                continue;
            }

            if !c.is_ascii() {
                utf16_chars.push(Utf16Char { start: curr_col, end: curr_col + c_len });
            }

            curr_col += c_len;
        }

        // Save any utf-16 characters seen in the last line
        if !utf16_chars.is_empty() {
            utf16_lines.insert(line, utf16_chars);
        }

        LineIndex { newlines, utf16_lines }
    }

    pub fn line_col(&self, offset: TextSize) -> LineCol {
        let line = self.newlines.upper_bound(&offset) - 1;
        let line_start_offset = self.newlines[line];
        let col = offset - line_start_offset;

        LineCol { line: line as u32, col_utf16: self.utf8_to_utf16_col(line as u32, col) as u32 }
    }

    pub fn offset(&self, line_col: LineCol) -> TextSize {
        //FIXME: return Result
        let col = self.utf16_to_utf8_col(line_col.line, line_col.col_utf16);
        self.newlines[line_col.line as usize] + col
    }

    pub fn lines(&self, range: TextRange) -> impl Iterator<Item = TextRange> + '_ {
        let lo = self.newlines.lower_bound(&range.start());
        let hi = self.newlines.upper_bound(&range.end());
        let all = iter::once(range.start())
            .chain(self.newlines[lo..hi].iter().copied())
            .chain(iter::once(range.end()));

        all.clone()
            .zip(all.skip(1))
            .map(|(lo, hi)| TextRange::new(lo, hi))
            .filter(|it| !it.is_empty())
    }

    fn utf8_to_utf16_col(&self, line: u32, col: TextSize) -> usize {
        let mut res: usize = col.into();
        if let Some(utf16_chars) = self.utf16_lines.get(&line) {
            for c in utf16_chars {
                if c.end <= col {
                    res -= usize::from(c.len()) - 1;
                } else {
                    // From here on, all utf16 characters come *after* the character we are mapping,
                    // so we don't need to take them into account
                    break;
                }
            }
        }
        res
    }

    fn utf16_to_utf8_col(&self, line: u32, mut col: u32) -> TextSize {
        if let Some(utf16_chars) = self.utf16_lines.get(&line) {
            for c in utf16_chars {
                if col > u32::from(c.start) {
                    col += u32::from(c.len()) - 1;
                } else {
                    // From here on, all utf16 characters come *after* the character we are mapping,
                    // so we don't need to take them into account
                    break;
                }
            }
        }

        col.into()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_line_index() {
        let text = "hello\nworld";
        let index = LineIndex::new(text);
        assert_eq!(index.line_col(0.into()), LineCol { line: 0, col_utf16: 0 });
        assert_eq!(index.line_col(1.into()), LineCol { line: 0, col_utf16: 1 });
        assert_eq!(index.line_col(5.into()), LineCol { line: 0, col_utf16: 5 });
        assert_eq!(index.line_col(6.into()), LineCol { line: 1, col_utf16: 0 });
        assert_eq!(index.line_col(7.into()), LineCol { line: 1, col_utf16: 1 });
        assert_eq!(index.line_col(8.into()), LineCol { line: 1, col_utf16: 2 });
        assert_eq!(index.line_col(10.into()), LineCol { line: 1, col_utf16: 4 });
        assert_eq!(index.line_col(11.into()), LineCol { line: 1, col_utf16: 5 });
        assert_eq!(index.line_col(12.into()), LineCol { line: 1, col_utf16: 6 });

        let text = "\nhello\nworld";
        let index = LineIndex::new(text);
        assert_eq!(index.line_col(0.into()), LineCol { line: 0, col_utf16: 0 });
        assert_eq!(index.line_col(1.into()), LineCol { line: 1, col_utf16: 0 });
        assert_eq!(index.line_col(2.into()), LineCol { line: 1, col_utf16: 1 });
        assert_eq!(index.line_col(6.into()), LineCol { line: 1, col_utf16: 5 });
        assert_eq!(index.line_col(7.into()), LineCol { line: 2, col_utf16: 0 });
    }

    #[test]
    fn test_char_len() {
        assert_eq!('メ'.len_utf8(), 3);
        assert_eq!('メ'.len_utf16(), 1);
    }

    #[test]
    fn test_empty_index() {
        let col_index = LineIndex::new(
            "
const C: char = 'x';
",
        );
        assert_eq!(col_index.utf16_lines.len(), 0);
    }

    #[test]
    fn test_single_char() {
        let col_index = LineIndex::new(
            "
const C: char = 'メ';
",
        );

        assert_eq!(col_index.utf16_lines.len(), 1);
        assert_eq!(col_index.utf16_lines[&1].len(), 1);
        assert_eq!(col_index.utf16_lines[&1][0], Utf16Char { start: 17.into(), end: 20.into() });

        // UTF-8 to UTF-16, no changes
        assert_eq!(col_index.utf8_to_utf16_col(1, 15.into()), 15);

        // UTF-8 to UTF-16
        assert_eq!(col_index.utf8_to_utf16_col(1, 22.into()), 20);

        // UTF-16 to UTF-8, no changes
        assert_eq!(col_index.utf16_to_utf8_col(1, 15), TextSize::from(15));

        // UTF-16 to UTF-8
        assert_eq!(col_index.utf16_to_utf8_col(1, 19), TextSize::from(21));
    }

    #[test]
    fn test_string() {
        let col_index = LineIndex::new(
            "
const C: char = \"メ メ\";
",
        );

        assert_eq!(col_index.utf16_lines.len(), 1);
        assert_eq!(col_index.utf16_lines[&1].len(), 2);
        assert_eq!(col_index.utf16_lines[&1][0], Utf16Char { start: 17.into(), end: 20.into() });
        assert_eq!(col_index.utf16_lines[&1][1], Utf16Char { start: 21.into(), end: 24.into() });

        // UTF-8 to UTF-16
        assert_eq!(col_index.utf8_to_utf16_col(1, 15.into()), 15);

        assert_eq!(col_index.utf8_to_utf16_col(1, 21.into()), 19);
        assert_eq!(col_index.utf8_to_utf16_col(1, 25.into()), 21);

        assert!(col_index.utf8_to_utf16_col(2, 15.into()) == 15);

        // UTF-16 to UTF-8
        assert_eq!(col_index.utf16_to_utf8_col(1, 15), TextSize::from(15));

        // メ UTF-8: 0xE3 0x83 0xA1, UTF-16: 0x30E1
        assert_eq!(col_index.utf16_to_utf8_col(1, 17), TextSize::from(17)); // first メ at 17..20
        assert_eq!(col_index.utf16_to_utf8_col(1, 18), TextSize::from(20)); // space
        assert_eq!(col_index.utf16_to_utf8_col(1, 19), TextSize::from(21)); // second メ at 21..24

        assert_eq!(col_index.utf16_to_utf8_col(2, 15), TextSize::from(15));
    }

    #[test]
    fn test_splitlines() {
        fn r(lo: u32, hi: u32) -> TextRange {
            TextRange::new(lo.into(), hi.into())
        }

        let text = "a\nbb\nccc\n";
        let line_index = LineIndex::new(text);

        let actual = line_index.lines(r(0, 9)).collect::<Vec<_>>();
        let expected = vec![r(0, 2), r(2, 5), r(5, 9)];
        assert_eq!(actual, expected);

        let text = "";
        let line_index = LineIndex::new(text);

        let actual = line_index.lines(r(0, 0)).collect::<Vec<_>>();
        let expected = vec![];
        assert_eq!(actual, expected);

        let text = "\n";
        let line_index = LineIndex::new(text);

        let actual = line_index.lines(r(0, 1)).collect::<Vec<_>>();
        let expected = vec![r(0, 1)];
        assert_eq!(actual, expected)
    }
}
<commit_msg>Fix column conversion for supplementary plane characters<commit_after>//! `LineIndex` maps flat `TextSize` offsets into `(Line, Column)`
//! representation.
use std::iter;

use ra_syntax::{TextRange, TextSize};
use rustc_hash::FxHashMap;
use superslice::Ext;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LineIndex {
    /// Offset the the beginning of each line, zero-based
    pub(crate) newlines: Vec<TextSize>,
    /// List of non-ASCII characters on each line
    pub(crate) utf16_lines: FxHashMap<u32, Vec<Utf16Char>>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct LineCol {
    /// Zero-based
    pub line: u32,
    /// Zero-based
    pub col_utf16: u32,
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub(crate) struct Utf16Char {
    /// Start offset of a character inside a line, zero-based
    pub(crate) start: TextSize,
    /// End offset of a character inside a line, zero-based
    pub(crate) end: TextSize,
}

impl Utf16Char {
    /// Returns the length in 8-bit UTF-8 code units.
    fn len(&self) -> TextSize {
        self.end - self.start
    }

    /// Returns the length in 16-bit UTF-16 code units.
    fn len_utf16(&self) -> usize {
        if self.len() == TextSize::from(4) {
            2
        } else {
            1
        }
    }
}

impl LineIndex {
    pub fn new(text: &str) -> LineIndex {
        let mut utf16_lines = FxHashMap::default();
        let mut utf16_chars = Vec::new();

        let mut newlines = vec![0.into()];
        let mut curr_row = 0.into();
        let mut curr_col = 0.into();
        let mut line = 0;
        for c in text.chars() {
            let c_len = TextSize::of(c);
            curr_row += c_len;
            if c == '\n' {
                newlines.push(curr_row);

                // Save any utf-16 characters seen in the previous line
                if !utf16_chars.is_empty() {
                    utf16_lines.insert(line, utf16_chars);
                    utf16_chars = Vec::new();
                }

                // Prepare for processing the next line
                curr_col = 0.into();
                line += 1;
                continue;
            }

            if !c.is_ascii() {
                utf16_chars.push(Utf16Char { start: curr_col, end: curr_col + c_len });
            }

            curr_col += c_len;
        }

        // Save any utf-16 characters seen in the last line
        if !utf16_chars.is_empty() {
            utf16_lines.insert(line, utf16_chars);
        }

        LineIndex { newlines, utf16_lines }
    }

    pub fn line_col(&self, offset: TextSize) -> LineCol {
        let line = self.newlines.upper_bound(&offset) - 1;
        let line_start_offset = self.newlines[line];
        let col = offset - line_start_offset;

        LineCol { line: line as u32, col_utf16: self.utf8_to_utf16_col(line as u32, col) as u32 }
    }

    pub fn offset(&self, line_col: LineCol) -> TextSize {
        //FIXME: return Result
        let col = self.utf16_to_utf8_col(line_col.line, line_col.col_utf16);
        self.newlines[line_col.line as usize] + col
    }

    pub fn lines(&self, range: TextRange) -> impl Iterator<Item = TextRange> + '_ {
        let lo = self.newlines.lower_bound(&range.start());
        let hi = self.newlines.upper_bound(&range.end());
        let all = iter::once(range.start())
            .chain(self.newlines[lo..hi].iter().copied())
            .chain(iter::once(range.end()));

        all.clone()
            .zip(all.skip(1))
            .map(|(lo, hi)| TextRange::new(lo, hi))
            .filter(|it| !it.is_empty())
    }

    fn utf8_to_utf16_col(&self, line: u32, col: TextSize) -> usize {
        let mut res: usize = col.into();
        if let Some(utf16_chars) = self.utf16_lines.get(&line) {
            for c in utf16_chars {
                if c.end <= col {
                    res -= usize::from(c.len()) - c.len_utf16();
                } else {
                    // From here on, all utf16 characters come *after* the character we are mapping,
                    // so we don't need to take them into account
                    break;
                }
            }
        }
        res
    }

    fn utf16_to_utf8_col(&self, line: u32, mut col: u32) -> TextSize {
        if let Some(utf16_chars) = self.utf16_lines.get(&line) {
            for c in utf16_chars {
                if col > u32::from(c.start) {
                    col += u32::from(c.len()) - c.len_utf16() as u32;
                } else {
                    // From here on, all utf16 characters come *after* the character we are mapping,
                    // so we don't need to take them into account
                    break;
                }
            }
        }

        col.into()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_line_index() {
        let text = "hello\nworld";
        let index = LineIndex::new(text);
        assert_eq!(index.line_col(0.into()), LineCol { line: 0, col_utf16: 0 });
        assert_eq!(index.line_col(1.into()), LineCol { line: 0, col_utf16: 1 });
        assert_eq!(index.line_col(5.into()), LineCol { line: 0, col_utf16: 5 });
        assert_eq!(index.line_col(6.into()), LineCol { line: 1, col_utf16: 0 });
        assert_eq!(index.line_col(7.into()), LineCol { line: 1, col_utf16: 1 });
        assert_eq!(index.line_col(8.into()), LineCol { line: 1, col_utf16: 2 });
        assert_eq!(index.line_col(10.into()), LineCol { line: 1, col_utf16: 4 });
        assert_eq!(index.line_col(11.into()), LineCol { line: 1, col_utf16: 5 });
        assert_eq!(index.line_col(12.into()), LineCol { line: 1, col_utf16: 6 });

        let text = "\nhello\nworld";
        let index = LineIndex::new(text);
        assert_eq!(index.line_col(0.into()), LineCol { line: 0, col_utf16: 0 });
        assert_eq!(index.line_col(1.into()), LineCol { line: 1, col_utf16: 0 });
        assert_eq!(index.line_col(2.into()), LineCol { line: 1, col_utf16: 1 });
        assert_eq!(index.line_col(6.into()), LineCol { line: 1, col_utf16: 5 });
        assert_eq!(index.line_col(7.into()), LineCol { line: 2, col_utf16: 0 });
    }

    #[test]
    fn test_char_len() {
        assert_eq!('メ'.len_utf8(), 3);
        assert_eq!('メ'.len_utf16(), 1);
    }

    #[test]
    fn test_empty_index() {
        let col_index = LineIndex::new(
            "
const C: char = 'x';
",
        );
        assert_eq!(col_index.utf16_lines.len(), 0);
    }

    #[test]
    fn test_single_char() {
        let col_index = LineIndex::new(
            "
const C: char = 'メ';
",
        );

        assert_eq!(col_index.utf16_lines.len(), 1);
        assert_eq!(col_index.utf16_lines[&1].len(), 1);
        assert_eq!(col_index.utf16_lines[&1][0], Utf16Char { start: 17.into(), end: 20.into() });

        // UTF-8 to UTF-16, no changes
        assert_eq!(col_index.utf8_to_utf16_col(1, 15.into()), 15);

        // UTF-8 to UTF-16
        assert_eq!(col_index.utf8_to_utf16_col(1, 22.into()), 20);

        // UTF-16 to UTF-8, no changes
        assert_eq!(col_index.utf16_to_utf8_col(1, 15), TextSize::from(15));

        // UTF-16 to UTF-8
        assert_eq!(col_index.utf16_to_utf8_col(1, 19), TextSize::from(21));

        let col_index = LineIndex::new("a𐐏b");
        assert_eq!(col_index.utf16_to_utf8_col(0, 3), TextSize::from(5));
    }

    #[test]
    fn test_string() {
        let col_index = LineIndex::new(
            "
const C: char = \"メ メ\";
",
        );

        assert_eq!(col_index.utf16_lines.len(), 1);
        assert_eq!(col_index.utf16_lines[&1].len(), 2);
        assert_eq!(col_index.utf16_lines[&1][0], Utf16Char { start: 17.into(), end: 20.into() });
        assert_eq!(col_index.utf16_lines[&1][1], Utf16Char { start: 21.into(), end: 24.into() });

        // UTF-8 to UTF-16
        assert_eq!(col_index.utf8_to_utf16_col(1, 15.into()), 15);

        assert_eq!(col_index.utf8_to_utf16_col(1, 21.into()), 19);
        assert_eq!(col_index.utf8_to_utf16_col(1, 25.into()), 21);

        assert!(col_index.utf8_to_utf16_col(2, 15.into()) == 15);

        // UTF-16 to UTF-8
        assert_eq!(col_index.utf16_to_utf8_col(1, 15), TextSize::from(15));

        // メ UTF-8: 0xE3 0x83 0xA1, UTF-16: 0x30E1
        assert_eq!(col_index.utf16_to_utf8_col(1, 17), TextSize::from(17)); // first メ at 17..20
        assert_eq!(col_index.utf16_to_utf8_col(1, 18), TextSize::from(20)); // space
        assert_eq!(col_index.utf16_to_utf8_col(1, 19), TextSize::from(21)); // second メ at 21..24

        assert_eq!(col_index.utf16_to_utf8_col(2, 15), TextSize::from(15));
    }

    #[test]
    fn test_splitlines() {
        fn r(lo: u32, hi: u32) -> TextRange {
            TextRange::new(lo.into(), hi.into())
        }

        let text = "a\nbb\nccc\n";
        let line_index = LineIndex::new(text);

        let actual = line_index.lines(r(0, 9)).collect::<Vec<_>>();
        let expected = vec![r(0, 2), r(2, 5), r(5, 9)];
        assert_eq!(actual, expected);

        let text = "";
        let line_index = LineIndex::new(text);

        let actual = line_index.lines(r(0, 0)).collect::<Vec<_>>();
        let expected = vec![];
        assert_eq!(actual, expected);

        let text = "\n";
        let line_index = LineIndex::new(text);

        let actual = line_index.lines(r(0, 1)).collect::<Vec<_>>();
        let expected = vec![r(0, 1)];
        assert_eq!(actual, expected)
    }
}
<|endoftext|>


<commit_before>//! Weave parsing

use crate::{header::Header, Error, NamingConvention, Result};
use flate2::read::GzDecoder;
use log::info;
use std::{
    cell::RefCell,
    fs::File,
    io::{BufRead, BufReader, Lines, Read},
    mem,
    rc::Rc,
};

/// A Sink is a place that a parsed weave can be sent to.  The insert/delete/end commands match
/// those in the weave file, and `plain` are the lines of data.  With each plain is a flag
/// indicating if that line should be included in the output (all lines are called, so that updates
/// can use this same code).  All methods return a result, with the Err value stopping the parse.
/// Note that the default implementations just return success, and ignore the result.
pub trait Sink {
    /// Begin an insert sequence for the given delta.
    fn insert(&mut self, _delta: usize) -> Result<()> {
        Ok(())
    }

    /// Begin a delete sequence.
    fn delete(&mut self, _delta: usize) -> Result<()> {
        Ok(())
    }

    /// End a previous insert or delete.
    fn end(&mut self, _delta: usize) -> Result<()> {
        Ok(())
    }

    /// A single line of plain text from the weave.  `keep` indicates if the line should be
    /// included in the requested delta.
    fn plain(&mut self, _text: &str, _keep: bool) -> Result<()> {
        Ok(())
    }
}

/// The PullParser returns the entries as nodes.  These are equivalent to
/// the values in Sink.
#[derive(Debug)]
pub enum Entry {
    /// Begin an insert sequence for the given delta.
    Insert { delta: usize },

    /// Begin a delete sequence.
    Delete { delta: usize },

    /// End a previous insert or delete.
    End { delta: usize },

    /// A single line of plaintext from the weave.  `keep` indicates if the
    /// line should be included in the requested delta.
    Plain { text: String, keep: bool },

    /// A control message.  Doesn't currently contain any data, which can be added later if needed.
    Control,
}

/// A Parser is used to process a weave file, extracting either everything, or only a specific
/// delta.
pub struct Parser<S: Sink, B> {
    /// The lines of the input.
    source: Lines<B>,

    /// The sink to be given each line record in the weave file.
    sink: Rc<RefCell<S>>,

    /// The desired delta to retrieve, which affects the parse_to call as well as the `keep`
    /// argument passed to the sink's `plain` call.
    delta: usize,

    /// The delta state is kept sorted with the newest (largest) delta at element 0.
    delta_state: Vec<OneDelta>,

    /// A pending input line kept from the last invocation.
    pending: Option<String>,

    /// Indicates that we are currently "keeping" lines.
    keeping: bool,

    /// The line number indicator.
    lineno: usize,

    /// The header extracted from the file.
    header: Header,
}

impl<S: Sink> Parser<S, BufReader<Box<dyn Read>>> {
    /// Construct a parser, based on the main file of the naming convention.
    pub fn new(
        naming: &dyn NamingConvention,
        sink: S,
        delta: usize,
    ) -> Result<Parser<S, BufReader<Box<dyn Read>>>> {
        let rd = if naming.is_compressed() {
            let fd = File::open(naming.main_file())?;
            Box::new(GzDecoder::new(fd)) as Box<dyn Read>
        } else {
            Box::new(File::open(naming.main_file())?) as Box<dyn Read>
        };
        let lines = BufReader::new(rd).lines();
        Parser::new_raw(lines, Rc::new(RefCell::new(sink)), delta)
    }
}

impl<S: Sink, B: BufRead> Parser<S, B> {
    /// Construct a new Parser, reading from the given Reader, giving records to the given Sink,
    /// and aiming for the specified `delta`.  This is not the intended constructor, normal users
    /// should use `new`.  (This is public, for testing).
    pub fn new_raw(
        mut source: Lines<B>,
        sink: Rc<RefCell<S>>,
        delta: usize,
    ) -> Result<Parser<S, B>> {
        if let Some(line) = source.next() {
            let line = line?;
            let header = Header::decode(&line)?;

            Ok(Parser {
                source,
                sink,
                delta,
                delta_state: vec![],
                pending: None,
                keeping: false,
                lineno: 0,
                header,
            })
        } else {
            Err(Error::EmptyWeave)
        }
    }

    /// Run the parser until we either reach the given line number, or the end of the weave.  Lines
    /// are numbered from 1, so calling with a lineno of zero will run the parser until the end of
    /// the input.  Returns Ok(0) for the end of input, Ok(n) for stopping at line n (which should
    /// always be the same as the passed in lineno, or Err if there is an error.
    pub fn parse_to(&mut self, lineno: usize) -> Result<usize> {
        // Handle any pending input line.
        if let Some(pending) = mem::replace(&mut self.pending, None) {
            self.sink.borrow_mut().plain(&pending, self.keeping)?;
        }

        loop {
            // Get the next input line, finishing if we are done.
            let line = match self.source.next() {
                None => return Ok(0),
                Some(line) => line?,
            };

            info!("line: {:?}", line);

            // Detect the first character, without borrowing.
            let textual = match line.bytes().next() {
                None => true,
                Some(ch) if ch != b'\x01' => true,
                _ => false,
            };

            if textual {
                // Textual line.  Count line numbers for the lines we're keeping.
                if self.keeping {
                    self.lineno += 1;
                    if self.lineno == lineno {
                        // This is the desired stopping point, hold onto this line, and return to
                        // the caller.
                        self.pending = Some(line);
                        return Ok(lineno);
                    }
                }

                // Otherwise, call the Sink, and continue.
                info!("textual: keeping={}", self.keeping);
                self.sink.borrow_mut().plain(&line, self.keeping)?;
                continue;
            }

            let linebytes = line.as_bytes();

            // At this point, all should be control lines.  Skip any that are too short.
            if linebytes.len() < 4 {
                continue;
            }

            // Ignore control lines other than the insert/delete/end lines.
            if linebytes[1] != b'I' && linebytes[1] != b'D' && linebytes[1] != b'E' {
                continue;
            }

            let this_delta: usize = line[3..].parse()?;

            match linebytes[1] {
                b'E' => {
                    self.sink.borrow_mut().end(this_delta)?;
                    self.pop(this_delta);
                }
                b'I' => {
                    self.sink.borrow_mut().insert(this_delta)?;

                    // Do this insert if this insert is at least as old as the requested delta.
                    if self.delta >= this_delta {
                        self.push(this_delta, StateMode::Keep);
                    } else {
                        self.push(this_delta, StateMode::Skip);
                    }
                }
                b'D' => {
                    self.sink.borrow_mut().delete(this_delta)?;

                    // Do this delete if this delete is newer than current.  If not, don't account
                    // for it.
                    if self.delta >= this_delta {
                        self.push(this_delta, StateMode::Skip);
                    } else {
                        self.push(this_delta, StateMode::Next);
                    }
                }
                _ => unreachable!(),
            }

            self.update_keep();
        }
    }

    /// Remove the given numbered state.
    fn pop(&mut self, delta: usize) {
        // The binary search is reversed, so the largest are first.
        let pos = match self
            .delta_state
            .binary_search_by(|ent| delta.cmp(&ent.delta))
        {
            Ok(pos) => pos,
            Err(_) => unreachable!(),
        };

        self.delta_state.remove(pos);
    }

    /// Add a new state.  It will be inserted in the proper place in the array, based on the delta
    /// number.
    fn push(&mut self, delta: usize, mode: StateMode) {
        match self
            .delta_state
            .binary_search_by(|ent| delta.cmp(&ent.delta))
        {
            Ok(_) => panic!("Duplicate state in push"),
            Err(pos) => self.delta_state.insert(pos, OneDelta { delta, mode }),
        }
    }

    /// Update the keep field, based on the current state.
    fn update_keep(&mut self) {
        info!("Update: {:?}", self.delta_state);
        for st in &self.delta_state {
            match st.mode {
                StateMode::Keep => {
                    self.keeping = true;
                    return;
                }
                StateMode::Skip => {
                    self.keeping = false;
                    return;
                }
                _ => (),
            }
        }

        // This shouldn't be reached if there are any more context lines, but we may get here when
        // we reach the end of the input.
        self.keeping = false;
    }

    /// Get the header read from this weave file.
    pub fn get_header(&self) -> &Header {
        &self.header
    }

    /// Consume the parser, returning the header.
    pub fn into_header(self) -> Header {
        self.header
    }

    /// Get a copy of the sink.
    pub fn get_sink(&self) -> Rc<RefCell<S>> {
        self.sink.clone()
    }
}

/*
/// A PullIterator returns entities in a weave file, extracting either
/// everything, or only a specific delta.
pub struct PullIterator<B> {
    /// The lines of the input.
    source: Lines<B>,

    /// The desired delta to retrieve.
    delta: usize,

    /// The delta state is kept sorted with the newest (largest) delta at
    /// element 0.
    delta_state: Vec<OneDelta>,

    /// Indicates we are currently keeping lines.
    keeping: bool,

    /// The current line number.
    lineno: usize,

    /// The header extracted from the file.
    header: Header,
}
*/

pub struct PullParser<B> {
    /// The lines of the input.
    source: Lines<B>,

    /// The desired delta to retrieve.
    delta: usize,

    /// The delta state is kept sorted with the newest (largest) delta at element 0.
    delta_state: Vec<OneDelta>,

    /// Indicates that we are currently "keeping" lines.
    keeping: bool,

    /// The header extracted from the file.
    header: Header,
}

impl PullParser<BufReader<Box<dyn Read>>> {
    /// Construct a parser, based on the main file of the naming
    /// convention.
    pub fn new(
        naming: &dyn NamingConvention,
        delta: usize,
    ) -> Result<PullParser<BufReader<Box<dyn Read>>>> {
        let rd = if naming.is_compressed() {
            let fd = File::open(naming.main_file())?;
            Box::new(GzDecoder::new(fd)) as Box<dyn Read>
        } else {
            Box::new(File::open(naming.main_file())?) as Box<dyn Read>
        };
        let lines = BufReader::new(rd).lines();
        PullParser::new_raw(lines, delta)
    }
}

impl<B: BufRead> PullParser<B> {
    /// Construct a new Parser, reading from the given Reader.  The parser
    /// will act as an iterator.  This is the intended constructor, normal
    /// users should use `new`.  (This is public for testing).
    pub fn new_raw(mut source: Lines<B>, delta: usize) -> Result<PullParser<B>> {
        if let Some(line) = source.next() {
            let line = line?;
            let header = Header::decode(&line)?;

            Ok(PullParser {
                source,
                delta,
                delta_state: vec![],
                keeping: false,
                header,
            })
        } else {
            Err(Error::EmptyWeave)
        }
    }

    /// Remove the given numbered state.
    fn pop(&mut self, delta: usize) {
        // The binary search is reversed, so the largest are first.
        let pos = match self
            .delta_state
            .binary_search_by(|ent| delta.cmp(&ent.delta))
        {
            Ok(pos) => pos,
            Err(_) => unreachable!(),
        };

        self.delta_state.remove(pos);
    }

    /// Add a new state.  It will be inserted in the proper place in the array, based on the delta
    /// number.
    fn push(&mut self, delta: usize, mode: StateMode) {
        match self
            .delta_state
            .binary_search_by(|ent| delta.cmp(&ent.delta))
        {
            Ok(_) => panic!("Duplicate state in push"),
            Err(pos) => self.delta_state.insert(pos, OneDelta { delta, mode }),
        }
    }

    /// Update the keep field, based on the current state.
    fn update_keep(&mut self) {
        info!("Update: {:?}", self.delta_state);
        for st in &self.delta_state {
            match st.mode {
                StateMode::Keep => {
                    self.keeping = true;
                    return;
                }
                StateMode::Skip => {
                    self.keeping = false;
                    return;
                }
                _ => (),
            }
        }

        // This shouldn't be reached if there are any more context lines, but we may get here when
        // we reach the end of the input.
        self.keeping = false;
    }

    /// Get the header read from this weave file.
    pub fn get_header(&self) -> &Header {
        &self.header
    }

    /// Consume the parser, returning the header.
    pub fn into_header(self) -> Header {
        self.header
    }
}

impl<B: BufRead> Iterator for PullParser<B> {
    type Item = Result<Entry>;

    fn next(&mut self) -> Option<Result<Entry>> {
        // At this level, there is a 1:1 correspondence between weave input
        // lines and those returned.
        let line = match self.source.next() {
            None => return None,
            Some(Ok(line)) => line,
            Some(Err(e)) => return Some(Err(From::from(e))),
        };

        info!("line: {:?}", line);

        // Detect the first character, without borrowing.
        let textual = match line.bytes().next() {
            None => true,
            Some(ch) if ch != b'\x01' => true,
            _ => false,
        };

        if textual {
            return Some(Ok(Entry::Plain {
                text: line,
                keep: self.keeping,
            }));
        }

        let linebytes = line.as_bytes();

        if linebytes.len() < 4 {
            return Some(Ok(Entry::Control));
        }

        if linebytes[1] != b'I' && linebytes[1] != b'D' && linebytes[1] != b'E' {
            return Some(Ok(Entry::Control));
        };

        // TODO: Don't panic, but fail.
        let this_delta: usize = line[3..].parse().unwrap();

        match linebytes[1] {
            b'E' => {
                self.pop(this_delta);
                self.update_keep();
                Some(Ok(Entry::End { delta: this_delta }))
            }
            b'I' => {
                if self.delta >= this_delta {
                    self.push(this_delta, StateMode::Keep);
                } else {
                    self.push(this_delta, StateMode::Skip);
                }
                self.update_keep();

                Some(Ok(Entry::Insert { delta: this_delta }))
            }
            b'D' => {
                if self.delta >= this_delta {
                    self.push(this_delta, StateMode::Skip);
                } else {
                    self.push(this_delta, StateMode::Next);
                }
                self.update_keep();

                Some(Ok(Entry::Delete { delta: this_delta }))
            }
            _ => unreachable!(),
        }
    }
}

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
enum StateMode {
    Keep,
    Skip,
    Next,
}

#[derive(Debug)]
struct OneDelta {
    delta: usize,
    mode: StateMode,
}
<commit_msg>weave: Simplified push parser<commit_after>//! Weave parsing

use crate::{header::Header, Error, NamingConvention, Result};
use flate2::read::GzDecoder;
use log::info;
use std::{
    cell::RefCell,
    fs::File,
    io::{BufRead, BufReader, Lines, Read},
    mem,
    rc::Rc,
};

/// A Sink is a place that a parsed weave can be sent to.  The insert/delete/end commands match
/// those in the weave file, and `plain` are the lines of data.  With each plain is a flag
/// indicating if that line should be included in the output (all lines are called, so that updates
/// can use this same code).  All methods return a result, with the Err value stopping the parse.
/// Note that the default implementations just return success, and ignore the result.
pub trait Sink {
    /// Begin an insert sequence for the given delta.
    fn insert(&mut self, _delta: usize) -> Result<()> {
        Ok(())
    }

    /// Begin a delete sequence.
    fn delete(&mut self, _delta: usize) -> Result<()> {
        Ok(())
    }

    /// End a previous insert or delete.
    fn end(&mut self, _delta: usize) -> Result<()> {
        Ok(())
    }

    /// A single line of plain text from the weave.  `keep` indicates if the line should be
    /// included in the requested delta.
    fn plain(&mut self, _text: &str, _keep: bool) -> Result<()> {
        Ok(())
    }
}

/// The PullParser returns the entries as nodes.  These are equivalent to
/// the values in Sink.
#[derive(Debug)]
pub enum Entry {
    /// Begin an insert sequence for the given delta.
    Insert { delta: usize },

    /// Begin a delete sequence.
    Delete { delta: usize },

    /// End a previous insert or delete.
    End { delta: usize },

    /// A single line of plaintext from the weave.  `keep` indicates if the
    /// line should be included in the requested delta.
    Plain { text: String, keep: bool },

    /// A control message.  Doesn't currently contain any data, which can be added later if needed.
    Control,
}

/// A Parser is used to process a weave file.  This is a wrapper around the pull parser that
/// invokes a push parser.
pub struct Parser<S: Sink, B> {
    /// The pull parser.
    pull: PullParser<B>,

    /// The sink to be given each line record in the weave file.
    sink: Rc<RefCell<S>>,

    /// A single pending line, kept from the last invocation.
    pending: Option<String>,

    /// Tracking the line number.
    lineno: usize,
}

impl<S: Sink> Parser<S, BufReader<Box<dyn Read>>> {
    /// Construct a parser, based on the main file of the naming convention.
    pub fn new(
        naming: &dyn NamingConvention,
        sink: S,
        delta: usize,
    ) -> Result<Parser<S, BufReader<Box<dyn Read>>>> {
        let rd = if naming.is_compressed() {
            let fd = File::open(naming.main_file())?;
            Box::new(GzDecoder::new(fd)) as Box<dyn Read>
        } else {
            Box::new(File::open(naming.main_file())?) as Box<dyn Read>
        };
        let lines = BufReader::new(rd).lines();
        Parser::new_raw(lines, Rc::new(RefCell::new(sink)), delta)
    }
}

impl<S: Sink, B: BufRead> Parser<S, B> {
    /// Construct a new Parser, reading from the given Reader, giving records to the given Sink,
    /// and aiming for the specified `delta`.  This is not the intended constructor, normal users
    /// should use `new`.  (This is public, for testing).
    pub fn new_raw(
        source: Lines<B>,
        sink: Rc<RefCell<S>>,
        delta: usize,
    ) -> Result<Parser<S, B>> {
        let pull = PullParser::new_raw(source, delta)?;
        Ok(Parser {
            pull,
            sink,
            pending: None,
            lineno: 0,
        })
    }

    /// Run the parser until we either reach the given line number, or the end of the weave.  Lines
    /// are numbered from 1, so calling with a lineno of zero will run the parser until the end of
    /// the input.  Returns Ok(0) for the end of input, Ok(n) for stopping at line n (which should
    /// always be the same as the passed in lineno, or Err if there is an error.
    pub fn parse_to(&mut self, lineno: usize) -> Result<usize> {
        // Handle any pending input line.  Pending lines only happen while keeping.
        if let Some(text) = mem::replace(&mut self.pending, None) {
            self.sink.borrow_mut().plain(&text, true)?;
        }

        loop {
            match self.pull.next() {
                Some(Ok(Entry::Plain { text, keep })) => {
                    if keep {
                        self.lineno += 1;
                        if self.lineno == lineno {
                            // This is the desired stopping point, hold onto this line, and return
                            // to the caller.
                            self.pending = Some(text);
                            return Ok(lineno);
                        }
                    }

                    self.sink.borrow_mut().plain(&text, keep)?;
                }
                Some(Ok(Entry::Insert { delta })) => {
                    self.sink.borrow_mut().insert(delta)?;
                }
                Some(Ok(Entry::Delete { delta })) => {
                    self.sink.borrow_mut().delete(delta)?;
                }
                Some(Ok(Entry::End { delta })) => {
                    self.sink.borrow_mut().end(delta)?;
                }
                Some(Ok(Entry::Control)) => (),
                Some(Err(err)) => {
                    return Err(err);
                }
                None => {
                    return Ok(0);
                }
            }
        }
    }

    /// Get the header read from this weave file.
    pub fn get_header(&self) -> &Header {
        &self.pull.header
    }

    /// Consume the parser, returning the header.
    pub fn into_header(self) -> Header {
        self.pull.into_header()
    }

    /// Get a copy of the sink.
    pub fn get_sink(&self) -> Rc<RefCell<S>> {
        self.sink.clone()
    }
}

/*
/// A PullIterator returns entities in a weave file, extracting either
/// everything, or only a specific delta.
pub struct PullIterator<B> {
    /// The lines of the input.
    source: Lines<B>,

    /// The desired delta to retrieve.
    delta: usize,

    /// The delta state is kept sorted with the newest (largest) delta at
    /// element 0.
    delta_state: Vec<OneDelta>,

    /// Indicates we are currently keeping lines.
    keeping: bool,

    /// The current line number.
    lineno: usize,

    /// The header extracted from the file.
    header: Header,
}
*/

pub struct PullParser<B> {
    /// The lines of the input.
    source: Lines<B>,

    /// The desired delta to retrieve.
    delta: usize,

    /// The delta state is kept sorted with the newest (largest) delta at element 0.
    delta_state: Vec<OneDelta>,

    /// Indicates that we are currently "keeping" lines.
    keeping: bool,

    /// The header extracted from the file.
    header: Header,
}

impl PullParser<BufReader<Box<dyn Read>>> {
    /// Construct a parser, based on the main file of the naming
    /// convention.
    pub fn new(
        naming: &dyn NamingConvention,
        delta: usize,
    ) -> Result<PullParser<BufReader<Box<dyn Read>>>> {
        let rd = if naming.is_compressed() {
            let fd = File::open(naming.main_file())?;
            Box::new(GzDecoder::new(fd)) as Box<dyn Read>
        } else {
            Box::new(File::open(naming.main_file())?) as Box<dyn Read>
        };
        let lines = BufReader::new(rd).lines();
        PullParser::new_raw(lines, delta)
    }
}

impl<B: BufRead> PullParser<B> {
    /// Construct a new Parser, reading from the given Reader.  The parser
    /// will act as an iterator.  This is the intended constructor, normal
    /// users should use `new`.  (This is public for testing).
    pub fn new_raw(mut source: Lines<B>, delta: usize) -> Result<PullParser<B>> {
        if let Some(line) = source.next() {
            let line = line?;
            let header = Header::decode(&line)?;

            Ok(PullParser {
                source,
                delta,
                delta_state: vec![],
                keeping: false,
                header,
            })
        } else {
            Err(Error::EmptyWeave)
        }
    }

    /// Remove the given numbered state.
    fn pop(&mut self, delta: usize) {
        // The binary search is reversed, so the largest are first.
        let pos = match self
            .delta_state
            .binary_search_by(|ent| delta.cmp(&ent.delta))
        {
            Ok(pos) => pos,
            Err(_) => unreachable!(),
        };

        self.delta_state.remove(pos);
    }

    /// Add a new state.  It will be inserted in the proper place in the array, based on the delta
    /// number.
    fn push(&mut self, delta: usize, mode: StateMode) {
        match self
            .delta_state
            .binary_search_by(|ent| delta.cmp(&ent.delta))
        {
            Ok(_) => panic!("Duplicate state in push"),
            Err(pos) => self.delta_state.insert(pos, OneDelta { delta, mode }),
        }
    }

    /// Update the keep field, based on the current state.
    fn update_keep(&mut self) {
        info!("Update: {:?}", self.delta_state);
        for st in &self.delta_state {
            match st.mode {
                StateMode::Keep => {
                    self.keeping = true;
                    return;
                }
                StateMode::Skip => {
                    self.keeping = false;
                    return;
                }
                _ => (),
            }
        }

        // This shouldn't be reached if there are any more context lines, but we may get here when
        // we reach the end of the input.
        self.keeping = false;
    }

    /// Get the header read from this weave file.
    pub fn get_header(&self) -> &Header {
        &self.header
    }

    /// Consume the parser, returning the header.
    pub fn into_header(self) -> Header {
        self.header
    }
}

impl<B: BufRead> Iterator for PullParser<B> {
    type Item = Result<Entry>;

    fn next(&mut self) -> Option<Result<Entry>> {
        // At this level, there is a 1:1 correspondence between weave input
        // lines and those returned.
        let line = match self.source.next() {
            None => return None,
            Some(Ok(line)) => line,
            Some(Err(e)) => return Some(Err(From::from(e))),
        };

        info!("line: {:?}", line);

        // Detect the first character, without borrowing.
        let textual = match line.bytes().next() {
            None => true,
            Some(ch) if ch != b'\x01' => true,
            _ => false,
        };

        if textual {
            return Some(Ok(Entry::Plain {
                text: line,
                keep: self.keeping,
            }));
        }

        let linebytes = line.as_bytes();

        if linebytes.len() < 4 {
            return Some(Ok(Entry::Control));
        }

        if linebytes[1] != b'I' && linebytes[1] != b'D' && linebytes[1] != b'E' {
            return Some(Ok(Entry::Control));
        };

        // TODO: Don't panic, but fail.
        let this_delta: usize = line[3..].parse().unwrap();

        match linebytes[1] {
            b'E' => {
                self.pop(this_delta);
                self.update_keep();
                Some(Ok(Entry::End { delta: this_delta }))
            }
            b'I' => {
                if self.delta >= this_delta {
                    self.push(this_delta, StateMode::Keep);
                } else {
                    self.push(this_delta, StateMode::Skip);
                }
                self.update_keep();

                Some(Ok(Entry::Insert { delta: this_delta }))
            }
            b'D' => {
                if self.delta >= this_delta {
                    self.push(this_delta, StateMode::Skip);
                } else {
                    self.push(this_delta, StateMode::Next);
                }
                self.update_keep();

                Some(Ok(Entry::Delete { delta: this_delta }))
            }
            _ => unreachable!(),
        }
    }
}

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
enum StateMode {
    Keep,
    Skip,
    Next,
}

#[derive(Debug)]
struct OneDelta {
    delta: usize,
    mode: StateMode,
}
<|endoftext|>