optimize operations with immediates for compound assignment ops

1 files changed, 68 insertions, 53 deletions

diff --git a/pez.c b/pez.c
index 1611770..2e30a75 100644
--- a/pez.c
+++ b/pez.c
@@ -2895,6 +2895,65 @@ commutate(enum op *op)
 static bool discard(Comp *cm);
 
 static bool
+binexprimm(Comp *cm, enum op op, int8_t imm)
+{
+    enum op iop = 0;
+    if (op == Osub) {
+        // X - k turns into X + -k, but if k == -128 we cannot encode 128
+        if (imm == -128) {
+            goto Bail;
+        }
+        imm = -imm;
+        goto Oadd;
+    } else if (op == Odiv && imm == 0) {
+        // the idiv fastpath assumes dividend is nonzero
+        goto Bail;
+    } else if (op == Omod && imm == 0) {
+        // the imod fastpath assumes dividend is nonzero
+        // always x % 0 = 0
+        cm->has_k = 0;
+        TRY(discard(cm) && compop(cm, Ozero));
+        return 1;
+    } else if (op == Omul) {
+        // try turn mul by powers of 2 into shifts
+        for (int exp = 0; exp < 8; ++exp) {
+            if (imm == 1 << exp) {
+                iop = Oishl;
+                imm = exp;
+                goto CompIop;
+            }
+        }
+    }
+    switch (op) {
+Oadd: case Oadd:  iop = Oiadd; break;
+    case Omul:  iop = Oimul; break;
+    case Odiv:  iop = Oidiv; break;
+    case Omod:  iop = Oimod; break;
+    case Oband: iop = Oiband; break;
+    case Obior: iop = Oibior; break;
+    case Obxor: iop = Oibxor; break;
+    case Oshl:  iop = Oishl; break;
+    case Oshra: iop = Oishra; break;
+    case Oshrl: iop = Oishrl; break;
+    case Olt:   iop = Oilt; break;
+    case Ole:   iop = Oile; break;
+    case Ogt:   iop = Oigt; break;
+    case Oge:   iop = Oige; break;
+    case Oeq:   iop = Oieq; break;
+    case One:   iop = Oine; break;
+    default: break;
+    }
+CompIop:
+    assert(iop);
+    cm->has_k = 0;
+    TRY(compop(cm, iop) && compbyte(cm, imm));
+    return 1;
+Bail:
+    TRY(compop(cm, op));
+    return 1;
+}
+
+static bool
 binexpr(Comp *cm, char okind, bool (*prev)(Comp *))
 {
     uint save = cm->code.len, save2, save3;
@@ -2911,7 +2970,7 @@ binexpr(Comp *cm, char okind, bool (*prev)(Comp *))
         Val lk = cm->k;
         bool has_lk = cm->has_k;
         int imm;
-        enum op opx = op, iop = 0;
+        enum op opx = op;
 
         TRY(flushconst(cm));
         save2 = cm->code.len;
@@ -2943,59 +3002,10 @@ binexpr(Comp *cm, char okind, bool (*prev)(Comp *))
         } else if (cm->has_k && isimm(&imm, cm->k)) {
             // try immediate variants of ops
         TryImm:
-            if (opx == Osub) {
-                // X - k turns into X + -k, but if k == -128 we cannot encode 128
-                if (imm == -128) {
-                    goto Bail;
-                }
-                imm = -imm;
-                goto Oadd;
-            } else if (opx == Odiv && imm == 0) {
-                // the idiv fastpath assumes dividend is nonzero
-                goto Bail;
-            } else if (opx == Omod && imm == 0) {
-                // the imod fastpath assumes dividend is nonzero
-                // always x % 0 = 0
-                cm->has_k = 0;
-                TRY(discard(cm) && compop(cm, Ozero));
-                continue;
-            } else if (opx == Omul) {
-                // try turn mul by powers of 2 into shifts
-                for (int exp = 0; exp < 8; ++exp) {
-                    if (imm == 1 << exp) {
-                        iop = Oishl;
-                        imm = exp;
-                        goto CompIop;
-                    }
-                }
-            }
-            switch (opx) {
-      Oadd: case Oadd:  iop = Oiadd; break;
-            case Omul:  iop = Oimul; break;
-            case Odiv:  iop = Oidiv; break;
-            case Omod:  iop = Oimod; break;
-            case Oband: iop = Oiband; break;
-            case Obior: iop = Oibior; break;
-            case Obxor: iop = Oibxor; break;
-            case Oshl:  iop = Oishl; break;
-            case Oshra: iop = Oishra; break;
-            case Oshrl: iop = Oishrl; break;
-            case Olt:   iop = Oilt; break;
-            case Ole:   iop = Oile; break;
-            case Ogt:   iop = Oigt; break;
-            case Oge:   iop = Oige; break;
-            case Oeq:   iop = Oieq; break;
-            case One:   iop = Oine; break;
-            default: break;
-            }
-        CompIop:
-            assert(iop);
-            cm->has_k = 0;
-            TRY(compop(cm, iop) && compbyte(cm, (int8_t)imm));
+            TRY(binexprimm(cm, opx, imm));
             continue;
         }
-    Bail:
-        TRY(compop(cm, opx));
+        TRY(compop(cm, op));
     } while (okind != 'C' && (op2 = getbinop(&kind, &chr, cm)) == op);
     if (op2 != op && kind != 0) {
         stashbinop(cm, op2, kind, chr);
@@ -3223,7 +3233,12 @@ setexpr(Comp *cm)
         TRY(condexpr(cm));
 
         if (binop != Oset) {
-            TRY(compop(cm, binop));
+            int imm;
+            if (cm->has_k && isimm(&imm, cm->k)) {
+                TRY(binexprimm(cm, binop, imm));
+            } else {
+                TRY(compop(cm, binop));
+            }
         }
 
         if (opcode == Oglobal) {