From 39141d961e55f56d540b3678f9f8fe94868f2652 Mon Sep 17 00:00:00 2001
From: Mikko Rasa <tdb@tdb.fi>
Date: Sun, 28 Feb 2021 13:19:25 +0200
Subject: [PATCH] Store functions definitions in Stage

---
 source/glsl/debug.cpp    |  3 +++
 source/glsl/generate.cpp | 26 ++++++++++++++++++++------
 source/glsl/generate.h   |  2 +-
 source/glsl/syntax.h     |  1 +
 source/glsl/visitor.cpp  |  7 +++++++
 source/glsl/visitor.h    |  1 +
 6 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/source/glsl/debug.cpp b/source/glsl/debug.cpp
index c29a60ce..b913f5b4 100644
--- a/source/glsl/debug.cpp
+++ b/source/glsl/debug.cpp
@@ -27,6 +27,9 @@ const std::string &DumpTree::apply(Stage &stage)
 			append(text);
 		}
 
+	for(std::map<string, FunctionDeclaration *>::const_iterator i=stage.functions.begin(); i!=stage.functions.end(); ++i)
+		append(format("Function: %%%d %s", get_label(*i->second), i->first));
+
 	last_branch();
 	stage.content.visit(*this);
 	return formatted;
diff --git a/source/glsl/generate.cpp b/source/glsl/generate.cpp
index 13b62aa3..ff863824 100644
--- a/source/glsl/generate.cpp
+++ b/source/glsl/generate.cpp
@@ -309,28 +309,42 @@ void VariableResolver::visit(InterfaceBlock &iface)
 }
 
 
+void FunctionResolver::apply(Stage &s)
+{
+	stage = &s;
+	s.content.visit(*this);
+}
+
 void FunctionResolver::visit(FunctionCall &call)
 {
-	map<string, vector<FunctionDeclaration *> >::iterator i = functions.find(call.name);
-	if(i!=functions.end())
-		call.declaration = i->second.back();
+	map<string, FunctionDeclaration *>::iterator i = stage->functions.find(call.name);
+	if(i!=stage->functions.end())
+		call.declaration = i->second;
 
 	TraversingVisitor::visit(call);
 }
 
 void FunctionResolver::visit(FunctionDeclaration &func)
 {
-	vector<FunctionDeclaration *> &decls = functions[func.name];
+	FunctionDeclaration *&stage_decl = stage->functions[func.name];
+	vector<FunctionDeclaration *> &decls = declarations[func.name];
 	if(func.definition==&func)
 	{
+		stage_decl = &func;
+
 		for(vector<FunctionDeclaration *>::iterator i=decls.begin(); i!=decls.end(); ++i)
 		{
 			(*i)->definition = func.definition;
 			(*i)->body.body.clear();
 		}
 	}
-	else if(!decls.empty() && decls.back()->definition)
-		func.definition = decls.back()->definition;
+	else
+	{
+		if(!stage_decl)
+			stage_decl = &func;
+
+		func.definition = stage_decl->definition;
+	}
 	decls.push_back(&func);
 
 	TraversingVisitor::visit(func);
diff --git a/source/glsl/generate.h b/source/glsl/generate.h
index dc90b6eb..4359a09f 100644
--- a/source/glsl/generate.h
+++ b/source/glsl/generate.h
@@ -93,7 +93,7 @@ class FunctionResolver: private TraversingVisitor
 {
 private:
 	Stage *stage;
-	std::map<std::string, std::vector<FunctionDeclaration *> > functions;
+	std::map<std::string, std::vector<FunctionDeclaration *> > declarations;
 
 public:
 	void apply(Stage &);
diff --git a/source/glsl/syntax.h b/source/glsl/syntax.h
index 55f47402..70ea50d1 100644
--- a/source/glsl/syntax.h
+++ b/source/glsl/syntax.h
@@ -410,6 +410,7 @@ struct Stage
 	Block content;
 	std::map<std::string, StructDeclaration *> types;
 	std::map<std::string, InterfaceBlock *> interface_blocks;
+	std::map<std::string, FunctionDeclaration *> functions;
 	std::map<std::string, unsigned> locations;
 	Features required_features;
 
diff --git a/source/glsl/visitor.cpp b/source/glsl/visitor.cpp
index fd25bf34..d8b18139 100644
--- a/source/glsl/visitor.cpp
+++ b/source/glsl/visitor.cpp
@@ -186,6 +186,13 @@ void NodeRemover::visit(InterfaceBlock &iface)
 	TraversingVisitor::visit(iface);
 }
 
+void NodeRemover::visit(FunctionDeclaration &func)
+{
+	if(to_remove->count(&func))
+		remove_from_map(stage->functions, func.name+func.signature, func);
+	TraversingVisitor::visit(func);
+}
+
 void NodeRemover::visit(Iteration &iter)
 {
 	if(to_remove->count(iter.init_statement.get()))
diff --git a/source/glsl/visitor.h b/source/glsl/visitor.h
index 4e48a1a1..83b310f6 100644
--- a/source/glsl/visitor.h
+++ b/source/glsl/visitor.h
@@ -106,6 +106,7 @@ private:
 	virtual void visit(StructDeclaration &);
 	virtual void visit(VariableDeclaration &);
 	virtual void visit(InterfaceBlock &);
+	virtual void visit(FunctionDeclaration &);
 	virtual void visit(Iteration &);
 };
 
-- 
2.43.0