1 files changed, 50 insertions, 11 deletions

diff --git a/src/bin/main.c b/src/bin/main.c
index 02b3d0a..481874c 100644
--- a/src/bin/main.c
+++ b/src/bin/main.c
@@ -23,6 +23,9 @@
 #include <errno.h>
 #include "aplwc_repl.h"
 
+static uint64_t alloc_count;
+static long memdbg;
+
 static void *
 my_alloc(size_t min, size_t max, size_t *actual)
 {
@@ -61,6 +64,16 @@ my_dealloc(void *p)
 	free(p);
 }
 
+static int
+hook_alloc(void)
+{
+	alloc_count++;
+	if((memdbg >= 0) && (alloc_count > ((uint64_t)memdbg)))
+		return 0;
+	else
+		return 1;
+}
+
 int
 main(int argc, char *argv[])
 {
@@ -84,9 +97,12 @@ main(int argc, char *argv[])
 	exit_code = 0;
 	line_number = 0;
 
+	memdbg = getenv_long("APLWC_ALLOC_FAIL_AFTER", -2);
+
 	memset(&fns, 0, sizeof(fns));
 	fns.alloc = my_alloc;
 	fns.dealloc = my_dealloc;
+	fns.hook_alloc = hook_alloc;
 
 	while((opt = getopt(argc, argv, "x")) != -1){
 		switch(opt){
@@ -96,7 +112,8 @@ main(int argc, char *argv[])
 		default:
 usage:
 			fprintf(stderr, "Usage: %s [-x] [file]\n", argv[0]);
-			return 1;
+			exit_code = 1;
+			goto end;
 		}
 	}
 
@@ -105,7 +122,8 @@ usage:
 		f = fopen(argv[optind], "r");
 		if(f == NULL){
 			fprintf(stderr, "Could not open file: %s: %s\n", argv[optind], strerror(errno));
-			return 1;
+			exit_code = 1;
+			goto end;
 		}
 	}else if(optind == argc)
 		f = stdin;
@@ -115,14 +133,16 @@ usage:
 	aplwc = aplwc_new(fns, &ret);
 	if(ret){
 		fprintf(stderr, "Couldn't create new aplwc instance: %s\n", aplwc_errstr_rc(ret));
-		return 1;
+		exit_code = 1;
+		goto end;
 	}
 
 	if(input_file){
 		if(fseek(f, 0, SEEK_END)){
 file_size:
 			fprintf(stderr, "Failed to determine file size\n");
-			return 1;
+			exit_code = 1;
+			goto end;
 		}
 		tmp = ftell(f);
 		if(tmp < 0)
@@ -132,11 +152,13 @@ file_size:
 		input = malloc(input_size);
 		if(input == NULL){
 			fprintf(stderr, "Failed to allocate space for file contents\n");
-			return 1;
+			exit_code = 1;
+			goto end;
 		}
 		if(fread(input, 1, input_size, f) < input_size){
 			fprintf(stderr, "Failed to read file contents\n");
-			return 1;
+			exit_code = 1;
+			goto end;
 		}
 		fclose(f);
 
@@ -156,14 +178,15 @@ file_size:
 		ret = aplwc_load_module(aplwc, "Main", src, src_size);
 		if(ret){
 			fprintf(stderr, "Couldn't load Main module from file %s: %s\n", argv[optind], aplwc_errstr_rc(ret));
-			return 1;
+			exit_code = 1;
+			goto end;
 		}
 		free(input);
 		input = NULL;
 	}
 
 	if(xflag)
-		goto exit;
+		goto end;
 
 	for(;;){
 		if(input)
@@ -183,10 +206,11 @@ file_size:
 			aplwc_read(aplwc, input, &line);
 			if(!aplwc_valid_line(&line)){
 line_error:
+				ret = aplwc_last_error(aplwc);
 				fprintf(stderr, "error on line %d: %s: %s\n", line_number, aplwc_errstr(aplwc), input);
 				if(input_file){
 					exit_code = 1;
-					goto exit;
+					goto end;
 				}
 				continue;
 			}
@@ -197,14 +221,29 @@ line_error:
 		}
 	}
 
-exit:
+end:
 	if(input)
 		free(input);
 	if(aplwc_valid_line(&line))
 		aplwc_free_line(&line);
 	if(aplwc_valid_value(&value))
 		aplwc_free_value(&value);
+	if(aplwc)
+		aplwc_destroy(aplwc);
+
+	if(memdbg == -1){
+		f = fdopen(3, "w");
+		if(f == NULL){
+			fprintf(stderr, "Failed to open fd 3\n");
+			exit_code = 1;
+		}else{
+			fprintf(f, "%ld\n", alloc_count);
+			fclose(f);
+		}
+	}
+
+	if((exit_code != 0) && (memdbg >= 0) && (ret == APLWC_OUT_OF_MEMORY))
+		exit_code = 0;
 
-	aplwc_destroy(aplwc);
 	return exit_code;
 }