From e21606144f3c3ced1a600e75c1c723135db4ac70 Mon Sep 17 00:00:00 2001
From: K1-R1 <77465250+K1-R1@users.noreply.github.com>
Date: Wed, 16 Oct 2024 09:55:03 +0100
Subject: [PATCH] fix: u256 log suggestions (#6645)

## Description

u256.log() now always recovers the flag register. Additionally the
method now returns 0u256 with invalid args when unsafe math panic is
disabled.

Closes #6644

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [x] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [x] If my change requires substantial documentation changes, I have
[requested support from the DevRel
team](https://github.com/FuelLabs/devrel-requests/issues/new/choose)
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [x] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.
---
 sway-lib-std/src/math.sw                               |  8 ++++++++
 .../test_programs/bytes_inline_tests/src/main.sw       |  2 +-
 .../test_programs/math_inline_tests/src/main.sw        | 10 +++++++++-
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/sway-lib-std/src/math.sw b/sway-lib-std/src/math.sw
index a52627693ff..6ffd4071865 100644
--- a/sway-lib-std/src/math.sw
+++ b/sway-lib-std/src/math.sw
@@ -339,10 +339,18 @@ impl Logarithm for u256 {
             assert(base >= 2);
             // Logarithm is undefined for 0
             assert(self != 0);
+        } else {
+            // Logarithm is undefined for bases less than 2
+            // Logarithm is undefined for 0
+            if (base < 2) || (self == 0) {
+                set_flags(flags);
+                return 0x00u256;
+            }
         }
 
         // Decimals rounded to 0
         if self < base {
+            set_flags(flags);
             return 0x00u256;
         }
 
diff --git a/test/src/in_language_tests/test_programs/bytes_inline_tests/src/main.sw b/test/src/in_language_tests/test_programs/bytes_inline_tests/src/main.sw
index 9a0b0eae392..65597943ef7 100644
--- a/test/src/in_language_tests/test_programs/bytes_inline_tests/src/main.sw
+++ b/test/src/in_language_tests/test_programs/bytes_inline_tests/src/main.sw
@@ -716,7 +716,7 @@ fn bytes_append() {
 
     assert(bytes2.len() == second_length);
     assert(!bytes2.is_empty());
-    
+
     assert(bytes2.get(0).unwrap() == d);
     assert(bytes2.get(1).unwrap() == e);
     assert(bytes2.get(2).unwrap() == f);
diff --git a/test/src/in_language_tests/test_programs/math_inline_tests/src/main.sw b/test/src/in_language_tests/test_programs/math_inline_tests/src/main.sw
index b294d52788b..c2641d71272 100644
--- a/test/src/in_language_tests/test_programs/math_inline_tests/src/main.sw
+++ b/test/src/in_language_tests/test_programs/math_inline_tests/src/main.sw
@@ -1,6 +1,6 @@
 library;
 
-use std::flags::disable_panic_on_overflow;
+use std::{flags::{disable_panic_on_overflow, disable_panic_on_unsafe_math}, registers::flags};
 
 #[test]
 fn math_root_u256() {
@@ -272,6 +272,7 @@ fn math_log_u8() {
 #[test]
 fn math_log_u256() {
     let max_u256 = u256::max();
+    let prior_flags = flags();
     assert(0x2u256.log(0x2u256) == 0x1u256);
     assert(0x1u256.log(0x3u256) == 0);
     assert(0x8u256.log(0x2u256) == 0x3u256);
@@ -279,6 +280,13 @@ fn math_log_u256() {
     assert(0x64u256.log(0x2u256) == 0x6u256);
     assert(0x64u256.log(0x9u256) == 0x2u256);
     assert(max_u256.log(0x2u256) == 0xffu256);
+    assert(prior_flags == flags());
+
+    let _ = disable_panic_on_unsafe_math();
+    let prior_flags = flags();
+    assert(0x1u256.log(0x1u256) == 0);
+    assert(0x0u256.log(0x3u256) == 0);
+    assert(prior_flags == flags());
 }
 
 #[test]