From f1c6ef565577ac322693255d764eea1f2cab9e77 Mon Sep 17 00:00:00 2001
From: Mikko Rasa <tdb@tdb.fi>
Date: Mon, 14 Nov 2016 17:46:35 +0200
Subject: [PATCH] Remove unused functions from the shader source

---
 source/programcompiler.cpp | 49 ++++++++++++++++++++++++++++++++++++++
 source/programcompiler.h   | 20 ++++++++++++++++
 source/programsyntax.cpp   |  1 +
 source/programsyntax.h     |  2 ++
 4 files changed, 72 insertions(+)

diff --git a/source/programcompiler.cpp b/source/programcompiler.cpp
index 9b326675..937639a4 100644
--- a/source/programcompiler.cpp
+++ b/source/programcompiler.cpp
@@ -159,6 +159,7 @@ void ProgramCompiler::generate(Stage &stage)
 {
 	inject_block(stage.content, module->shared.content);
 
+	apply<FunctionResolver>(stage);
 	apply<VariableResolver>(stage);
 	apply<InterfaceGenerator>(stage);
 	apply<VariableResolver>(stage);
@@ -171,6 +172,8 @@ bool ProgramCompiler::optimize(Stage &stage)
 	apply<VariableResolver>(stage);
 
 	set<Node *> unused = apply<UnusedVariableLocator>(stage);
+	set<Node *> unused2 = apply<UnusedFunctionLocator>(stage);
+	unused.insert(unused2.begin(), unused2.end());
 	apply<NodeRemover>(stage, unused);
 
 	return !unused.empty();
@@ -616,6 +619,34 @@ void ProgramCompiler::VariableResolver::visit(InterfaceBlock &iface)
 }
 
 
+void ProgramCompiler::FunctionResolver::visit(FunctionCall &call)
+{
+	map<string, vector<FunctionDeclaration *> >::iterator i = functions.find(call.name);
+	if(i!=functions.end())
+		call.declaration = i->second.back();
+
+	TraversingVisitor::visit(call);
+}
+
+void ProgramCompiler::FunctionResolver::visit(FunctionDeclaration &func)
+{
+	vector<FunctionDeclaration *> &decls = functions[func.name];
+	if(func.definition)
+	{
+		for(vector<FunctionDeclaration *>::iterator i=decls.begin(); i!=decls.end(); ++i)
+			(*i)->definition = func.definition;
+		decls.clear();
+		decls.push_back(&func);
+	}
+	else if(!decls.empty() && decls.back()->definition)
+		func.definition = decls.back()->definition;
+	else
+		decls.push_back(&func);
+
+	TraversingVisitor::visit(func);
+}
+
+
 ProgramCompiler::InterfaceGenerator::InterfaceGenerator():
 	scope_level(0),
 	remove_node(false)
@@ -1198,6 +1229,24 @@ void ProgramCompiler::UnusedVariableLocator::visit(Iteration &iter)
 }
 
 
+void ProgramCompiler::UnusedFunctionLocator::visit(FunctionCall &call)
+{
+	TraversingVisitor::visit(call);
+
+	unused_nodes.erase(call.declaration);
+	if(call.declaration && call.declaration->definition!=call.declaration)
+		used_definitions.insert(call.declaration->definition);
+}
+
+void ProgramCompiler::UnusedFunctionLocator::visit(FunctionDeclaration &func)
+{
+	TraversingVisitor::visit(func);
+
+	if(func.name!="main" && !used_definitions.count(&func))
+		unused_nodes.insert(&func);
+}
+
+
 ProgramCompiler::NodeRemover::NodeRemover(const set<Node *> &r):
 	to_remove(r)
 { }
diff --git a/source/programcompiler.h b/source/programcompiler.h
index f585ba92..8681d363 100644
--- a/source/programcompiler.h
+++ b/source/programcompiler.h
@@ -108,6 +108,14 @@ private:
 		virtual void visit(ProgramSyntax::InterfaceBlock &);
 	};
 
+	struct FunctionResolver: Visitor
+	{
+		std::map<std::string, std::vector<ProgramSyntax::FunctionDeclaration *> > functions;
+
+		virtual void visit(ProgramSyntax::FunctionCall &);
+		virtual void visit(ProgramSyntax::FunctionDeclaration &);
+	};
+
 	struct InterfaceGenerator: Visitor
 	{
 		std::string in_prefix;
@@ -207,6 +215,18 @@ private:
 		virtual void visit(ProgramSyntax::Iteration &);
 	};
 
+	struct UnusedFunctionLocator: Visitor
+	{
+		typedef std::set<ProgramSyntax::Node *> ResultType;
+
+		std::set<ProgramSyntax::Node *> unused_nodes;
+		std::set<ProgramSyntax::FunctionDeclaration *> used_definitions;
+
+		const ResultType &get_result() const { return unused_nodes; }
+		virtual void visit(ProgramSyntax::FunctionCall &);
+		virtual void visit(ProgramSyntax::FunctionDeclaration &);
+	};
+
 	struct NodeRemover: Visitor
 	{
 		std::set<ProgramSyntax::Node *> to_remove;
diff --git a/source/programsyntax.cpp b/source/programsyntax.cpp
index f87488dc..098130f8 100644
--- a/source/programsyntax.cpp
+++ b/source/programsyntax.cpp
@@ -72,6 +72,7 @@ void Assignment::visit(NodeVisitor &visitor)
 
 
 FunctionCall::FunctionCall():
+	declaration(0),
 	constructor(false)
 { }
 
diff --git a/source/programsyntax.h b/source/programsyntax.h
index 9545760c..08968332 100644
--- a/source/programsyntax.h
+++ b/source/programsyntax.h
@@ -52,6 +52,7 @@ public:
 
 struct StructDeclaration;
 struct VariableDeclaration;
+struct FunctionDeclaration;
 
 struct Block: Node
 {
@@ -145,6 +146,7 @@ struct Assignment: BinaryExpression
 struct FunctionCall: Expression
 {
 	std::string name;
+	FunctionDeclaration *declaration;
 	bool constructor;
 	std::vector<NodePtr<Expression> > arguments;
 
-- 
2.45.2