constcombine pass optimizing unary and binary operators (FuelLabs#4780)

## Description

Closes FuelLabs#4708.

This PR does NOT change the IR generation for expressions. Instead, it
implements constant folding for unary and binary operators. The idea is
that IR generation will be kept simple and stable, and all optimizations
are pushed onto the opt passes.

To better test this, this PR also changes some implementation details on
IR generation tests:

1 - `filecheck` directives are now collected and tested in groups. This
allows a new "::check-ir-optimized::", that can be configured with the
passes you want to test, or just fallback to "o1".

2 - `filechecker` explanations are now pretty-printed, and can be
printed when `verbose=true`. This allows the test to be checked more
easily.


![image](https://github.com/FuelLabs/sway/assets/83425/75d0f628-b271-458d-8c33-7d4d47af789a)

![image](https://github.com/FuelLabs/sway/assets/83425/9591a3d5-b0b4-4623-8a7c-1a34f5f96a99)

`breaking` label because of this change:
https://github.com/FuelLabs/sway/pull/4780/files#diff-b40fc3999876b0ff447de112b508250c7d0e7993f9023db29ddce8590d9b2286L5

## 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] 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.

---------

Co-authored-by: Anton Trunov <[email protected]>

sway-ir/src/optimize.rs

@@ -35,3 +35,89 @@ pub mod simplify_cfg;
 pub use simplify_cfg::*;
 
 mod target_fuel;
+
+#[cfg(test)]
+pub mod tests {
+    use crate::{PassGroup, PassManager};
+    use sway_types::SourceEngine;
+
+    /// This function parses the IR text representation and run the specified optimizers passes.
+    /// Then, depending on the `expected` parameter it checks if the IR was optimized or not.
+    ///
+    /// This comparison is done by capturing all instructions with metadata "!0".
+    ///
+    /// For example:
+    ///
+    /// ```rust, ignore
+    /// assert_optimization(
+    ///     &["constcombine"],
+    ///     "entry fn main() -> u64 {
+    ///        entry():
+    ///             l = const u64 1
+    ///             r = const u64 2
+    ///             result = add l, r, !0
+    ///             ret u64 result
+    ///     }",
+    ///     ["const u64 3"],
+    /// );
+    /// ```
+    pub(crate) fn assert_optimization<'a>(
+        passes: &[&'static str],
+        body: &str,
+        expected: Option<impl IntoIterator<Item = &'a str>>,
+    ) {
+        let source_engine = SourceEngine::default();
+        let mut context = crate::parse(
+            &format!(
+                "script {{
+                {body}
+            }}
+
+            !0 = \"a.sw\""
+            ),
+            &source_engine,
+        )
+        .unwrap();
+
+        let mut pass_manager = PassManager::default();
+        crate::register_known_passes(&mut pass_manager);
+
+        let mut group = PassGroup::default();
+        for pass in passes {
+            group.append_pass(pass);
+        }
+
+        let modified = pass_manager.run(&mut context, &group).unwrap();
+        assert_eq!(expected.is_some(), modified);
+
+        let Some(expected) = expected else {
+            return;
+        };
+
+        let actual = context
+            .to_string()
+            .lines()
+            .filter_map(|x| {
+                if x.contains(", !0") {
+                    Some(format!("{}\n", x.trim()))
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<String>>();
+
+        assert!(!actual.is_empty());
+
+        let mut expected_matches = actual.len();
+
+        for (actual, expected) in actual.iter().zip(expected) {
+            if !actual.contains(expected) {
+                panic!("error: {actual:?} {expected:?}");
+            } else {
+                expected_matches -= 1;
+            }
+        }
+
+        assert_eq!(expected_matches, 0);
+    }
+}

sway-ir/src/optimize/constants.rs

@@ -43,6 +43,16 @@ pub fn combine_constants(
             continue;
         }
 
+        if combine_binary_op(context, &function) {
+            modified = true;
+            continue;
+        }
+
+        if combine_unary_op(context, &function) {
+            modified = true;
+            continue;
+        }
+
         // Other passes here... always continue to the top if pass returns true.
         break;
     }
@@ -143,3 +153,239 @@ fn combine_cmp(context: &mut Context, function: &Function) -> bool {
         true
     })
 }
+
+fn combine_binary_op(context: &mut Context, function: &Function) -> bool {
+    let candidate = function
+        .instruction_iter(context)
+        .find_map(
+            |(block, inst_val)| match &context.values[inst_val.0].value {
+                ValueDatum::Instruction(Instruction::BinaryOp { op, arg1, arg2 })
+                    if arg1.is_constant(context) && arg2.is_constant(context) =>
+                {
+                    let val1 = arg1.get_constant(context).unwrap();
+                    let val2 = arg2.get_constant(context).unwrap();
+                    let v = match op {
+                        crate::BinaryOpKind::Add => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => l.checked_add(*r),
+                            _ => None,
+                        },
+                        crate::BinaryOpKind::Sub => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => l.checked_sub(*r),
+                            _ => None,
+                        },
+                        crate::BinaryOpKind::Mul => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => l.checked_mul(*r),
+                            _ => None,
+                        },
+                        crate::BinaryOpKind::Div => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => l.checked_div(*r),
+                            _ => None,
+                        },
+                        crate::BinaryOpKind::And => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => Some(l & r),
+                            _ => None,
+                        },
+                        crate::BinaryOpKind::Or => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => Some(l | r),
+                            _ => None,
+                        },
+                        crate::BinaryOpKind::Xor => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => Some(l ^ r),
+                            _ => None,
+                        },
+                        crate::BinaryOpKind::Mod => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => Some(l % r),
+                            _ => None,
+                        },
+                        crate::BinaryOpKind::Rsh => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => {
+                                u32::try_from(*r).ok().and_then(|r| l.checked_shr(r))
+                            }
+                            _ => None,
+                        },
+                        crate::BinaryOpKind::Lsh => match (&val1.value, &val2.value) {
+                            (ConstantValue::Uint(l), ConstantValue::Uint(r)) => {
+                                u32::try_from(*r).ok().and_then(|r| l.checked_shl(r))
+                            }
+                            _ => None,
+                        },
+                    };
+
+                    v.map(|v| {
+                        (
+                            inst_val,
+                            block,
+                            Constant {
+                                ty: val1.ty,
+                                value: ConstantValue::Uint(v),
+                            },
+                        )
+                    })
+                }
+                _ => None,
+            },
+        );
+
+    // Replace this binary op instruction with a constant.
+    candidate.map_or(false, |(inst_val, block, new_value)| {
+        inst_val.replace(context, ValueDatum::Constant(new_value));
+        block.remove_instruction(context, inst_val);
+        true
+    })
+}
+
+fn combine_unary_op(context: &mut Context, function: &Function) -> bool {
+    let candidate = function
+        .instruction_iter(context)
+        .find_map(
+            |(block, inst_val)| match &context.values[inst_val.0].value {
+                ValueDatum::Instruction(Instruction::UnaryOp { op, arg })
+                    if arg.is_constant(context) =>
+                {
+                    let val = arg.get_constant(context).unwrap();
+                    match op {
+                        crate::UnaryOpKind::Not => match &val.value {
+                            ConstantValue::Uint(v) => {
+                                val.ty.get_uint_width(context).and_then(|width| {
+                                    let max = match width {
+                                        8 => u8::MAX as u64,
+                                        16 => u16::MAX as u64,
+                                        32 => u32::MAX as u64,
+                                        64 => u64::MAX,
+                                        _ => return None,
+                                    };
+                                    Some((
+                                        inst_val,
+                                        block,
+                                        Constant {
+                                            ty: val.ty,
+                                            value: ConstantValue::Uint((!v) & max),
+                                        },
+                                    ))
+                                })
+                            }
+                            _ => None,
+                        },
+                    }
+                }
+                _ => None,
+            },
+        );
+
+    // Replace this unary op instruction with a constant.
+    candidate.map_or(false, |(inst_val, block, new_value)| {
+        inst_val.replace(context, ValueDatum::Constant(new_value));
+        block.remove_instruction(context, inst_val);
+        true
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::optimize::tests::*;
+
+    fn assert_operator(opcode: &str, l: &str, r: Option<&str>, result: Option<&str>) {
+        let expected = result.map(|result| format!("v0 = const u64 {result}"));
+        let expected = expected.as_ref().map(|x| vec![x.as_str()]);
+        let body = format!(
+            "
+    entry fn main() -> u64 {{
+        entry():
+        l = const u64 {l}
+        {r_inst}
+        result = {opcode} l, {result_inst} !0
+        ret u64 result
+    }}
+",
+            r_inst = r.map_or("".into(), |r| format!("r = const u64 {r}")),
+            result_inst = r.map_or("", |_| " r,")
+        );
+        assert_optimization(&["constcombine"], &body, expected);
+    }
+
+    #[test]
+    fn unary_op_are_optimized() {
+        assert_operator("not", &u64::MAX.to_string(), None, Some("0"));
+    }
+
+    #[test]
+    fn binary_op_are_optimized() {
+        assert_operator("add", "1", Some("1"), Some("2"));
+        assert_operator("sub", "1", Some("1"), Some("0"));
+        assert_operator("mul", "2", Some("2"), Some("4"));
+        assert_operator("div", "10", Some("5"), Some("2"));
+        assert_operator("mod", "12", Some("5"), Some("2"));
+        assert_operator("rsh", "16", Some("1"), Some("8"));
+        assert_operator("lsh", "16", Some("1"), Some("32"));
+
+        assert_operator(
+            "and",
+            &0x00FFF.to_string(),
+            Some(&0xFFF00.to_string()),
+            Some(&0xF00.to_string()),
+        );
+        assert_operator(
+            "or",
+            &0x00FFF.to_string(),
+            Some(&0xFFF00.to_string()),
+            Some(&0xFFFFF.to_string()),
+        );
+
+        assert_operator(
+            "xor",
+            &0x00FFF.to_string(),
+            Some(&0xFFF00.to_string()),
+            Some(&0xFF0FF.to_string()),
+        );
+    }
+
+    #[test]
+    fn binary_op_are_not_optimized() {
+        assert_operator("add", &u64::MAX.to_string(), Some("1"), None);
+        assert_operator("sub", "0", Some("1"), None);
+        assert_operator("mul", &u64::MAX.to_string(), Some("2"), None);
+        assert_operator("div", "1", Some("0"), None);
+
+        assert_operator("rsh", "1", Some("64"), None);
+        assert_operator("lsh", "1", Some("64"), None);
+    }
+
+    #[test]
+    fn ok_chain_optimization() {
+        // Unary operator
+
+        // `sub 1` is used to guarantee that the assert string is unique
+        assert_optimization(
+            &["constcombine"],
+            "
+        entry fn main() -> u64 {
+            entry():
+            a = const u64 18446744073709551615
+            b = not a, !0
+            c = not b, !0
+            d = const u64 1
+            result = sub c, d, !0
+            ret u64 result
+        }
+    ",
+            Some(["const u64 18446744073709551614"]),
+        );
+
+        // Binary Operators
+        assert_optimization(
+            &["constcombine"],
+            "
+        entry fn main() -> u64 {
+            entry():
+            l0 = const u64 1
+            r0 = const u64 2
+            l1 = add l0, r0, !0
+            r1 = const u64 3
+            result = add l1, r1, !0
+            ret u64 result
+        }
+    ",
+            Some(["const u64 6"]),
+        );
+    }
+}

sway-ir/src/parser.rs

@@ -897,7 +897,7 @@ mod ir_builder {
             fn_decl: IrAstFnDecl,
         ) -> Result<(), IrError> {
             let convert_md_idx = |opt_md_idx: &Option<MdIdxRef>| {
-                opt_md_idx.map(|mdi| self.md_map.get(&mdi).copied().unwrap())
+                opt_md_idx.and_then(|mdi| self.md_map.get(&mdi).copied())
             };
             let args: Vec<(String, Type, Option<MetadataIndex>)> = fn_decl
                 .args

test/src/e2e_vm_tests/test_programs/should_pass/require_contract_deployment/call_basic_storage/src/main.sw

@@ -2,7 +2,7 @@ script;
 use basic_storage_abi::{BasicStorage, Quad};
 
 fn main() -> u64 {
-    let addr = abi(BasicStorage, 0x6a550eadf838642881db8c66f40c1fc8442ee625212f3f2943850b2c12d12275);
+    let addr = abi(BasicStorage, 0xf2cabd05603d8f2b7eb273e2e763e60e3b88dd6405de5cb732bc7c498f13f213);
     let key = 0x0fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff;
     let value = 4242;