Create VariableExpression struct for referencing inlined variables.

In this CL, the result expression has been updated to use
VariableExpression; followup CLs will update the VariableRewriteMap to
leverage it as well.

This is intended to allow the inliner to inline simple expressions (such
as literals or swizzles) directly if they are not written to, instead of
making an extraneous copy.

Change-Id: I050057d8c3e940e5e44c22fde2f4bc37bb4c6754
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/318576
Auto-Submit: John Stiles <johnstiles@google.com>
Commit-Queue: Ethan Nicholas <ethannicholas@google.com>
Reviewed-by: Ethan Nicholas <ethannicholas@google.com>
diff --git a/src/sksl/SkSLInliner.cpp b/src/sksl/SkSLInliner.cpp
index 023c1a0..2eb7758 100644
--- a/src/sksl/SkSLInliner.cpp
+++ b/src/sksl/SkSLInliner.cpp
@@ -390,12 +390,12 @@
 std::unique_ptr<Statement> Inliner::inlineStatement(int offset,
                                                     VariableRewriteMap* varMap,
                                                     SymbolTable* symbolTableForStatement,
-                                                    const Variable* returnVar,
+                                                    const VariableExpression& resultExpr,
                                                     bool haveEarlyReturns,
                                                     const Statement& statement) {
     auto stmt = [&](const std::unique_ptr<Statement>& s) -> std::unique_ptr<Statement> {
         if (s) {
-            return this->inlineStatement(offset, varMap, symbolTableForStatement, returnVar,
+            return this->inlineStatement(offset, varMap, symbolTableForStatement, resultExpr,
                                          haveEarlyReturns, *s);
         }
         return nullptr;
@@ -451,14 +451,13 @@
         case Statement::Kind::kReturn: {
             const ReturnStatement& r = statement.as<ReturnStatement>();
             if (r.fExpression) {
-                auto assignment = std::make_unique<ExpressionStatement>(
-                        std::make_unique<BinaryExpression>(
-                            offset,
-                            std::make_unique<VariableReference>(offset, *returnVar,
-                                                                VariableReference::kWrite_RefKind),
-                            Token::Kind::TK_EQ,
-                            expr(r.fExpression),
-                            &returnVar->type()));
+                auto assignment =
+                        std::make_unique<ExpressionStatement>(std::make_unique<BinaryExpression>(
+                                offset,
+                                resultExpr.cloneWithRefKind(VariableReference::kWrite_RefKind),
+                                Token::Kind::TK_EQ,
+                                expr(r.fExpression),
+                                &resultExpr.type()));
                 if (haveEarlyReturns) {
                     std::vector<std::unique_ptr<Statement>> block;
                     block.push_back(std::move(assignment));
@@ -613,11 +612,13 @@
     };
 
     // Create a variable to hold the result in the extra statements (excepting void).
-    const Variable* resultVar = nullptr;
+    VariableExpression resultExpr;
     if (function.fDeclaration.fReturnType != *fContext->fVoid_Type) {
         std::unique_ptr<Expression> noInitialValue;
-        resultVar = makeInlineVar(String(function.fDeclaration.fName),
-                                  &function.fDeclaration.fReturnType, Modifiers{}, &noInitialValue);
+        const Variable* var = makeInlineVar(String(function.fDeclaration.fName),
+                                            &function.fDeclaration.fReturnType,
+                                            Modifiers{}, &noInitialValue);
+        resultExpr.fInnerVariable = std::make_unique<VariableReference>(offset, *var);
     }
 
     // Create variables in the extra statements to hold the arguments, and assign the arguments to
@@ -645,7 +646,7 @@
     inlineBlock->fStatements.reserve(body.fStatements.size());
     for (const std::unique_ptr<Statement>& stmt : body.fStatements) {
         inlineBlock->fStatements.push_back(this->inlineStatement(
-                offset, &varMap, symbolTableForCall, resultVar, hasEarlyReturn, *stmt));
+                offset, &varMap, symbolTableForCall, resultExpr, hasEarlyReturn, *stmt));
     }
     if (hasEarlyReturn) {
         // Since we output to backends that don't have a goto statement (which would normally be
@@ -685,7 +686,10 @@
 
     if (function.fDeclaration.fReturnType != *fContext->fVoid_Type) {
         // Return a reference to the result variable as our replacement expression.
-        inlinedCall.fReplacementExpr = std::make_unique<VariableReference>(offset, *resultVar);
+        resultExpr.fInnerVariable->setRefKind(VariableReference::kRead_RefKind);
+        inlinedCall.fReplacementExpr = resultExpr.fOuterExpression
+                                               ? std::move(resultExpr.fOuterExpression)
+                                               : std::move(resultExpr.fInnerVariable);
     } else {
         // It's a void function, so it doesn't actually result in anything, but we have to return
         // something non-null as a standin.
diff --git a/src/sksl/SkSLInliner.h b/src/sksl/SkSLInliner.h
index 484f731..819f46b 100644
--- a/src/sksl/SkSLInliner.h
+++ b/src/sksl/SkSLInliner.h
@@ -12,6 +12,7 @@
 #include <unordered_map>
 
 #include "src/sksl/ir/SkSLProgram.h"
+#include "src/sksl/ir/SkSLVariableReference.h"
 
 namespace SkSL {
 
@@ -56,6 +57,23 @@
     bool analyze(Program& program);
 
 private:
+    struct VariableExpression {
+        std::unique_ptr<Expression> cloneWithRefKind(VariableReference::RefKind refKind) const {
+            fInnerVariable->setRefKind(refKind);
+            return fOuterExpression ? fOuterExpression->clone() : fInnerVariable->clone();
+        }
+        const Type& type() const {
+            return fOuterExpression ? fOuterExpression->type() : fInnerVariable->type();
+        }
+
+        // All VariableExpressions are expected to bottom out at a variable.
+        std::unique_ptr<VariableReference> fInnerVariable;
+
+        // fOuterExpression applies an expression to the fInnerVariable (e.g. a Swizzle). This can
+        // be null if the variable needs no adjustment.
+        std::unique_ptr<Expression> fOuterExpression;
+    };
+
     using VariableRewriteMap = std::unordered_map<const Variable*, const Variable*>;
 
     String uniqueNameForInlineVar(const String& baseName, SymbolTable* symbolTable);
@@ -66,7 +84,7 @@
     std::unique_ptr<Statement> inlineStatement(int offset,
                                                VariableRewriteMap* varMap,
                                                SymbolTable* symbolTableForStatement,
-                                               const Variable* returnVar,
+                                               const VariableExpression& resultExpr,
                                                bool haveEarlyReturns,
                                                const Statement& statement);