1 files changed, 26 insertions, 5 deletions

diff --git a/src/go-flex.skl b/src/go-flex.skl
index 431cde8..e9ce1f1 100755
--- a/src/go-flex.skl
+++ b/src/go-flex.skl
@@ -194,9 +194,11 @@ m4_define([[YY_SC_TO_UI]], [[((YY_CHAR)($1))]])
 /* Action number for EOF rule of a given start state. */
 m4_define([[YY_STATE_EOF]], [[YY_END_OF_BUFFER + $1 + 1]])
 
-/* The state buf must be large enough to hold one state per character in the main buffer.
+/* The state buf must be large enough to hold one state per character in the main buffer,
+ * plus the start state, plus the two end-of-buffer byte states.
  */
-m4_define([[YY_STATE_BUF_SIZE]], [[((flexInputBufferSize + 2) * sizeof(yyStateType))]])
+m4_define([[YY_STATE_BUF_EXTRA_SPACE]], [[3]])
+m4_define([[YY_STATE_BUF_SIZE]], [[(YY_BUF_SIZE + YY_STATE_BUF_EXTRA_SPACE)]])
 
 const bool FLEX_DEBUG = m4_ifdef([[M4_MODE_DEBUG]], [[true]], [[false]]);
 
@@ -1444,6 +1446,8 @@ bool yyatbol(FlexLexer *yyscanner) {
  */
 void yy_switch_to_buffer(yybuffer newBuffer, FlexLexer *yyscanner)
 {
+	size_t newSize = 0;
+	yyStateType *newStateBuf = 0;
 
 	/* TODO. We should be able to replace this entire function body
 	 * with
@@ -1463,6 +1467,23 @@ void yy_switch_to_buffer(yybuffer newBuffer, FlexLexer *yyscanner)
 
 	yyscanner->yyBufferStack[yyscanner->yyBufferStackTop] = newBuffer;
 	yy_load_buffer_state(yyscanner);
+	
+m4_ifdef( [[M4_MODE_USES_REJECT]], [[
+	/* Ensure the reject state buffer is large enough.
+	 */
+	if ( yyscanner->yyStateBufMax < (yy_current_buffer(yyscanner)->yyInputBufSize + YY_STATE_BUF_EXTRA_SPACE) ) {
+		newSize = yy_current_buffer(yyscanner)->yyInputBufSize + YY_STATE_BUF_EXTRA_SPACE;
+		newStateBuf = (yyStateType *)yyrealloc( yyscanner->yyStateBuf, (new_size * sizeof(yyStateType)), yyscanner );
+		
+		if ( newStateBuf == NULL ) {
+			YY_FATAL_ERROR( "out of dynamic memory in yylex()" );
+		}
+		else {
+			yyscanner->yyStateBuf = newStateBuf;
+			yyscanner->yyStateBufMax = newSize;
+		}
+	}
+]] )
 
 	/* We don't actually know whether we did this switch during
 	 * EOF (yywrap()) processing, but the only time this flag
@@ -1861,14 +1882,14 @@ m4_ifdef([[M4_MODE_USES_REJECT]],
 			}
 		}
 		
-		if ( ! yyscanner->yyStateBuf ) {
+		if ( yyscanner->yyStateBuf == NULL ) {
 			yypanic( "out of dynamic memory in yylex()", yyscanner );
 		}
 		else if ( yy_current_buffer(yyscanner) == NULL ) {
 			yyscanner->yyStateBufMax = YY_STATE_BUF_SIZE;
 		}
-		if (yyscanner->yyStateBuf == NULL) {
-			yypanic("out of dynamic memory in yylex()", yyscanner);
+		else {
+			yyscanner->yyStateBufMax = (yy_current_buffer(yyscanner)->yyInputBufSize + YY_STATE_BUF_EXTRA_SPACE);
 		}
 ]])