From b7a23276688be1188335e859cdd171899a160165 Mon Sep 17 00:00:00 2001
From: Zachary Snow <zach@zachjs.com>
Date: Sat, 30 Apr 2022 20:19:02 -0600
Subject: [PATCH] simple for loop elaboration applies in more cases

---
 CHANGELOG.md                  |  1 +
 src/Convert/Jump.hs           | 44 +++++++++++++++++------------------
 test/core/simple_loop_jump.sv | 10 ++++++++
 test/core/simple_loop_jump.v  |  8 +++++++
 4 files changed, 41 insertions(+), 22 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9f79f23..0b52be2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@
 * Added conversion for `do` `while` loops
 * Added support for passing through DPI imports and exports
 * Added support for passing through functions with output ports
+* Extended applicability of simplified Yosys-compatible `for` loop elaboration
 
 ### Other Enhancements
 
diff --git a/src/Convert/Jump.hs b/src/Convert/Jump.hs
index bcd8042..85e64fb 100644
--- a/src/Convert/Jump.hs
+++ b/src/Convert/Jump.hs
@@ -121,12 +121,25 @@ convertStmts stmts = do
 pattern SimpleLoopInits :: Identifier -> [(LHS, Expr)]
 pattern SimpleLoopInits var <- [(LHSIdent var, _)]
 
-pattern SimpleLoopGuard :: Identifier -> Expr
-pattern SimpleLoopGuard var <- BinOp _ (Ident var) _
-
 pattern SimpleLoopIncrs :: Identifier -> [(LHS, AsgnOp, Expr)]
 pattern SimpleLoopIncrs var <- [(LHSIdent var, _, _)]
 
+-- check if an expression contains a reference to the given identifier
+usesIdent :: Identifier -> Expr -> Writer Any ()
+usesIdent x (Ident y)
+    | x == y = tell $ Any True
+    | otherwise = return ()
+usesIdent x expr =
+    collectSinglyNestedExprsM (usesIdent x) expr
+
+-- identifies loops which could likely be statically unrolled, and so may
+-- benefit from avoiding complicating the loop guard with jump state, assuming
+-- the guard and incrementation do not have side effects
+simpleLoopVar :: [(LHS, Expr)] -> Expr -> [(LHS, AsgnOp, Expr)] -> Identifier
+simpleLoopVar (SimpleLoopInits var1) comp (SimpleLoopIncrs var3)
+    | var1 == var3, getAny $ execWriter $ usesIdent var1 comp = var1
+simpleLoopVar _ _ _ = ""
+
 -- rewrites the given statement, and returns the type of any unfinished jump
 convertStmt :: Stmt -> State Info Stmt
 
@@ -140,20 +153,12 @@ convertStmt (Block Par x decls stmts) = do
 
 convertStmt (Block Seq ""
     decls@[CommentDecl{}, Variable Local _ var0 [] Nil]
-    [ comment@CommentStmt{}
-    , For
-        inits@(SimpleLoopInits var1)
-        comp@(SimpleLoopGuard var2)
-        incr@(SimpleLoopIncrs var3)
-        stmt
-    ]) =
-    convertLoop localInfo loop comp incr stmt
+    [comment@CommentStmt{}, For inits comp incr stmt])
+    | var1@(_ : _) <- simpleLoopVar inits comp incr, var0 == var1 =
+    convertLoop (Just var1) loop comp incr stmt
         >>= return . Block Seq "" decls . (comment :) . pure
     where
         loop c i s = For inits c i s
-        localInfo = if var0 /= var1 || var1 /= var2 || var2 /= var3
-                        then Nothing
-                        else Just ""
 
 convertStmt (Block Seq x decls stmts) =
     step stmts >>= return . Block Seq x decls
@@ -192,16 +197,11 @@ convertStmt (Case unique kw expr cases) = do
     modify $ \s -> s { sHasJump = hasJump }
     return $ Case unique kw expr cases'
 
-convertStmt (For
-    inits@(SimpleLoopInits var1)
-    comp@(SimpleLoopGuard var2)
-    incr@(SimpleLoopIncrs var3) stmt) =
-    convertLoop localInfo loop comp incr stmt
+convertStmt (For inits comp incr stmt)
+    | var@(_ : _) <- simpleLoopVar inits comp incr =
+    convertLoop (Just var) loop comp incr stmt
     where
         loop c i s = For inits c i s
-        localInfo = if var1 /= var2 || var2 /= var3
-                        then Nothing
-                        else Just var1
 convertStmt (For inits comp incr stmt) =
     convertLoop Nothing loop comp incr stmt
     where loop c i s = For inits c i s
diff --git a/test/core/simple_loop_jump.sv b/test/core/simple_loop_jump.sv
index f21395e..a5f4b29 100644
--- a/test/core/simple_loop_jump.sv
+++ b/test/core/simple_loop_jump.sv
@@ -44,6 +44,15 @@ module top;
                 return k * 3;
         k = k * 2 + 1;
     endfunction
+    function automatic integer l;
+        input integer inp;
+        l = 1;
+        for (integer idx = inp; 0 < idx; idx--) begin
+            if (l == 32)
+                break;
+            l = l * 2;
+        end
+    endfunction
     integer i;
     initial
         for (i = 0; i < 10; i = i + 1) begin
@@ -52,5 +61,6 @@ module top;
             $display("h(%0d) = %0d", i, h(i));
             $display("j(%0d) = %0d", i, j(i));
             $display("k(%0d) = %0d", i, k(i));
+            $display("l(%0d) = %0d", i, l(i));
         end
 endmodule
diff --git a/test/core/simple_loop_jump.v b/test/core/simple_loop_jump.v
index 3535809..ae3f66f 100644
--- a/test/core/simple_loop_jump.v
+++ b/test/core/simple_loop_jump.v
@@ -34,6 +34,13 @@ module top;
         else
             k = inp * 2 + 1;
     endfunction
+    function automatic integer l;
+        input integer inp;
+        if (inp > 5)
+            l = 32;
+        else
+            l = 2 ** inp;
+    endfunction
     integer i;
     initial
         for (i = 0; i < 10; i = i + 1) begin
@@ -42,5 +49,6 @@ module top;
             $display("h(%0d) = %0d", i, h(i));
             $display("j(%0d) = %0d", i, j(i));
             $display("k(%0d) = %0d", i, k(i));
+            $display("l(%0d) = %0d", i, l(i));
         end
 endmodule