Require LoopUnrollInfo in ForStatement::Make.

Previously, callers could pass a null LoopUnrollInfo into
ForStatement::Make and it would just try to make a new one. Now, we
honor whatever is passed in as-is.

This required some fixes in the Inliner and Rehydrator, which relied on
the recreate-at-will behavior. However, it should improve the
performance of inlining a For loop, as we no longer need to recompute
the unroll info from scratch; we now fix up the existing unroll info.

Change-Id: Ie6d6be2094a70b58ff07b13e33e8fc91f74ab796
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/456468
Commit-Queue: John Stiles <johnstiles@google.com>
Commit-Queue: Brian Osman <brianosman@google.com>
Auto-Submit: John Stiles <johnstiles@google.com>
Reviewed-by: Brian Osman <brianosman@google.com>

src/sksl/SkSLInliner.cpp

@@ -217,6 +217,24 @@ public:
 
 }  // namespace
 
+const Variable* Inliner::RemapVariable(const Variable* variable,
+                                       const VariableRewriteMap* varMap) {
+    auto iter = varMap->find(variable);
+    if (iter == varMap->end()) {
+        SkDEBUGFAILF("rewrite map does not contain variable '%.*s'",
+                     (int)variable->name().size(), variable->name().data());
+        return variable;
+    }
+    Expression* expr = iter->second.get();
+    SkASSERT(expr);
+    if (!expr->is<VariableReference>()) {
+        SkDEBUGFAILF("rewrite map contains non-variable replacement for '%.*s'",
+                     (int)variable->name().size(), variable->name().data());
+        return variable;
+    }
+    return expr->as<VariableReference>().variable();
+}
+
 Inliner::ReturnComplexity Inliner::GetReturnComplexity(const FunctionDefinition& funcDef) {
     int returnsAtEndOfControlFlow = count_returns_at_end_of_control_flow(funcDef);
     CountReturnsWithLimit counter{funcDef, returnsAtEndOfControlFlow + 1};
@@ -489,10 +507,17 @@ std::unique_ptr<Statement> Inliner::inlineStatement(int line,
             // need to ensure initializer is evaluated first so that we've already remapped its
             // declarations by the time we evaluate test & next
             std::unique_ptr<Statement> initializer = stmt(f.initializer());
-            // We can't reuse the unroll info from the original for loop, because it uses a
-            // different induction variable. Ours is a clone.
+
+            std::unique_ptr<LoopUnrollInfo> unrollInfo;
+            if (f.unrollInfo()) {
+                // The for loop's unroll-info points to the Variable in the initializer as the
+                // index. This variable has been rewritten into a clone by the inliner, so we need
+                // to update the loop-unroll info to point to the clone.
+                unrollInfo = std::make_unique<LoopUnrollInfo>(*f.unrollInfo());
+                unrollInfo->fIndex = RemapVariable(unrollInfo->fIndex, varMap);
+            }
             return ForStatement::Make(*fContext, line, std::move(initializer), expr(f.test()),
-                                      expr(f.next()), stmt(f.statement()), /*unrollInfo=*/nullptr,
+                                      expr(f.next()), stmt(f.statement()), std::move(unrollInfo),
                                       SymbolTable::WrapIfBuiltin(f.symbols()));
         }
         case Statement::Kind::kIf: {

src/sksl/SkSLInliner.h

@@ -73,6 +73,13 @@ private:
                                                const Statement& statement,
                                                bool isBuiltinCode);
 
+    /**
+     * Searches the rewrite map for an rewritten Variable* for the passed-in one. Asserts if the
+     * rewrite map doesn't contain the variable, or contains a different type of expression.
+     */
+    static const Variable* RemapVariable(const Variable* variable,
+                                         const VariableRewriteMap* varMap);
+
     /** Determines if a given function has multiple and/or early returns. */
     static ReturnComplexity GetReturnComplexity(const FunctionDefinition& funcDef);
 

src/sksl/SkSLRehydrator.cpp

@@ -13,6 +13,7 @@
 #include "include/private/SkSLModifiers.h"
 #include "include/private/SkSLProgramElement.h"
 #include "include/private/SkSLStatement.h"
+#include "src/sksl/SkSLAnalysis.h"
 #include "src/sksl/ir/SkSLBinaryExpression.h"
 #include "src/sksl/ir/SkSLBreakStatement.h"
 #include "src/sksl/ir/SkSLConstructor.h"
@@ -340,9 +341,12 @@ std::unique_ptr<Statement> Rehydrator::statement() {
             std::unique_ptr<Expression> next = this->expression();
             std::unique_ptr<Statement> body = this->statement();
             std::shared_ptr<SymbolTable> symbols = this->symbolTable();
+            std::unique_ptr<LoopUnrollInfo> unrollInfo =
+                    Analysis::GetLoopUnrollInfo(/*line=*/-1, initializer.get(), test.get(),
+                                                next.get(), body.get(), /*errors=*/nullptr);
             return ForStatement::Make(fContext, /*line=*/-1, std::move(initializer),
                                       std::move(test), std::move(next), std::move(body),
-                                      /*unrollInfo=*/nullptr, std::move(symbols));
+                                      std::move(unrollInfo), std::move(symbols));
         }
         case Rehydrator::kIf_Command: {
             bool isStatic = this->readU8();

src/sksl/ir/SkSLForStatement.cpp

@@ -164,13 +164,7 @@ std::unique_ptr<Statement> ForStatement::Make(const Context& context, int line,
              is_vardecl_block_initializer(initializer.get()));
     SkASSERT(!test || test->type() == *context.fTypes.fBool);
     SkASSERT(!Analysis::DetectVarDeclarationWithoutScope(*statement));
-
-    // If the caller didn't provide us with unroll info, we can compute it here if needed.
-    if (!unrollInfo && context.fConfig->strictES2Mode()) {
-        unrollInfo = Analysis::GetLoopUnrollInfo(line, initializer.get(), test.get(),
-                                                 next.get(), statement.get(), /*errors=*/nullptr);
-        SkASSERT(unrollInfo);
-    }
+    SkASSERT(unrollInfo || !context.fConfig->strictES2Mode());
 
     return std::make_unique<ForStatement>(line, std::move(initializer), std::move(test),
                                           std::move(next), std::move(statement),