1 files changed, 112 insertions, 0 deletions
diff --git a/src/3rdparty/angle/src/compiler/translator/RunAtTheEndOfShader.cpp b/src/3rdparty/angle/src/compiler/translator/RunAtTheEndOfShader.cpp
new file mode 100644
index 0000000000..3c4209c539
--- /dev/null
+++ b/src/3rdparty/angle/src/compiler/translator/RunAtTheEndOfShader.cpp
@@ -0,0 +1,112 @@
+//
+// Copyright (c) 2017 The ANGLE Project Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+//
+// RunAtTheEndOfShader.cpp: Add code to be run at the end of the shader. In case main() contains a
+// return statement, this is done by replacing the main() function with another function that calls
+// the old main, like this:
+//
+// void main() { body }
+// =>
+// void main0() { body }
+// void main()
+// {
+//     main0();
+//     codeToRun
+// }
+//
+// This way the code will get run even if the return statement inside main is executed.
+//
+
+#include "compiler/translator/RunAtTheEndOfShader.h"
+
+#include "compiler/translator/FindMain.h"
+#include "compiler/translator/IntermNode.h"
+#include "compiler/translator/IntermNode_util.h"
+#include "compiler/translator/IntermTraverse.h"
+#include "compiler/translator/SymbolTable.h"
+
+namespace sh
+{
+
+namespace
+{
+
+class ContainsReturnTraverser : public TIntermTraverser
+{
+  public:
+    ContainsReturnTraverser() : TIntermTraverser(true, false, false), mContainsReturn(false) {}
+
+    bool visitBranch(Visit visit, TIntermBranch *node) override
+    {
+        if (node->getFlowOp() == EOpReturn)
+        {
+            mContainsReturn = true;
+        }
+        return false;
+    }
+
+    bool containsReturn() { return mContainsReturn; }
+
+  private:
+    bool mContainsReturn;
+};
+
+bool ContainsReturn(TIntermNode *node)
+{
+    ContainsReturnTraverser traverser;
+    node->traverse(&traverser);
+    return traverser.containsReturn();
+}
+
+void WrapMainAndAppend(TIntermBlock *root,
+                       TIntermFunctionDefinition *main,
+                       TIntermNode *codeToRun,
+                       TSymbolTable *symbolTable)
+{
+    // Replace main() with main0() with the same body.
+    TSymbolUniqueId oldMainId(symbolTable);
+    std::stringstream oldMainName;
+    oldMainName << "main" << oldMainId.get();
+    TIntermFunctionDefinition *oldMain = CreateInternalFunctionDefinitionNode(
+        TType(EbtVoid), oldMainName.str().c_str(), main->getBody(), oldMainId);
+
+    bool replaced = root->replaceChildNode(main, oldMain);
+    ASSERT(replaced);
+
+    // void main()
+    TIntermFunctionPrototype *newMainProto = new TIntermFunctionPrototype(
+        TType(EbtVoid), main->getFunctionPrototype()->getFunctionSymbolInfo()->getId());
+    newMainProto->getFunctionSymbolInfo()->setName("main");
+
+    // {
+    //     main0();
+    //     codeToRun
+    // }
+    TIntermBlock *newMainBody     = new TIntermBlock();
+    TIntermAggregate *oldMainCall = CreateInternalFunctionCallNode(
+        TType(EbtVoid), oldMainName.str().c_str(), oldMainId, new TIntermSequence());
+    newMainBody->appendStatement(oldMainCall);
+    newMainBody->appendStatement(codeToRun);
+
+    // Add the new main() to the root node.
+    TIntermFunctionDefinition *newMain = new TIntermFunctionDefinition(newMainProto, newMainBody);
+    root->appendStatement(newMain);
+}
+
+}  // anonymous namespace
+
+void RunAtTheEndOfShader(TIntermBlock *root, TIntermNode *codeToRun, TSymbolTable *symbolTable)
+{
+    TIntermFunctionDefinition *main = FindMain(root);
+    if (!ContainsReturn(main))
+    {
+        main->getBody()->appendStatement(codeToRun);
+        return;
+    }
+
+    WrapMainAndAppend(root, main, codeToRun, symbolTable);
+}
+
+}  // namespace sh