Josephus permutation - follow up Follow up of this question

Changes:


The pop_min function now panics if the tree is empty.
Moved the size attribute to the Tree struct.
Applied most of the suggestions given by the answers.
Added some other functions.
Refactored some functions.
Moved everything to the same file.




extern crate itertools;
use itertools::Unfold;
use std::cmp::Ordering;


pub fn permutation(size: u32, m: u32) -&gt; Box&lt;Iterator&lt;Item = u32&gt;&gt; {
    let mut tree = Tree::new(1, size);
    let x = Unfold::new(1, move |a| {
        *a = (*a + m - 2) % tree.size + 1;
        Some(tree.pop_rank(*a))
    });
    Box::new(x.take(size as usize))
}

#[derive(Debug)]
struct Node {
    data: u32,
    left: Tree,
    right: Tree,
}

#[derive(Debug)]
struct Tree {
    size: u32,
    root: Option&lt;Box&lt;Node&gt;&gt;,
}

impl Tree {
    fn new(from: u32, to: u32) -&gt; Tree {
        if from &gt; to {
            return Tree { root: None, size: 0, };
        }
        let mid = from + (to - from) / 2;
        let node = Node {
            data: mid,
            left: Tree::new(from, mid - 1),
            right: Tree::new(mid + 1, to),
        };

        Tree {
            size: 1 + node.left.size + node.right.size,
            root: Some(Box::new(node)),
        }
    }

    fn find_rank(&amp;mut self, rank: u32) -&gt; &amp;mut Tree {
        let r = self.root
                    .as_mut()
                    .expect("rank out of range")
                    .left
                    .size + 1;

        self.size -= 1;
        match rank.cmp(&amp;r) {
            Ordering::Equal =&gt; self,
            Ordering::Less =&gt; self.as_mut().left.find_rank(rank),
            Ordering::Greater =&gt; self.as_mut().right.find_rank(rank - r),
        }
    }

    fn as_mut(&amp;mut self) -&gt; &amp;mut Box&lt;Node&gt; {
        self.root.as_mut().unwrap()
    }

    fn as_ref(&amp;self) -&gt; &amp;Box&lt;Node&gt; {
        self.root.as_ref().unwrap()
    }

    fn empty(&amp;self) -&gt; bool {
        self.root.is_none()
    }

    fn delete_node(&amp;mut self) {
        *self = self.root
                    .take()
                    .map(|mut x| {
                        if x.left.empty() {
                            x.right
                        } else if x.right.empty() {
                            x.left
                        } else {
                            x.data = x.right.pop_min();
                            Tree {
                                root: Some(x),
                                size: self.size,
                            }
                        }
                    })
                    .expect("Can't delete None");
    }

    fn pop_min(&amp;mut self) -&gt; u32 {

        if self.empty() {
            panic!("underflow");
        }

        if self.as_ref().left.empty() {
            let data = self.as_ref().data;
            *self = self.root.take().unwrap().right;
            data
        } else {
            self.size -= 1;
            self.as_mut().left.pop_min()
        }
    }

    fn pop_rank(&amp;mut self, rank: u32) -&gt; u32 {
        let ranked = self.find_rank(rank);
        let data = ranked.as_ref().data;
        ranked.delete_node();
        data
    }
}
This mostly looks OK, but I really don't like

*self = self.root
            .take()
            .map(|mut x| {
                if x.left.empty() {
                    x.right
                } else if x.right.empty() {
                    x.left
                } else {
                    x.data = x.right.pop_min();
                    Tree {
                        root: Some(x),
                        size: self.size,
                    }
                }
            })
            .expect("Can't delete None");


You can shift up the expect to get

let mut node = self.root.take().expect("Can't delete None");
*self =
    if node.left.empty() {
        node.right
    } else if node.right.empty() {
        node.left
    } else {
        node.data = node.right.pop_min();
        Tree {
            root: Some(node),
            size: self.size,
        }
    };


which is already much nicer. Then deal with the error and naming:

fn delete_root_node(&amp;mut self) {
    let mut root = self.root.take().expect("Empty tree has no root");
    *self =
        if root.left.empty() {
            root.right
        } else if root.right.empty() {
            root.left
        } else {
            root.data = root.right.pop_min();
            Tree {
                root: Some(root),
                size: self.size,
            }
        };
}


Then I'd flatten it back out

fn delete_root_node(&amp;mut self) {
    let mut root = self.root.take().expect("Empty tree has no root");

    if root.left.empty() {
        *self = root.right
    } else if root.right.empty() {
        *self = root.left
    } else {
        root.data = root.right.pop_min();
        self.root = Some(root);
    };
}


Rust's tools are great, but they shouldn't be abused. In this case, a simple imperative style reads much nicer than the semi-functional transformations, so should be preferred.

More rightward drift can be dealt with in find_rank:

fn find_rank(&amp;mut self, rank: u32) -&gt; &amp;mut Tree {
    let right_offset = self.as_mut().left.size + 1;

    self.size -= 1;
    match rank.cmp(&amp;right_offset) {
        Ordering::Equal =&gt; self,
        Ordering::Less =&gt; self.as_mut().left.find_rank(rank),
        Ordering::Greater =&gt; self.as_mut().right.find_rank(rank - right_offset),
    }
}


Now, find_rank is doing something really horrible:

self.size -= 1;


This breaks the Tree's invariants. Much better would be to call this pop_rank and move pop_rank inside of it. You can make this a bit nicer by having pop_root return a u32 and just depending on that, though.

Then we have a nice chain of functions

fn pop_rank(&amp;mut self, rank: u32) -&gt; u32 {
    let left_size = self.as_mut().left.size;
    self.size -= 1;

    match rank.cmp(&amp;(1 + left_size)) {
        Ordering::Equal   =&gt; {
            self.size += 1;
            self.pop_root()
        },
        Ordering::Less    =&gt; self.as_mut().left.pop_rank(rank),
        Ordering::Greater =&gt; self.as_mut().right.pop_rank(rank - 1 - left_size),
    }
}

fn pop_root(&amp;mut self) -&gt; u32 {
    let mut root = self.root.take().expect("Empty tree has no root");
    let data = root.data;

    if root.left.empty() {
        *self = root.right
    } else if root.right.empty() {
        *self = root.left
    } else {
        root.data = root.right.pop_min();
        self.root = Some(root);
        self.size -= 1;
    };

    data
}

fn pop_min(&amp;mut self) -&gt; u32 {
    if self.as_ref().left.empty() {
        let data = self.as_ref().data;
        *self = self.root.take().unwrap().right;
        data
    } else {
        self.size -= 1;
        self.as_mut().left.pop_min()
    }
}


The restoring of self.size in pop_rank before calling pop_root might slightly increase computational burden but it means that each function only relies on the data structure's basic invariants to work properly.

Note that pop_min doesn't actually need to panic! as that's captured by self.as_ref's auto-panic.

as_ref and as_mut are misleading names since they actually give handles to root; better would be get_root and get_root_mut. I'd also swap pop_min for pop_first for clarity.

I'd also change m to step in permutation.



I've been transfixed by this problem for a while now, so I'm going to introduce my thoughts about it. Very little of what follows has anything to do with your implementation, but hopefully someone at least finds it a little interesting.

It's worth noting that Box-heavy structures tend to be slow. A much more efficient tree structure to work with could be a Vec&lt;Vec&lt;usize&gt;&gt; like so:

tree[0]               13
tree[1]        8               5
tree[2]    4       4       4       1
tree[3]  2   2   2   2   2   2   1
tree[4] 1 1 1 1 1 1 1 1 1 1 1 1 1


tree[0..3] are used to guide you to an index in the bottom. "Using" the fifth value would result in the tree modified like so

tree[0]               12
tree[1]        7               5
tree[2]    4       3       4       1
tree[3]  2   2   1   2   2   2   1
tree[4] 1 1 1 1 0 1 1 1 1 1 1 1 1


Then it's worth noting that you never actually need to read the right branches on traversal, so those can be discarded:

tree[0]               13
tree[1]        8               .
tree[2]    4       .       4       .
tree[3]  2   .   2   .   2   .   1
tree[4] 1 . 1 . 1 . 1 . 1 . 1 . 1


Then move the size into the data structure itself, since that needs to be read frequently in the algorithm and not actually during traversal.

tree[0]        8               .
tree[1]    4       .       4       .
tree[2]  2   .   2   .   2   .   1
tree[3] 1 . 1 . 1 . 1 . 1 . 1 . 1


I've found that just this gives a noticeable but underwhelming speed-up factor of 3-4x; I was expecting with all the Boxing for that to be closer to 10. The problem seems to be that the branches are still unpredictable and the behaviour is still cache-thrashing; the main improvement came not from the change in format but from not writing to the right branch (not the eliding step but just not writing to it). Plus, allocators are likely to put all those boxed values next to each other, and the lack of space efficiency doesn't really hurt when doing psuedo-random access anyway.

Then note that with something like a (14345577, 15347)-permutation, \$\log 14345577\$ is much greater than \$\log 15347\$. You can use this to your advantage by ignoring the top of the tree:

tree[ 0]             8388608 ...
tree[ 1]           4194304 ...
...
tree[ 8]       32768 ...
tree[ 9]     16384 ...
...
tree[22]   2 ...
tree[23] 1 ...


You can actually discard tree[0..8]; jumps are normally going to be no greater than steps of 15347 times the "sparsity" of the data. When the sparsity gets so large that you are taking jumps that large, you can "recompact" the whole array in \$O(n)\$ time. Since you only need to do this when density drops somewhere below 50%, this is a relatively cheap way of moving from a total \$O(n \log n)\$ to \$O(n \log m)\$ time complexity (resizing each time capacity halves results in only \$O(n)\$ extra work).

Note that the top level then becomes a bunch of roots into sub-trees, not a single tree. Thus you can't ignore the right-nodes, or you'll be losing the size of the tree. Conceptually, those right nodes are calculated from the node above, so if there's no node above you aren't able to drop the value.

An extreme example would be a (10000000, 3)-permutation, where the entire tree is clearly overkill. This tree trimming works well there.

Then it's worth looking at those bottom rows, to see if we can do something more intelligent about them:

tree[X]                            64   ...
...
tree[X]    4       4       4       4    ...
tree[X]  2   .   2   .   2   .   2   .  ...
tree[X] 1 . 1 . 1 . 1 . 1 . 1 . 1 . 1 . ...


Despite the elided right branches, this is still a ton of data to express 64 bits of data. In fact, it uses 63 times that amount of data even after elision. Much more effective is to have an actual 64 bit integer.

0b1111111111111111111111111...


This can still be efficiently binary searched by utilizing a specialized routine, bit shifts, masks and u64::count_ones. In fact, on newer x86 processors you don't even have to; Jukka Suomela points out that it's as simple as

pub fn position_nth_set_bit(value: u64, mut bit: u32) -&gt; u32 {
    // Holy smokes
    bmi_pdep_64(1 &lt;&lt; bit, value).trailing_zeros()
}


Sadly I think right now I need to use asm!, but that's no show-stopper. If pdep isn't available, the binary search with count_ones is still pretty fast - but it's certainly not nearly as fast.

In fact, this also makes the compacting step that we need to elide the top of the tree much simpler; the bit-set acts as a mask over our values that we can just iterate.



Let's quickly calculate what the overhead is. Our entire tree looks like

roots           x              x
tree[X]     x       .       x
tree[X]   x   .   x   .   x   .
tree[X]  x . x . x . x . x . x .
masks    bbbbbbbbbbbbbbbbbbbbbbbb
values   VVVVVVVVVVVVVVVVVVVVVVVV  ← We don't need this until after compacting


values costs a usize per output value, but can actually be ignored until the first compaction as before then it's an identity map (values[k] == k at first).

masks introduces one further bit per value.

If you look at tree and roots the right way, you can see you can "slide" the xs down the right branches to fill exactly one row (I'll re-use this property later). The bottom-most tree[X] is at the 128 level (since masks.count_ones() works at the 64 level), so that's a usize of bits spent per 128 values.

Ergo we have 1-2 extra bits per value, and only half of the values actually need to be stored. So our space complexity is pretty close to n * usize::BITS / 2. If m is expected to approach n in size or space matters more than the move from \$O(n \log n)\$ to \$O(n \log m)\$, you can instead use masks without trimming the top of the tree, so no values array is needed, to get a space cost of roughly just 2 * n bits.

A pleasing hybrid might be obtained by making tree a little bit taller, so one has to resize slightly less frequently. Two extra layers means two extra branches per iteration but delays resizing until the density is about 1/8, leaving memory consumption of about n * usize::BITS / 8 while still having \$O(n \log m)\$ operations. Plus, doing the first resize early hurts as before it values is not used and the later resizes are smaller. On the other hand, a few extra branches is expensive when the branches are unpredictable, so it actually makes sense to keep the tree shallow and do a linear search over the roots. It's worth noting that resizing can make the tree smaller and isn't actually expensive, so the trade to do it later isn't so clear-cut. The precise parameters are tunable.

This is starting to feel optimal, since \$\log m\$ is about the smallest cost for a jump one can normally reasonably expect and we're approaching a small number of bits of overhead. However, we might be able to make the data structure more optimal. One possible improvement is to take the "compact" view of the tree we used to calculate its size: slide each value down the "right-hand slide":

roots           a              b
tree[X]     c       ↘       d
tree[X]   e   ↘   f   ↘   g   ↘
tree[X]  x ↘ x ↘ x ↘ x ↘ x ↘ x ↘
masks    bbbbbbbbbbbbbbbbbbbbbbbb
values   VVVVVVVVVVVVVVVVVVVVVVVV


to get

tree     x e x c x f x a x g x d · · · b
masks    bbbbbbbbbbbbbbbbbbbbbbbb
values   VVVVVVVVVVVVVVVVVVVVVVVV


Note that traversal is now just a binary search over a sub-range of the array. This might look tough to traverse, but it's actually a fairly mechanical transformation.

Generating this is more tough, but it does have the lovely property that each modulo is

1, 2, 1, 4, 1, 2, 1, 8, 1, 2, 1, 4, 1, 2, ...


which you can take the \$\log_2\$ of to get

0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, ...


OEIS helpfully adds that,


  a(n) is the number of 0's at the end of n when n is written in base 2


and then you note that n &gt;&gt; (n.trailing_zeros() + 1) gives how many times you've seen this value before, which lets you deal with the leftover parts.

Note that the modulo is for a "cropped" root, and so you need to take the min against the top-level modulo you're using, and you also need to add 7 before then to multiply all the modulos by 128 (because the bottom groups of 64 are dealt with in masks).



With my implementation, for very large \$m\$ (same order of magnitude as \$n\$), the tree pruning is entirely ineffective (and is just pointless overhead), so the only performance gains come from the change in how iteration happens, the move to iteration over a contiguous vector and the masks array removing the last few layers of the tree. The move to a contiguous vector is actually rather ineffective as a speed enhancer since the Boxes are probably contiguous anyway and the major overhead is in the unpredictable branching. There's roughly a factor 6 improvement over your code here.

For "medium"-sized \$m\$ (eg. \$n\$ is 14345577, \$m\$ is 15347), tree pruning comes into play and speed improves 3-fold. Your code does not, so the speed advantage increases. There's roughly a factor 20 improvement over your code here.

For small \$m\$ (eg. 3), the tree pruning pushes further increases speed significantly, resulting in massive decreases in time. Further, the BMI instructions at the bottom of the tree end up improving speed significantly. However, your code also becomes very predictable, since branches repeat frequently, so the cost there actually falls too and the overall relative speed improvement is not as great as the time reduction alone would imply. There's roughly a factor 32 improvement over your code here.



Honestly I was hoping for more, but I can see why the improvements were only this good. It should be possible to improve this a lot by using a B-tree instead of a binary tree, by effectively replacing

roots           x              x
tree[X]     x       .       x
tree[X]   x   .   x   .   x   .
tree[X]  x . x . x . x . x . x .
masks    bbbbbbbbbbbbbbbbbbbbbbbb
values   VVVVVVVVVVVVVVVVVVVVVVVV


with

roots       x       x       x       x
tree[X]       .       .       .
tree[X]  x x x . x x x . x x x .
masks    bbbbbbbbbbbbbbbbbbbbbbbb
values   VVVVVVVVVVVVVVVVVVVVVVVV


which should improve branch predictability by using intermittent unrolled linear searches. However, the reduction from doing this the "obvious" way is slight as you still end up with poorly predictable branches over spread-out locations.

Ideally, one should reorganize the values as

roots                      x x x x
tree[X]  x x x|x x x|x x x|
masks    bbbbbbbbbbbbbbbbbbbbbbbb
values   VVVVVVVVVVVVVVVVVVVVVVVV


instead.

Another oversight I've been purposefully ignoring is SIMD. Small SIMD vectors for usize-sized data makes this relatively weak, but lower down the tree a u16x8 vector is likely extremely useful, and basically makes an 8-wide B-tree optimal. Newer AVX instructions would do far better, and allow the whole tree to be heavily vectorized. I'm going to continue to ignore this, though, 'cause SIMD with Rust is not ideal right now.