Struggling with closures and lifetimes in Rust

Question

I'm trying to port a little benchmark from F# to Rust. The F# code looks like this:

let inline iterNeighbors f (i, j) =
  f (i-1, j)
  f (i+1, j)
  f (i, j-1)
  f (i, j+1)

let rec nthLoop n (s1: HashSet<_>) (s2: HashSet<_>) =
  match n with
  | 0 -> s1
  | n ->
      let s0 = HashSet(HashIdentity.Structural)
      let add p =
        if not(s1.Contains p || s2.Contains p) then
          ignore(s0.Add p)
      Seq.iter (fun p -> iterNeighbors add p) s1
      nthLoop (n-1) s0 s1

let nth n p =
  nthLoop n (HashSet([p], HashIdentity.Structural)) (HashSet(HashIdentity.Structural))

(nth 2000 (0, 0)).Count

It computes the nth-nearest neighbor shells from an initial vertex in a potentially infinite graph. I used something similar during my PhD to study amorphous materials.

I've spent many hours trying and failing to port this to Rust. I have managed to get one version working but only by manually inlining the closure and converting the recursion into a loop with local mutables (yuk!).

I tried writing the iterNeighbors function like this:

use std::collections::HashSet;

fn iterNeighbors(f: &F, (i, j): (i32, i32)) -> ()
where
    F: Fn((i32, i32)) -> (),
{
    f((i - 1, j));
    f((i + 1, j));
    f((i, j - 1));
    f((i, j + 1));
}

I think that is a function that accepts a closure (that itself accepts a pair and returns unit) and a pair and returns unit. I seem to have to double bracket things: is that correct?

I tried writing a recursive version like this:

fn nthLoop(n: i32, s1: HashSet<(i32, i32)>, s2: HashSet<(i32, i32)>) -> HashSet<(i32, i32)> {
    if n == 0 {
        return &s1;
    } else {
        let mut s0 = HashSet::new();
        for &p in s1 {
            if !(s1.contains(&p) || s2.contains(&p)) {
                s0.insert(p);
            }
        }
        return &nthLoop(n - 1, s0, s1);
    }
}

Note that I haven't even bothered with the call to iterNeighbors yet.

I think I'm struggling to get the lifetimes of the arguments correct because they are rotated in the recursive call. How should I annotate the lifetimes if I want s2 to be deallocated just before the returns and I want s1 to survive either when returned or into the recursive call?

The caller would look something like this:

fn nth<'a>(n: i32, p: (i32, i32)) -> &'a HashSet<(i32, i32)> {
    let s0 = HashSet::new();
    let mut s1 = HashSet::new();
    s1.insert(p);
    return &nthLoop(n, &s1, s0);
}

I gave up on that and wrote it as a while loop with mutable locals instead:

fn nth<'a>(n: i32, p: (i32, i32)) -> HashSet<(i32, i32)> {
    let mut n = n;
    let mut s0 = HashSet::new();
    let mut s1 = HashSet::new();
    let mut s2 = HashSet::new();
    s1.insert(p);
    while n > 0 {
        for &p in &s1 {
            let add = &|p| {
                if !(s1.contains(&p) || s2.contains(&p)) {
                    s0.insert(p);
                }
            };
            iterNeighbors(&add, p);
        }
        std::mem::swap(&mut s0, &mut s1);
        std::mem::swap(&mut s0, &mut s2);
        s0.clear();
        n -= 1;
    }
    return s1;
}

This works if I inline the closure by hand, but I cannot figure out how to invoke the closure. Ideally, I'd like static dispatch here.

The main function is then:

fn main() {
    let s = nth(2000, (0, 0));
    println!("{}", s.len());
}

So... what am I doing wrong? :-)

Also, I only used HashSet in the F# because I assume Rust doesn't provide a purely functional Set with efficient set-theoretic operations (union, intersection and difference). Am I correct in assuming that?

Struggling with closures and lifetimes in Rust

2 Answers