Fix subtraction optimization

crates/rue-compiler/src/compile/expr/prefix.rs

@@ -1,7 +1,7 @@
 use std::borrow::Cow;
 
 use log::debug;
-use rue_ast::{AstNode, AstPrefixExpr};
+use rue_ast::{AstExpr, AstNode, AstPrefixExpr};
 use rue_diagnostic::DiagnosticKind;
 use rue_hir::{Hir, UnaryOp, Value};
 use rue_lir::{atom_bigint, bigint_atom};
@@ -51,6 +51,13 @@ pub fn compile_prefix_expr(ctx: &mut Compiler, prefix: &AstPrefixExpr) -> Value
             }
         }
         T![-] => {
+            if let AstExpr::PrefixExpr(inner) = &expr
+                && let Some(op) = inner.op()
+                && op.kind() == T![-]
+            {
+                ctx.diagnostic(prefix.syntax(), DiagnosticKind::DoubleNegationWarning);
+            }
+
             if ctx.is_assignable(value.ty, ctx.builtins().types.int) {
                 let ty = if let Some(Atoms::Restricted(restrictions)) =
                     rue_types::extract_atoms(ctx.types_mut(), value.ty, true)

crates/rue-diagnostic/src/kind.rs

@@ -230,6 +230,11 @@ pub enum DiagnosticKind {
 
     #[error("Cannot end path in `super`")]
     SuperAtEnd,
+
+    #[error(
+        "Use of `--` in an expression will perform a double negation, but could be confused with the 'prefix decrement' operator in other languages"
+    )]
+    DoubleNegationWarning,
 }
 
 impl DiagnosticKind {
@@ -309,7 +314,8 @@ impl DiagnosticKind {
             | Self::AlwaysFalseCondition
             | Self::AlwaysTrueCondition
             | Self::UnusedStatementValue
-            | Self::UnusedGlobImport => DiagnosticSeverity::Warning,
+            | Self::UnusedGlobImport
+            | Self::DoubleNegationWarning => DiagnosticSeverity::Warning,
         }
     }
 }

crates/rue-lir/src/optimize/ops.rs

@@ -1,3 +1,5 @@
+use std::num::Saturating;
+
 use id_arena::Arena;
 use num_bigint::{BigInt, Sign};
 use num_integer::Integer;
@@ -142,67 +144,96 @@ pub fn opt_add(arena: &mut Arena<Lir>, args: Vec<LirId>) -> LirId {
     arena.alloc(Lir::Add(result))
 }
 
-// If the value is an atom, we can add it to a sum to subtract from
-// We can collapse nested subs
-// If the value is a raise, we can return it directly
+// Split the subtraction into minuends and subtrahends and sum them together,
+// then subtract the subtrahend from the minuend
 pub fn opt_sub(arena: &mut Arena<Lir>, args: Vec<LirId>) -> LirId {
+    let mut minuend_value = BigInt::from(0);
+    let mut other_minuends = Vec::new();
+    let mut other_subtrahends = Vec::new();
+
     let mut args = ArgList::new(args);
-    let mut result = Vec::new();
-    let mut first = None;
-    let mut sum = BigInt::from(0);
+    let mut minuend_count = Saturating(1usize);
 
     while let Some(arg) = args.next() {
         match arena[arg].clone() {
             Lir::Atom(atom) => {
-                if let Some(first_lir) = first {
-                    if let Lir::Atom(first_atom) = arena[first_lir].clone() {
-                        first = Some(arena.alloc(Lir::Atom(bigint_atom(
-                            atom_bigint(first_atom) - atom_bigint(atom),
-                        ))));
-                    } else {
-                        sum += atom_bigint(atom);
-                    }
+                let value = atom_bigint(atom);
+
+                if minuend_count > Saturating(0) {
+                    minuend_value += value;
                 } else {
-                    first = Some(arg);
+                    // Fold subtrahend values into the minuend as an optimization
+                    minuend_value -= value;
                 }
             }
-            Lir::Sub(items) => {
+            Lir::Add(items) => {
+                if minuend_count > Saturating(0) {
+                    minuend_count += items.len();
+                }
+
                 args.prepend(items);
             }
+            Lir::Sub(items) => {
+                if !items.is_empty() {
+                    if minuend_count > Saturating(0) {
+                        // Just unroll the inner subtraction into the outer subtraction
+                        // The first item is the new minuend, the rest are the subtrahends
+                        args.prepend(items);
+                        minuend_count += 1;
+                    } else {
+                        // The minuend of the inner subtraction is a subtrahend of the outer subtraction
+                        args.prepend(vec![items[0]]);
+
+                        // The subtrahends of the inner subtraction are minuends of the outer subtraction
+                        // One extra because minuend_count will be subtracted at the end of this iteration
+                        minuend_count += items.len();
+                        args.prepend(items.iter().skip(1).copied().collect());
+                    }
+                }
+            }
             Lir::Raise(_) => return arg,
             _ => {
-                if first.is_none() {
-                    first = Some(arg);
-                    continue;
+                if minuend_count > Saturating(0) {
+                    other_minuends.push(arg);
+                } else {
+                    other_subtrahends.push(arg);
                 }
-                result.push(arg);
             }
         }
-    }
 
-    if let Some(first) = first {
-        result.insert(0, first);
+        minuend_count -= 1;
     }
 
-    if sum != BigInt::from(0) {
-        result.push(arena.alloc(Lir::Atom(bigint_atom(sum))));
+    if other_minuends.is_empty() && other_subtrahends.is_empty() {
+        return arena.alloc(Lir::Atom(bigint_atom(minuend_value)));
     }
 
-    if result.is_empty() {
-        return arena.alloc(Lir::Atom(vec![]));
+    if minuend_value > BigInt::from(0)
+        || (minuend_value != BigInt::from(0) && !other_subtrahends.is_empty())
+    {
+        other_minuends.insert(0, arena.alloc(Lir::Atom(bigint_atom(minuend_value))));
+    } else if minuend_value < BigInt::from(0) {
+        other_subtrahends.insert(0, arena.alloc(Lir::Atom(bigint_atom(-minuend_value))));
     }
 
-    if result.len() == 1 && matches!(arena[result[0]], Lir::Atom(_)) {
-        return result[0];
+    let minuend = if other_minuends.len() == 1
+        && (matches!(arena[other_minuends[0]], Lir::Atom(_)) || !other_subtrahends.is_empty())
+    {
+        other_minuends[0]
+    } else if other_minuends.is_empty() {
+        arena.alloc(Lir::Atom(vec![]))
+    } else {
+        arena.alloc(Lir::Add(other_minuends))
+    };
+
+    if other_subtrahends.is_empty() {
+        return minuend;
     }
 
-    if result.len() == 2
-        && let Lir::Atom(first) = arena[result[0]].clone()
-        && let Lir::Atom(second) = arena[result[1]].clone()
-    {
-        return arena.alloc(Lir::Atom(bigint_atom(
-            atom_bigint(first) - atom_bigint(second),
-        )));
+    let mut result = vec![minuend];
+
+    for subtrahend in other_subtrahends {
+        result.push(subtrahend);
     }
 
     arena.alloc(Lir::Sub(result))

tests/optimizer/constant_folding.yaml

@@ -52,7 +52,7 @@ tests:
   byte_cost: 60000
   total_cost: 60020
 - name: nest_sub_add_ref
-  program: (- (+ 2 (q . 200)) 2)
+  program: (- (+ (q . 200) 2) 2)
   debug_program: (- (+ 2 (q . 200)) 2)
   solution: (100)
   output: '200'

tests/optimizer/nested_subtraction.rue

@@ -0,0 +1,29 @@
+fn double_negate(num: Int) -> Int {
+    --num
+}
+
+fn subtract(a: Int, b: Int) -> Int {
+    a - b
+}
+
+test fn double_negate_cases() {
+    assert double_negate(100) == 100;
+    assert double_negate(-100) == -100;
+}
+
+test fn double_negate_literal() {
+    assert --100 == 100;
+    assert --(-100) == -100;
+}
+
+test fn nested_subtraction() {
+    assert subtract(100, 200) - 300 == -400;
+    assert subtract(100, subtract(200, 300)) == 200;
+    assert (100 - 200) - 300 == -400;
+    assert 100 - (200 - 300) == 200;
+    assert 100 - 200 - 300 == -400;
+}
+
+test fn negative_right_logical_shift() {
+    assert 10 >>> -1 == 20;
+}

tests/optimizer/nested_subtraction.yaml

@@ -0,0 +1,33 @@
+tests:
+- name: double_negate_cases
+  program: (a (q 2 (i (= (a 2 (q . 100)) (q . 100)) (q 2 (i (= (a 2 (q . -100)) (q . -100)) (q) (q 8)) 1) (q 8)) 1) (c (q 16 1) ()))
+  debug_program: (a (q 2 (i (= (a 2 (q . 100)) (q . 100)) (q 2 (i (= (a 2 (- () (q . 100))) (- () (q . 100))) (q) (q 8 (q . "assertion failed at nested_subtraction.rue:11:5"))) 1) (q 8 (q . "assertion failed at nested_subtraction.rue:10:5"))) 1) (c (q 17 () (- () 1)) ()))
+  output: ()
+  runtime_cost: 2240
+  byte_cost: 1140000
+  total_cost: 1142240
+- name: double_negate_literal
+  program: (a (i (= (q . 100) (q . 100)) (q 2 (i (= (q . -100) (q . -100)) (q) (q 8)) 1) (q 8)) 1)
+  debug_program: (a (i (= (- () (- () (q . 100))) (q . 100)) (q 2 (i (= (- () (- () (- () (q . 100)))) (- () (q . 100))) (q) (q 8 (q . "assertion failed at nested_subtraction.rue:16:5"))) 1) (q 8 (q . "assertion failed at nested_subtraction.rue:15:5"))) 1)
+  output: ()
+  runtime_cost: 782
+  byte_cost: 756000
+  total_cost: 756782
+- name: nested_subtraction
+  program: (a (q 2 (i (= (- (a 2 (c (q . 100) (q . 200))) (q . 300)) (q . -400)) (q 2 (i (= (a 2 (c (q . 100) (a 2 (c (q . 200) (q . 300))))) (q . 200)) (q 2 (i (= (q . -400) (q . -400)) (q 2 (i (= (q . 200) (q . 200)) (q 2 (i (= (q . -400) (q . -400)) (q) (q 8)) 1) (q 8)) 1) (q 8)) 1) (q 8)) 1) (q 8)) 1) (c (q 17 2 3) ()))
+  debug_program: (a (q 2 (i (= (- (a 2 (c (q . 100) (q . 200))) (q . 300)) (- () (q . 400))) (q 2 (i (= (a 2 (c (q . 100) (a 2 (c (q . 200) (q . 300))))) (q . 200)) (q 2 (i (= (- (- (q . 100) (q . 200)) (q . 300)) (- () (q . 400))) (q 2 (i (= (- (q . 100) (- (q . 200) (q . 300))) (q . 200)) (q 2 (i (= (- (- (q . 100) (q . 200)) (q . 300)) (- () (q . 400))) (q) (q 8 (q . "assertion failed at nested_subtraction.rue:24:5"))) 1) (q 8 (q . "assertion failed at nested_subtraction.rue:23:5"))) 1) (q 8 (q . "assertion failed at nested_subtraction.rue:22:5"))) 1) (q 8 (q . "assertion failed at nested_subtraction.rue:21:5"))) 1) (q 8 (q . "assertion failed at nested_subtraction.rue:20:5"))) 1) (c (q 17 2 3) ()))
+  output: ()
+  runtime_cost: 6119
+  byte_cost: 2964000
+  total_cost: 2970119
+- name: negative_right_logical_shift
+  program: (a (i (= (lsh (q . 10) (q . 1)) (q . 20)) (q) (q 8)) 1)
+  debug_program: (a (i (= (lsh (q . 10) (- () (- () (q . 1)))) (q . 20)) (q) (q 8 (q . "assertion failed at nested_subtraction.rue:28:5"))) 1)
+  output: ()
+  runtime_cost: 727
+  byte_cost: 468000
+  total_cost: 468727
+diagnostics:
+- Use of `--` in an expression will perform a double negation, but could be confused with the 'prefix decrement' operator in other languages at nested_subtraction.rue:2:5
+- Use of `--` in an expression will perform a double negation, but could be confused with the 'prefix decrement' operator in other languages at nested_subtraction.rue:15:12
+- Use of `--` in an expression will perform a double negation, but could be confused with the 'prefix decrement' operator in other languages at nested_subtraction.rue:16:12