Better codegen for array repeats in Rust -- fixes #137

[msprotz]

No description provided.

[msprotz]

@franziskuskiefer I would like you to try this out on mlkem and mldsa before merging, so I'll wait until we get libcrux back to green before merging this, I'm mildly concerned with regressions

please also note that in the case of nested arrays, there will still be a long series of assignments -- if this is a problem I left a TODO in the code that indicates what to tweak to improve that situation, too

[msprotz]

Note that this patch still improves things: previously,

  17   │ fn init2() -> [[u8; 32]; 32] {
  18   │     [[0u8; 32]; 32]
  19   │ }

would generate 1024 assignments -- now it only generates 32, but with 32 distinct buffers on the stack...

[msprotz]

ok looks like I need to revamp the array repeat desugaring phase

[msprotz]

As far as I can tell, this generates a rather pleasant diff on branch protz_refresh of libcrux -- see my latest commit there to see the effect of this PR.

@franziskuskiefer if the changes are agreeable to you, please go ahead and merge... thanks

[franziskuskiefer]

Thanks for tackling this!

This improves a bunch of cases.

But there's at least one it makes worse.

let mut tmp_stack = [[0i32; 263], [0i32; 263], [0i32; 263], [0i32; 263]];

(Maybe this should be different in Rust already, but that's a different question.)
used to generate

int32_t tmp_stack[4U][263U] = {{0U}};

Now it generates

int32_t repeat_expression[263U] = {0U};
int32_t tmp_stack[4U][263U] = {
    {(int32_t)0, ...},
    {(int32_t)0, ...},
    {(int32_t)0, ...},
    repeat_expression};

Which is not only worse, but won't compile.

[msprotz]

Ah good catch. I'll add this as a test-case, I didn't have it locally. Thank you

[msprotz]

I ended up revamping the handling of nested arrays and simplified pretty awful code. (I actually vividly remember writing this code in a rush and thinking I wasn't sure I was doing the right thing...)

You'll need FStarLang/karamel#530

The consequence should now be uniformly better code. For your specific example I do not yet generate a perfect nested initializer (i.e. int32_t tmp_stack[4U][263U] = { 0 }) but the code is pretty good. Let me know if you'd like me to push a little more on code quality.

For instance, if you have arrays nested three levels deep (i.e. arrays of arrays of arrays) there is an opportunity to introduce a for-loop in the code I just touched, rather than statically generating a slew of assignments. With the big cleanup, this should be pretty easy to do, so if this is something you need, let me know. Cheers.

[msprotz]

I was able to confirm that on my branch protz_refresh of mlkem, I see no regression (i.e. the code is still as improved as before), and this passes the tests locally.

[franziskuskiefer]

Thanks!
I'm seeing another issue though. Looks like there's a counter that's not counting anymore.

Here's the diff.

-  memcpy(ret[1U], copy_of_out1, (size_t)200U * sizeof(uint8_t));
-  memcpy(ret[2U], copy_of_out2, (size_t)200U * sizeof(uint8_t));
-  memcpy(ret[3U], uu____3, (size_t)200U * sizeof(uint8_t));
+  memcpy(ret[0U], copy_of_out1, (size_t)200U * sizeof(uint8_t));
+  memcpy(ret[0U], copy_of_out2, (size_t)200U * sizeof(uint8_t));
+  memcpy(ret[0U], uu____3, (size_t)200U * sizeof(uint8_t));
 }

The Rust code for this is

let mut out0 = [0u8; 200];
let mut out1 = [0u8; 200];
let mut out2 = [0u8; 200];
let mut out3 = [0u8; 200];
// ...
[out0, out1, out2, out3]

[franziskuskiefer]

Thanks. Looks like this is good to go now.