Distribution::sample_iter could use owned values instead of references to be more general

[huonw]

Distribution::sample_iter and the DistIter type currently have signature:

fn sample_iter<'a, R>(&'a self, rng: &'a mut R) -> DistIter<'a, Self, R, T> { ... }

struct DistIter<'a, D: 'a, R: 'a, T> { ... }

The lifetimes and borrowing mean that it is impossible to construct a distribution or RNG inside a function and return the sample_iter, e.g. if one is trying to abstract some iterator pipeline

fn ten_dice_rolls_other_than_five<'a, R: Rng>(rng: &'a mut R) -> impl Iterator<i32> + 'a {
    Uniform::new_inclusive(1, 6)
        .sample_iter(rng)
        .filter(|x| *x != 5)
        .take(10)
}

However, these lifetimes aren't entirely necessary, because of the impls of Distribution for &D and RngCore for &mut R:

impl<'a, T, D: Distribution<T>> Distribution<T> for &'a D
impl<'a, R> RngCore for &'a mut R where R: RngCore + ?Sized

(Side note: the Distribution impl could use D: ?Sized like RngCore, to allow &Distribution<T> trait objects to impl Distribution<T>.)

These impls mean that the signature for sample_iter and definition of DistIter could be simplified to:

fn sample_iter<R>(self, rng: R) -> DistIter<Self, R, T> { ... }

struct DistIter<D, R, T> { ... }

The impls above specifically mean that if someone wants to pass a borrowed rng they still can, and similarly for a borrowed distribution. This also applies to Rng::sample_iter (both parameters should become owned).

See https://play.rust-lang.org/?gist=c050e777ed23fb40e320ba0eb340d031&version=stable&mode=debug&edition=2015 for a playground demonstrating the problem with the old API and how the new API fixes it without breaking (much) code.

It also shows that a by_ref function similar to Iterator::by_ref (and Read and Write) may be useful with this sort of API, to use as a postfix &/&mut:

trait RngCore {
    ...
    fn by_ref(&mut self) -> &mut Self { self }
}
trait Distribution<T> {
    ... 
    fn by_ref(&self) -> &Self { self }
}

[dhardy]

Hello @huonw, nice to see you've taken an interest in this project again!

We actually had a discussion on this a while back: #287. There's also some notes on #244.

IIRC the real problem with fn sample<R: Rng>(&self, rng: R) -> T is the lack of user-space reborrowing (see rust-lang/rfcs#1403 and rust-lang/rfcs#2364): dist.sample(&mut rng) is fine, but if rng: &mut Rng then dist.sample(rng) doesn't work because of the limited reborrowing.

That said, we could still switch to passing rng that way, but see #244.

As for sample_iter taking self by value, this seems less problematic, because distributions are less often passed as trait objects anyway. sample could have the same adjustment.

There are potential niggles though: does code-gen cause unnecessary copying in this case? And with non-Copy distributions, (&dist).sample(..) must be used to avoid consuming dist. Example:

let dist = Poisson::new(2.0);
let mut rng = ... ;
dist.sample(&mut rng);  // does this copy dist?
dist.sample(&mut rng);  // again?

let dist = WeightedIndex::new(..);  // is Clone but not Copy
dist.sample(&mut rng);  // consumes dist
// dist.sample(&mut rng);

[huonw]

Hello @huonw, nice to see you've taken an interest in this project again!

I've always been interested, and am so glad you've and the rest of the team have made it so much better than what it was. :)

I'm now in a position to interact with open source projects more than zero, but this is definitely your project still, and please don't think you have to defer to me just because I did some work on it years ago.

We actually had a discussion on this a while back: #287.

I'm afraid I'm not too sure how much of that discussion applies here: sample itself is different because it isn't returning an object involving Self or the rng.

IIRC the real problem with fn sample<R: Rng>(&self, rng: R) -> T is the lack of user-space reborrowing (see rust-lang/rfcs#1403 and rust-lang/rfcs#2364): dist.sample(&mut rng) is fine, but if rng: &mut Rng then dist.sample(rng) doesn't work because of the limited reborrowing.

Yeah, using an inferred generic param does mean that the &mut isn't reborrowed automatically, which is unfortunate, requiring people write sample(&mut *rng), sample(&mut rng) or sample(rng.by_ref()) as you say.

However, from my perspective, it is even more unfortunate to make the borrows compulsory to work around a small (hopefully temporary) language deficiency. Rust generally tends towards providing the most general APIs and then also providing glue APIs that allow them to work more nicely for specific cases (and, adds language features (like the reborrowing RFCs) once the pain points of using the general APIs in practice are understood).

For instance, all the iterator adaptors take self and their other parameters by value, even ones like take where one might want to do something like let first10 = iter.take(10).collect(); let rest = iter.collect(). Instead of forcing take to borrow the underlying iterator and thus not work for returning iterator chains, Iterator provides by_ref so the first instance can be written iter.by_ref().take(10).

(From my perspective, the semantic split between sample and sample_iter is quite similar to the split between Iterator by-ref functions like next & collect and by-value functions like take & map, and that seems like a good place from which to take API-design cues.)

There are potential niggles though: does code-gen cause unnecessary copying in this case?

It could... but, I suspect for most code any unnecessary byte copies won't be noticeable, since they'll just occur during construction/movement of the DistIter type, and not during iteration, which uses &mut.

Additionally, conversely, does the & form cause unnecessary pointer dereferences and inhibit optimizations? Compilers generally have a much easier time of things when there's plain values/no pointers, as that allows them to more easily do things like breaking up a struct into its individual fields and reasoning about each of them independently. As soon as an address is taken, this is much harder as the compiler first has to prove that the address is unnecessary (and it's even worse for pointers-to-pointers, like one gets with self inside the Iterator::next call).

Finally, I suspect the common case for this API will be taking the Rng either as &mut (it will be some global instance/parameter that is passed everywhere as &mut R) or as thread_rng, either of which is a single pointer, and only the (usually-small) Distribution by value.

And with non-Copy distributions, (&dist).sample(..) must be used to avoid consuming dist.

Yes, that's what by_ref is for (see line 81 in the playground link too). I do think it's slightly unintuitive the first few times it is used, but it's now a common API trend in Rust as mentioned at the end of the issue.

[dhardy]

I'm now in a position to interact with open source projects more than zero, but this is definitely your project still, and please don't think you have to defer to me just because I did some work on it years ago.

Glad to hear you have more flexibility now! As far as I'm concerned, this isn't my project, but just one I've driven for a while because it wasn't getting much input and because I had a few ideas. But @pitdicker probably gave as much input as I did and several others have been involved too. I just invited you as a collaborator.

[dhardy]

Additionally, conversely, does the & form cause unnecessary pointer dereferences and inhibit optimizations?

Rust has strong no-alias rules which put it ahead of most languages. There may be other applicable optimisations though.

Okay, seems to me you have a good argument here. 👍

@vks any comments?

[burdges]

I think this sounds reasonable, but doc comment examples should use the (&dist).sample_iter(&mut rng) form to sample twice, or maybe use reborrowing from variables, ala

let dist = &dist;
let rng = &mut rng;
let x = dist.sample_iter(rng).collect();
let y = dist.sample_iter(rng).collect();

There is an orthogonal concern that cryptographic Rngs like ChaCha should never be Copy because duplicating one accidentally tends to be insecure. It's not common for sampling distributions to require a CSPRNG, so this API change does not make this concern significantly more serious.

We do however sample from distributions using CSPRNGs in lattice based cryptography, anonymity tools, and byzantine consensus algorithm, ala proof-of-stake, so probably CryptoRng should warn against Copy.

We cannot impose CryptoRng: Drop of course because &mut R: CryptoRng if R: CryptoRng, and CryptoRngs like OsRng could employ interior mutability to provide a safe Copy interface.

Also, I expect the same reasoning applied to all Rngs because duplicating one accidentally tends to produce unwanted statistical dependence, so maybe all Rng itself should warn against Copy.

[huonw]

Rust has strong no-alias rules which put it ahead of most languages. There may be other applicable optimisations though.

These are definitely good and work well for some cases, but in practice the benefits don't apply to every situation. In particular, I believe rustc currently only really informs LLVM of pointer-aliasing information for pointers that are function arguments (at the LLVM level), which means it doesn't work for pointers loaded from memory (including through other function arguments).

For instance, in the following, x (or r.x) and y (or r.y) should never alias. This means that loads from x can be moved past writes to y, and thus, in a perfect world, all of the following should have a single load from x. However, that optimization currently can't be applied when x and y themselves have to be loaded from memory.

// All of these are testing if the `a` load can be
// moved past the `10` write to turn the return value
// into `2 * a`

// take two pointers as arguments directly
pub fn args(x: &i32, y: &mut i32) -> i32 {
    let a = *x;
    *y = 10;
    a + *x
}

pub struct Refs<'a> {
    x: &'a i32,
    y: &'a mut i32,
}
// take the pointers as a by-value struct (which
// gets exploded into individual pointer arguments
// like `args`)
pub fn struct_val(r: Refs) -> i32 {
    let a = *r.x;
    *r.y = 10;
    a + *r.x
}

// now take the struct by reference, so the pointers
// need to be loaded from memory
pub fn struct_ref(r: &mut Refs) -> i32 {
    let x = r.x;
    let y = &mut *r.y;

    let a = *r.x;
    *r.y = 10;
    a + *r.x
}

example::struct_val:
        mov     eax, dword ptr [rdi] ; let a = *r.x;
        mov     dword ptr [rsi], 10  ; *r.y = 10;
        add     eax, eax             ; a + a (optimisation worked!)
        ret

example::struct_ref:
        mov     rcx, qword ptr [rdi]     ; let x = r.x;
        mov     rdx, qword ptr [rdi + 8] ; let y = &mut *r.y;
        mov     eax, dword ptr [rcx]     ; let a = *x;
        mov     dword ptr [rdx], 10      ; *y = 10;
        add     eax, dword ptr [rcx]     ; a + *x (failed!)
        ret

example::args:
        jmp     example::struct_val@PLT ; this function is identical to struct_val, they're merged

https://rust.godbolt.org/z/HezTT4