First Rust program - squares which are sum of squares

Question

I've been programming in C++ for over a decade and am thinking of learning Rust. Finding a number whose square is the sum of three different positive squares in at least two disjoint ways, which I used in test code in a C++ program, I wrote a Rust program to find out what other numbers' squares can be written as the sum of three different positive squares. Here's the program; how can it be improved?

fn main()
{
  for a in 1..101
  {
    let mut b=1;
    while a*a-b*b>=0
    {
      let mut c=1;
      while a*a-b*b-c*c>=0 && c<b
      {
        let mut d=1;
        while a*a-b*b-c*c-d*d>=0 && d<c
        {
          if a*a-b*b-c*c-d*d==0
          {
            println!("{}²={}²+{}²+{}²",a,b,c,d);
          }
          d+=1;
        }
        c+=1;
      }
      b+=1;
    }
  }
}

Answer 1

How can it be improved? Well, it all depends on your goals with the code. I'll assume the goal is to make it prettier and more in line with typical Rust style.

First of all, while loops are rare in Rust code and you should replace them with for. This is because while works on your hand-rolled logic, whereas for .. in .. works with external iteration that is composable and idiomatic Rust.

I see you have one for loop already. Let's turn all others into for.

You always start with 1, so the first iterator will be (1..). We keep taking elements of the range as long as a condition is satisfied, so the next element of the iteration pipeline will be .take_while(|&b| ...). This Rust standard library method best fits the situation.

fn main()
{
  for a in 1..101
  {
    for b in (1..).take_while(|&b| a*a-b*b >= 0)
    {
      for c in (1..b).take_while(|&c| a*a-b*b-c*c>=0)
      {
        for d in (1..c).take_while(|&d| a*a-b*b-c*c-d*d>=0)
        {
          if a*a-b*b-c*c-d*d==0
          {
            println!("{}²={}²+{}²+{}²",a,b,c,d);
          }
        }
      }
    }
  }
}

This code is much more idiomatic already. Next up, I noticed that the intent of the two inner loops is to iterate until we hit the number above, so let's make two simple changes:

• for c in (1..b).take_while(|&c| a*a-b*b-c*c>=0)
• for d in (1..c).take_while(|&d| a*a-b*b-c*c-d*d>=0)

For readability, we can Don't-Repeat-Yourself by making a tiny function.

fn square_diff(a: i32, b: i32, c: i32, d: i32) -> i32 { a*a-b*b-c*c-d*d }

For all instances of a*a-b*b or a*a-b*b-c*c or ..., we replace them with square_diff(a, b, 0, 0) etc.

One more observation: (1..).take_while(|b| square_diff(a, b, 0, 0) >= 0) is equivalent to 1..=a.

Next up, we don't do C-style indentation in Rust. We put a single space after every comma in arguments. The result so far is:

fn main() {
  for a in 1..101 {
    for b in 1..=a {
      for c in (1..b).take_while(|&c| square_diff(a, b, c, 0) >= 0) {
        for d in (1..c).take_while(|&d| square_diff(a, b, c, d) >= 0) {
          if square_diff(a, b, c, d) == 0 {
            println!("{}²={}²+{}²+{}²", a, b, c, d);
          }
        }
      }
    }
  }
}

fn square_diff(a: i32, b: i32, c: i32, d: i32) -> i32 { a*a - b*b - c*c - d*d }

I tried to test whether these abstractions really reduce to zero-cost, but alas, Rust@godbolt failed me. Benchmarking is left as a small exercise.

Good luck on your Rust journeys.

Answer 2

The main readability improvement I see is using for and ranges for all of your loops — this is debatable but I see it as highlighting that the overall structure of your program is about iterating over numbers one by one (there is no 'skipping' where you write c += b or anything like that).

Secondarily, I reformatted the code with rustfmt so as to follow standard style.

fn main() {
    for a in 1..=100 {
        for b in (1..).take_while(|b| a * a - b * b >= 0) {
            for c in (1..b).take_while(|c| a * a - b * b - c * c >= 0) {
                for d in (1..c).take_while(|d| a * a - b * b - c * c - d * d >= 0) {
                    if a * a - b * b - c * c - d * d == 0 {
                        println!("{}²={}²+{}²+{}²", a, b, c, d);
                    }
                }
            }
        }
    }
}

I chose to use take_while to express the stopping conditions that weren't simply upper limits of the numeric ranges, but this is equivalent to explicitly stopping the loop such as by

        ...
        for b in 1.. {
            if !(a * a - b * b >= 0) {
                break;
            }
            ...

After that, the next thing that's bothering me is all the repeated subexpressions. Arguably, keeping them as they are makes the program mathematically clearer, but I think it will further help readability of the algorithm if the expressions are less repeated.

fn main() {
    for a in 1..=100 {
        let target = a * a;
        for b in (1..).take_while(|b| target - b * b >= 0) {
            let target = target - b * b;
            for c in (1..b).take_while(|c| target - c * c >= 0) {
                let target = target - c * c;
                for d in (1..c).take_while(|d| target - d * d >= 0) {
                    if target - d * d == 0 {
                        println!("{}²={}²+{}²+{}²", a, b, c, d);
                    }
                }
            }
        }
    }
}

Now it is clear that the program is doing the same search three times nested. This would be a perfectly fine time to stop and say we've made it readable.

But what if we'd like to not actually have to write the search three times? We can do that by moving the repeated algorithm into a function that produces an iterator.

fn squares_subtracted(minuend: i32, stop_before: i32) -> impl Iterator<Item = (i32, i32)> {
    (1..stop_before)
        .map(move |root| {
            let result = minuend - root * root;
            (root, result)
        })
        .take_while(|(_, result)| *result >= 0)
}

fn main() {
    for a in 1..=100 {
        for (b, target) in squares_subtracted(a * a, a) {
            for (c, target) in squares_subtracted(target, b) {
                for (d, target) in squares_subtracted(target, c) {
                    if target == 0 {
                        println!("{}²={}²+{}²+{}²", a, b, c, d);
                    }
                }
            }
        }
    }
}

We could go further and refactor into a form where instead of nested loops, there's one iterator adapter applied three times. However, I don't think that would actually serve readability (or maintainability); it would only make sense if you are working in a larger context of mathematical brute-force searches in general, where you might want to use the same operation in several different combinations.