179 lines · 5.3 KB · Rust
RawDownload
// Copyright 2023 The Jujutsu Authors//// Licensed under the Apache License, Version 2.0 (the "License");// you may not use this file except in compliance with the License.// You may obtain a copy of the License at//// https://www.apache.org/licenses/LICENSE-2.0//// Unless required by applicable law or agreed to in writing, software// distributed under the License is distributed on an "AS IS" BASIS,// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.// See the License for the specific language governing permissions and// limitations under the License. use std::collections::BinaryHeap;use std::mem; #[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]pub(super) struct RevWalkWorkItem<P, T> {    pub pos: P,    pub value: T,} #[derive(Clone)]pub(super) struct RevWalkQueue<P, T> {    items: BinaryHeap<RevWalkWorkItem<P, T>>,    // Optionally keep the greatest item out of the heap, so pop() + push() of    // the greatest item won't have to rebalance the heap.    scratch_item: Option<RevWalkWorkItem<P, T>>,    min_pos: P,} impl<P: Ord, T: Ord> RevWalkQueue<P, T> {    pub fn with_min_pos(min_pos: P) -> Self {        Self {            items: BinaryHeap::new(),            scratch_item: None,            min_pos,        }    }     #[cfg_attr(not(test), allow(dead_code))]    pub fn len(&self) -> usize {        self.items.len() + self.scratch_item.is_some() as usize    }     pub fn push(&mut self, pos: P, value: T) {        if pos < self.min_pos {            return;        }        self.push_item(RevWalkWorkItem { pos, value });    }     fn push_item(&mut self, new: RevWalkWorkItem<P, T>) {        if let Some(next) = self.scratch_item.as_mut() {            if new < *next {                self.items.push(new);            } else {                let next = mem::replace(next, new);                self.items.push(next);            }        } else if let Some(next) = self.items.peek() {            if new < *next {                // items[0] could be moved to scratch_item, but simply leave                // scratch_item empty.                self.items.push(new);            } else {                self.scratch_item = Some(new);            }        } else {            self.scratch_item = Some(new);        }    }     pub fn extend(&mut self, positions: impl IntoIterator<Item = P>, value: T)    where        T: Clone,    {        for pos in positions {            self.push(pos, value.clone());        }    }     pub fn peek(&self) -> Option<&RevWalkWorkItem<P, T>> {        self.scratch_item.as_ref().or_else(|| self.items.peek())    }     pub fn pop(&mut self) -> Option<RevWalkWorkItem<P, T>> {        let next = self.scratch_item.take().or_else(|| self.items.pop())?;        Some(next)    }     pub fn pop_if(        &mut self,        predicate: impl FnOnce(&RevWalkWorkItem<P, T>) -> bool,    ) -> Option<RevWalkWorkItem<P, T>> {        let next = self.peek()?;        predicate(next).then(|| self.pop().unwrap())    }     pub fn pop_eq(&mut self, pos: &P) -> Option<RevWalkWorkItem<P, T>> {        self.pop_if(|next| next.pos == *pos)    }     pub fn skip_while_eq(&mut self, pos: &P) {        while self.pop_eq(pos).is_some() {            continue;        }    }} #[cfg(test)]mod tests {    use assert_matches::assert_matches;     use super::*;     #[test]    fn test_push_pop_in_forward_order() {        let mut queue: RevWalkQueue<u32, ()> = RevWalkQueue::with_min_pos(0);         queue.push(0, ());        assert!(queue.scratch_item.is_some());        assert_eq!(queue.items.len(), 0);         queue.push(1, ());        assert!(queue.scratch_item.is_some());        assert_eq!(queue.items.len(), 1);         assert_matches!(queue.pop(), Some(RevWalkWorkItem { pos: 1, .. }));        assert!(queue.scratch_item.is_none());        assert_eq!(queue.items.len(), 1);         queue.push(2, ());        assert!(queue.scratch_item.is_some());        assert_eq!(queue.items.len(), 1);         assert_matches!(queue.pop(), Some(RevWalkWorkItem { pos: 2, .. }));        assert!(queue.scratch_item.is_none());        assert_eq!(queue.items.len(), 1);         assert_matches!(queue.pop(), Some(RevWalkWorkItem { pos: 0, .. }));        assert!(queue.scratch_item.is_none());        assert_eq!(queue.items.len(), 0);         assert_matches!(queue.pop(), None);    }     #[test]    fn test_push_pop_in_reverse_order() {        let mut queue: RevWalkQueue<u32, ()> = RevWalkQueue::with_min_pos(0);         queue.push(2, ());        assert!(queue.scratch_item.is_some());        assert_eq!(queue.items.len(), 0);         queue.push(1, ());        assert!(queue.scratch_item.is_some());        assert_eq!(queue.items.len(), 1);         assert_matches!(queue.pop(), Some(RevWalkWorkItem { pos: 2, .. }));        assert!(queue.scratch_item.is_none());        assert_eq!(queue.items.len(), 1);         queue.push(0, ());        assert!(queue.scratch_item.is_none());        assert_eq!(queue.items.len(), 2);         assert_matches!(queue.pop(), Some(RevWalkWorkItem { pos: 1, .. }));        assert!(queue.scratch_item.is_none());        assert_eq!(queue.items.len(), 1);         assert_matches!(queue.pop(), Some(RevWalkWorkItem { pos: 0, .. }));        assert!(queue.scratch_item.is_none());        assert_eq!(queue.items.len(), 0);         assert_matches!(queue.pop(), None);    }}