diff --git a/src/Shader/PixelProgram.cpp b/src/Shader/PixelProgram.cpp
index 632de99..1644c06 100644
--- a/src/Shader/PixelProgram.cpp
+++ b/src/Shader/PixelProgram.cpp
@@ -37,7 +37,7 @@
 		ifFalseBlock.resize(limits.ifs);
 		loopRepTestBlock.resize(limits.loops);
 		loopRepEndBlock.resize(limits.loops);
-		labelBlock.resize(limits.functions);
+		labelBlock.resize(limits.maxLabel + 1);
 		isConditionalIf.resize(limits.ifs);
 
 		loopDepth = -1;
diff --git a/src/Shader/Shader.cpp b/src/Shader/Shader.cpp
index 3528b1e..fe0c0ac 100644
--- a/src/Shader/Shader.cpp
+++ b/src/Shader/Shader.cpp
@@ -1957,32 +1957,47 @@
 			bool reachable; // Is this function reachable?
 		};
 
-		// Begin by doing a pass over the instructions to identify all the
-		// labels that denote the start of a function.
 		std::unordered_map<FunctionID, FunctionInfo> functions;
 
-		for(auto &inst : instruction)
-		{
-			if (inst->isCall())
-			{
-				FunctionID id = inst->dst.label;
-				ASSERT(id != MAIN_ID); // If this fires, we're going to have to represent main with something else.
-				functions[id] = FunctionInfo();
-			}
-		}
+		uint32_t maxLabel = 0; // Highest label found
 
 		// Add a definition for the main entry point.
 		// This starts at the beginning of the instructions and does not have
 		// its own label.
 		functions[MAIN_ID] = FunctionInfo();
 		functions[MAIN_ID].reachable = true;
+
+		// Begin by doing a pass over the instructions to identify all the
+		// functions. These start with a label and end with a ret. Note that
+		// functions can have labels within them.
 		FunctionID currentFunc = MAIN_ID;
+		for(auto &inst : instruction)
+		{
+			switch (inst->opcode)
+			{
+				case OPCODE_LABEL:
+					if (currentFunc == INVALID_ID)
+					{
+						// Start of a function.
+						FunctionID id = inst->dst.label;
+						ASSERT(id != MAIN_ID); // If this fires, we're going to have to represent main with something else.
+						functions[id] = FunctionInfo();
+					}
+					break;
+				case OPCODE_RET:
+					currentFunc = INVALID_ID;
+					break;
+				default:
+					break;
+			}
+		}
 
 		// Limits for the currently analyzed function.
 		FunctionLimits currentLimits;
 
 		// Now loop over the instructions gathering the limits of each of the
 		// functions.
+		currentFunc = MAIN_ID;
 		for(size_t i = 0; i < instruction.size(); i++)
 		{
 			const auto& inst = instruction[i];
@@ -1990,15 +2005,14 @@
 			{
 				case OPCODE_LABEL:
 				{
-					FunctionID id = inst->dst.label;
-					auto it = functions.find(id);
-					if (it == functions.end())
+					maxLabel = std::max(maxLabel, inst->dst.label);
+					if (currentFunc == INVALID_ID)
 					{
-						// Label does not start a new function.
-						continue;
+						// Start of a function.
+						FunctionID id = inst->dst.label;
+						ASSERT(functions.find(id) != functions.end()); // Sanity check
+						currentFunc = id;
 					}
-					ASSERT(currentFunc == INVALID_ID); // Functions overlap
-					currentFunc = id;
 					break;
 				}
 				case OPCODE_CALL:
@@ -2006,6 +2020,7 @@
 				{
 					ASSERT(currentFunc != INVALID_ID);
 					FunctionID id = inst->dst.label;
+					ASSERT(functions.find(id) != functions.end());
 					functions[currentFunc].calls.emplace(id);
 					functions[id].reachable = true;
 					break;
@@ -2067,6 +2082,9 @@
 		// stripped in earlier stages). Unreachable functions may be code
 		// generated, but their own limits are not considered below, potentially
 		// causing OOB indexing in later stages.
+		// If we ever find cases where there are unreachable functions, we can
+		// replace this assert with NO-OPing or stripping out the dead
+		// functions.
 		for (auto it : functions) { ASSERT(it.second.reachable); }
 
 		// We have now gathered all the information about each of the functions
@@ -2097,6 +2115,6 @@
 		};
 
 		limits = traverse(MAIN_ID);
-		limits.functions = (uint32_t)functions.size();
+		limits.maxLabel = maxLabel;
 	}
 }
diff --git a/src/Shader/Shader.hpp b/src/Shader/Shader.hpp
index 64a86cd..b9d0188 100644
--- a/src/Shader/Shader.hpp
+++ b/src/Shader/Shader.hpp
@@ -561,7 +561,7 @@
 			uint32_t loops = 0; // maximum nested loop and reps.
 			uint32_t ifs = 0; // maximum nested if statements.
 			uint32_t stack = 0; // maximum call depth.
-			uint32_t functions = 0; // total number of functions.
+			uint32_t maxLabel = 0; // highest label in use.
 		};
 
 		Shader();
diff --git a/src/Shader/VertexProgram.cpp b/src/Shader/VertexProgram.cpp
index 25c1b0e..67307ef 100644
--- a/src/Shader/VertexProgram.cpp
+++ b/src/Shader/VertexProgram.cpp
@@ -36,7 +36,7 @@
 		ifFalseBlock.resize(limits.ifs);
 		loopRepTestBlock.resize(limits.loops);
 		loopRepEndBlock.resize(limits.loops);
-		labelBlock.resize(limits.functions);
+		labelBlock.resize(limits.maxLabel + 1);
 		isConditionalIf.resize(limits.ifs);
 
 		loopDepth = -1;