Allow improveIf of con/undet when only one constructor

src/comp/IExpand.hs

@@ -4398,6 +4398,42 @@ improveIf f t cnd thn@(IAps chr@(ICon _ (ICPrim _ PrimChr)) ts1 [chr_thn])
   let chrArgType = iGetType chr_thn
   (e', _) <- improveIf f chrArgType cnd chr_thn chr_els
   return (IAps chr ts1 [e'], True)
+
+-- Push if improvement inside constructors when one arm is undefined
+-- and the type has only one constructor
+--
+-- Further down, a general improveIf rule optimizes 'if c e _' to just 'e'.
+-- But that can cause poor code generation for if-else chains returning
+-- the constructors of a union type, so an earlier improveIf rule catches
+-- that situation (before the general rule can apply).
+-- However, if there is only one constructor, we do want an optimization to apply,
+-- so we put that here, prior to the blocking rule.
+--
+improveIf f t cnd thn@(IAps (ICon i1 c1@(ICCon {})) ts1 es1)
+                  els@(ICon i2 (ICUndet { iuKind = u }))
+  | numCon (conTagInfo c1) == 1
+  = do
+      when doTraceIf $ traceM ("improveIf ICCon/ICUndet triggered" ++ ppReadable (cnd,thn,els))
+      let realConType = itInst (iConType c1) ts1
+          (argTypes, _) = itGetArrows realConType
+      when (length argTypes /= length es1) $ internalError ("improveIf Con/Undet:" ++ ppReadable (argTypes, es1))
+      let mkUndet t = icUndetAt (getIdPosition i2) t u
+      (es', bs) <- mapAndUnzipM (\(t, e1) -> improveIf f t cnd e1 (mkUndet t)) (zip argTypes es1)
+      -- unambiguous improvement because the ICCon has propagated out
+      return ((IAps (ICon i1 c1) ts1 es'), True)
+improveIf f t cnd thn@(ICon i1 (ICUndet { iuKind = u }))
+                  els@(IAps (ICon i2 c2@(ICCon {})) ts2 es2)
+  | numCon (conTagInfo c2) == 1
+  = do
+      when doTraceIf $ traceM ("improveIf ICCon/ICUndet triggered" ++ ppReadable (cnd,thn,els))
+      let realConType = itInst (iConType c2) ts2
+          (argTypes, _) = itGetArrows realConType
+      when (length argTypes /= length es2) $ internalError ("improveIf Con/Undet:" ++ ppReadable (argTypes, es2))
+      let mkUndet t = icUndetAt (getIdPosition i1) t u
+      (es', bs) <- mapAndUnzipM (\(t, e2) -> improveIf f t cnd (mkUndet t) e2) (zip argTypes es2)
+      -- unambiguous improvement because the ICCon has propagated out
+      return ((IAps (ICon i2 c2) ts2 es'), True)
+
 -- Do not "optimize" constructors against undefined values because this can remove
 -- the conditions required to optimize chains of ifs like these:
 -- if (x == 0) 0 else if (x == 1) 1 else ... back to just x

testsuite/bsc.evaluator/opt/ImproveIf_ConUndet_OneCon.bsv

@@ -0,0 +1,74 @@
+// This example comes from GitHub Issue #742
+//
+// This is a simple example that should compile without a lot of
+// unfolding steps or heap space.  However, if the evaluator is
+// missing an optimization for this:
+//
+//     if (cond)
+//       then Ctor e1 e2 ...
+//       else _
+//
+// the compiler will take time and memory until eventually exiting
+// with an error message about max unfolding steps reached or
+// (if the max steps is increased) that stack space was exhausted.
+//
+// For better code generation of "pack of unpack" for union types, we
+// don't want the evaluator to optimize away the don't-care like this:
+//
+//     ==> Ctor e1 e2 ...
+//
+// However, when the type only has one constructor, such as Vector in
+// this example, we can optimize the expression to this:
+//
+//     ==> Ctor
+//           (if (cond) then e1 else _)
+//           (if (cond) then e2 else _)
+//           ...
+//
+// With this optimization, the example will successfully compile
+// (quickly, without many unfolding steps or stack space usage).
+//
+
+import Vector::*;
+
+`define Q_SIZE 8
+typedef Bit#(2) T;
+
+(*synthesize*)
+module mkImproveIf_ConUndet_OneCon ();
+
+  Vector#(`Q_SIZE, Reg#(T))           vec_data  <- replicateM(mkRegU);
+  Reg #(Maybe#(Vector#(`Q_SIZE, T)))  rg_1   <- mkRegU;
+  Reg #(Maybe#(Vector#(`Q_SIZE, T)))  rg_2   <- mkRegU;
+  Reg #(Vector#(`Q_SIZE, Bool))       rg_3   <- mkRegU;
+  Reg #(Maybe#(Vector#(`Q_SIZE, T)))  rg_4   <- mkRegU;
+  Reg #(Vector#(`Q_SIZE, Bool))       rg_5   <- mkRegU;
+
+  rule r;
+    Vector#(`Q_SIZE, T) tmp_data = readVReg (vec_data);
+
+    if (rg_1 matches tagged Valid .d1)
+      tmp_data  = d1;
+
+    if (rg_2 matches tagged Valid .d2)
+      begin
+        for (int i = 0 ; i< `Q_SIZE ;i = i+1)
+          begin
+            if (rg_3[i] == True)
+              tmp_data[i] = d2[i];
+            end
+      end
+
+    if (rg_4 matches tagged Valid .d4)
+      begin
+        for (int i = 0 ; i< `Q_SIZE ;i = i+1)
+          begin
+            if (rg_5[i] == True)
+              tmp_data[i] = d4[i];
+          end
+      end
+
+    writeVReg(vec_data,tmp_data);
+  endrule
+
+endmodule

testsuite/bsc.evaluator/opt/opt.exp

@@ -63,3 +63,13 @@ if { $vtest == 1 } {
     find_n_strings sysCompareSameExpr.v {rSLE} 1
     find_n_strings sysCompareSameExpr.v {rSGE} 1
 }
+
+# Test that the evaluator has an 'improveIf' optimization for
+# ICCon/ICUndet when the type has only one constructor
+# (GitHub Issue #742)
+#
+# Without the optimization, there is excessive unfolding steps
+# and stack usage
+#
+compile_verilog_pass ImproveIf_ConUndet_OneCon.bsv {} \
+    {-steps-max-intervals 10 -steps-warn-interval 100000 +RTS -K10m -RTS}