Experiment with incremental timsort inspired by std::slice::sort

Author: juntyr

necsim/impls/no-std/src/lib.rs

@@ -12,6 +12,8 @@
 #![feature(associated_type_bounds)]
 #![feature(const_float_bits_conv)]
 #![feature(core_intrinsics)]
+#![feature(ptr_sub_ptr)]
+#![feature(is_sorted)]
 #![allow(incomplete_features)]
 #![feature(specialization)]
 

necsim/impls/no-std/src/parallelisation/independent/monolithic/mod.rs

@@ -115,7 +115,7 @@ pub fn simulate<
         P::IsLive,
         R,
         P,
-    >>::WaterLevelReporter::new(event_slice.get(), local_partition);
+    >>::WaterLevelReporter::new(event_slice.get(), local_partition, 100_000); // TODO: configure here as well
     let mut min_spec_samples = dedup_cache.construct(slow_lineages.len());
 
     let mut total_steps = 0_u64;

necsim/impls/no-std/src/parallelisation/independent/monolithic/reporter/live.rs

@@ -1,5 +1,5 @@
-use alloc::vec::Vec;
-use core::{fmt, marker::PhantomData};
+use alloc::{vec::Vec, collections::VecDeque};
+use core::{fmt, marker::PhantomData, ops::ControlFlow};
 use necsim_core_bond::NonNegativeF64;
 
 use necsim_core::{
@@ -12,10 +12,21 @@ use necsim_partitioning_core::LocalPartition;
 
 use super::WaterLevelReporterProxy;
 
+#[derive(Clone, Copy)]
+struct Run {
+    start: usize,
+    len: usize,
+}
+
 #[allow(clippy::module_name_repetitions)]
 pub struct LiveWaterLevelReporterProxy<'l, 'p, R: Reporter, P: LocalPartition<'p, R>> {
     water_level: NonNegativeF64,
     slow_events: Vec<PackedEvent>,
+    tmp_events: Vec<PackedEvent>,
+    run: Run,
+    runs: Vec<Run>,
+    overflow: VecDeque<Run>,
+    sort_batch_size: usize,
     fast_events: Vec<PackedEvent>,
 
     local_partition: &'l mut P,
@@ -36,6 +47,8 @@ impl<'l, 'p, R: Reporter, P: LocalPartition<'p, R>> fmt::Debug
 
         fmt.debug_struct(stringify!(LiveWaterLevelReporterProxy))
             .field("water_level", &self.water_level)
+            .field("runs", &self.runs.len())
+            .field("overflow", &self.overflow.len())
             .field("slow_events", &EventBufferLen(self.slow_events.len()))
             .field("fast_events", &EventBufferLen(self.fast_events.len()))
             .finish()
@@ -46,35 +59,248 @@ impl<'l, 'p, R: Reporter, P: LocalPartition<'p, R>> Reporter
     for LiveWaterLevelReporterProxy<'l, 'p, R, P>
 {
     impl_report!(speciation(&mut self, speciation: MaybeUsed<R::ReportSpeciation>) {
-        if speciation.event_time < self.water_level {
-            self.slow_events.push(speciation.clone().into());
+        let speciation: PackedEvent = speciation.clone().into();
+
+        if speciation.event_time() < self.water_level {
+            let new_run = self.run.len > self.sort_batch_size; // self.slow_events.last().map_or(true, |prev| prev > &speciation);
+
+            if new_run {
+                let old_run = core::mem::replace(&mut self.run, Run {
+                    start: self.slow_events.len(),
+                    len: 1,
+                });
+                self.overflow.push_back(old_run);
+            } else {
+                self.run.len += 1;
+            }
+
+            self.slow_events.push(speciation);
         } else {
-            self.fast_events.push(speciation.clone().into());
+            self.fast_events.push(speciation);
         }
     });
 
     impl_report!(dispersal(&mut self, dispersal: MaybeUsed<R::ReportDispersal>) {
-        if dispersal.event_time < self.water_level {
-            self.slow_events.push(dispersal.clone().into());
+        let dispersal: PackedEvent = dispersal.clone().into();
+
+        if dispersal.event_time() < self.water_level {
+            let new_run = self.run.len > self.sort_batch_size; // self.slow_events.last().map_or(true, |prev| prev > &dispersal);
+
+            if new_run {
+                let old_run = core::mem::replace(&mut self.run, Run {
+                    start: self.slow_events.len(),
+                    len: 1,
+                });
+                self.overflow.push_back(old_run);
+            } else {
+                self.run.len += 1;
+            }
+
+            self.slow_events.push(dispersal);
         } else {
-            self.fast_events.push(dispersal.clone().into());
+            self.fast_events.push(dispersal);
         }
     });
 
     impl_report!(progress(&mut self, _progress: Ignored) {});
 }
 
+impl<'l, 'p, R: Reporter, P: LocalPartition<'p, R>> LiveWaterLevelReporterProxy<'l, 'p, R, P> {
+    #[inline]
+    fn collapse(&self, force_merge: bool) -> Option<usize> {
+        let n = self.runs.len();
+        if n >= 2
+            && (force_merge
+                || self.runs[n - 2].len <= self.runs[n - 1].len
+                || (n >= 3 && self.runs[n - 3].len <= self.runs[n - 2].len + self.runs[n - 1].len)
+                || (n >= 4 && self.runs[n - 4].len <= self.runs[n - 3].len + self.runs[n - 2].len))
+        {
+            if n >= 3 && self.runs[n - 3].len < self.runs[n - 1].len { Some(n - 3) } else { Some(n - 2) }
+        } else {
+            None
+        }
+    }
+
+    unsafe fn merge<T, F>(v: &mut [T], mid: usize, buf: *mut T, is_less: &mut F)
+    where
+        F: FnMut(&T, &T) -> bool,
+    {
+        unsafe fn get_and_increment<T>(ptr: &mut *mut T) -> *mut T {
+            let old = *ptr;
+            *ptr = unsafe { ptr.add(1) };
+            old
+        }
+
+        unsafe fn decrement_and_get<T>(ptr: &mut *mut T) -> *mut T {
+            *ptr = unsafe { ptr.sub(1) };
+            *ptr
+        }
+
+        // When dropped, copies the range `start..end` into `dest..`.
+        struct MergeHole<T> {
+            start: *mut T,
+            end: *mut T,
+            dest: *mut T,
+        }
+
+        impl<T> Drop for MergeHole<T> {
+            fn drop(&mut self) {
+                // `T` is not a zero-sized type, and these are pointers into a slice's elements.
+                unsafe {
+                    let len = self.end.sub_ptr(self.start);
+                    core::ptr::copy_nonoverlapping(self.start, self.dest, len);
+                }
+            }
+        }
+
+        let len = v.len();
+        let v = v.as_mut_ptr();
+        let (v_mid, v_end) = unsafe { (v.add(mid), v.add(len)) };
+
+        // The merge process first copies the shorter run into `buf`. Then it traces the newly copied
+        // run and the longer run forwards (or backwards), comparing their next unconsumed elements and
+        // copying the lesser (or greater) one into `v`.
+        //
+        // As soon as the shorter run is fully consumed, the process is done. If the longer run gets
+        // consumed first, then we must copy whatever is left of the shorter run into the remaining
+        // hole in `v`.
+        //
+        // Intermediate state of the process is always tracked by `hole`, which serves two purposes:
+        // 1. Protects integrity of `v` from panics in `is_less`.
+        // 2. Fills the remaining hole in `v` if the longer run gets consumed first.
+        //
+        // Panic safety:
+        //
+        // If `is_less` panics at any point during the process, `hole` will get dropped and fill the
+        // hole in `v` with the unconsumed range in `buf`, thus ensuring that `v` still holds every
+        // object it initially held exactly once.
+        let mut hole;
+
+        if mid <= len - mid {
+            // The left run is shorter.
+            unsafe {
+                core::ptr::copy_nonoverlapping(v, buf, mid);
+                hole = MergeHole { start: buf, end: buf.add(mid), dest: v };
+            }
+
+            // Initially, these pointers point to the beginnings of their arrays.
+            let left = &mut hole.start;
+            let mut right = v_mid;
+            let out = &mut hole.dest;
+
+            while *left < hole.end && right < v_end {
+                // Consume the lesser side.
+                // If equal, prefer the left run to maintain stability.
+                unsafe {
+                    let to_copy = if is_less(&*right, &**left) {
+                        get_and_increment(&mut right)
+                    } else {
+                        get_and_increment(left)
+                    };
+                    core::ptr::copy_nonoverlapping(to_copy, get_and_increment(out), 1);
+                }
+            }
+        } else {
+            // The right run is shorter.
+            unsafe {
+                core::ptr::copy_nonoverlapping(v_mid, buf, len - mid);
+                hole = MergeHole { start: buf, end: buf.add(len - mid), dest: v_mid };
+            }
+
+            // Initially, these pointers point past the ends of their arrays.
+            let left = &mut hole.dest;
+            let right = &mut hole.end;
+            let mut out = v_end;
+
+            while v < *left && buf < *right {
+                // Consume the greater side.
+                // If equal, prefer the right run to maintain stability.
+                unsafe {
+                    let to_copy = if is_less(&*right.sub(1), &*left.sub(1)) {
+                        decrement_and_get(left)
+                    } else {
+                        decrement_and_get(right)
+                    };
+                    core::ptr::copy_nonoverlapping(to_copy, decrement_and_get(&mut out), 1);
+                }
+            }
+        }
+        // Finally, `hole` gets dropped. If the shorter run was not fully consumed, whatever remains of
+        // it will now be copied into the hole in `v`.
+    }
+
+    fn sort_slow_events_step(&mut self, force_merge: bool) -> ControlFlow<()> {
+        let r = loop {
+            if let Some(r) = self.collapse(force_merge && self.overflow.is_empty()) {
+                break r;
+            }
+
+            let next_run = match self.overflow.pop_front() {
+                Some(next_run) => next_run,
+                None if self.run.len > 0 => core::mem::replace(&mut self.run, Run { start: self.slow_events.len(), len: 0 }),
+                None => return ControlFlow::Break(()),
+            };
+
+            // let Some(mut next_run) = self.overflow.pop_front() else {
+            //     return ControlFlow::Break(());
+            // };
+
+            // if next_run.len < self.sort_batch_size {
+            //     while next_run.len < self.sort_batch_size {
+            //         let Some(extra_run) = self.overflow.pop_front() else {
+            //             break;
+            //         };
+            //         next_run.len += extra_run.len;
+            //     }
+            //     self.slow_events[next_run.start..next_run.start+next_run.len].sort_unstable();
+            // }
+
+            self.slow_events[next_run.start..next_run.start+next_run.len].sort_unstable();
+
+            self.runs.push(next_run);
+        };
+
+        let left = self.runs[r];
+        let right = self.runs[r + 1];
+
+        let min_len = left.len.min(right.len);
+
+        if min_len > self.tmp_events.capacity() {
+            self.tmp_events.reserve(min_len - self.tmp_events.capacity());
+        }
+
+        unsafe {
+            Self::merge(
+                &mut self.slow_events[left.start..right.start + right.len],
+                left.len,
+                self.tmp_events.as_mut_ptr(),
+                &mut core::cmp::PartialOrd::lt,
+            );
+        }
+
+        self.runs[r] = Run { start: left.start, len: left.len + right.len };
+        self.runs.remove(r + 1);
+
+        ControlFlow::Continue(())
+    }
+}
+
 #[contract_trait]
 impl<'l, 'p, R: Reporter, P: LocalPartition<'p, R>> WaterLevelReporterProxy<'l, 'p, R, P>
     for LiveWaterLevelReporterProxy<'l, 'p, R, P>
 {
-    fn new(capacity: usize, local_partition: &'l mut P) -> Self {
+    fn new(capacity: usize, local_partition: &'l mut P, sort_batch_size: usize) -> Self {
         info!("Events will be reported using the live water-level algorithm ...");
 
         Self {
             water_level: NonNegativeF64::zero(),
             slow_events: Vec::with_capacity(capacity),
+            tmp_events: Vec::with_capacity(capacity),
+            run: Run { start: 0, len: 0 },
+            runs: Vec::new(),
+            overflow: VecDeque::new(),
             fast_events: Vec::with_capacity(capacity),
+            sort_batch_size,
 
             local_partition,
             _marker: PhantomData::<(&'p (), R)>,
@@ -85,9 +311,22 @@ impl<'l, 'p, R: Reporter, P: LocalPartition<'p, R>> WaterLevelReporterProxy<'l,
         self.water_level
     }
 
+    fn partial_sort_step(&mut self) -> ControlFlow<()> {
+        self.sort_slow_events_step(false)
+    }
+
     fn advance_water_level(&mut self, water_level: NonNegativeF64) {
+        let mut i = 0;
+
         // Report all events below the water level in sorted order
-        self.slow_events.sort();
+        while let ControlFlow::Continue(()) = self.sort_slow_events_step(true) {
+            if (i % 100) == 0 {
+                info!("{:?}", self);
+            }
+            i += 1;
+        }
+
+        debug_assert!(self.slow_events.is_sorted());
 
         for event in self.slow_events.drain(..) {
             match event.into() {
@@ -104,14 +343,28 @@ impl<'l, 'p, R: Reporter, P: LocalPartition<'p, R>> WaterLevelReporterProxy<'l,
             }
         }
 
+        self.runs.clear();
+        self.run = Run { start: 0, len: 0 };
+
         // Update the water level
         self.water_level = water_level;
 
         // Move fast events below the new water level into slow events
-        self.slow_events.extend(
-            self.fast_events
-                .drain_filter(|event| event.event_time() < water_level),
-        );
+        for event in self.fast_events.drain_filter(|event| event.event_time() < water_level) {
+            let new_run = self.run.len > self.sort_batch_size; // self.slow_events.last().map_or(true, |prev| prev > &event);
+
+            if new_run {
+                let old_run = core::mem::replace(&mut self.run, Run {
+                    start: self.slow_events.len(),
+                    len: 1,
+                });
+                self.overflow.push_back(old_run);
+            } else {
+                self.run.len += 1;
+            }
+
+            self.slow_events.push(event);
+        }
     }
 
     fn local_partition(&mut self) -> &mut P {
@@ -124,7 +377,9 @@ impl<'l, 'p, R: Reporter, P: LocalPartition<'p, R>> Drop
 {
     fn drop(&mut self) {
         // Report all events below the water level in sorted order
-        self.slow_events.sort();
+        while let ControlFlow::Continue(()) = self.sort_slow_events_step(true) {}
+
+        debug_assert!(self.slow_events.is_sorted());
 
         for event in self.slow_events.drain(..) {
             match event.into() {
@@ -141,6 +396,9 @@ impl<'l, 'p, R: Reporter, P: LocalPartition<'p, R>> Drop
             }
         }
 
+        self.runs.clear();
+        self.run = Run { start: 0, len: 0 };
+
         // Report all events above the water level in sorted order
         self.fast_events.sort();