Lift recursive function to for loop - rust

Question

I have a simple recursive function that uses borrows in Rust. I want to convert it into a for loop. I have made a trimmed down example.
enum Counter {
Val(usize),
}
fn count(n: &Counter) -> Result<usize, String> {
match n {
Counter::Val(a) => {
match a {
0..=9 => {
let b = Counter::Val(a+1);
Ok(count(&b)? + 1)
},
10 => Ok(0),
_ => Err("Out of range".to_string()),
}
},
}
}
fn main() {
print!("{:?}\n", count(&Counter::Val(8)));
}
If I convert without being careful, I end up with
enum Counter {
Val(usize),
}
fn count(n: &Counter) -> Result<usize, String> {
enum State<'a> {
Start(&'a Counter),
Val0(Counter),
}
let mut stack = vec![State::Start(n)];
let mut ret = None;
loop {
match (ret, stack.pop()) {
(Some(r), None) => return Ok(r),
(None, Some(State::Start(n))) => {
match n {
Counter::Val(a) => {
match a {
0..=9 => {
let b = Counter::Val(a+1);
stack.push(State::Val0(b));
let br = match &stack[stack.len()-1] {
State::Val0(a) => a,
_ => return Err("Wrong state type".to_string()),
};
stack.push(State::Start(br));
ret = None;
},
10 => { ret = Some(0); },
_ => return Err("Out of range".to_string()),
}
},
}
},
(Some(r), Some(State::Val0(_))) => {
ret = Some(r + 1);
},
_ => return Err("Invalid state".to_string()),
}
}
}
fn main() {
print!("{:?}\n", count(&Counter::Val(8)));
}
This is correctly rejected by the borrow checker.
error[E0502]: cannot borrow `stack` as mutable because it is also borrowed as immutable
--> a1.rs:15:17
|
15 | match (ret, stack.pop()) {
| ^^^^^^^^^^^ mutable borrow occurs here
...
23 | stack.push(State::Val0(b));
| -------------------------- immutable borrow later used here
24 | let br = match &stack[stack.len()-1] {
| ----- immutable borrow occurs here
error[E0502]: cannot borrow `stack` as mutable because it is also borrowed as immutable
--> a1.rs:23:17
|
23 | stack.push(State::Val0(b));
| ^^^^^^^^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
24 | let br = match &stack[stack.len()-1] {
| -----
| |
| immutable borrow occurs here
| immutable borrow later used here
error[E0502]: cannot borrow `stack` as mutable because it is also borrowed as immutable
--> a1.rs:28:17
|
24 | let br = match &stack[stack.len()-1] {
| ----- immutable borrow occurs here
...
28 | stack.push(State::Start(br));
| ^^^^^^----^^^^^^^^^^^^^^^^^^
| | |
| | immutable borrow later used by call
| mutable borrow occurs here
The stack might need to grow, leading to a reallocation, invalidating all of the borrows. So, this program is incorrect. To account for that, I can use Box to
have pointers to the heap, that won't change.
However, this still gets a similar error. To get
it to compile, I need to add an unsafe to
circumvent the borrow checker.
enum Counter {
Val(usize),
}
fn count(n: &Counter) -> Result<usize, String> {
enum State<'a> {
Start(&'a Counter),
Val0(Box<Counter>),
}
let mut stack = vec![State::Start(n)];
let mut ret = None;
loop {
match (ret, stack.pop()) {
(Some(r), None) => return Ok(r),
(None, Some(State::Start(n))) => {
match n {
Counter::Val(a) => {
match a {
0..=9 => {
let b = Counter::Val(a+1);
stack.push(State::Val0(Box::new(b)));
let br = match &stack[stack.len()-1] {
State::Val0(a) => &**a,
_ => return Err("Wrong state type".to_string()),
};
let br0 = unsafe { &*(br as *const Counter) };
stack.push(State::Start(br0));
ret = None;
},
10 => { ret = Some(0); },
_ => return Err("Out of range".to_string()),
}
},
}
},
(Some(r), Some(State::Val0(_))) => {
ret = Some(r + 1);
},
_ => return Err("Invalid state".to_string()),
}
}
}
fn main() {
print!("{:?}\n", count(&Counter::Val(8)));
}
To avoid using unsafe, I could also use a separate data stack.
enum Counter {
Val(usize),
}
fn count(n: &Counter) -> Result<usize, String> {
enum State {
Start(usize),
Val0,
}
let mut data = vec![];
let mut stack = vec![State::Start(0)];
let mut ret = None;
loop {
match (ret, stack.pop()) {
(Some(r), None) => return Ok(r),
(None, Some(State::Start(nr))) => {
let n0 = if nr == 0 {
n
} else {
&data[nr-1]
};
match n0 {
Counter::Val(a) => {
match a {
0..=9 => {
let b = Counter::Val(a+1);
data.push(b);
stack.push(State::Val0);
stack.push(State::Start(data.len()));
ret = None;
},
10 => { ret = Some(0); },
_ => return Err("Out of range".to_string()),
}
},
}
},
(Some(r), Some(State::Val0)) => {
data.pop();
ret = Some(r + 1);
},
_ => return Err("Invalid state".to_string()),
}
}
}
fn main() {
print!("{:?}\n", count(&Counter::Val(8)));
}
This also works.
My questions are:
That the borrow checker works for the recursive form, is there a way to explain the stack form to it (the Box version)?
Is there a reason to prefer the unsafe with the Box form over the data stack form? In both cases, the borrow checker has been circumvented, so they are both equally correct, as far as I can tell.
Is there a way to wrap a data structure around this, perhaps using unsafe internally, but exposing a safe exterior, that could allow writing this function with keeping the benefit of borrow checking?
Is there an idiomatic way to write the for loop form of the function?